henry 發(fā)自 凹非寺
量子位 | 公眾號(hào) QbitAI

全員本科生！

剛剛，何愷明攜本科生“軍團(tuán)”又放出一篇新論文。

繼去年探索直接從像素預(yù)測(cè)圖像的JiT架構(gòu)后，團(tuán)隊(duì)這次又把這套“刪繁就簡(jiǎn)”的思路擴(kuò)展到了文生圖領(lǐng)域，推出全新工作：

MiniT2I

在今天動(dòng)輒數(shù)十億參數(shù)、海量圖文數(shù)據(jù)訓(xùn)練文生圖模型的背景下，MiniT2I選擇了另一條路。

它基于全新的MM-JiT架構(gòu)，直接在像素空間進(jìn)行擴(kuò)散生成，同時(shí)盡可能壓縮模型復(fù)雜度和訓(xùn)練成本。

最終，僅用258M參數(shù)，就實(shí)現(xiàn)了不錯(cuò)的文生圖效果。

更關(guān)鍵的是，整個(gè)訓(xùn)練成本只相當(dāng)于一次標(biāo)準(zhǔn)ImageNet實(shí)驗(yàn)。

這是怎么做到的？

從JiT到MM-JiT

整體看來(lái)，MM-JiT是愷明組之前論文「Back to Basics」在T2I（文本生成圖像）方向上的延伸。

Back to Basics中，愷明和他的博后黎天鴻提出了JiT架構(gòu)，Just image Transformers。

JiT的核心主張是：拋開(kāi)VAE編解碼器，直接在像素空間預(yù)測(cè)干凈圖像（x-prediction），而不是像傳統(tǒng)擴(kuò)散模型那樣預(yù)測(cè)噪聲。

這樣做的好處是，整個(gè)生成流程更加直接，符合流形假設(shè)以及“從像素出發(fā)”的第一性原理。

不過(guò)，當(dāng)時(shí)的JiT主要針對(duì)類別條件生成（class-conditional generation），任務(wù)范圍相對(duì)有限，模型只能根據(jù)ImageNet的類別標(biāo)簽生成對(duì)應(yīng)圖像。

然而，真實(shí)的圖像生成任務(wù)往往不限于ImageNet的1000個(gè)固定類別，而是需要理解并遵循開(kāi)放的文本Prompt

問(wèn)題也隨之而來(lái)。一旦從類別生成擴(kuò)展到文生圖，訓(xùn)練成本往往會(huì)迅速攀升。

無(wú)論是SD3、FLUX.1-dev還是DALL·E 3，背后都依賴多階段訓(xùn)練流程、龐大的文本編碼器以及海量數(shù)據(jù)資源。

對(duì)于大多數(shù)學(xué)術(shù)團(tuán)隊(duì)而言，從零開(kāi)始訓(xùn)練一個(gè)完整的文生圖模型，幾乎是一項(xiàng)難以承擔(dān)的工程。

于是，MiniT2I應(yīng)運(yùn)而生。

它試圖回答一個(gè)更現(xiàn)實(shí)的問(wèn)題：

如果只用接近ImageNet訓(xùn)練規(guī)模的計(jì)算資源，能不能也做出效果不錯(cuò)的文生圖模型？

答案是，可以。

研究發(fā)現(xiàn)，當(dāng)文本首先被預(yù)訓(xùn)練語(yǔ)言模型編碼為語(yǔ)義表示后，對(duì)于生成模型而言，文本條件本質(zhì)上只是另一種形式的上下文條件。

換句話說(shuō)，文生圖或許并沒(méi)有想象中那么特殊。

在模型架構(gòu)、訓(xùn)練計(jì)算量，甚至所需數(shù)據(jù)規(guī)模上，它與類別條件生成的差距遠(yuǎn)沒(méi)有業(yè)界普遍認(rèn)為的那么大。

如果這個(gè)判斷成立，那么一個(gè)很自然的問(wèn)題就出現(xiàn)了：

既然類別條件生成已經(jīng)能用JiT這樣的極簡(jiǎn)架構(gòu)完成，那么文生圖任務(wù)里那些復(fù)雜的模塊，究竟哪些是真正必要的？

MM-JiT給出的答案是：把它們一個(gè)個(gè)刪掉，再看模型還能不能工作。

MM-JiT：刪繁就簡(jiǎn)

對(duì)于上面這個(gè)問(wèn)題，MiniT2I項(xiàng)目負(fù)責(zé)人王銜邦在X上的總結(jié)非常精煉：

我們的原則很簡(jiǎn)單，能去掉的全去掉。起點(diǎn)是像素空間、標(biāo)準(zhǔn)的T5-Large編碼器，以及一個(gè)采用x-prediction的簡(jiǎn)潔多模態(tài)骨干MM-JiT。

