降低擴散模型生成的計算成本，性能還保持在高水平！

最新研究提出一種用于極低位差分量化的混合精度量化方法。

△圖1 W2A6位寬下基線和MPQ-DM生成的樣本可視化

目前，擴散模型在生成任務中受到了廣泛關注。然而，昂貴的計算成本阻礙了擴散模型在資源受限場景中的應用。

量化是一種實用的解決方案，通過減少參數的位寬來顯著節省存儲和計算。然而，現有的擴散模型量化方法仍然會導致性能嚴重下降，特別是在極低的位寬（2-4位）下。

性能的主要下降來自低位量化時激活值的顯著離散化。太少的激活候選對離群值顯著權重通道量化不友好，離散特征阻礙了擴散模型在不同時間步長上的穩定學習。

為此，來自中國科學院計算技術研究所、蘇黎世聯邦理工學院、北京交通大學的研究人員提出了由離群值驅動的混合量化（OMD）和時間平滑關系蒸餾（TRD）組成的極低位擴散模型（MPQ-DM）的混合精度量化。

△圖2 MPQ-DM概述

MPQ-DM主要依賴于兩種技術：

• 為減輕異常值權重信道引起的量化誤差，團隊提出了一種異常值驅動的混合量化（OMQ）技術，該技術使用峰度來量化異常值顯著信道，并應用優化的層內混合精度比特寬度分配，以在目標效率范圍內恢復精度性能。
• 為了魯棒地學習跨越時間步長的表示，團隊在量化擴散模型與其全精度對應物之間構建了一個時間平滑關系蒸餾（TRD）方案，將離散和連續的潛勢轉移到統一的關系空間，以減少表示的不一致性。

綜合實驗表明，與SOTA量化方法相比，MPQ-DM在極低的比特寬度下實現了顯著的精度增益。與基線相比，在W2A4設置下，MPQ-DM的FID降低了58%，而所有其他方法甚至無法生成正常的圖像。

相關工作

擴散模型

擴散模型通過逐步向數據分布x0～q（x）添加噪聲來執行前向采樣過程。在DDPM中，擴散模型的前向噪聲添加過程是馬爾可夫鏈，其形式為：

其中α_t=1?β_t , β_t 是時間相關的時間表。擴散模型通過對隨機采樣的高斯噪聲 x_T～N(0, I)應用去噪過程來生成高質量圖像，形式如下：

擴散模型量化

訓練后量化（PTQ）和量化感知訓練（QAT）是模型量化的兩種主要方法。

常用的QAT方法，如LSQ和擴散模型Q-dm和Binarydm的方法，確保了模型在極低比特寬度甚至二進制量化下的性能，但與PTQ方法相比，它們需要大量額外的訓練時間，從而導致更大的訓練負擔。對擴散模型PTQ4DM和Q-diffusion的PTQ方法進行了初步探索。

以下工作PTQ-D，TFMQ-DM，APQ-DM和QuEST在量化誤差方向、瞬態特征、校準數據和校準模塊方面進行了改進。量化后的擴散模型的性能得到了進一步的提升。

然而，基于PTQ的方法在極低的比特寬度下性能嚴重下降。為了結合QAT的優點并減少所需的訓練時間，EfficientDM使用LoRA方法對量化擴散模型進行微調。然而，這兩種有效的量化方法都不能保證差分模型在低位下的性能。因此，本文的重點是最大限度地提高極低位量化擴散模型的性能。

方法

模型量化

模型量化將模型權重和激活映射到低位整數值，以減少內存占用并加速推理。對于浮動向量x_f，量化過程可以公式化為：

其中表示整數中的量化向量，?·?是舍入函數，clip(·)是將值限制在[0，2N?1]范圍內的函數，s是比例因子，z是量化零點。l和u分別是量化閾值的下限和上限。它們由x_f和目標位寬決定。相反，為了將低位整數量化向量恢復到全精度表示，反量化過程公式為：

其中是用于正向處理的去量化向量。

離群值指導的混合精度量化

