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

夢(mèng)晨 2024-06-23 12:03:11 來(lái)源：量子位

夢(mèng)晨 發(fā)自 凹非寺
量子位 | 公眾號(hào) QbitAI

何愷明入職MIT副教授后，首次帶隊(duì)的新作來(lái)了！

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

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

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

鄧明揚(yáng)，IMO、IOI雙料奧賽金牌得主，在競(jìng)賽圈人稱“乖神”。

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

接下來(lái)具體介紹一下，這篇論文研究了什么。

借鑒擴(kuò)散模型，大改自回歸生成

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

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

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

并且從另一個(gè)方面來(lái)看，矢量量化方法還會(huì)帶來(lái)一系列麻煩：

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

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

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

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

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

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

這個(gè)去噪網(wǎng)絡(luò)和自回歸模型是端到端聯(lián)合訓(xùn)練的，鏈?zhǔn)椒▌t直接把損失傳給自回歸模型，使其學(xué)會(huì)輸出最佳的條件z。

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

其中掩碼自回歸（MAR）模型，可以在任意隨機(jī)位置同時(shí)預(yù)測(cè)多個(gè)token，同時(shí)還能和擴(kuò)散損失完美配合。

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

通過(guò)消除矢量量化，團(tuán)隊(duì)訓(xùn)練的圖像生成模型獲得了強(qiáng)大的結(jié)果，同時(shí)享受序列建模的速度優(yōu)勢(shì)。

論文在AR、MAR的各種變體上做了大量實(shí)驗(yàn)，結(jié)果表明擴(kuò)散損失比交叉熵?fù)p失穩(wěn)定帶來(lái)2-3倍的提升。

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

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

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

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

當(dāng)然，這個(gè)方法還有進(jìn)一步探索的空間，團(tuán)隊(duì)提出，目前在在某些復(fù)雜的幾何圖形理解任務(wù)上還有待提高。

何愷明團(tuán)隊(duì)都有誰(shuí)

最后再來(lái)介紹一下即將或可能加入何愷明課題組的團(tuán)隊(duì)成員。。

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

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

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

目前鄧明揚(yáng)的研究重點(diǎn)是機(jī)器學(xué)習(xí)，特別是理解和推進(jìn)生成式基礎(chǔ)模型，包括擴(kuò)散模型和大型語(yǔ)言模型。

不過(guò)他的個(gè)人主頁(yè)上還沒(méi)有透露下一步計(jì)劃。

One More Thing

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

現(xiàn)在，何愷明在AI4S方向的參與的首篇論文也來(lái)了：強(qiáng)化學(xué)習(xí)+量子物理學(xué)方向。

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

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

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

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