擴(kuò)散模型低位量化突破！有效擴(kuò)散量化的極限推向2-4位，W2A4位寬下FID降低58%，超越SOTA方法

中科院提出混合精度量化

原標(biāo)題：擴(kuò)散模型低位量化突破！有效擴(kuò)散量化的極限推向2-4位，W2A4位寬下FID降低58%，超越SOTA方法
文章來(lái)源：量子位
內(nèi)容字?jǐn)?shù)：11110字

極低位擴(kuò)散模型量化：MPQ-DM方法詳解

本文總結(jié)了MPQ-DM團(tuán)隊(duì)發(fā)表在arXiv上的最新研究，該研究提出了一種混合精度量化方法MPQ-DM，用于降低擴(kuò)散模型的計(jì)算成本，同時(shí)保持高水平的性能。現(xiàn)有擴(kuò)散模型量化方法在極低位寬（2-4位）下性能嚴(yán)重下降，主要原因是激活值的顯著離散化。MPQ-DM通過(guò)結(jié)合離群值驅(qū)動(dòng)的混合量化(OMQ)和時(shí)間平滑關(guān)系蒸餾(TRD)兩種技術(shù)來(lái)解決這個(gè)問(wèn)題。

1. 擴(kuò)散模型與量化

擴(kuò)散模型通過(guò)逐步添加噪聲然后去除噪聲來(lái)生成圖像。然而，其計(jì)算成本很高。量化通過(guò)降低參數(shù)的位寬來(lái)節(jié)省存儲(chǔ)和計(jì)算資源，但現(xiàn)有方法在低位寬下性能下降嚴(yán)重。

2. MPQ-DM的核心技術(shù)

MPQ-DM的核心在于OMQ和TRD兩項(xiàng)技術(shù)：

1. 離群值驅(qū)動(dòng)的混合量化 (OMQ): 該方法利用峰度(kurtosis)來(lái)識(shí)別權(quán)重通道中的離群值。離群值顯著的通道被分配更高的位寬，而其他通道分配較低的位寬，從而在保持平均位寬不變的情況下提高精度。這種層內(nèi)混合精度量化避免了傳統(tǒng)逐層分配方法的不足。
2. 時(shí)間平滑關(guān)系蒸餾 (TRD): 為了解決低位量化導(dǎo)致的特征表示不一致性，TRD在量化模型和全精度模型之間構(gòu)建了一個(gè)時(shí)間平滑的關(guān)系蒸餾方案。它通過(guò)融合多個(gè)連續(xù)時(shí)間步長(zhǎng)的中間特征，并使用KL散度來(lái)衡量特征相似性分布之間的差異，而不是直接比較數(shù)值，從而提高了模型的魯棒性。

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

實(shí)驗(yàn)結(jié)果表明，MPQ-DM在LSUN-Bedrooms、LSUN-Churches和ImageNet數(shù)據(jù)集上，以及Stable Diffusion模型上都顯著優(yōu)于現(xiàn)有方法。尤其是在極低位寬(例如W2A4)下，MPQ-DM取得了巨大的性能提升，而其他方法甚至無(wú)法生成正常的圖像。消融實(shí)驗(yàn)也驗(yàn)證了OMQ和TRD的有效性。

4. 主要發(fā)現(xiàn)

研究發(fā)現(xiàn)：

• 層內(nèi)混合精度量化比逐層量化更有效。
• 基于峰度的離群值選擇方法優(yōu)于隨機(jī)選擇方法。
• 關(guān)系蒸餾比直接數(shù)值對(duì)齊更適合處理離散和連續(xù)特征的差異。

5. 總結(jié)

MPQ-DM通過(guò)巧妙地結(jié)合OMQ和TRD，有效地解決了極低位擴(kuò)散模型量化中的性能下降問(wèn)題，為擴(kuò)散模型在資源受限場(chǎng)景中的應(yīng)用提供了新的可能性。該方法在多個(gè)數(shù)據(jù)集和模型上的優(yōu)異表現(xiàn)，以及全面的消融實(shí)驗(yàn)結(jié)果，都證明了其有效性和先進(jìn)性。

聯(lián)系作者

文章來(lái)源：量子位
作者微信：
作者簡(jiǎn)介：追蹤人工智能新趨勢(shì)，關(guān)注科技行業(yè)新突破