比Stable Diffusion便宜118倍！1890美元訓(xùn)出11.6億參數(shù)高質(zhì)量文生圖模型
訓(xùn)練一個(gè)擴(kuò)散模型要多少錢？
之前最便宜的方法（Wuerstchen）用了28400美元，而像Stable Diffusion這樣的模型還要再貴一個(gè)數(shù)量級(jí)。
大模型時(shí)代，一般人根本玩不起。想要各種文生小姐姐，還得靠廠商們負(fù)重前行
為了降低這龐大的開銷，研究者們嘗試了各種方案。
比如，原始的擴(kuò)散模型從噪聲到圖像大約需要1000步，目前已經(jīng)被減少到20步左右，甚至更少。
當(dāng)擴(kuò)散模型中的基礎(chǔ)模塊逐漸由Unet（CNN）替換為DiT（Transformer）之后，一些根據(jù)Transformer特性來(lái)做的優(yōu)化也跟了上來(lái)。
比如量化，比如跳過(guò)Attention中的一些冗余計(jì)算，比如pipeline。
而近日，來(lái)自加州大學(xué)爾灣分校等機(jī)構(gòu)的研究人員，把「省錢」這個(gè)目標(biāo)直接向前推進(jìn)了一大步：
論文地址：https://arxiv.org/abs/2407.15811
——從頭開始訓(xùn)練一個(gè)11.6億參數(shù)的擴(kuò)散模型，只需要1890美元！
對(duì)比SOTA有了一個(gè)數(shù)量級(jí)的提升，讓普通人也看到了能摸一摸預(yù)訓(xùn)練的希望。
更重要的是，降低成本的技術(shù)并沒有影響模型的性能，11.6億個(gè)參數(shù)給出了下面這樣非常不錯(cuò)的效果。
除了觀感，模型的數(shù)據(jù)指標(biāo)也很優(yōu)秀，比如下表給出的FID分?jǐn)?shù)，非常接近Stable Diffusion 1.5和DALL·E 2。
相比之下，Wuerstchen的降成本方案則導(dǎo)致自己的考試分?jǐn)?shù)不甚理想。
省錢的秘訣
抱著「Stretching Each Dollar」的目標(biāo)，研究人員從擴(kuò)散模型的基礎(chǔ)模塊DiT入手。
首先，序列長(zhǎng)度是Transformer計(jì)算成本的大敵，需要除掉。
對(duì)于圖像來(lái)說(shuō)，就需要在不影響性能的情況下，盡量減少參加計(jì)算的patch數(shù)量（同時(shí)也減少了內(nèi)存開銷）。
減少圖像切塊數(shù)可以有兩種方式，一是增大每塊的尺寸，二是干掉一部分patch（mask）。
因?yàn)榍罢邥?huì)顯著降低模型性能，所以我們考慮進(jìn)行mask的方式。
最樸素的mask（Naive token masking）類似于卷積UNet中隨機(jī)裁剪的訓(xùn)練，但允許對(duì)圖像的非連續(xù)區(qū)域進(jìn)行訓(xùn)練。
而之前最先進(jìn)的方法（MaskDiT），在輸出之前增加了一個(gè)恢復(fù)重建的結(jié)構(gòu)，通過(guò)額外的損失函數(shù)來(lái)訓(xùn)練，希望通過(guò)學(xué)習(xí)彌補(bǔ)丟掉的信息。
這兩種mask都為了降低計(jì)算成本，在一開始就丟棄了大部分patch，信息的損失顯著降低了Transformer的整體性能，即使MaskDiT試圖彌補(bǔ)，也只是獲得了不太多的改進(jìn)。
——丟掉信息不可取，那么怎樣才能減小輸入又不丟信息呢？
延遲掩蔽
本文提出了一種延遲掩蔽策略（deferred masking strategy），在mask之前使用混合器（patch-mixer）進(jìn)行預(yù)處理，把被丟棄patch的信息嵌入到幸存的patch中，從而顯著減少高mask帶來(lái)的性能下降。
在本架構(gòu)中，patch-mixer是通過(guò)注意力層和前饋層的組合來(lái)實(shí)現(xiàn)的，使用二進(jìn)制掩碼進(jìn)行mask，整個(gè)模型的損失函數(shù)為：
與MaskDiT相比，這里不需要額外的損失函數(shù)，整體設(shè)計(jì)和訓(xùn)練更加簡(jiǎn)單。
而混合器本身是個(gè)非常輕量的結(jié)構(gòu)，符合省錢的標(biāo)準(zhǔn)。
微調(diào)
由于非常高的掩蔽比（masking ratio）會(huì)顯著降低擴(kuò)散模型學(xué)習(xí)圖像中全局結(jié)構(gòu)的能力，并引入訓(xùn)練到測(cè)試的分布偏移，所以作者在預(yù)訓(xùn)練（mask）后進(jìn)行了小幅度的微調(diào)（unmask）。
另外，微調(diào)還可以減輕由于使用mask而產(chǎn)生的任何不良生成偽影。
MoE和分層擴(kuò)展
MoE能夠增加模型的參數(shù)和表達(dá)能力，而不會(huì)顯著增加訓(xùn)練成本。
作者使用基于專家選擇路由的簡(jiǎn)化MoE層，每個(gè)專家確定路由到它的token，而不需要任何額外的輔助損失函數(shù)來(lái)平衡專家之間的負(fù)載。
此外，作者還考慮了分層縮放方法，線性增加Transformer塊的寬度（即注意力層和前饋層中的隱藏層尺寸）。
由于視覺模型中的更深層傾向于學(xué)習(xí)更復(fù)雜的特征，因此在更深層中使用更多的參數(shù)將帶來(lái)更好的性能。
實(shí)驗(yàn)設(shè)置
作者使用兩種DiT的變體：DiT-Tiny/2和DiT-Xl/2，patch大小為2。
使用具有余弦學(xué)習(xí)率衰減和高權(quán)重衰減的AdamW優(yōu)化器訓(xùn)練所有模型。
模型前端使用Stable-Diffusion-XL模型中的四通道變分自動(dòng)編碼器（VAE）來(lái)提取圖像特征，另外還測(cè)試了最新的16通道VAE在大規(guī)模訓(xùn)練（省錢版）中的性能。
作者使用EDM框架作為所有擴(kuò)散模型的統(tǒng)一訓(xùn)練設(shè)置，使用FID以及CLIP分?jǐn)?shù)來(lái)衡量圖像生成模型的性能。
文本編碼器選擇了最常用的CLIP模型，盡管T5-xxl這種較大的模型在文本合成等具有挑戰(zhàn)性的任務(wù)上表現(xiàn)更好，但為了省錢的目標(biāo)，這里沒有采用。
訓(xùn)練數(shù)據(jù)集
使用三個(gè)真實(shí)圖像數(shù)據(jù)集（Conceptual Captions、Segment Anything、TextCaps），包含2200萬(wàn)個(gè)圖像文本對(duì)。
由于SA1B不提供真實(shí)的字幕，這里使用LLaVA模型生成的合成字幕。作者還在大規(guī)模訓(xùn)練中添加了兩個(gè)包含1500萬(wàn)個(gè)圖像文本對(duì)的合成圖像數(shù)據(jù)集：JourneyDB和DiffusionDB。
對(duì)于小規(guī)模消融，研究人員通過(guò)從較大的COYO-700M數(shù)據(jù)集中對(duì)10個(gè)CIFAR-10類的圖像進(jìn)行二次采樣，構(gòu)建了一個(gè)名為cifar-captions的文本到圖像數(shù)據(jù)集。
評(píng)估
使用DiT-Tiny/2模型和cifar-captions數(shù)據(jù)集（256×256分辨率）進(jìn)行所有評(píng)估實(shí)驗(yàn)。
對(duì)每個(gè)模型進(jìn)行60K優(yōu)化步驟的訓(xùn)練，并使用AdamW優(yōu)化器和指數(shù)移動(dòng)平均值（最后10K步平滑系數(shù)為0.995）。
延遲掩蔽
實(shí)驗(yàn)的基線選擇我們上面提到的Naive masking，而本文的延遲掩蔽則加入一個(gè)輕量的patch-mixer，參數(shù)量小于主干網(wǎng)絡(luò)的10%。
一般來(lái)說(shuō)，丟掉的patch越多（高masking ratio），模型的性能會(huì)越差，比如MaskDiT在超過(guò)50%后表現(xiàn)大幅下降。
這里的對(duì)比實(shí)驗(yàn)采用默認(rèn)的超參數(shù)（學(xué)習(xí)率1.6×10e-4、0.01的權(quán)重衰減和余弦學(xué)習(xí)率）來(lái)訓(xùn)練兩個(gè)模型。
上圖的結(jié)果顯示了延遲屏蔽方法在FID、Clip-FID和Clip score三個(gè)指標(biāo)上都獲得了提升。
并且，與基線的性能差距隨著掩蔽率的增加而擴(kuò)大。在掩蔽率為75%的情況下，樸素掩蔽會(huì)將FID分?jǐn)?shù)降低至 16.5，而本文的方法則達(dá)到5.03，更接近于無(wú)掩蔽時(shí)的FID分?jǐn)?shù)（3.79）。
超參數(shù)
沿著訓(xùn)練LLM的一般思路，這里比較兩個(gè)任務(wù)的超參數(shù)選擇。
首先，在前饋層中，SwiGLU激活函數(shù)優(yōu)于GELU。其次，較高的權(quán)重衰減會(huì)帶來(lái)更好的圖像生成性能。
另外，與LLM訓(xùn)練不同的是，當(dāng)對(duì)AdamW二階矩 (β) 使用更高的運(yùn)行平均系數(shù)時(shí)，本文的擴(kuò)散模型可以達(dá)到更好的性能。
最后，作者發(fā)現(xiàn)使用少量的訓(xùn)練步驟，而將學(xué)習(xí)率增加到最大可能值（直到訓(xùn)練不穩(wěn)定）也顯著提高了圖像生成性能。
混合器的設(shè)計(jì)
大力出奇跡一般都是對(duì)的，作者也觀察到使用更大的patch-mixer后，模型性能得到持續(xù)改善。
然而，本著省錢的目的，這里還是選擇使用小型的混合器。
作者將噪聲分布修改為 (?0.6, 1.2)，這改善了字幕和生成圖像之間的對(duì)齊。
如下圖所示，在75% masking ratio下，作者還研究了采用不同patch大小所帶來(lái)的影響。
當(dāng)連續(xù)區(qū)域變多（patch變大）時(shí)，模型的性能會(huì)下降，因此保留隨機(jī)屏蔽每個(gè)patch的原始策略。
分層縮放
這個(gè)實(shí)驗(yàn)訓(xùn)練了DiT-Tiny架構(gòu)的兩種變體，一種具有恒定寬度，另一種采用分層縮放的結(jié)構(gòu)。
兩種方法都使用Naive masking，并調(diào)整Transformer的尺寸，保證兩種情況下的模型算力相同，同時(shí)執(zhí)行相同的訓(xùn)練步驟和訓(xùn)練時(shí)間。
由上表結(jié)果可知發(fā)現(xiàn)，在所有三個(gè)性能指標(biāo)上，分層縮放方法都優(yōu)于基線的恒定寬度方法，這表明分層縮放方法更適合DiT的掩蔽訓(xùn)練。
本文轉(zhuǎn)自 新智元 ，作者：新智元
