現(xiàn)階段，視覺生成模型擅長創(chuàng)建逼真的視覺內(nèi)容，然而從頭開始訓練這些模型的成本和工作量仍然很高。比如 Stable Diffusion 2.1 花費了 200000 個 A100 GPU 小時。即使研究者使用最先進的方法，也需要在 8×H100 GPU 上訓練一個多月的時間。

此外，訓練大模型也對數(shù)據(jù)集提出了挑戰(zhàn)，這些數(shù)據(jù)基本以億為單位，同樣給訓練模型帶來挑戰(zhàn)。

高昂的訓練成本和對數(shù)據(jù)集的要求為大規(guī)模擴散模型的開發(fā)造成了難以逾越的障礙。

現(xiàn)在，來自 Sony AI 等機構(gòu)的研究者僅僅花了 1890 美元，就訓練了一個不錯的擴散模型， 具有 11.6 億參數(shù)的稀疏 transformer。

• 論文地址：https://arxiv.org/pdf/2407.15811
• 論文標題：Stretching Each Dollar: Diffusion Training from Scratch on a Micro-Budget
• 項目（即將發(fā)布）：https://github.com/SonyResearch/micro_diffusion

具體而言，在這項工作中，作者通過開發(fā)一種低成本端到端的 pipeline 用于文本到圖像擴散模型，使得訓練成本比 SOTA 模型降低了一個數(shù)量級還多，同時還不需要訪問數(shù)十億張訓練圖像或?qū)Ｓ袛?shù)據(jù)集。

作者考慮了基于視覺 transformer 的潛在擴散模型進行文本到圖像生成，主要原因是這種方式設(shè)計簡單，并且應用廣泛。為了降低計算成本，作者利用了 transformer 計算開銷與輸入序列大?。疵繌垐D像的 patch 數(shù)量）的強依賴關(guān)系。

本文的主要目標是在訓練過程中減少 transformer 處理每張圖像的有效 patch 數(shù)。通過在 transformer 的輸入層隨機掩蔽（mask）掉部分 token，可以輕松實現(xiàn)這一目標。

然而，現(xiàn)有的掩蔽方法無法在不大幅降低性能的情況下將掩蔽率擴展到 50% 以上，特別是在高掩蔽率下，很大一部分輸入 patch 完全不會被擴散 transformer 觀察到。

為了減輕掩蔽造成的性能大幅下降，作者提出了一種延遲掩蔽（deferred masking）策略，其中所有 patch 都由輕量級 patch 混合器（patch-mixer）進行預處理，然后再傳輸?shù)綌U散 transformer。Patch 混合器包含擴散 transformer 中參數(shù)數(shù)量的一小部分。

與 naive 掩蔽方法相比，在 patch mixing 處理之后進行掩蔽允許未掩蔽的 patch 保留有關(guān)整個圖像的語義信息，并能夠在非常高的掩蔽率下可靠地訓練擴散 transformer，同時與現(xiàn)有的最先進掩蔽相比不會產(chǎn)生額外的計算成本。 

作者還證明了在相同的計算預算下，延遲掩蔽策略比縮小模型規(guī)模（即減小模型大?。崿F(xiàn)了更好的性能。最后，作者結(jié)合 Transformer 架構(gòu)的最新進展，例如逐層縮放、使用 MoE 的稀疏 Transformer，以提高大規(guī)模訓練的性能。

作者提出的低成本訓練 pipeline 減少了實驗開銷。除了使用真實圖像，作者還考慮在訓練數(shù)據(jù)集中組合其他合成圖像。組合數(shù)據(jù)集僅包含 3700 萬張圖像，比大多數(shù)現(xiàn)有的大型模型所需的數(shù)據(jù)量少得多。

在這個組合數(shù)據(jù)集上，作者以 1890 美元的成本訓練了一個 11.6 億參數(shù)的稀疏 transformer，并在 COCO 數(shù)據(jù)集上的零樣本生成中實現(xiàn)了 12.7 FID。

