近年來，文本生成圖像的技術(shù)不斷突破，但隨著模型規(guī)模的擴(kuò)大，計算成本也隨之急劇上升。

為此，英偉達(dá)聯(lián)合MIT、清華、北大等機(jī)構(gòu)的研究人員提出了一種高效可擴(kuò)展的線性擴(kuò)散Transformer——SANA，在大幅降低計算需求的情況下，還能保持有競爭力的性能。

SANA1.5在此基礎(chǔ)上，聚焦了兩個關(guān)鍵問題：

1. 線性擴(kuò)散Transformer的可擴(kuò)展性如何？
2. 在擴(kuò)展大規(guī)模線性DiT時，怎樣降低訓(xùn)練成本？

論文鏈接：https://arxiv.org/pdf/2501.18427

SANA 1.5：高效模型擴(kuò)展三大創(chuàng)新

SANA 1.5在SANA 1.0（已被ICLR 2025接收）的基礎(chǔ)上，有三項關(guān)鍵創(chuàng)新。

首先，研究者提出了一種高效的模型增長策略，使得SANA可以從1.6B（20層）擴(kuò)展到4.8B（60層）參數(shù)，同時顯著減少計算資源消耗，并結(jié)合了一種節(jié)省內(nèi)存的8位優(yōu)化器。

與傳統(tǒng)的從頭開始訓(xùn)練大模型不同，通過有策略地初始化額外模塊，可以讓大模型保留小模型的先驗知識。與從頭訓(xùn)練相比，這種方法能減少60%的訓(xùn)練時間。

其二，引入了模型深度剪枝技術(shù)，實現(xiàn)了高效的模型壓縮。通過識別并保留關(guān)鍵的塊，實現(xiàn)高效的模型壓縮，然后通過微調(diào)快速恢復(fù)模型質(zhì)量，實現(xiàn)靈活的模型配置。

其三，研究者提出了一種推理期間擴(kuò)展策略，引入了重復(fù)采樣策略，使得SANA在推理時通過計算而非參數(shù)擴(kuò)展，使小模型也能達(dá)到大模型的生成質(zhì)量。

通過生成多個樣本，并利用基于視覺語言模型（VLM）的選擇機(jī)制，將GenEval分?jǐn)?shù)從0.72提升至0.80。

與從頭開始訓(xùn)練大模型不同，研究者首先將一個包含N個Transformer層的基礎(chǔ)模型擴(kuò)展到N+M層（在實驗中，N=20，M=40），同時保留其學(xué)到的知識。

在推理階段，采用兩種互補(bǔ)的方法，實現(xiàn)高效部署：

• 模型深度剪枝機(jī)制：識別并保留關(guān)鍵的Transformer塊，從而在小的微調(diào)成本下，實現(xiàn)靈活的模型配置。
• 推理時擴(kuò)展策略：借助重復(fù)采樣和VLM引導(dǎo)選擇，在計算資源和模型容量之間權(quán)衡。

同時，內(nèi)存高效CAME-8bit優(yōu)化器讓單個消費級GPU上微調(diào)十億級別的模型成為可能。

下圖展示了這些組件如何在不同的計算資源預(yù)算下協(xié)同工作，實現(xiàn)高效擴(kuò)展。

模型增長

研究者提出一種高效的模型增長策略，目的是對預(yù)訓(xùn)練的DiT模型進(jìn)行擴(kuò)展，把它從??層增加到??+??層，同時保留模型已經(jīng)學(xué)到的知識。

研究過程中，探索了三種初始化策略，最終選定部分保留初始化方法。這是因為該方法既簡單又穩(wěn)定。

在這個策略里，預(yù)訓(xùn)練的N層繼續(xù)發(fā)揮特征提取的作用，而新增加的M層一開始是隨機(jī)初始化，從恒等映射起步，慢慢學(xué)習(xí)優(yōu)化特征表示。

實驗結(jié)果顯示，與循環(huán)擴(kuò)展和塊擴(kuò)展策略相比，這種部分保留初始化方法在訓(xùn)練時的動態(tài)表現(xiàn)最為穩(wěn)定。

模型剪枝

本文提出了一種模型深度剪枝方法，能高效地將大模型壓縮成各種較小的配置，同時保持模型質(zhì)量。

受Minitron啟發(fā)，通過輸入輸出相似性模式分析塊的重要性：

這里的表示第i個transformer的第t個token。

模型的頭部和尾部塊的重要性較高，而中間層的輸入和輸出特征相似性較高，表明這些層主要用于逐步優(yōu)化生成的結(jié)果。根據(jù)排序后的塊重要性，對transformer塊進(jìn)行剪枝。

剪枝會逐步削弱高頻細(xì)節(jié)，因為，在剪枝后進(jìn)一步微調(diào)模型，以彌補(bǔ)信息損失。

