新智元報道

編輯：LRST 好困

【新智元導讀】拯救4bit擴散模型精度，僅需時間特征維護——以超低精度量化技術重塑圖像內容生成！

近日，北航、莫納什、UT Austin聯(lián)合推出了時間特征維護的擴散模型低精度無損量化方法TFMQ-DM。

不僅以4bit的權重大小實現(xiàn)了目前無損條件下最極限的擴散模型訓練后壓縮，同時還實現(xiàn)了超過2.38倍真實硬件加速。

這一發(fā)現(xiàn)再次將Diffusion壓縮推向全新的高度。

目前，該工作已被CVPR 2024高分接收，并被接收為Highlight Poster (Top 2.8%)。

論文地址：https://arxiv.org/pdf/2311.16503

代碼地址：https://github.com/ModelTC/TFMQ-DM

擴散模型由于引入了時間變量t，于是在模型中注入了時序信息，這正是擴散模型不同于以往傳統(tǒng)視覺模型的一大特征。

同時該變量也通過將時序特征融入模型去控制去噪過程。研究人員首次定義了時間特征emb，同時發(fā)現(xiàn)現(xiàn)有量化算法對于這些特征將產(chǎn)生嚴重擾動，從而破壞圖片生成質量：

1. 時間特征擾動：研究人員發(fā)現(xiàn)量化導致了明顯的時間特征誤差，并將這種特征錯誤內的擾動現(xiàn)象稱為時間特征擾動；

2. 時間信息失配：時間特征擾動改變了原始嵌入的時間信息。具體來說，旨在對應于時間步長。然而，由于存在顯著的誤差，量化模型的不再準確地與相關聯(lián)，傾向于與對應的時間特征更為接近，這導致了該方法所說的時間信息不匹配；

3. 去噪軌跡偏離：時間信息不匹配傳遞了錯誤的時間信息，因此導致圖像在去噪軌跡中對應的時間位置發(fā)生了偏差，最終導致圖片不再按原軌跡去噪：

（Diffusion中的時間特征擾動）

擾動誘因分析

研究人員發(fā)現(xiàn)該擾動主要由以下兩個原因造成：

1. 不合適的重建對象：已有量化重建方法并未直接優(yōu)化時間特征，同時時間特征將會受到有限的校準數(shù)據(jù)影響產(chǎn)生過擬合現(xiàn)象，如下圖Prev所示，其中Freeze代表凍結相關量化參數(shù)；

2. 忽略了時間特征相關模塊中的有限激活：由于輸入是有限整數(shù)，因此產(chǎn)生時間特征的模塊將僅產(chǎn)生有限且隨時間變化的激活，而已有量化策略均考慮分布層級優(yōu)化，忽略了對于此類有限激活的擬合近似。

（Diffusion中時間特征相關模塊有限激活）

時間特征維護-——TFMQ-DM

基于以上的誘因分析研究人員提出了如下時間特征維護策略，在低bit量化下完美的保證了Diffusion模型的時間特征精度與圖像生成質量。

（TFMQDM整體壓縮框架）

1. 時間信息塊：將時間特征生成相關模塊進行整合得到時間信息塊（見框架圖）：

2. 時間信息感知重建：基于時間信息塊，研究人員提出了時間信息感知重建 (TIAR) 來應對第一個誘因。在重構過程中，該塊的優(yōu)化目標如下：

3. 有限集校準：為了解決第二個誘因中有限集內激活范圍寬泛的挑戰(zhàn)，研究人員提出了有限集校準 (FSC) 用于激活量化。這個策略為所有時間信息塊內的每個激活使用T組量化參數(shù)，例如激活x的量化參數(shù)可為。在時間步長為t時，x的量化函數(shù)可以表示為：

其中sT, zt分別為量化縮放因子和零偏移。

實驗結果

研究團隊在DDIM，LDM以及Stable Diffusion系列模型上驗證了無條件生成、分類條件生成、文本條件以及多種先進采樣器生成下TFMQ-DM框架性能。

TFMQ-DM在平均4或8 bit權重，8或32 bit激活時，所有評價指標均超過Q-Diffusion，PTQD等已有最先進方法，且在大部分場景下，該方案使用4bit權重量化超越了已有技術在8-bit權重甚至于全精度權重的模型的性能。

結果表明，TFMQ-DM率先在4bit權重、8bit權重下實現(xiàn)了對于 Diffusion的無損壓縮加速，推動了Diffusion無訓練量化的邊界。

（LDM系列無條件生成對比結果）

（左：LDM系列分類條件生成對比結果；右：DDIM系列無條件生成對比結果）

具體來說，在 CelebA-HQ 256 × 256 數(shù)據(jù)集上，與當前最先進的方法相比，該團隊的方法在 w4a8 設置下將FID與sFID分別降低了 6.71和 6.60（越低越好）。

值得注意的是，現(xiàn)有方法，無論是4 bit還是8bit權重量化，在人臉數(shù)據(jù)集如 CelebA-HQ 256 × 256 和 FFHQ 256 × 256 上與 FP 模型相比都顯示出明顯的性能下降，而 TFMQ-DM 與全精度模型相比幾乎沒有性能損失。

（Stable Diffusion系列文本條件生成對比結果）

（左：PLMS系列先進采樣器無條件生成對比結果右：DPM++系列先進采樣器無條件生成對比結果）

此外，對于當下最流行的文本條件生成類模型Stable Diffusion，TFMQ-DM在 w8a8 設置下的 FID 和在 w4a32 設置下的 sFID 甚至略低于全精度模型。

然而，雖然現(xiàn)有的指標無法充分評估生成圖像的語義一致性以及物體細節(jié)，該團隊提出的方法產(chǎn)生了更高質量的圖像（見后文），具有更真實的細節(jié)，更好地展示了語義信息。

由于現(xiàn)有指標并不能完全反映生成圖像效果的優(yōu)劣，該團隊研究人員提供了大量可視化效果對比圖，體現(xiàn)出更加細膩和準確的生成質量：

（LDM上無條件圖像生成效果圖）

（LDM上分類條件圖像生成效果圖）

（Stable Diffusion上文本條件圖像生成效果圖，左提示詞：“A digital illustration of the Babel tower, detailed, trending in artstation, fantasy vivid colors”。右提示詞：“A beautiful castle beside a waterfall in the woods”。）

除測試集精度和可視化效果外，團隊人員還在Intel? Xeon? Gold 6248R 處理器上驗證框架的推理加速效果，相比原始浮點模型實現(xiàn)了 2.38 倍的顯著速度提升。

（Stable Diffusion在CPU上真實加速）

總結

基于時間特征維護的校準量化可確保 Diffusion的量化參數(shù)準確保留原始時間信息。廣泛的實驗證明，TFMQ-DM 在 DDIM、LDM及Stable-Diffusion 系列中實現(xiàn)了令人信服的精度提升，尤其是在w4a8等極低比特設置下表現(xiàn)出了明顯優(yōu)勢；同時該方案也實現(xiàn)了真實硬件部署加速。

其次TFMQ-DM 具有良好的兼容性，可與各種Diffusion量化框架無縫集成。總的來說，TFMQ-DM 提供的顯著量化精度提升與其對于硬件的高度友好，有助于在資源受限的情況下進行實際部署，進一步促進了文生圖模型的更廣泛普及和應用。

TFMQ-DM 具有顯著的兼容性，可與各種Diffusion量化框架無縫集成。TFMQ-DM 提供的顯著量化精度，有助于在資源受限的情況下進行實際部署。

參考資料：

https://arxiv.org/pdf/2311.16503