新智元報(bào)道

編輯：LRST 好困

【新智元導(dǎo)讀】拯救4bit擴(kuò)散模型精度，僅需時(shí)間特征維護(hù)——以超低精度量化技術(shù)重塑圖像內(nèi)容生成！

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

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

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

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

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

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

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

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

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

2. 時(shí)間信息失配：時(shí)間特征擾動(dòng)改變了原始嵌入的時(shí)間信息。具體來(lái)說(shuō)，旨在對(duì)應(yīng)于時(shí)間步長(zhǎng)。然而，由于存在顯著的誤差，量化模型的不再準(zhǔn)確地與相關(guān)聯(lián)，傾向于與對(duì)應(yīng)的時(shí)間特征更為接近，這導(dǎo)致了該方法所說(shuō)的時(shí)間信息不匹配；

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

（Diffusion中的時(shí)間特征擾動(dòng)）

擾動(dòng)誘因分析

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

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

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

（Diffusion中時(shí)間特征相關(guān)模塊有限激活）

時(shí)間特征維護(hù)-——TFMQ-DM

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

（TFMQDM整體壓縮框架）

1. 時(shí)間信息塊：將時(shí)間特征生成相關(guān)模塊進(jìn)行整合得到時(shí)間信息塊（見(jiàn)框架圖）：

2. 時(shí)間信息感知重建：基于時(shí)間信息塊，研究人員提出了時(shí)間信息感知重建 (TIAR) 來(lái)應(yīng)對(duì)第一個(gè)誘因。在重構(gòu)過(guò)程中，該塊的優(yōu)化目標(biāo)如下：

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

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

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

研究團(tuán)隊(duì)在DDIM，LDM以及Stable Diffusion系列模型上驗(yàn)證了無(wú)條件生成、分類條件生成、文本條件以及多種先進(jìn)采樣器生成下TFMQ-DM框架性能。

TFMQ-DM在平均4或8 bit權(quán)重，8或32 bit激活時(shí)，所有評(píng)價(jià)指標(biāo)均超過(guò)Q-Diffusion，PTQD等已有最先進(jìn)方法，且在大部分場(chǎng)景下，該方案使用4bit權(quán)重量化超越了已有技術(shù)在8-bit權(quán)重甚至于全精度權(quán)重的模型的性能。

結(jié)果表明，TFMQ-DM率先在4bit權(quán)重、8bit權(quán)重下實(shí)現(xiàn)了對(duì)于 Diffusion的無(wú)損壓縮加速，推動(dòng)了Diffusion無(wú)訓(xùn)練量化的邊界。

（LDM系列無(wú)條件生成對(duì)比結(jié)果）

（左：LDM系列分類條件生成對(duì)比結(jié)果；右：DDIM系列無(wú)條件生成對(duì)比結(jié)果）

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

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

（Stable Diffusion系列文本條件生成對(duì)比結(jié)果）

（左：PLMS系列先進(jìn)采樣器無(wú)條件生成對(duì)比結(jié)果右：DPM++系列先進(jìn)采樣器無(wú)條件生成對(duì)比結(jié)果）

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

然而，雖然現(xiàn)有的指標(biāo)無(wú)法充分評(píng)估生成圖像的語(yǔ)義一致性以及物體細(xì)節(jié)，該團(tuán)隊(duì)提出的方法產(chǎn)生了更高質(zhì)量的圖像（見(jiàn)后文），具有更真實(shí)的細(xì)節(jié)，更好地展示了語(yǔ)義信息。

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

（LDM上無(wú)條件圖像生成效果圖）

（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”。）

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

（Stable Diffusion在CPU上真實(shí)加速）

總結(jié)

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

其次TFMQ-DM 具有良好的兼容性，可與各種Diffusion量化框架無(wú)縫集成?？偟膩?lái)說(shuō)，TFMQ-DM 提供的顯著量化精度提升與其對(duì)于硬件的高度友好，有助于在資源受限的情況下進(jìn)行實(shí)際部署，進(jìn)一步促進(jìn)了文生圖模型的更廣泛普及和應(yīng)用。

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

參考資料：

https://arxiv.org/pdf/2311.16503