使用與大模型相同的訓(xùn)練損失來監(jiān)督剪枝后的模型。剪枝模型的適配過程非常簡單，僅需100步微調(diào)，剪枝后的1.6B參數(shù)模型就能達(dá)到與完整的4.8B參數(shù)模型相近的質(zhì)量，并且優(yōu)于SANA 1.0的1.6B模型。

推理時擴(kuò)展

SANA 1.5經(jīng)過充分訓(xùn)練，在高效擴(kuò)展的基礎(chǔ)上，生成能力有了顯著提升。受LLM推理時擴(kuò)展的啟發(fā)，研究者也想通過這種方式，讓SANA 1.5表現(xiàn)得更好。

對SANA和很多擴(kuò)散模型來說，增加去噪步數(shù)是一種常見的推理時擴(kuò)展方法。但實際上，這個方法不太理想。一方面，新增的去噪步驟沒辦法修正之前出現(xiàn)的錯誤；另一方面，生成質(zhì)量很快就會達(dá)到瓶頸。

相較而言，增加采樣次數(shù)是更有潛力的方向。

研究者用視覺語言模型（VLM）來判斷生成圖像和文本提示是否匹配。他們以NVILA-2B為基礎(chǔ)模型，專門制作了一個數(shù)據(jù)集對其進(jìn)行微調(diào)。

微調(diào)后的VLM能自動比較并評價生成的圖像，經(jīng)過多輪篩選，選出排名top-N的候選圖像。這不僅確保了評選結(jié)果的可靠性，還能有效過濾與文本提示不匹配的圖像。

模型增長、模型深度剪枝和推理擴(kuò)展，構(gòu)成了一個高效的模型擴(kuò)展框架。三種方法協(xié)同配合，證明了精心設(shè)計的優(yōu)化策略，遠(yuǎn)比單純增加參數(shù)更有效。

• 模型增長策略探索了更大的優(yōu)化空間，挖掘出更優(yōu)質(zhì)的特征表示。
• 模型深度剪枝精準(zhǔn)識別并保留了關(guān)鍵特征，從而實現(xiàn)高效部署。
• 推理時間擴(kuò)展表明，當(dāng)模型容量有限時，借助額外的推理時間和計算資源，能讓模型達(dá)到與大模型相似甚至更好的效果。

為了實現(xiàn)大模型的高效訓(xùn)練與微調(diào)，研究者對CAME進(jìn)行擴(kuò)展，引入按塊8位量化，從而實現(xiàn)CAME-8bit優(yōu)化器。

CAME-8bit相比AdamW-32bit減少了約8倍的內(nèi)存使用，同時保持訓(xùn)練的穩(wěn)定性。

該優(yōu)化器不僅在預(yù)訓(xùn)練階段效果顯著，在單GPU微調(diào)場景中更是意義非凡。用RTX 4090這樣的消費級GPU，就能輕松微調(diào)SANA 4.8B。

研究揭示了高效擴(kuò)展不僅僅依賴于增加模型容量。通過充分利用小模型的知識，并設(shè)計模型的增長-剪枝，更高的生成質(zhì)量并不一定需要更大的模型。

SANA 1.5 評估結(jié)果

實驗表明，SANA 1.5的訓(xùn)練收斂速度比傳統(tǒng)方法（擴(kuò)大規(guī)模并從頭開始訓(xùn)練）快2.5倍。

訓(xùn)練擴(kuò)展策略將GenEval分?jǐn)?shù)從0.66提升至0.72，并通過推理擴(kuò)展將其進(jìn)一步提高至0.80，在GenEval基準(zhǔn)測試中達(dá)到了最先進(jìn)的性能。

模型增長

將SANA-4.8B與當(dāng)前最先進(jìn)的文本生成圖像方法進(jìn)行了比較，結(jié)果如表所示。

從SANA-1.6B到4.8B的擴(kuò)展帶來了顯著的改進(jìn)：GenEval得分提升0.06（從0.66增加到0.72），F(xiàn)ID降低0.34（從5.76降至5.42），DPG得分提升0.2（從84.8增加到85.0）。

和當(dāng)前最先進(jìn)的方法相比，SANA-4.8B模型的參數(shù)數(shù)量少很多，卻能達(dá)到和大模型一樣甚至更好的效果。

SANA-4.8B的GenEval得分為0.72，接近Playground v3的0.76。

在運(yùn)行速度上，SANA-4.8B的延遲比FLUX-dev（23.0秒）低5.5倍；吞吐量為0.26樣本/秒，是FLUX-dev（0.04樣本/秒）的6.5倍，這使得SANA-4.8B在實際應(yīng)用中更具優(yōu)勢。

模型剪枝

為了和SANA 1.0（1.6B）公平比較，此次訓(xùn)練的SANA 1.5（4.8B）模型，沒有用高質(zhì)量數(shù)據(jù)做監(jiān)督微調(diào)。

所有結(jié)果都是針對512×512尺寸的圖像評估得出的。經(jīng)過修剪和微調(diào)的模型，僅用較低的計算成本，得分就達(dá)到了0.672，超過了從頭訓(xùn)練模型的0.664。

