何愷明新作再戰(zhàn)AI生成：入職MIT后首次帶隊，奧賽雙料金牌得主鄧明揚參與

夢晨 2024-06-23 12:03:11 來源：量子位

夢晨 發(fā)自 凹非寺
量子位 | 公眾號 QbitAI

何愷明入職MIT副教授后，首次帶隊的新作來了！

讓自回歸模型拋棄矢量量化，使用連續(xù)值生成圖像。并借鑒擴散模型的思想，提出Diffusion Loss。

他加入MIT后，此前也參與過另外幾篇CV方向的論文，不過都是和MIT教授Wojciech Matusik團隊等合作的。

這次何愷明自己帶隊，參與者中還出現(xiàn)一個熟悉的名字：

鄧明揚，IMO、IOI雙料奧賽金牌得主，在競賽圈人稱“乖神”。

目前鄧明揚MIT本科在讀，按入學時間推算現(xiàn)在剛好大四，所以也有不少網(wǎng)友猜測他如果繼續(xù)在MIT讀博可能會加入何愷明團隊。

接下來具體介紹一下，這篇論文研究了什么。

借鑒擴散模型，大改自回歸生成

傳統(tǒng)觀點認為，圖像生成的自回歸模型通常伴隨著矢量量化（Vector Quantization），比如DALL·E一代就使用了經(jīng)典的VQ-VAE方法。

但團隊觀察到，自回歸生成的本質(zhì)是根據(jù)先前的值預測下一個token，這其實與值是離散還是連續(xù)沒啥必然聯(lián)系啊。

關(guān)鍵是要對token的概率分布進行建模，只要該概率分布可以通過損失函數(shù)來測量并用于從中抽取樣本就行。

并且從另一個方面來看，矢量量化方法還會帶來一系列麻煩：

• 需要一個離散的token詞表，需要精心設(shè)計量化的目標函數(shù)，訓練困難，對梯度近似策略很敏感
• 量化誤差會帶來信息損失，導致還原圖像質(zhì)量打折
• 離散token適合建模分類分布，有表達能力上的局限

那么有什么更好的替代方法？

何愷明團隊選擇在損失函數(shù)上動刀，借鑒近年大火的擴散模型的思想，提出Diffusion Loss，消除了離散tokenizer的必要性。

如此一來，在連續(xù)值空間中應(yīng)用自回歸模型生成圖像就可行了。

具體來說，它讓自回歸模型輸出一個潛變量z作為條件，去訓練一個小型的去噪MLP網(wǎng)絡(luò)。

通過反向擴散過程，這個小網(wǎng)絡(luò)就學會了如何根據(jù)z去采樣生成連續(xù)值的token x。擴散的過程天然能建模任意復雜的分布，所以沒有類別分布的局限。

這個去噪網(wǎng)絡(luò)和自回歸模型是端到端聯(lián)合訓練的，鏈式法則直接把損失傳給自回歸模型，使其學會輸出最佳的條件z。

這篇工作的另一個亮點，是各種自回歸模型的變體都適用。它統(tǒng)一了標準的自回歸AR、隨機順序的AR、以及何愷明擅長的掩碼方法。

其中掩碼自回歸（MAR）模型，可以在任意隨機位置同時預測多個token，同時還能和擴散損失完美配合。

在這個統(tǒng)一的框架下，所有變體要么逐個token預測，要么并行預測一批token，但本質(zhì)上都是在已知token的基礎(chǔ)上去預測未知token，都是廣義的自回歸模型，所以擴散損失都能適用。

通過消除矢量量化，團隊訓練的圖像生成模型獲得了強大的結(jié)果，同時享受序列建模的速度優(yōu)勢。

論文在AR、MAR的各種變體上做了大量實驗，結(jié)果表明擴散損失比交叉熵損失穩(wěn)定帶來2-3倍的提升。

與其他領(lǐng)先模型一比也毫不遜色，小模型都能做到1.98的FID分數(shù)，大模型更是創(chuàng)下了1.55的SOTA。

而且它生成256×256圖像速度也很快，不到0.3秒一張。這得益于自回歸生成本來就很快，比擴散模型少采樣很多步，再加上去噪網(wǎng)絡(luò)又很小。

最后總結(jié)一下，這項工作通過自回歸建模token間的相關(guān)性，再搭配擴散過程對每個token的分布進行建模。

這也有別于普通的潛空間擴散模型中用單個大擴散模型對所有token的聯(lián)合分布建模，而是做局部擴散，在效果、速度和靈活性上都展現(xiàn)出了巨大的潛力。

當然，這個方法還有進一步探索的空間，團隊提出，目前在在某些復雜的幾何圖形理解任務(wù)上還有待提高。

何愷明團隊都有誰

最后再來介紹一下即將或可能加入何愷明課題組的團隊成員。。

Tianhong LI（黎天鴻），清華姚班校友，MIT博士生在讀，將于2024年9月加入何愷明的課題組，擔任博士后。

Mingyang Deng（鄧明揚），MIT本科數(shù)學和計算機科學專業(yè)在讀。

他在高一獲得IMO金牌，高三獲得IOI金牌，是競賽圈為數(shù)不多的雙料金牌得主，也是IOI歷史上第三位滿分選手。

目前鄧明揚的研究重點是機器學習，特別是理解和推進生成式基礎(chǔ)模型，包括擴散模型和大型語言模型。

不過他的個人主頁上還沒有透露下一步計劃。

One More Thing

何愷明當初在MIT的求職演講備受關(guān)注，其中提到未來工作方向會是AI for Science，還引起圈內(nèi)一陣熱議。

現(xiàn)在，何愷明在AI4S方向的參與的首篇論文也來了：強化學習+量子物理學方向。

把Transformer模型用在了動態(tài)異構(gòu)量子資源調(diào)度問題上，利用自注意力機制處理量子比特對的序列信息。并在概率性環(huán)境中訓練強化學習代理，提供動態(tài)實時調(diào)度指導，最終顯著提升了量子系統(tǒng)性能，比基于規(guī)則的方法提高了3倍以上。

這樣一來，何愷明在自己的成名領(lǐng)域CV和探索新領(lǐng)域AI4S上都沒耽誤，兩開花，兩開花。

論文：
https://arxiv.org/abs/2406.11838

參考鏈接：
[1]https://www.tianhongli.me
[2]https://lambertae.github.io
[3]https://arxiv.org/abs/2405.16380