近日，來自加州大學爾灣分校等機構的研究人員利用延遲遮罩、MoE、分層擴展等策略，將擴散模型的訓練成本降到了1890美元。這比之前最便宜的方法(Wuerstchen)的28400美元還要低一個數量級，而像Stable Diffusion這樣的模型的成本則更高。
研究者們嘗試了各種方案來降低這龐大的開銷。比如，原始的擴散模型從噪聲到圖像大約需要1000步，目前已經被減少到20步左右，甚至更少。當擴散模型中的基礎模組逐漸由Unet(CNN)替換為DiT(Transformer)之後，一些根據Transformer特性來做的優化也跟了上來，如量化、跳過Attention中的一些冗餘計算、pipeline等。
本研究從頭開始訓練一個11.6億參數的擴散模型，只需要1890美元。這對比SOTA有了一個數量級的提升，讓普通人也看到了能摸一摸預訓練的希望。更重要的是，降低成本的技術並沒有影響模型的性能，11.6億個參數給出了非常不錯的效果。除了觀感，模型的數據指標也很優秀，比如FID分數非常接近Stable Diffusion 1.5和DALL·E 2。
省錢的秘訣主要包括:
-
延遲遮罩策略：在mask之前使用混合器(patch-mixer)進行預處理，把被丟棄patch的信息嵌入到倖存的patch中，從而顯著減少高mask帶來的性能下降。
-
微調：在預訓練(mask)後進行小幅度的微調(unmask)，以減輕由於使用mask而產生的不良生成偽影。
-
MoE和分層擴展：使用基於專家選擇路由的簡化MoE層，增加模型的參數和表達能力，而不會顯著增加訓練成本。同時考慮了分層縮放方法，線性增加Transformer塊的寬度。
實驗設置方面，作者使用兩種DiT的變體：DiT-Tiny/2和DiT-Xl/2，patch大小為2。使用具有餘弦學習率衰減和高權重衰減的AdamW優化器訓練所有模型。模型前端使用Stable-Diffusion-XL模型中的四通道變分自動編碼器(VAE)來提取圖像特徵，另外還測試了最新的16通道VAE在大規模訓練(省錢版)中的性能。
作者使用EDM框架作為所有擴散模型的統一訓練設置，使用FID以及CLIP分數來衡量圖像生成模型的性能。文本編碼器選擇了最常用的CLIP模型。
訓練數據集使用了三個真實圖像數據集(Conceptual Captions、Segment Anything、TextCaps)，包含2200萬個圖像文本對。由於SA1B不提供真實的字幕，這裡使用LLaVA模型生成的合成字幕。