降低擴(kuò)散模型生成的計(jì)算成本，性能還保持在高水平！

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

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

目前，擴(kuò)散模型在生成任務(wù)中受到了廣泛關(guān)注。然而，昂貴的計(jì)算成本阻礙了擴(kuò)散模型在資源受限場景中的應(yīng)用。

量化是一種實(shí)用的解決方案，通過減少參數(shù)的位寬來顯著節(jié)省存儲和計(jì)算。然而，現(xiàn)有的擴(kuò)散模型量化方法仍然會導(dǎo)致性能嚴(yán)重下降，特別是在極低的位寬（2-4位）下。

性能的主要下降來自低位量化時(shí)激活值的顯著離散化。太少的激活候選對離群值顯著權(quán)重通道量化不友好，離散特征阻礙了擴(kuò)散模型在不同時(shí)間步長上的穩(wěn)定學(xué)習(xí)。

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

△圖2 MPQ-DM概述

MPQ-DM主要依賴于兩種技術(shù)：

• 為減輕異常值權(quán)重信道引起的量化誤差，團(tuán)隊(duì)提出了一種異常值驅(qū)動的混合量化（OMQ）技術(shù)，該技術(shù)使用峰度來量化異常值顯著信道，并應(yīng)用優(yōu)化的層內(nèi)混合精度比特寬度分配，以在目標(biāo)效率范圍內(nèi)恢復(fù)精度性能。
• 為了魯棒地學(xué)習(xí)跨越時(shí)間步長的表示，團(tuán)隊(duì)在量化擴(kuò)散模型與其全精度對應(yīng)物之間構(gòu)建了一個時(shí)間平滑關(guān)系蒸餾（TRD）方案，將離散和連續(xù)的潛勢轉(zhuǎn)移到統(tǒng)一的關(guān)系空間，以減少表示的不一致性。

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

相關(guān)工作

擴(kuò)散模型

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

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

擴(kuò)散模型量化

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

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

以下工作PTQ-D，TFMQ-DM，APQ-DM和QuEST在量化誤差方向、瞬態(tài)特征、校準(zhǔn)數(shù)據(jù)和校準(zhǔn)模塊方面進(jìn)行了改進(jìn)。量化后的擴(kuò)散模型的性能得到了進(jìn)一步的提升。

然而，基于PTQ的方法在極低的比特寬度下性能嚴(yán)重下降。為了結(jié)合QAT的優(yōu)點(diǎn)并減少所需的訓(xùn)練時(shí)間，EfficientDM使用LoRA方法對量化擴(kuò)散模型進(jìn)行微調(diào)。然而，這兩種有效的量化方法都不能保證差分模型在低位下的性能。因此，本文的重點(diǎn)是最大限度地提高極低位量化擴(kuò)散模型的性能。

方法

模型量化

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

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

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

離群值指導(dǎo)的混合精度量化

擴(kuò)散模型的一些研究發(fā)現(xiàn)，存在一些具有顯著權(quán)重值的通道，這些通道對量化誤差和模型性能也至關(guān)重要。

團(tuán)隊(duì)在圖3a中可視化了擴(kuò)散模型的一層權(quán)重分布，在圖3b中可視化了其對應(yīng)的異常顯著通道分布?？梢园l(fā)現(xiàn)，不僅不同通道的權(quán)重在數(shù)值范圍內(nèi)存在較大差距，而且一些權(quán)重通道也存在嚴(yán)重的異常值。

△圖3 LDM-4 ImageNet 256x256的權(quán)重分布

在圖3c和圖3d中，團(tuán)隊(duì)可視化了兩個通道的權(quán)重分布。

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

這種誤差在極低的位寬（2-4位）下會大大放大，因此團(tuán)隊(duì)自然希望為這些通道分配更多的位寬，以確保它們的性能。然而，在特定通道中的這種異?，F(xiàn)象可能發(fā)生在模型的各個層。

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

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

峰度κ可用于量化實(shí)值變量概率分布的“尾性”。這自然與團(tuán)隊(duì)的目標(biāo)相一致，因?yàn)楫惓Ｖ祵惓５拈L尾添加到正態(tài)分布中，如上所述。如圖3c和圖3d所示，與沒有異常值的正態(tài)分布相比，具有更顯著異常值現(xiàn)象的通道具有更高的κ值。因此，團(tuán)隊(duì)使用峰度κ作為量化不同權(quán)重通道難度的指標(biāo)，為每層權(quán)重W_f找到最佳比特分配方法。

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

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

其中N是目標(biāo)平均比特寬度，n是權(quán)重中的通道數(shù)。 表示量化和去量化過程。

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

為了加速優(yōu)化過程，團(tuán)隊(duì)將k個通道設(shè)置為一個搜索組，離群顯著通道的搜索區(qū)域約束在中。團(tuán)隊(duì)根據(jù)經(jīng)驗(yàn)將k設(shè)置為，因此搜索時(shí)間是5倍，這對于層比特分配來說足夠有效。

時(shí)間步平滑的關(guān)系蒸餾

對于擴(kuò)散優(yōu)化，現(xiàn)有的方法，如EfficientDM，通過對齊全精度（FP）模型和量化模型的輸出來優(yōu)化量化模型的訓(xùn)練參數(shù):

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

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

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

在極低位量化中，僅對齊模型的最終輸出是不夠的，因?yàn)榱炕Ｐ偷谋磉_(dá)能力嚴(yán)重不足。團(tuán)隊(duì)通常將量化模型的最終投影層設(shè)置為8位，并將投影層之前的層設(shè)置為低位。這樣，團(tuán)隊(duì)就不會直接感知到量化信息嚴(yán)重缺失的部分。因此，團(tuán)隊(duì)可以通過在最后一個項(xiàng)目層之前提取模型特征層，為極低位量化模型提供更細(xì)粒度的指導(dǎo)

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

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

