填補領域空白！TerDiT：首次探索大規模DiT模型量化問題 精華
論文鏈接：https://arxiv.org/pdf/2405.14854
項目鏈接：https://github.com/Lucky-Lance/TerDiT
最近在大規模預訓練的文本到圖像擴散模型方面的發展顯著提高了高保真圖像的生成能力，特別是基于transformer架構的擴散模型（DiTs）的出現。在這些擴散模型中，擴散transformer展示了卓越的圖像生成能力，降低了FID分數并提高了可擴展性。然而，由于其龐大的參數數量，部署大規模的DiT模型可能會非常昂貴。盡管現有研究已經探索了擴散模型的高效部署技術，如模型量化，但關于基于DiT模型的研究仍然很少。為了解決這一研究空白，本文提出了TerDiT，一種面向量化感知訓練（QAT）和高效部署的基于transformer的三值化擴散模型方案。本文專注于DiT網絡的三值化化，并將模型規模從600M擴展到4.2B。本文的工作為大規模DiT模型的高效部署策略探索做出了貢獻，證明了從頭訓練極低比特的擴散transformer模型的可行性，同時在圖像生成能力上保持與全精度模型相當的競爭力。
介紹
大規模預訓練的文本到圖像擴散模型的進步已經成功生成了復雜且高度保真于輸入條件的圖像。值得注意的是，基于transformer架構的擴散模型（DiTs）的出現代表了該研究領域的一個重要進展。與其他擴散模型相比，擴散transformer展示了在更高計算量下實現更低FID分數的能力。最新的研究突出了擴散transformer架構在圖像生成方面的顯著能力，例如在Stable Diffusion 3方法中展示的成果，以及在視頻生成方面的出色表現，如Sora所展示的工作。
鑒于擴散transformer模型的出色表現，研究人員現在越來越多地深入研究這些視覺模型的擴展規律，這與大語言模型相似。例如，Stable Diffusion 3提供了一系列參數規模的訓練DiT模型，從8億到80億。此外，有研究人員推測Sora可能擁有大約30億個參數。由于這些模型擁有龐大的參數數量，部署這些DiT模型往往成本高昂，尤其是在某些終端設備上。
為了應對部署難題，最近已經有一些關于高效部署擴散模型的研究工作，其中大多數集中在模型量化方面。然而，據本文所知，目前的研究仍然存在兩個主要缺陷。首先，雖然量化基于 U-Net 的擴散模型已經受到了很多關注，但對于基于transformer的擴散模型的量化方法探索仍然非常有限。其次，目前文獻中的大多數主流方法主要依賴于后訓練量化 (PTQ) 技術來進行模型量化，這會導致不可接受的性能下降，特別是在極低比特寬度（例如2比特和1比特）下。然而，神經網絡的極低比特量化非常重要，因為它可以顯著減少部署所需的計算資源，尤其對于具有巨大參數規模的模型。在本文的研究過程中，本文發現目前還沒有研究考慮 DiT 模型的極低比特量化。
為了解決這些缺陷，本文提出利用量化感知訓練（QAT）技術對大規模 DiT 模型進行極低比特量化。在大規模模型領域，低比特 QAT 方法已在大語言模型（LLM）領域進行了討論。最近的研究表明，從頭開始訓練具有極低比特參數（例如二進制和三進制）的大語言模型，也可以達到與全精度模型相當的競爭性能。這表明大規模模型中仍然存在顯著的精度冗余，并且暗示了 QAT 方案對于大規模 DiT 模型的可行性。
在本文中，本文主要關注三值權重網絡，并提供了 TerDiT，這是本文所知的首個用于 DiT 的量化方案。本文的方法實現了三值擴散transformer模型的量化感知訓練（僅限權重）和高效部署。與 LLM 和 CNN 中線性層的簡單量化不同，本文發現直接對 DiT 模塊中的 adaLN 模塊進行權重三值化，會導致歸一化層中的大尺寸尺度和偏移值（由于權重量化和梯度近似），這與全精度模型相比，導致收斂速度較慢和模型性能較差。因此，本文提出了一種 adaLN 的變體，通過在 adaLN 模塊的三值線性層之后應用 RMS Norm，有效地緩解了這一訓練問題。
通過這種修改，本文將三值 DiT 模型的參數規模從 600M（DiT-XL/2的規模）擴展到 4.2B（Large-DiT-4.2B的規模），發現具有更多參數的模型能夠收斂到更好的結果。本文進一步采用現有的2-bit CUDA 內核來部署訓練后的三值 DiT 模型，使模型checkpoint 大小減少了十倍以上，推理內存消耗減少了約六倍，同時實現了具有競爭力的生成質量。主要貢獻總結如下：
- 受低位 LLMs 量化感知訓練方案的啟發，本文研究了針對三值 DiT 模型的 QAT 方法，并引入了 DiT 特定的改進以獲得更好的訓練效果，這在 DiT 文獻中尚未被探索。
- 本文將三值DiT模型的參數規模從600M擴展到4.2B，并基于現有的2-bit CUDA內核在GPU上部署了訓練后的三值DiT模型，使得4.2B DiT模型的推理內存消耗小于3GB。
- 與全精度模型在ImageNet基準測試（圖像生成）中的對比評估結果展示了本文提出的TerDiT方案的有效性。
本文的研究是首次嘗試探索DiT模型的量化問題。本文專注于量化感知訓練和大規模三值DiT模型的高效部署，為未來研究在極低比特精度下部署DiT模型提供了寶貴的見解。
相關工作
擴散模型。 近年來，擴散模型因其生成高質量圖像的能力和多種應用潛力而受到了廣泛關注。擴散模型的概念最早由提出，該研究提出了一種通過學習逆向擴散過程的生成模型。這項工作為該領域的后續研究奠定了基礎。[1]進一步擴展了這一思想，提出了去噪擴散概率模型（DDPMs），這類模型已成為圖像生成任務中的熱門選擇。DDPMs被應用于廣泛的領域，包括無條件圖像生成、圖像修復和圖像超分辨率。此外，擴散模型還被用于文本到圖像的合成，如DALL-E模型和Imagen模型所展示的那樣，這些模型展示了擴散模型從文本描述中生成高度逼真和多樣化圖像的能力。進一步地，擴散模型還被擴展到其他模態，如音頻合成和視頻生成，展示了其在多模態應用中的多樣性和潛力。
擴散模型的量化。 近年來，為了提高擴散模型的效率，研究人員對擴散模型的量化進行了研究。后訓練量化（PTQ）方法，如[9, 11, 13, 14, 15]中提出的方法，在量化時間和數據使用方面具有優勢。然而，當這些方法應用于低比特設置時，通常會導致性能不佳。為了解決這個問題，[31]提出了將量化感知低秩適配器（QALoRA）與PTQ方法相結合，從而提高評估結果。作為PTQ的替代方案，量化感知訓練（QAT）方法專門用于低比特擴散模型的量化。盡管這些QAT方法有效，但目前僅限于小規模的基于U-Net的擴散模型，這揭示了在大規模DiT模型上應用QAT的研究空白。進一步探索適用于大規模DiT模型的極低比特寬度QAT技術，可能會帶來更大的效率提升，并在資源受限的環境中有效部署擴散模型。
三值權重網絡。 三值權重網絡作為一種內存高效和計算高效的網絡結構，已經引起了廣泛關注，其在推理內存使用上的顯著減少潛力尤為突出。在專用硬件的支持下，三值權重網絡還可以提供顯著的計算加速。在量化方法中，三值權重網絡備受關注，主要有兩種方法：僅權重量化和權重-激活量化。在僅權重量化中，如[35]所述，僅對權重進行三值量化。而權重-激活量化，如[36, 37]所述，則對權重和激活值同時進行三值量化。近期研究表明，三值權重網絡在訓練大語言模型方面具有可行性，并且其結果可與全精度模型相媲美。基于這些進展，本文的工作首次引入了針對三值DiT模型的量化感知訓練和高效部署方案。通過在DiT模型中利用三值量化的優勢，本文旨在推動效率的極限，并使強大的擴散模型在資源受限的環境中得以部署，從而為實際應用開辟新的可能性。
TerDiT
TerDiT，這是一個用于進行僅權重量化感知訓練和高效部署大規模三值DiT模型的框架。本文首先簡要回顧擴散transformer（DiT）模型。然后，基于之前開源的Large-DiT，闡述了量化函數和量化感知訓練方案，并進行特定于QAT的模型結構改進以優化網絡訓練，并介紹了三值部署方案。
擴散transformer模型
擴散transformer（Diffusion Transformer）。擴散transformer（DiT）是一種架構，它用操作潛在patches的transformer替代了擴散模型中常用的U-Net骨干結構。類似于下圖2（C）中展示的視覺transformer（ViT）架構，DiT首先將空間輸入劃分為一系列tokens，然后通過一系列transformer塊（下圖2（B））進行去噪處理。為了處理額外的條件信息（例如噪聲時間步t、類別標簽l、自然語言輸入），DiT利用自適應歸一化模塊（adaLNZero）將這些額外的條件輸入插入到transformer塊中。在最后一個transformer塊之后，應用標準線性解碼器來預測最終的噪聲和協方差。DiT模型的訓練方式與基于U-Net的擴散模型相同。
DiT中的AdaLN模塊。DiT與傳統ViT的主要區別在于需要注入條件信息以進行圖像生成。DiT在每個transformer塊中使用零初始化的自適應層歸一化（adaLN-Zero）模塊，如上圖2（B）紅色部分所示，該模塊根據輸入條件c計算維度級的縮放和偏移值。
AdaLN 是 DiT 模型中的一個重要組件，其效果已被證明優于交叉注意力和上下文條件方法。在 DiT 架構中，AdaLN 模塊集成了一個包含大量參數的 MLP 層，占模型總參數的約 10% 到 20%。在 TerDiT 的訓練過程中，本文觀察到直接對該模塊進行權重三值化會導致不理想的訓練結果。
模型量化
如上文所示，理解DiT模型的擴展規律越來越受到關注，這對于開發和優化大語言模型（LLM）至關重要。在最近的探索中，Large-DiT成功地將模型參數從600M擴展到7B，結合了LLaMA和DiT的方法。結果表明，參數擴展可以潛在地提升模型性能，并加快標簽條件的ImageNet生成任務的收斂速度。受此啟發，本文提出進一步研究DiT模型的三值化，這可以緩解部署大規模DiT模型相關的挑戰。在本小節中，本文介紹量化函數和量化感知訓練方案。
填補領域空白！TerDiT：首次探索大規模DiT模型量化問題-AI.x社區
填補領域空白！TerDiT：首次探索大規模DiT模型量化問題-AI.x社區
TerDiT 是一種僅對權重進行量化的方案，本文不對激活值進行量化。
量化感知訓練方案。 基于上述設計的量化函數，本文從頭開始訓練一個 DiT 模型，利用直接傳遞估計器（STE），允許梯度通過不可微分的網絡組件傳播。在整個訓練過程中，本文保留網絡的全精度參數。對于每一步訓練，通過前向傳播中的三值量化函數從全精度參數計算出三值權重，并在反向傳播中將三值權重的梯度直接應用于全精度參數進行參數更新。
然而，本文發現收斂速度非常慢。即使經過多次訓練迭代，損失值也無法降低到合理范圍。本文認為這個問題可能源于三值線性層通常會導致較大的激活值，并提出在接下來的小節中通過針對 QAT（量化感知訓練）特定的模型結構改進來解決這個問題。
QAT特定模型結構改進
填補領域空白！TerDiT：首次探索大規模DiT模型量化問題-AI.x社區
通過對三值線性層的輸出應用層歸一化，可以緩解由三值線性權重帶來的大激活值問題。本文在三值線性層之后添加了一個RMS歸一化層（類似于LLaMA），并獲得了激活值分布（如下圖3左側所示）。在這種情況下，激活值在通過歸一化層后被縮放到一個合理范圍，從而導致更穩定的訓練行為。這一觀察結果也與[17]中的結論一致，其中在每個量化線性層的激活量化之前應用了層歸一化函數。
RMS歸一化的AdaLN模塊。 基于上述見解，本文分析了DiT模型以改進QAT特定的模型結構。在標準的ViT Transformer塊中，層歸一化應用于每個自注意力層和前饋層。DiT塊中的自注意力層和前饋層也是如此，這有助于適當地縮放激活值范圍。然而，由于上文中介紹的AdaLN模塊的存在，DiT塊與傳統的Transformer塊有所不同。值得注意的是，這個模塊沒有應用層歸一化。在全精度訓練的情況下，缺乏層歸一化并不會產生顯著影響。然而，對于三值DiT網絡來說，其缺失可能會導致adaLN（歸一化）模塊中的維度尺度和偏移值過大，從而對模型訓練產生不良影響。為了解決這個問題，本文在每個三值DiT塊的AdaLN模塊的MLP層之后引入了RMS歸一化：
最終的TerDiT模型結構如上圖2(A)所示。這個小的修改可以帶來更快的收斂速度和更低的訓練損失，從而在定量和定性評估中取得更好的結果。為了更好地展示這一效果，在原文附錄中分析了模型訓練后引入或不引入RMS Norm的實際激活分布。
部署方案
在訓練了DiT模型之后，本文發現目前沒有有效的開源三值網絡部署解決方案。在這種情況下，本文使用2位實現來部署訓練好的網絡。具體來說，本文使用文獻[44]提供的??pack_2bit_u8()???函數，將三值線性權重打包成??int8???值（4個三值數打包成一個??int8???數）。在DiT模型的推斷過程中，本文即時調用相應的??unpack_2bit_u8()??函數，將打包的2位數字恢復為浮點數值，然后進行后續計算。添加解包操作會減慢推斷過程，但本文相信，隨著對模型三值化研究的深入，將會有更多硬件支持來加速推斷過程。
實驗
在本節中，本文進行了一系列實驗來評估本文提出的TerDiT。本文展示了主要的評估結果，進行了部署效率比較，并說明了RMS Normalized adaLN模塊的有效性。本文的DiT實現基于開源代碼Large-DiT-ImageNet4。本文分別對具有600M（DiT-XL/2的大小）和4.2B（Large-DiT-4.2B的大小）參數的三值DiT模型進行了實驗。
主要評價結果
本文在本小節中提供了TerDiT模型的定量和定性評估結果。據本文所知，目前尚無關于擴散transformer模型量化的已發表工作，因此本文主要在本小節中將其與具有代表性的全精度擴散模型進行比較。
關于TerDiT基線的備注。 據本文所知，目前仍沒有研究DiT模型量化的工作。除了在本小節中與全精度模型進行比較外，本文還在其他小節中建立了一些基線進行比較。對于QAT基線，本文直接訓練了在Sec. 4.3中的adaLN模塊中沒有RMS Norm的三值DiT模型。為了與現有的PTQ方法進行比較，本文對預訓練模型進行了4位權重量化，使用與TerDiT相同的一組參數，結果發現它們無法生成可視的圖像。
實驗設置。 按照原始DiT論文的評估設置，本文在ImageNet數據集上訓練了600M和4.2B的三值DiT模型。由于計算資源的限制，本文在256×256分辨率下訓練和評估模型，但本文認為評估結果已經具有很強的代表性。本文將TerDiT與一系列全精度擴散模型進行比較，并報告FID、sFID、Inception Score、Precision和Recall（50k生成圖像），參考[48]。本文還提供了訓練階段的總圖像數量（百萬），如[23]所示，以進一步了解不同生成模型的收斂速度。
填補領域空白！TerDiT：首次探索大規模DiT模型量化問題-AI.x社區
定量結果分析。 評估結果列在下表1中。TerDiT是針對DiT模型的QAT方案，因此在所有全精度模型中，本文特別關注DiT-XL/2（675M）和Large-DiT-4.2B。在沒有分類器自由指導的情況下，TerDiT-4.2B在測試結果上與DiT-XL/2非常相似（使用的訓練圖像數量要少得多）。在有分類器自由指導（cfg=1.5）的情況下，TerDiT-4.2B-G的表現優于LDM-G，同時與兩個全精度DiT結構模型相比僅帶來了非常輕微的性能下降。此外，TerDiT-4.2B-G的評估結果優于TerDiT-600M-G，這表明參數更多的模型在量化后會帶來更小的性能下降。
為了直觀地展示TerDiT的有效性，本文在下圖4中展示了一些定性比較結果，涉及TerDiT-4.2B、DiT-XL/2和Large-DiT4.2B。從視覺感知的角度來看，TerDiT生成的圖像與全精度模型生成的圖像之間沒有顯著差異。
部署效率對比
部署效率的提升是本文提出TerDiT方案的動機。在本小節中，本文對TerDiT-600M/4.2B、DiT-XL/2和Large-DiT-4.2B進行了比較，以討論TerDiT在實際部署中所能帶來的效率提升。下表2展示了四種DiT模型的checkpoint 大小。本文還記錄了在單個A100-80G GPU上，總擴散采樣循環（步數=250）的內存使用情況和推理時間。
從表格中可以看出，TerDiT大大減少了checkpoint 大小和內存使用。4.2B三值化DiT模型的checkpoint 大小和內存使用顯著小于Large-DiT-4.2B，甚至比DiT-XL/2還要小。這為在終端設備（如手機）上部署模型帶來了顯著優勢。盡管由于需要解包操作，本文觀察到推理速度較慢，但本文相信，隨著未來硬件支持的提升，三值化權重網絡的計算優勢將會得到充分展示。
RMS歸一化AdaLN模塊的討論
TerDiT對DiT模型結構的主要修改是在adaLN模塊中的MLP之后增加了RMS Norm。在這一部分，本文與基線三值化模型進行比較，以展示RMS Norm對訓練過程和訓練結果的影響。
實驗設置。 本文在ImageNet數據集上以256×256分辨率訓練具有600M和4.2B參數的三值化DiT模型。對于每種參數規模，本文訓練了兩個模型，一個在adaLN模塊中使用了RMS Norm，另一個則沒有（本文的基線模型）。本文記錄了訓練過程中的損失曲線，并每100k訓練步測量一次FID-50k分數（cfg=1.5）。
填補領域空白！TerDiT：首次探索大規模DiT模型量化問題-AI.x社區
結果分析： 訓練損失和評估得分分別顯示在下圖5和下圖6中。如圖所示，使用RMS Normalized adaLN模塊進行訓練將導致更快的收斂速度和更低的FID分數。另一個觀察結果是，參數更多的模型相比參數較少的模型能實現更快且更好的訓練效果。這在一定程度上也反映了三值化DiT模型的擴展規律。
討論和未來展望
本文在成功的大語言模型低比特訓練方法的基礎上，提出了針對大規模三值化DiT模型的量化感知訓練（QAT）和高效部署方法。在ImageNet數據集（256×256）上的競爭性評估結果證明了從頭開始訓練大型三值化DiT模型的可行性，同時實現了與全精度模型相當的結果。據本文所知，這是首個關于DiT模型量化的研究。
雖然本文認為這項工作為DiT模型的低比特量化提供了有價值的見解，但它仍然存在一些局限性。首先，訓練三值化DiT比全精度網絡更不穩定且耗時。在本文中，盡管本文討論了通過添加歸一化方法來使訓練更穩定，但相較于訓練全精度網絡（如Large-DiT-4.2B），它仍然更耗時，這將在更廣泛的背景下導致模型訓練期間二氧化碳排放量的增加。
其次，由于計算資源的限制，本文沒有進行ImageNet 512×512實驗，也沒有進行文本到圖像生成任務的實驗。然而，本文相信ImageNet 256×256基準上的評估結果已經相當具有代表性。剩余的任務將留待本文未來的工作中進行。本文希望本文的工作可以減少圖像生成模型的部署需求，并能激勵社區在未來加入本文，共同促進這一研究領域的更廣泛發展。
本文轉自 AI生成未來 ，作者：Xudong Lu等
