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