這套思路的第一刀，砍向了VAE。

眾所周知，當(dāng)前主流文生圖模型大多采用潛在擴(kuò)散（Latent Diffusion）路線：

先通過(guò)VAE把圖像壓縮到低維潛空間，再在潛空間里完成擴(kuò)散生成，最后解碼回像素。

這樣做的好處是顯著降低計(jì)算量，但代價(jià)也很明顯——

VAE會(huì)帶來(lái)重建誤差和偽影，同時(shí)還額外增加了一套編解碼器的訓(xùn)練流程。

針對(duì)這一問(wèn)題，在前作JiT中，團(tuán)隊(duì)已經(jīng)證明，至少在ImageNet任務(wù)上，直接在像素空間建模并不存在所謂的“不可逾越瓶頸”。

那么在文生圖任務(wù)里，VAE是否真的不可替代？

團(tuán)隊(duì)決定直接把它刪掉試試。

MiniT2I將擴(kuò)散過(guò)程重新搬回像素空間，希望驗(yàn)證一個(gè)看似反常識(shí)的判斷：直接在像素空間擴(kuò)散，不僅完全可行，而且未必比潛空間路線更貴。

實(shí)驗(yàn)表明，傳統(tǒng)潛空間模型單次前向傳播需要1379 GFLOPs，而徹底擺脫VAE之后，MiniT2I的計(jì)算開(kāi)銷僅為265 GFLOPs，直接降低了約80%。

刪掉VAE之后，團(tuán)隊(duì)又把目光轉(zhuǎn)向了模型架構(gòu)本身。

前作JiT面向的是ImageNet分類條件生成，因此采用標(biāo)準(zhǔn)DiT，并通過(guò)AdaLN-Zero注入類別標(biāo)簽和時(shí)間步信息。

但到了開(kāi)放式文生圖任務(wù)，最自然的參考對(duì)象就變成了SD3采用的MM-DiT。

在團(tuán)隊(duì)看來(lái)，MM-DiT身上仍然掛著不少“歷史包袱”。

其中最典型的就是AdaLN機(jī)制。模型會(huì)把時(shí)間步和池化后的文本特征轉(zhuǎn)換成縮放、偏移和門控參數(shù)，并注入到每一層網(wǎng)絡(luò)中。

MM-JiT的選擇則相當(dāng)激進(jìn)：直接把AdaLN整個(gè)刪掉。

理由也很簡(jiǎn)單——擴(kuò)散模型當(dāng)前所處的噪聲水平，其實(shí)已經(jīng)包含在加噪后的輸入 z_t 里。

換句話說(shuō)，模型完全可以自己推斷當(dāng)前處于擴(kuò)散過(guò)程的哪個(gè)階段，并不需要額外開(kāi)一條通道專門傳遞時(shí)間步信息。

于是，條件信息只通過(guò)聯(lián)合注意力這一條路徑進(jìn)入模型，整個(gè)骨干網(wǎng)絡(luò)也回歸到更接近標(biāo)準(zhǔn)Pre-Norm Transformer的形式。

與此同時(shí)，團(tuán)隊(duì)只額外增加了兩個(gè)Text Adapter Block，放在聯(lián)合注意力之前，讓凍結(jié)的T5文本特征先完成一次適配，再與圖像Token交互。

實(shí)驗(yàn)結(jié)果再次驗(yàn)證了團(tuán)隊(duì)的判斷。

參數(shù)量幾乎保持不變，依舊只有260M，但模型性能卻一路提升：

FID從18.7（MM-DiT像素空間基線），提升到17.4（加入Text Adapter），最終達(dá)到13.7（移除AdaLN后的MM-JiT）。

訓(xùn)練與實(shí)驗(yàn)

在具體實(shí)現(xiàn)上，MiniT2I基于流匹配（Flow Matching）框架，網(wǎng)絡(luò)直接預(yù)測(cè)干凈圖像，并在速度空間計(jì)算損失。

訓(xùn)練分為兩個(gè)階段：首先在CC12M上預(yù)訓(xùn)練25萬(wàn)步，學(xué)習(xí)基礎(chǔ)視覺(jué)分布；隨后在12萬(wàn)張高質(zhì)量合成圖像上微調(diào)4萬(wàn)步，進(jìn)一步提升Prompt遵循能力。

結(jié)果證明，這套極簡(jiǎn)設(shè)計(jì)并沒(méi)有犧牲性能。

B/16版本總參數(shù)量不到600M，在GenEval上達(dá)到0.87、DPG-Bench達(dá)到84.2，超過(guò)了多款參數(shù)規(guī)模數(shù)倍于自身的像素空間文生圖模型。

