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

全員本科生！

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

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

MiniT2I

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

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

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

更關(guān)鍵的是，整個訓練成本只相當于一次標準ImageNet實驗。

這是怎么做到的？

從JiT到MM-JiT

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

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

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

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

不過，當時的JiT主要針對類別條件生成（class-conditional generation），任務范圍相對有限，模型只能根據(jù)ImageNet的類別標簽生成對應圖像。

然而，真實的圖像生成任務往往不限于ImageNet的1000個固定類別，而是需要理解并遵循開放的文本Prompt

問題也隨之而來。一旦從類別生成擴展到文生圖，訓練成本往往會迅速攀升。

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

對于大多數(shù)學術(shù)團隊而言，從零開始訓練一個完整的文生圖模型，幾乎是一項難以承擔的工程。

于是，MiniT2I應運而生。

它試圖回答一個更現(xiàn)實的問題：

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

答案是，可以。

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

換句話說，文生圖或許并沒有想象中那么特殊。

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

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

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

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

MM-JiT：刪繁就簡

對于上面這個問題，MiniT2I項目負責人王銜邦在X上的總結(jié)非常精煉：

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

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

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

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

這樣做的好處是顯著降低計算量，但代價也很明顯——

VAE會帶來重建誤差和偽影，同時還額外增加了一套編解碼器的訓練流程。

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

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

團隊決定直接把它刪掉試試。

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

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

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

前作JiT面向的是ImageNet分類條件生成，因此采用標準DiT，并通過AdaLN-Zero注入類別標簽和時間步信息。

但到了開放式文生圖任務，最自然的參考對象就變成了SD3采用的MM-DiT。

在團隊看來，MM-DiT身上仍然掛著不少“歷史包袱”。

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

MM-JiT的選擇則相當激進：直接把AdaLN整個刪掉。

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

換句話說，模型完全可以自己推斷當前處于擴散過程的哪個階段，并不需要額外開一條通道專門傳遞時間步信息。

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

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

實驗結(jié)果再次驗證了團隊的判斷。

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

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

訓練與實驗

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

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

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

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

更重要的是，完成這一訓練僅需約等于一次標準ImageNet實驗的算力預算——8張H100，大約3天。

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

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

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

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

作者介紹

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

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

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

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

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

更早之前，他還在2021年和2022年斬獲全國信息學奧林匹克競賽的銀牌。

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

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

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

核心貢獻者陸伊煬(Yiyang Lu)則來自清華大學姚班，目前大二，在MIT CSAIL實習，導師正是何愷明。

高中時期，他是物理競賽生，曾以江蘇省第一、全國第九的成績獲得第39屆全國中學生物理競賽（CPhO）金牌。

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

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

2024年，他在第36屆國際信息學奧林匹克競賽（IOI）中奪冠，并以600分滿分成為當屆唯一滿分選手。

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

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

他曾擔任中國國家隊隊長，在第56屆國際化學奧林匹克競賽（IChO 2024）中獲得金牌。

最后再簡單介紹一下何愷明。

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

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

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

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