推理時擴(kuò)展

將推理擴(kuò)展應(yīng)用于SANA 1.5（4.8B）模型，并在GenEval基準(zhǔn)上與其他大型圖像生成模型進(jìn)行了比較。

通過從2048張生成的圖像中選擇樣本，經(jīng)過推理擴(kuò)展的模型在整體準(zhǔn)確率上比單張圖像生成提高了8%，在「顏色」「位置」和「歸屬」子任務(wù)上提升明顯。

不僅如此，借助推理時擴(kuò)展，SANA 1.5（4.8B）模型的整體準(zhǔn)確率比Playground v3 (24B）高4%。

結(jié)果表明，即使模型容量有限，提高推理效率，也能提升模型生成圖像的質(zhì)量和準(zhǔn)確性。

SANA：超高效文生圖

在這里介紹一下之前的SANA工作。

SANA是一個超高效的文本生成圖像框架，能生成高達(dá)4096×4096分辨率的圖像，不僅畫質(zhì)清晰，還能讓圖像和輸入文本精準(zhǔn)匹配，而且生成速度超快，在筆記本電腦的GPU上就能運(yùn)行。

SANA為何如此強(qiáng)大？這得益于它的創(chuàng)新設(shè)計：

• 深度壓縮自動編碼器：傳統(tǒng)自動編碼器壓縮圖像的能力有限，一般只能壓縮8倍。而SANA的自動編碼器能達(dá)到32倍壓縮，大大減少了潛在tokens數(shù)量，計算效率也就更高了。
• 線性DiT：SANA用線性注意力替換了DiT中的標(biāo)準(zhǔn)注意力。在處理高分辨率圖像時，速度更快，還不會降低圖像質(zhì)量。
• 僅解碼文本編碼器：SANA不用T5做文本編碼器了，而是采用現(xiàn)代化的小型僅解碼大模型。同時，通過上下文學(xué)習(xí)，設(shè)計出更貼合實際需求的指令，讓生成的圖像和輸入文本對應(yīng)得更好。
• 高效訓(xùn)練與采樣：SANA提出了Flow-DPM-Solver方法，減少了采樣步驟。再配合高效的字幕標(biāo)注與選取，讓模型更快收斂。

經(jīng)過這些優(yōu)化，SANA-0.6B表現(xiàn)十分出色。

它生成圖像的質(zhì)量和像Flux-12B這樣的現(xiàn)代大型擴(kuò)散模型差不多，但模型體積縮小了20倍，數(shù)據(jù)處理能力卻提升了100倍以上。

SANA-0.6B運(yùn)行要求不高，在只有16GB顯存的筆記本GPU上就能運(yùn)行，生成一張1024×1024分辨率的圖像，用時不到1秒。

這意味著，創(chuàng)作者們用普通的筆記本電腦，就能輕松制作高質(zhì)量圖像，大大降低了內(nèi)容創(chuàng)作的成本。

研究者提出新的深度壓縮自動編碼器，將壓縮比例提升到32倍，和壓縮比例為8倍的自動編碼器相比，F(xiàn)32自動編碼器生成的潛在tokens減少了16倍。

這一改進(jìn)對于高效訓(xùn)練和超高分辨率圖像生成，至關(guān)重要。

研究者提出一種全新的線性DiT，用線性注意力替代傳統(tǒng)的二次復(fù)雜度注意力，將計算復(fù)雜度從原本的O(N2) 降低至O(N)。另一方面，在MLP層引入3×3深度可分卷積，增強(qiáng)潛在tokens的局部信息。

在生成效果上，線性注意力與傳統(tǒng)注意力相當(dāng)，在生成4K圖像時，推理延遲降低了1.7倍。Mix-FFN結(jié)構(gòu)讓模型無需位置編碼，也能生成高質(zhì)量圖像，這讓它成為首個無需位置嵌入的DiT變體。

在文本編碼器的選擇上，研究者選用了僅解碼的小型大語言模型Gemma，以此提升對提示詞的理解與推理能力。相較于CLIP和T5，Gemma在文本理解和指令執(zhí)行方面表現(xiàn)更為出色。

為充分發(fā)揮Gemma的優(yōu)勢，研究者優(yōu)化訓(xùn)練穩(wěn)定性，設(shè)計復(fù)雜人類指令，借助Gemma的上下文學(xué)習(xí)能力，進(jìn)一步提高了圖像與文本的匹配質(zhì)量。

研究者提出一種自動標(biāo)注與訓(xùn)練策略，借助多個視覺語言模型（VLM）生成多樣化的重新描述文本。然后，運(yùn)用基于CLIPScore的策略，篩選出CLIPScore較高的描述，以此增強(qiáng)模型的收斂性和對齊效果。

在推理環(huán)節(jié)，相較于Flow-Euler-Solver，F(xiàn)low-DPM-Solver將推理步驟從28-50步縮減至14-20步，不僅提升了速度，生成效果也更為出色。