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