擴散模型的一些研究發現，存在一些具有顯著權重值的通道，這些通道對量化誤差和模型性能也至關重要。

團隊在圖3a中可視化了擴散模型的一層權重分布，在圖3b中可視化了其對應的異常顯著通道分布。可以發現，不僅不同通道的權重在數值范圍內存在較大差距，而且一些權重通道也存在嚴重的異常值。

△圖3 LDM-4 ImageNet 256x256的權重分布

在圖3c和圖3d中，團隊可視化了兩個通道的權重分布。

總體權重分布可以被視為正態分布加上異常值分布。對于模型量化，正態分布的直接量化得到了很好的研究，但少量異常值的存在可能會導致一些異常值被箝位或占據部分比特寬度，從而難以量化主正態分布，導致量化誤差。

這種誤差在極低的位寬（2-4位）下會大大放大，因此團隊自然希望為這些通道分配更多的位寬，以確保它們的性能。然而，在特定通道中的這種異常現象可能發生在模型的各個層。

因此，使用傳統方法將不同的比特分配給不同的模型層進行混合精度量化顯然不是一個好的選擇。但是由于權重的量化是逐信道的，團隊可以在層內的不同信道之間重新分配量化比特寬度來解決這個問題，即層內混合精度量化。

為了確定具體的比特寬度分配方法，需要量化輸出的重要性。

峰度κ可用于量化實值變量概率分布的“尾性”。這自然與團隊的目標相一致，因為異常值將異常的長尾添加到正態分布中，如上所述。如圖3c和圖3d所示，與沒有異常值的正態分布相比，具有更顯著異常值現象的通道具有更高的κ值。因此，團隊使用峰度κ作為量化不同權重通道難度的指標，為每層權重W_f找到最佳比特分配方法。

團隊希望在混合比特寬度分配量化后，模型輸出盡可能與全精度模型保持一致。由于團隊的目標是模型輸出Y=XW_T，并且異常值是權重中的最大值或最小值，因此團隊可以在量化之前使用矩陣乘法的特性來降低異常值的顯著程度，而不會造成損失，如下所示：

其中Y、X和W表示輸出激活、輸入激活和模型權重。δ是一個通道平滑因子，通過將異常值縮放到更接近正態分布來平衡權重和激活的量化難度。經過預縮放后，離群顯著通道更平滑以進行量化。團隊為每個權重通道計算κ，并相應地對其進行排序。然后，團隊使用以下優化公式來確定異常突出和非客戶端通道，如下所示：

其中N是目標平均比特寬度，n是權重中的通道數。 表示量化和去量化過程。

團隊采用通道混合比特量化為，其中c_i表示第i個信道的比特寬度。C表示信道共享相同位寬的集合，|C|表示集合C中的信道數量。例如，對于3位量化，團隊將一些異常顯著信道分配給4位，并將相同數量的異常非客戶端信道重新分配給2位。這樣，層權重的平均總位寬仍然是3位，而不添加任何額外參數。

為了加速優化過程，團隊將k個通道設置為一個搜索組，離群顯著通道的搜索區域約束在中。團隊根據經驗將k設置為，因此搜索時間是5倍，這對于層比特分配來說足夠有效。

時間步平滑的關系蒸餾

對于擴散優化，現有的方法，如EfficientDM，通過對齊全精度（FP）模型和量化模型的輸出來優化量化模型的訓練參數:

其中θ_f和θ_q分別表示FP模型和量化模型，x_t是通過用FP模型迭代T?t步對高斯噪聲x_T～N(0, I)進行去噪而獲得的。

△圖4 從timesetp（左）T、（左中）T-1、（右中）T-2、（右）T-10的特征圖可視化

△圖5 （左）FP模型、（中）訓練好的量化模型、（右）未訓練量化模型的不同激活圖的可視化

在極低位量化中，僅對齊模型的最終輸出是不夠的，因為量化模型的表達能力嚴重不足。團隊通常將量化模型的最終投影層設置為8位，并將投影層之前的層設置為低位。這樣，團隊就不會直接感知到量化信息嚴重缺失的部分。因此，團隊可以通過在最后一個項目層之前提取模型特征層，為極低位量化模型提供更細粒度的指導