更重要的是，完成這一訓(xùn)練僅需約等于一次標(biāo)準(zhǔn)ImageNet實(shí)驗(yàn)的算力預(yù)算——8張H100，大約3天。

即便與工業(yè)級(jí)模型相比，MM-JiT也展現(xiàn)出不俗競(jìng)爭(zhēng)力。

在PRISM-Bench上，L/16版本取得62.4分，而FLUX.1-dev為68.5分。具體來(lái)看，模型在風(fēng)格表現(xiàn)和開(kāi)放想象力兩個(gè)維度甚至超過(guò)了FLUX；

短板方面，則主要集中在文字渲染和命名實(shí)體生成，這與公開(kāi)訓(xùn)練數(shù)據(jù)覆蓋范圍有限有關(guān)。

（注：具體實(shí)驗(yàn)設(shè)置可參考文末博客鏈接）

作者介紹

這篇工作最值得聊的，除了技術(shù)本身，還有背后的作者們。

整篇論文一共六位作者。除了何愷明之外，其余五位都還是本科生。

而且，這些年輕面孔并不是第一次出現(xiàn)在論文作者欄里。在何愷明團(tuán)隊(duì)此前的多篇工作中，他們都已經(jīng)開(kāi)始嶄露頭角。

項(xiàng)目負(fù)責(zé)人王銜邦(Xianbang Wang)目前是MIT大一本科生，去年剛從人大附中畢業(yè)。

2024年，他代表中國(guó)隊(duì)參加第65屆國(guó)際數(shù)學(xué)奧林匹克競(jìng)賽(IMO)，拿下金牌。

更早之前，他還在2021年和2022年斬獲全國(guó)信息學(xué)奧林匹克競(jìng)賽的銀牌。

在這項(xiàng)工作之前，他已經(jīng)是何愷明團(tuán)隊(duì)Bidirectional Normalizing Flow論文的共同第一作者。

另一位核心貢獻(xiàn)者趙瀚宏(Hanhong Zhao)，目前是MIT大二學(xué)生，曾獲得國(guó)際物理奧林匹克競(jìng)賽（IPhO）金牌。

不久前引發(fā)關(guān)注的ELF（連續(xù)擴(kuò)散語(yǔ)言模型）論文中，趙瀚宏也是作者之一。

核心貢獻(xiàn)者陸伊煬(Yiyang Lu)則來(lái)自清華大學(xué)姚班，目前大二，在MIT CSAIL實(shí)習(xí)，導(dǎo)師正是何愷明。

高中時(shí)期，他是物理競(jìng)賽生，曾以江蘇省第一、全國(guó)第九的成績(jī)獲得第39屆全國(guó)中學(xué)生物理競(jìng)賽（CPhO）金牌。

此前，他已經(jīng)與何愷明合作完成Bidirectional Normalizing Flow、Pixel Mean Flow等工作，在ELF論文中同樣名列作者名單。

周康陽(yáng)（Kangyang Zhou）也是MIT本科生(Class of 2029)，背景更偏信息學(xué)方向。

2024年，他在第36屆國(guó)際信息學(xué)奧林匹克競(jìng)賽（IOI）中奪冠，并以600分滿分成為當(dāng)屆唯一滿分選手。

更早的2023年，他以全國(guó)信息學(xué)奧林匹克競(jìng)賽（NOI）金牌第一名的成績(jī)?nèi)脒x國(guó)家集訓(xùn)隊(duì)，領(lǐng)先第二名55分。今年，他還作為MIT代表隊(duì)成員獲得ICPC 2026北美錦標(biāo)賽冠軍。

馬麟瑞（Linrui Ma）同樣畢業(yè)于人大附中，目前在MIT就讀本科。

他曾擔(dān)任中國(guó)國(guó)家隊(duì)隊(duì)長(zhǎng)，在第56屆國(guó)際化學(xué)奧林匹克競(jìng)賽（IChO 2024）中獲得金牌。

最后再簡(jiǎn)單介紹一下何愷明。

目前，他是MIT EECS終身副教授，同時(shí)兼任Google DeepMind杰出科學(xué)家。

他是深度學(xué)習(xí)、計(jì)算機(jī)視覺(jué)一系列重要工作，如ResNet、Faster R-CNN、Mask R-CNN、MoCo、MAE的作者。其中，ResNet是21世紀(jì)被引用次數(shù)最多的論文。

某種程度上說(shuō)，這篇論文最有意思的地方，不只是提出了一個(gè)新方法，更像是一群剛剛走出奧賽賽場(chǎng)的年輕人，已經(jīng)開(kāi)始站上AI研究最前沿的舞臺(tái)。

[1] https://peppaking8.github.io/#/post/minit2i