比Stable Diffusion便宜118倍,1890美元訓出11.6億參數高質量文生圖模型
近日,來自加州大學爾灣分校等機構的研究人員,利用延遲掩蔽、MoE、分層擴展等策略,將擴散模型的訓練成本降到了1890美元。
訓練一個擴散模型要多少錢?
之前最便宜的方法(Wuerstchen)用了28400美元,而像Stable Diffusion這樣的模型還要再貴一個數量級。
大模型時代,一般人根本玩不起。想要各種文生小姐姐,還得靠廠商們負重前行
為了降低這龐大的開銷,研究者們嘗試了各種方案。
![](http://n.sinaimg.cn/spider20240812/214/w819h195/20240812/5b1b-f9351d3392657bd9122d2c848b6f526d.jpg)
比如,原始的擴散模型從噪聲到圖像大約需要1000步,目前已經被減少到20步左右,甚至更少。
當擴散模型中的基礎模塊逐漸由Unet(CNN)替換為DiT(Transformer)之後,一些根據Transformer特性來做的優化也跟了上來。
![](http://n.sinaimg.cn/spider20240812/12/w1080h532/20240812/cece-9213ee7be4c689b4998efa069c5c8cf0.jpg)
比如量化,比如跳過Attention中的一些冗餘計算,比如pipeline。
而近日,來自加州大學爾灣分校等機構的研究人員,把「慳錢」這個目標直接向前推進了一大步:
![論文地址:https://arxiv.org/abs/2407.15811](http://n.sinaimg.cn/spider20240812/456/w959h297/20240812/0fa8-ccd1bbe246ff0081b3af41587352408a.jpg)
——從頭開始訓練一個11.6億參數的擴散模型,只需要1890美元!
對比SOTA有了一個數量級的提升,讓普通人也看到了能摸一摸預訓練的希望。
更重要的是,降低成本的技術並沒有影響模型的性能,11.6億個參數給出了下面這樣非常不錯的效果。
![](http://n.sinaimg.cn/spider20240812/368/w891h277/20240812/0a51-50a5201ab82b970537879575963af1a0.jpg)
![](http://n.sinaimg.cn/spider20240812/232/w1080h752/20240812/20f1-90bf6e709b3221a66f11023a61d77c93.jpg)
除了觀感,模型的數據指標也很優秀,比如下表給出的FID分數,非常接近Stable Diffusion 1.5和DALL·E 2。
相比之下,Wuerstchen的降成本方案則導致自己的考試分數不甚理想。
![](http://n.sinaimg.cn/spider20240812/426/w401h825/20240812/1d8e-b4631132ed7dfa9aba69b5b171195698.jpg)
慳錢的秘訣
抱著「Stretching Each Dollar」的目標,研究人員從擴散模型的基礎模塊DiT入手。
首先,序列長度是Transformer計算成本的大敵,需要除掉。
對於圖像來說,就需要在不影響性能的情況下,儘量減少參加計算的patch數量(同時也減少了內存開銷)。
![](http://n.sinaimg.cn/spider20240812/412/w850h362/20240812/2ff1-55db2dbc1c01fa31fa7e9fd8ff90f762.jpg)
減少圖像切塊數可以有兩種方式,一是增大每塊的尺寸,二是幹掉一部分patch(mask)。
![](http://n.sinaimg.cn/spider20240812/552/w883h469/20240812/8685-ee441dc58228fa1dc7566234687af6b4.jpg)
因為前者會顯著降低模型性能,所以我們考慮進行mask的方式。
最樸素的mask(Naive token masking)類似於卷積UNet中隨機裁剪的訓練,但允許對圖像的非連續區域進行訓練。
![](http://n.sinaimg.cn/spider20240812/312/w653h459/20240812/e7e8-b8794d9ee15f472c1f7ee136322ea255.jpg)
而之前最先進的方法(MaskDiT),在輸出之前增加了一個恢復重建的結構,通過額外的損失函數來訓練,希望通過學習彌補丟掉的信息。
![](http://n.sinaimg.cn/spider20240812/431/w1080h151/20240812/c494-13344c94426c1e0cd2959f2b7aa2f568.jpg)
這兩種mask都為了降低計算成本,在一開始就丟棄了大部分patch,信息的損失顯著降低了Transformer的整體性能,即使MaskDiT試圖彌補,也只是獲得了不太多的改進。
——丟掉信息不可取,那麼怎樣才能減小輸入又不丟信息呢?
延遲掩蔽
本文提出了一種延遲掩蔽策略(deferred masking strategy),在mask之前使用混合器(patch-mixer)進行預處理,把被丟棄patch的信息嵌入到倖存的patch中,從而顯著減少高mask帶來的性能下降。
![](http://n.sinaimg.cn/spider20240812/704/w771h733/20240812/ebf8-5c807eb5cff141a01651ed6826dacd10.jpg)
在本架構中,patch-mixer是通過注意力層和前饋層的組合來實現的,使用二進製掩碼進行mask,整個模型的損失函數為:
![](http://n.sinaimg.cn/spider20240812/246/w997h49/20240812/8363-602d9363fd4335eb4536a8a28eb6b5a0.jpg)
與MaskDiT相比,這裏不需要額外的損失函數,整體設計和訓練更加簡單。
而混合器本身是個非常輕量的結構,符合慳錢的標準。
微調
由於非常高的掩蔽比(masking ratio)會顯著降低擴散模型學習圖像中全局結構的能力,並引入訓練到測試的分佈偏移,所以作者在預訓練(mask)後進行了小幅度的微調(unmask)。
另外,微調還可以減輕由於使用mask而產生的任何不良生成偽影。
MoE和分層擴展
MoE能夠增加模型的參數和表達能力,而不會顯著增加訓練成本。
作者使用基於專家選擇路由的簡化MoE層,每個專家確定路由到它的token,而不需要任何額外的輔助損失函數來平衡專家之間的負載。
![](http://n.sinaimg.cn/spider20240812/106/w811h895/20240812/a46a-a43469a0d60e50a9a4036da28eb321a3.jpg)
此外,作者還考慮了分層縮放方法,線性增加Transformer塊的寬度(即注意力層和前饋層中的隱藏層尺寸)。
由於視覺模型中的更深層傾向於學習更複雜的特徵,因此在更深層中使用更多的參數將帶來更好的性能。
實驗設置
作者使用兩種DiT的變體:DiT-Tiny/2和DiT-Xl/2,patch大小為2。
使用具有餘弦學習率衰減和高權重衰減的AdamW優化器訓練所有模型。
![](http://n.sinaimg.cn/spider20240812/791/w1080h511/20240812/c650-264c4f125901143cc7f6dd7a03c6a153.jpg)
模型前端使用Stable-Diffusion-XL模型中的四通道變分自動編碼器(VAE)來提取圖像特徵,另外還測試了最新的16通道VAE在大規模訓練(慳錢版)中的性能。
![](http://n.sinaimg.cn/spider20240812/731/w1080h451/20240812/eff3-71a7f90a31e1cfd455df053a4f0d140a.jpg)
作者使用EDM框架作為所有擴散模型的統一訓練設置,使用FID以及CLIP分數來衡量圖像生成模型的性能。
文本編碼器選擇了最常用的CLIP模型,儘管T5-xxl這種較大的模型在文本合成等具有挑戰性的任務上表現更好,但為了慳錢的目標,這裏沒有採用。
訓練數據集
使用三個真實圖像數據集(Conceptual Captions、Segment Anything、TextCaps),包含2200萬個圖像文本對。
由於SA1B不提供真實的字幕,這裏使用LLaVA模型生成的合成字幕。作者還在大規模訓練中添加了兩個包含1500萬個圖像文本對的合成圖像數據集:JourneyDB和DiffusionDB。
對於小規模消融,研究人員通過從較大的COYO-700M數據集中對10個CIFAR-10類的圖像進行二次采樣,構建了一個名為cifar-captions的文本到圖像數據集。
評估
使用DiT-Tiny/2模型和cifar-captions數據集(256×256解像度)進行所有評估實驗。
對每個模型進行60K優化步驟的訓練,並使用AdamW優化器和指數移動平均值(最後10K步平滑係數為0.995)。
延遲掩蔽
實驗的基線選擇我們上面提到的Naive masking,而本文的延遲掩蔽則加入一個輕量的patch-mixer,參數量小於主幹網絡的10%。
一般來說,丟掉的patch越多(高masking ratio),模型的性能會越差,比如MaskDiT在超過50%後表現大幅下降。
這裏的對比實驗採用預設的超參數(學習率1.6×10e-4、0.01的權重衰減和餘弦學習率)來訓練兩個模型。
![](http://n.sinaimg.cn/spider20240812/434/w993h241/20240812/0ec8-0e5d68587488e5891c72de043c2ee324.jpg)
上圖的結果顯示了延遲屏蔽方法在FID、Clip-FID和Clip score三個指標上都獲得了提升。
並且,與基線的性能差距隨著掩蔽率的增加而擴大。在掩蔽率為75%的情況下,樸素掩蔽會將FID分數降低至 16.5,而本文的方法則達到5.03,更接近於無掩蔽時的FID分數(3.79)。
超參數
沿著訓練LLM的一般思路,這裏比較兩個任務的超參數選擇。
首先,在前饋層中,SwiGLU激活函數優於GELU。其次,較高的權重衰減會帶來更好的圖像生成性能。
![](http://n.sinaimg.cn/spider20240812/16/w603h213/20240812/7164-6c1b06a3d72b6864f5e712c478118699.jpg)
另外,與LLM訓練不同的是,當對AdamW二階矩 (β) 使用更高的運行平均係數時,本文的擴散模型可以達到更好的性能。
最後,作者發現使用少量的訓練步驟,而將學習率增加到最大可能值(直到訓練不穩定)也顯著提高了圖像生成性能。
混合器的設計
大力出奇蹟一般都是對的,作者也觀察到使用更大的patch-mixer後,模型性能得到持續改善。
然而,本著慳錢的目的,這裏還是選擇使用小型的混合器。
作者將噪聲分佈修改為 (−0.6, 1.2),這改善了字幕和生成圖像之間的對齊。
如下圖所示,在75% masking ratio下,作者還研究了採用不同patch大小所帶來的影響。
![](http://n.sinaimg.cn/spider20240812/20/w629h191/20240812/2c25-1f6a3a71306c89dab1207e06c5eac577.jpg)
當連續區域變多(patch變大)時,模型的性能會下降,因此保留隨機屏蔽每個patch的原始策略。
分層縮放
這個實驗訓練了DiT-Tiny架構的兩種變體,一種具有恒定寬度,另一種採用分層縮放的結構。
兩種方法都使用Naive masking,並調整Transformer的尺寸,保證兩種情況下的模型算力相同,同時執行相同的訓練步驟和訓練時間。
![](http://n.sinaimg.cn/spider20240812/558/w467h91/20240812/8b94-ed1c6992561423d8b92895ed4a767657.jpg)
由上表結果可知發現,在所有三個性能指標上,分層縮放方法都優於基線的恒定寬度方法,這表明分層縮放方法更適合DiT的掩蔽訓練。
參考資料:
https://arxiv.org/abs/2407.15811
本文來自微信公眾號「新智元」,作者:新智元,36氪經授權發佈。