值得注意的是，本文訓練的模型實現(xiàn)了具有競爭力的 FID 和高質(zhì)量生成，同時成本僅為 stable diffusion 模型的 1/118 ，是目前最先進的方法（成本為 28,400 美元）的 1/15。

方法介紹

為了大幅降低計算成本，patch 掩蔽要求在輸入主干 transformer 之前丟棄大部分輸入 patch，從而使 transformer 無法獲得被掩蔽 patch 的信息。高掩蔽率（例如 75% 的掩蔽率）會顯著降低 transformer 的整體性能。即使使用 MaskDiT，也只能觀察到它比 naive 掩蔽有微弱的改善，因為這種方法也會在輸入層本身丟棄大部分圖像 patch。

延遲掩蔽，保留所有 patch 的語義信息

由于高掩蔽率會去除圖像中大部分有價值的學習信號，作者不禁要問，是否有必要在輸入層進行掩蔽？只要計算成本不變，這就只是一種設(shè)計選擇，而不是根本限制。事實上，作者發(fā)現(xiàn)了一種明顯更好的掩蔽策略，其成本與現(xiàn)有的 MaskDiT 方法幾乎相同。由于 patch 來自擴散 Transformer 中的非重疊圖像區(qū)域，每個 patch 嵌入都不會嵌入圖像中其他 patch 的任何信息。因此，作者的目標是在掩蔽之前對 patch 嵌入進行預處理，使未被掩蔽的 patch 能夠嵌入整個圖像的信息。他們將預處理模塊稱為 patch-mixer。

使用 patch-mixer 訓練擴散 transformer

作者認為，patch-mixer 是任何一種能夠融合單個 patch 嵌入的神經(jīng)架構(gòu)。在 transformer 模型中，這一目標自然可以通過注意力層和前饋層的組合來實現(xiàn)。因此，作者使用一個僅由幾個層組成的輕量級 transformer 作為 patch-mixer。輸入序列 token 經(jīng) patch-mixer 處理后，他們將對其進行掩蔽（圖 2e）。

圖 2：壓縮 patch 序列以降低計算成本。由于擴散 transformer 的訓練成本與序列大?。?patch 數(shù)量）成正比，因此最好能在不降低性能的情況下縮減序列大小。這可以通過以下方法實現(xiàn)：b) 使用更大的 patch；c) 隨機簡單（naive）掩蔽一部分 patch；或者 d) 使用 MaskDiT，該方法結(jié)合了 naive 掩蔽和額外的自動編碼目標。作者發(fā)現(xiàn)這三種方法都會導致圖像生成性能顯著下降，尤其是在高掩蔽率的情況下。為了緩解這一問題，他們提出了一種直接的延遲掩蔽策略，即在 patch-mixer 處理完 patch 后再對其進行掩蔽。除了使用 patch-mixer 之外，他們的方法在所有方面都類似于 naive 掩蔽。與 MaskDiT 相比，他們的方法無需優(yōu)化任何替代目標，計算成本幾乎相同。

假定掩碼為二進制掩碼 m，作者使用以下?lián)p失函數(shù)來訓練模型：    

其中，M_? 是 patch-mixer 模型，F(xiàn)_θ 是主干 transformer。請注意，與 MaskDiT 相比，本文提出的方法還簡化了整體設(shè)計，不需要額外的損失函數(shù)，也不需要在訓練過程中在兩個損失之間進行相應的超參數(shù)調(diào)優(yōu)。在推理過程中，該方法不掩蔽任何 patch。

未掩蔽微調(diào)

由于極高的掩蔽率會大大降低擴散模型學習圖像全局結(jié)構(gòu)的能力，并在序列大小上引入訓練 - 測試分布偏移，因此作者考慮在掩蔽預訓練后進行少量的未掩蔽微調(diào)。微調(diào)還可以減輕由于使用 patch 掩蔽而產(chǎn)生的任何生成瑕疵。因此，在以前的工作中，恢復因掩蔽而急劇下降的性能至關(guān)重要，尤其是在采樣中使用無分類器引導時。然而，作者認為這并不是完全必要的，因為即使在掩蔽預訓練中，他們的方法也能達到與基線未掩蔽預訓練相當?shù)男阅?。作者只在大?guī)模訓練中使用這種方法，以減輕由于高度 patch 掩蔽而產(chǎn)生的任何未知 - 未知生成瑕疵。