這種離散化特征在數(shù)值上顯示出高度的不真實(shí)性，直接提取它效果不佳，甚至?xí)蓴_量化模型的正常訓(xùn)練。此外，由于擴(kuò)散模型的獨(dú)特迭代去噪過程，特征的輸出將繼續(xù)與去噪相加，進(jìn)一步放大這種不真實(shí)性。

通常，團(tuán)隊(duì)可以使用額外的投影頭將它們映射到均勻空間或?qū)ζ溥M(jìn)行數(shù)值正則化。然而，投影頭會帶來額外的訓(xùn)練參數(shù)，進(jìn)行正則化并不能解決迭代計(jì)算中的誤差累積問題。團(tuán)隊(duì)發(fā)現(xiàn)擴(kuò)散模型的特征與時(shí)間步長高度相關(guān)，并在某些時(shí)間步長中表現(xiàn)出相似性。

在圖4中，團(tuán)隊(duì)發(fā)現(xiàn)連續(xù)時(shí)間步長上的特征高度相似，而相距甚遠(yuǎn)的時(shí)間步長則截然不同。特征的相似性表明，去噪軌跡中連續(xù)時(shí)間步長的差異很小，因此團(tuán)隊(duì)可以通過融合多個連續(xù)步驟的中間特征來緩解不同時(shí)間步長的特征不真實(shí)性。

而不是強(qiáng)迫高度離散化的量化模型嚴(yán)格學(xué)習(xí)FP模型每一步的去噪軌跡，而是學(xué)習(xí)連續(xù)多步的去噪聲軌跡。

因此，團(tuán)隊(duì)使用N個連續(xù)步驟的中間特征作為平滑特征表示進(jìn)行提取。對于T時(shí)間步長優(yōu)化，團(tuán)隊(duì)將蒸餾公式改寫為

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

在圖5a中，團(tuán)隊(duì)可視化了FP模型、訓(xùn)練好的量化模型和未訓(xùn)練的量化模型之間的特征圖。

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

因此，建議使用關(guān)系對齊來代替嚴(yán)格的學(xué)習(xí)數(shù)值對齊和之間的特征相似關(guān)系。

在圖5b，團(tuán)隊(duì)將離散潛空間和連續(xù)潛空間之間的數(shù)值關(guān)系轉(zhuǎn)化為每個空間內(nèi)的特征相似關(guān)系，將蒸餾目標(biāo)統(tǒng)一到特征相似空間中。

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

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

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

團(tuán)隊(duì)在這里使用KL散度而不是L2損失。因?yàn)長2損失只能感知單個表示，而KL散度可以感知整個特征圖的信息?？傮w優(yōu)化目標(biāo)為

實(shí)驗(yàn)

實(shí)驗(yàn)設(shè)置

團(tuán)隊(duì)在常用的數(shù)據(jù)集LSUN-Bedrooms 256×256、LSUN-Churchs 256×256和ImageNet 256×256上進(jìn)行了LDM模型上無條件和有條件圖像生成任務(wù)的實(shí)驗(yàn)。

還對Stable Diffusion進(jìn)行了文本到圖像的生成任務(wù)，使用IS、FID、sFID和Precision來評估LDM的性能。對于Stable Diffusion，團(tuán)隊(duì)使用CLIP評分進(jìn)行評估。

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

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

實(shí)驗(yàn)結(jié)果

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

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

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

△表 2

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

△表 3

類條件生成

團(tuán)隊(duì)以LDM-4為重點(diǎn)，在ImageNet 256×256數(shù)據(jù)集上進(jìn)行了條件生成實(shí)驗(yàn)。

表1中的結(jié)果顯示，MPQ-DM在所有比特設(shè)置上都大大優(yōu)于現(xiàn)有方法。MPQ-DM通常比分層方法HAWQ-v3表現(xiàn)更好，這證明了在層內(nèi)進(jìn)行混合精度量化的必要性。

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

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

△表4

無條件的生成

團(tuán)隊(duì)在LDM-4的LSUN臥室數(shù)據(jù)集和LDM-8的LSUN教堂數(shù)據(jù)集上進(jìn)行了無條件生成實(shí)驗(yàn)，分辨率為256×256。在表2和表3中，MPQ-DM在所有比特設(shè)置下仍然優(yōu)于所有其他現(xiàn)有方法。

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

文本到圖像生成

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

消融實(shí)驗(yàn)

組件研究

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

團(tuán)隊(duì)提出的OMQ從量化的角度解決了現(xiàn)有的逐層比特定位方法，在總平均比特寬度不變的情況下，將更多的比特寬度定位到層內(nèi)具有顯著異?，F(xiàn)象的信道。

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

通過量化和優(yōu)化兩個方面的并行改進(jìn)，MPQ-DM實(shí)現(xiàn)了最先進(jìn)的量化性能。

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

異常值選擇方法研究

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

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

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

異常值選擇函數(shù)的消融研究：

△表6 抽取10k個樣本進(jìn)行評估

蒸餾指標(biāo)研究

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

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

這表明離散特征和連續(xù)特征不能直接通過數(shù)值很好地對齊，從而導(dǎo)致負(fù)優(yōu)化。然而，團(tuán)隊(duì)提出的關(guān)系蒸餾可以將所有特征轉(zhuǎn)移到統(tǒng)一的相似性空間中。這打破了兩個潛在空間之間的差異，提高了模型性能。

蒸餾度量的消融研究：

△表7 抽取10k個樣本進(jìn)行評估

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