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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

方法介紹

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

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

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

使用 patch-mixer 訓(xùn)練擴(kuò)散 transformer

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

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

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

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

未掩蔽微調(diào)

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

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

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

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

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

實(shí)驗(yàn)

實(shí)驗(yàn)采用擴(kuò)散 Transformer（DiT）兩個(gè)變體 DiT-Tiny/2 和 DiT-Xl/2。

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

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

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

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

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

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

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