其中F_f和F_q分別表示FP模型和量化模型的最后一個項目層之前的特征圖。D是測量兩個特征圖距離的度量。

然而，極低位量化的特征呈現出高度的離散化（例如，4位激活量化只有16個值）。

這種離散化特征在數值上顯示出高度的不真實性，直接提取它效果不佳，甚至會干擾量化模型的正常訓練。此外，由于擴散模型的獨特迭代去噪過程，特征的輸出將繼續與去噪相加，進一步放大這種不真實性。

通常，團隊可以使用額外的投影頭將它們映射到均勻空間或對其進行數值正則化。然而，投影頭會帶來額外的訓練參數，進行正則化并不能解決迭代計算中的誤差累積問題。團隊發現擴散模型的特征與時間步長高度相關，并在某些時間步長中表現出相似性。

在圖4中，團隊發現連續時間步長上的特征高度相似，而相距甚遠的時間步長則截然不同。特征的相似性表明，去噪軌跡中連續時間步長的差異很小，因此團隊可以通過融合多個連續步驟的中間特征來緩解不同時間步長的特征不真實性。

而不是強迫高度離散化的量化模型嚴格學習FP模型每一步的去噪軌跡，而是學習連續多步的去噪聲軌跡。

因此，團隊使用N個連續步驟的中間特征作為平滑特征表示進行提取。對于T時間步長優化，團隊將蒸餾公式改寫為

其中F_T?t,f 和F_T?t,q 分別表示FP模型和量化模型在時間步長T-t中的最后一個特征圖。

在圖5a中，團隊可視化了FP模型、訓練好的量化模型和未訓練的量化模型之間的特征圖。

盡管時間平滑特征提高了量化特征的魯棒性，但訓練有素的模型和FP模型之間的數值表達式仍然存在不匹配。這被歸咎于量化特征的離散潛空間和FP特征的連續潛空間之間的差異。因此，任何指標，如L2損耗數字對齊和無法避免這種不同空間之間的差異。

因此，建議使用關系對齊來代替嚴格的學習數值對齊和之間的特征相似關系。

在圖5b，團隊將離散潛空間和連續潛空間之間的數值關系轉化為每個空間內的特征相似關系，將蒸餾目標統一到特征相似空間中。

這成功地解決了數值不匹配的問題。

形式上說，對于，為了簡化寫作，團隊對其進行了重寫其中s=h×w。對于第i個特征表示，團隊可以計算它的余弦相似性分布因此，關系蒸餾度量為

其中和表示時間平滑特征相似度分布分別來自FP模型和量化模型。D_kl(·||·)代表兩個分布之間的Kullback-Leibler（KL）散度。

團隊在這里使用KL散度而不是L2損失。因為L2損失只能感知單個表示，而KL散度可以感知整個特征圖的信息。總體優化目標為

實驗

實驗設置

團隊在常用的數據集LSUN-Bedrooms 256×256、LSUN-Churchs 256×256和ImageNet 256×256上進行了LDM模型上無條件和有條件圖像生成任務的實驗。

還對Stable Diffusion進行了文本到圖像的生成任務，使用IS、FID、sFID和Precision來評估LDM的性能。對于Stable Diffusion，團隊使用CLIP評分進行評估。

為了應對2bit量化下的極端壓縮性能下降，團隊在OMQ的搜索過程中為2bit額外分配了10%的信道，稱為MPQ-DM+。

與FP模型相比，模型尺寸僅增加了0.6%。團隊將MPQ-DM與基線EfficientDM和分層混合精度HAWQ-v3以及其他具有相似耗時的基于PTQ的方法進行了比較。

實驗結果

ImageNet 256×256上全量化LDM-4模型的性能比較：

△表1 最佳結果以粗體顯示，次佳結果以下劃線顯示

LDM-4模型在LSUN臥室256×256上的無條件圖像生成結果：

