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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

模型增長(zhǎng)

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

研究過(guò)程中，探索了三種初始化策略，最終選定部分保留初始化方法。這是因?yàn)樵摲椒群?jiǎn)單又穩(wěn)定。

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

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

模型剪枝

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

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

這里的表示第i個(gè)transformer的第t個(gè)token。

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

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

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

推理時(shí)擴(kuò)展

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

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

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

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

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

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

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

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

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

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

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

SANA 1.5 評(píng)估結(jié)果

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

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

模型增長(zhǎng)

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

從SANA-1.6B到4.8B的擴(kuò)展帶來(lái)了顯著的改進(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在實(shí)際應(yīng)用中更具優(yōu)勢(shì)。

模型剪枝

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

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

推理時(shí)擴(kuò)展

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

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

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

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

SANA：超高效文生圖

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

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

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

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

經(jīng)過(guò)這些優(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分辨率的圖像，用時(shí)不到1秒。

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

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

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

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

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

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

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

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

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