利用 MoE 和 layer-wise scaling 改進主干 transformer 架構(gòu)

作者還利用 transformer 架構(gòu)設(shè)計方面的創(chuàng)新，在計算限制條件下提高了模型的性能。

他們使用混合專家層，因為它們在不顯著增加訓練成本的情況下增加了模型的參數(shù)和表現(xiàn)力。他們使用基于專家選擇路由的簡化 MoE 層，每個專家決定路由給它的 token，因為它不需要任何額外的輔助損失函數(shù)來平衡專家間的負載。他們還考慮了 layer-wise scaling，該方法最近被證明在大型語言模型中優(yōu)于典型 transformer。該方法線性增加 transformer 塊的寬度，即注意力層和前饋層的隱藏層維度。因此，網(wǎng)絡(luò)中較深的層比較早的層被分配了更多的參數(shù)。作者認為，由于視覺模型中的較深層往往能學習到更復雜的特征，因此在較深層使用更高的參數(shù)會帶來更好的性能。作者在圖 3 中描述了他們提出的擴散 Transformer 的整體架構(gòu)。

圖 3：本文提出的擴散 transformer 的整體架構(gòu)。作者在骨干 transformer 模型中加入了一個輕量級的 patch-mixer，它可以在輸入圖像中的所有 patch 被掩蔽之前對其進行處理。根據(jù)當前的研究成果，作者使用注意力層處理 caption 嵌入，然后再將其用于調(diào)節(jié)。他們使用正弦嵌入來表示時間步長。他們的模型只對未掩蔽的 patch 進行去噪處理，因此只對這些 patch 計算擴散損失（論文中的公式 3）。他們對主干 transformer 進行了修改，在單個層上使用了 layer-wise scaling，并在交替 transformer 塊中使用了混合專家層。

實驗

實驗采用擴散 Transformer（DiT）兩個變體 DiT-Tiny/2 和 DiT-Xl/2。

如圖 4 所示，延遲掩蔽方法在多個指標中都實現(xiàn)了更好的性能。此外，隨著掩蔽率的增加，性能差距會擴大。例如，在 75% 的掩蔽率下，naive 掩蔽會將 FID 得分降低到 16.5（越低越好），而本文方法可以達到 5.03，更接近沒有掩蔽的 FID 得分 3.79。

表 1 表明 layer-wise scaling 方法在擴散 transformer 的掩蔽訓練中具有更好的擬合效果。

比較不同的掩蔽策略。作者首先將本文方法與使用較大 patch 的策略進行比較。將 patch 大小從 2 增加到 4，相當于 75% 的 patch 掩蔽。與延遲掩蔽相比，其他方法表現(xiàn)不佳，分別僅達到 9.38、6.31 和 26.70 FID、Clip-FID 和 Clip-score。相比之下，延遲掩蔽分別達到 7.09、4.10 和 28.24 FID、Clip-FID 和 Clip-score。

下圖為延遲掩蔽 vs. 模型縮小以減少訓練成本的比較。在掩蔽率達到 75% 之前，作者發(fā)現(xiàn)延遲掩蔽在至少三個指標中的兩個方面優(yōu)于網(wǎng)絡(luò)縮小。但是，在極高的掩蔽率下，延遲掩蔽往往會實現(xiàn)較低的性能。這可能是因為在這些比率下掩蔽的信息損失太高導致的。

表 5 提供了有關(guān)模型訓練超參數(shù)的詳細信息。訓練過程分兩個階段。

計算成本。表 2 提供了每個訓練階段的計算成本明細，包括訓練 FLOP 和經(jīng)濟成本。第 1 階段和第 2 階段訓練分別消耗了總計算成本的 56% 和 44%。模型在 8×H100 GPU 集群上的總時鐘訓練時間為 2.6 天，相當于在 8×A100 GPU 集群上為 6.6 天。

了解更多結(jié)果，請參考原論文。