△表 2

LDM-8模型在LSUN Churches 256×256上的無條件圖像生成結果：

△表 3

類條件生成

團隊以LDM-4為重點，在ImageNet 256×256數據集上進行了條件生成實驗。

表1中的結果顯示，MPQ-DM在所有比特設置上都大大優于現有方法。MPQ-DM通常比分層方法HAWQ-v3表現更好，這證明了在層內進行混合精度量化的必要性。

MPQ-DM W3A4型號在FID上甚至超過了FP型號。在W2A4設置中，基于PTQ的方法無法生成圖像，而EfficientDM的性能較差。MPQ-DM大大改善了基線，FID顯著降低了27.86。MPQ-DM+僅使用4.8 MB的額外模型大小，就進一步導致FID減少9.48。

使用10k COCO2014驗證集提示的Stable Diffusion v1.4的文本到圖像生成結果（512×512）：

△表4

無條件的生成

團隊在LDM-4的LSUN臥室數據集和LDM-8的LSUN教堂數據集上進行了無條件生成實驗，分辨率為256×256。在表2和表3中，MPQ-DM在所有比特設置下仍然優于所有其他現有方法。

對于LSUN臥室數據集，與基線相比，W3A4的FID降低了5.59，W2A6降低了8.53，W2A4甚至降低了12.81。在W2A4設置下，團隊是第一種將sFID降至20以下的方法，與基線相比，sFID降低了6.12。

文本到圖像生成

團隊在512×512分辨率的Stable Diffusion v1.4模型上，對隨機選擇的10k COCO2014驗證集提示進行了文本到圖像生成實驗。在表4中，團隊的方法比基線和SOTA PTQ方法具有更好的性能。在W3A4和W2A6設置中，團隊實現了超過0.3的CLIP評分改善。MPQ-DM+在CLIP評分上進一步提高了1.79，只增加了10.2 MB的模型大小。

消融實驗

組件研究

在表5中，團隊對LDM-4 ImageNet 256×256模型進行了全面的消融研究，以評估每個擬議組件的有效性。

團隊提出的OMQ從量化的角度解決了現有的逐層比特定位方法，在總平均比特寬度不變的情況下，將更多的比特寬度定位到層內具有顯著異常現象的信道。

這種層內混合精度量化方法大大提高了基線的性能，增益IS增加了58.88。此外，TSD從模型優化的角度提高了蒸餾過程的魯棒性，也取得了一定的改進。

通過量化和優化兩個方面的并行改進，MPQ-DM實現了最先進的量化性能。

△表5 對提出方法的消融研究

異常值選擇方法研究

在表6中，團隊研究了LDM-4 ImageNet 256×256模型上混合精度量化中的不同異常顯著信道選擇方法。團隊發現，即使隨機選擇一些通道到更高或更低的比特，也會帶來一定的性能提升。

這表明，在極低比特量化中，增加某些信道的比特帶來的增益遠遠大于減少某些信道帶來的影響，這證明了混合量化的必要性。雖然在隨機或從權重的頭部和尾部選擇通道方面有一些收獲，但基于峰度的異常值選擇方法實現了最顯著的性能提升。

這表明峰度選擇了性能提高最顯著的異常顯著信道，這證明了基于峰度的信道選擇方法的有效性。

異常值選擇函數的消融研究：

△表6 抽取10k個樣本進行評估

蒸餾指標研究

在表7中，團隊研究了LDM-4 ImageNet模型上方程9中使用的不同蒸餾度量。

團隊與未蒸餾的模型進行對比。使用L2損耗來對齊和僅顯示sFID幾乎沒有改善，但降低了FID和精度。

這表明離散特征和連續特征不能直接通過數值很好地對齊，從而導致負優化。然而，團隊提出的關系蒸餾可以將所有特征轉移到統一的相似性空間中。這打破了兩個潛在空間之間的差異，提高了模型性能。

蒸餾度量的消融研究：

△表7 抽取10k個樣本進行評估

論文鏈接：https://arxiv.org/pdf/2412.11549