突破瓶頸！北航ETH等首次將擴散模型完全量化至1bit，28倍存儲節(jié)省+52.7倍效率提升

BiDM團隊 投稿

量子位 | 公眾號 QbitAI

將擴散模型量化到1比特極限，又有新SOTA了！

來自北航、ETH等機構的研究人員提出了一種名為BiDM的新方法，首次將擴散模型（DMs）的權重和激活完全二值化。

具體而言，作者們從時間和空間的角度對擴散模型進行了優(yōu)化：

從時間角度來看，引入了“時間步友好二值結(jié)構”（TBS），通過可學習的激活二值化器和跨時間步特征連接來應對DMs高度時間步相關的激活特征。

從空間角度來看，提出了“空間分塊蒸餾”（SPD），目標是解決二值化特征匹配的困難，特別關注圖像生成任務和噪聲估計網(wǎng)絡的空間局部性。

實驗結(jié)果顯示，W1A1 BiDM在LSUN-Bedrooms 256×256上的LDM-4模型上取得了22.74的FID分數(shù)，遠遠超越了當前狀態(tài)的最先進通用二值化方法的59.44分，并實現(xiàn)了高達28倍的存儲節(jié)省和52.7倍的操作效率提升。

下面具體來看。

現(xiàn)有擴散模型1位量化主要集中在「權重」

目前為止，擴散模型由于其高質(zhì)量和多樣化的生成能力，在圖像、語音和視頻等多個領域引起了極大的關注和應用。它可以通過多達1000步的去噪步驟，從隨機噪聲生成數(shù)據(jù)。

不過，雖然一些加速采樣方法能夠有效減少生成任務所需的步驟數(shù)量，但每個時間步的昂貴浮點計算仍然限制了該模型在資源受限場景中的廣泛應用。

因此，對擴散模型的壓縮成為其更廣泛應用的關鍵步驟，現(xiàn)有的壓縮方法主要包括量化、蒸餾、剪枝等。這些壓縮方法的目標是在保持準確性的同時減少存儲和計算成本。

其中，量化被認為是一種非常有效的模型壓縮技術，通過將權重和/或激活量化為低位整數(shù)或二值化，實現(xiàn)緊湊存儲和推理中的高效計算。

所以，已有一些研究將量化應用于擴散模型，以在保持生成質(zhì)量的同時實現(xiàn)模型的壓縮和加速。

而1位量化，即二值化，能夠最大限度地節(jié)省模型的存儲空間，并且在卷積神經(jīng)網(wǎng)絡（CNN）等判別模型中表現(xiàn)良好。此外，當權重和激活同時量化為1位時，如完全二值化，使用類似XNOR和bitcount的高效位運算可以替代矩陣乘法，達到最高效的加速效果。

一些現(xiàn)有的工作嘗試將擴散模型量化為1位，但它們的探索主要集中在權重上，離完全二值化仍有較大距離。

實際上，對于擴散模型這樣的生成模型，完全二值化權重和激活的影響是災難性的：

a)?作為生成模型，擴散模型的豐富中間表示與時間步密切相關，而高度動態(tài)的激活范圍在使用二值化權重和激活時受到嚴重限制；

b)?像擴散模型這樣的生成模型通常需要輸出完整的圖像，但高度離散的參數(shù)和特征空間使得在訓練過程中很難與真實值匹配。離散空間中的優(yōu)化難度以及與時間步動態(tài)相關的表示能力不足，導致二值化擴散模型難以收斂，甚至在優(yōu)化過程中崩潰。

提出BiDM：完全二值化「權重」和「激活」

面對上述不足，作者們提出了BiDM。通過完全二值化權重和激活，將擴散模型推向極限的壓縮和加速。

概括而言，BiDM目標是解決擴散模型激活特征、模型結(jié)構以及生成任務的獨特需求，克服完全二值化帶來的挑戰(zhàn)。

它包含兩項創(chuàng)新技術：

從時間角度出發(fā)，作者們觀察到擴散模型的激活特征與時間步高度相關。

因此引入了“時間步友好二值結(jié)構”（TBS），通過可學習的激活二值量化器匹配擴散模型的動態(tài)激活范圍，并設計跨時間步的特征連接，利用相鄰時間步之間的特征相似性，增強二值模型的表示能力。

而從空間角度出發(fā)，作者們注意到生成任務中的空間局部性以及擴散模型使用的基于卷積的U-Net結(jié)構。

于是提出了“空間分塊蒸餾”（SPD），引入一個全精度模型作為監(jiān)督，通過對分塊的自注意模仿，專注于局部特征，更好地引導二值擴散模型的優(yōu)化方向。

△具有時間步友好二值結(jié)構和空間分塊蒸餾的BiDM

廣泛的實驗表明，與現(xiàn)有的SOTA完全二值化方法相比，BiDM在保持相同推理效率的同時顯著提高了精度，在各種評價指標上超越了所有現(xiàn)有基線方法。

具體來說，在像素空間擴散模型中，BiDM是唯一一種將IS提高到5.18的方法，接近全精度模型的水平，比最佳基線方法高出0.95。

在LDM 中，BiDM將LSUN-Bedrooms的FID從SOTA方法的59.44降低到了令人印象深刻的22.74，同時節(jié)省了28.0倍的存儲空間和52.7倍的OPs。

作為第一種完全二值化的擴散模型方法，大量生成的樣本也證明BiDM是目前唯一一種能夠生成可接受的完全二值化DM圖像的方法，從而使DM在低資源場景中得到有效應用。

具體咋實現(xiàn)的

二值擴散模型基線

擴散模型：給定一個數(shù)據(jù)分布0~（0），前向過程生成一個隨機變量序列 ? ｛1，…，｝，該序列通過轉(zhuǎn)移核（|-1）生成，通常是高斯擾動，其表達式為：

其中 ?（0，1）是噪聲調(diào)度。高斯轉(zhuǎn)移核允許作者們邊緣化聯(lián)合分布，因此利用：=1-和

，可以通過采樣一個高斯向量~N（0，I）并應用變換

輕松獲得的樣本。

反向過程旨在通過去除噪聲生成樣本，使用可學習的轉(zhuǎn)移核（-1|）近似不可用的條件分布q（-1|），其表達式為：

均值

和方差

可以使用重參數(shù)化技巧得到：

其中是具有可學習參數(shù)的函數(shù)逼近，用于根據(jù)預測。

對于擴散模型的訓練，通常使用變分下界的簡化變體作為損失函數(shù)，以提高樣本質(zhì)量，其表達式為：

由于 U-Net 能夠融合低層次和高維特征，它已成為擴散模型的主流骨干。U-Net 的輸入輸出塊可以表示為

和

，其中較小m對應的塊是更低層次的。跳躍連接將低層次信息從Dm(·）傳播到Um(·），因此Um接收到的輸入表示為：

二值化：量化通過將權重和激活離散化為低位寬，壓縮并加速噪聲估計模型。在二值化擴散模型的基線中，權重W被二值化為1位：

其中符號函數(shù)將W限制為+1或-1，并設0為閾值。是浮點標量，初始化為

（n表示權重數(shù)量），并在訓練過程中進行學習。

同時，激活通常通過簡單的BNN量化器進行量化：

當權重和激活同時被量化為1位時，去噪模型的計算可以由XNOR和bitcount運算代替，從而實現(xiàn)顯著的壓縮和加速。

時間步友好二值結(jié)構

在詳細介紹所提方法之前，作者們先總結(jié)對擴散模型（DMs）屬性的觀察：

下圖中， (a) 全精度DDIM模型在CIFAR-10上第4層卷積層的激活范圍隨去噪步驟變化。(b) 在LSUN-Bedrooms數(shù)據(jù)集上，全精度LDM-4模型在每一步的輸出特征與前一步相似。

觀察1：激活范圍在長期時間步中變化顯著，但激活特征在短期相鄰時間步中相似。

之前的研究，如TDQ和Q-DM，已經(jīng)表明，DMs的激活分布在去噪過程中高度依賴于時間步，表現(xiàn)為相鄰時間步之間的相似性，而遠距離時間步之間的差異較大，如圖2(a)所示。

因此，在所有時間步中應用固定的縮放因子會導致激活范圍的嚴重失真。除了分布范圍之外，Deepcache強調(diào)了連續(xù)時間步之間高維特征的顯著時間一致性，如圖2(b)所示。

這些現(xiàn)象促使作者們重新審視現(xiàn)有的二值化結(jié)構。

二值化，尤其是權重和激活的完全二值化，與4位等低位量化相比，會導致激活范圍和精度的更大損失。這使得生成豐富激活特征變得更加困難。激活范圍和輸出特征的不足嚴重損害了像DMs這樣豐富表示的生成模型。

因此，采用更靈活的激活范圍的二值量化器，并通過利用其特征輸出來增強模型的整體表達能力，是在完全二值化后提高其生成能力的關鍵策略。

作者們首先關注長期時間步之間的差異。

大多數(shù)現(xiàn)有的激活量化器，如BNN和Bi-Real，直接將激活量化為｛+1，-1｝，如式(7)所示。這種方法嚴重擾亂了激活特征，負面影響了生成模型的表達能力。一些改進的激活二值量化器，如XNOR++，采用了可訓練的縮放因子k：

其中K的形式可以是向量或多個向量的乘積，但在推理過程中保持為常量值。盡管這種方法部分恢復了激活的特征表達，但它與高度相關的時間步不匹配，可能仍會導致顯著的性能損失。

作者們將注意力轉(zhuǎn)向原始的XNOR，它采用動態(tài)計算的均值來構造激活二值量化器。其對二維卷積的運算可以表示為：

其中

，

，

，

。

表示二維過濾器，

和

分別表示帶和不帶乘法的卷積操作。這種方法自然保留了激活特征的范圍，并隨著不同時間步的輸入范圍動態(tài)調(diào)整。

然而，由于DM特征的豐富表達，局部激活在模塊通過前后的范圍內(nèi)表現(xiàn)出不一致，表明預定的k值不能有效恢復激活表示。

因此，作者們使k可調(diào)整，并允許其在訓練過程中進行學習，以自適應匹配激活在前后發(fā)生的范圍變化。作者們提出的可學習的微小卷積k的梯度計算過程如下：

值得注意的是，使k可學習不會增加任何額外的推理負擔。計算成本保持不變，允許進行高效的二值運算。

另一方面，作者們專注于相鄰時間步之間的相似性。

Deepcache直接提取高維特征作為緩存，以跳過U-Net中的大量深度計算，實現(xiàn)了顯著的推理加速。其過程表示為：

然而，這種方法不適用于二值化擴散模型，因為二值網(wǎng)絡的每個輸出的信息內(nèi)容非常有限。對于本質(zhì)上已經(jīng)實現(xiàn)了顯著壓縮和加速但表達能力有限的二值擴散模型，作者們預計相鄰時間步之間的特征相似性將增強二值表示，從而彌補表達挑戰(zhàn)。

作者們構建了一個跨時間步的信息增強連接，以使用前一時間步的特征豐富當前時間步的表達。其過程可以表示為：

其中

是一個可學習的縮放因子。正如圖2(b)所示，不同塊和時間步中的高維特征相似性各異。

因此，作者們設置了多個獨立的值，以允許模型在訓練過程中更有效地自適應學習。

總結(jié)來說，時間步友好的二值結(jié)構（TBS）包括在輸入平均和跨時間步連接后應用于縮放因子的可學習微卷積。它們的組合效果能夠適應擴散模型在長期時間步中的激活范圍變化，并利用相鄰時間步之間高維特征的相似性來增強信息表示。

從誤差減少的角度來看，圖3中展示了TBS的可視化。

首先，作者們將基線方法下二值擴散模型的輸出抽象為向量Bt-1。

縮放因子的不匹配會在它與全精度模型的輸出向量Ft-1之間產(chǎn)生顯著的長度差異。

通過使用提出的縮放因子和可學習的微卷積，Bt-1被擴展為Lt-1。Lt-1更接近Ft-1，但仍與全精度模型存在方向上的差異。

跨時間步的連接進一步結(jié)合了前一個時間步的輸出Ft、Bt和Lt。相鄰時間步之間的高維特征相似性意味著Ft-1和Ft之間的差距相對較小，從而促進了Lt-1和Lt的結(jié)合。

最后，作者們通過應用TBS得到二值化擴散模型的輸出，表達式為Tt-1=(1-)·Lt-1+·Lt，這一輸出最接近全精度模型的輸出Ft-1。

空間分塊蒸餾

由于生成模型的特性，擴散模型的優(yōu)化過程與以往的判別模型有著不同的特征：

觀察2：傳統(tǒng)的蒸餾難以引導完全二值化的擴散模型與全精度模型對齊，而擴散模型在生成任務中的特征在空間上表現(xiàn)出局部性。

在以往的實踐中，在量化模型的訓練過程中添加蒸餾損失是常見的做法。由于二值模型的數(shù)值空間有限，直接使用簡單的損失進行優(yōu)化會導致調(diào)整梯度更新方向的困難，使學習變得具有挑戰(zhàn)性。因此，向中間特征添加蒸餾損失可以更好地指導模型的局部和全局優(yōu)化過程。

然而，作為生成模型，擴散模型的高度豐富的特征表示使得二值模型極難細致地模仿全精度模型。

盡管原始擴散模型訓練中使用的L2損失與擴散過程中的高斯噪聲對齊，但并不適用于中間特征的蒸餾匹配。在常規(guī)蒸餾過程中，常用的L2損失往往優(yōu)先優(yōu)化差異較大的像素，從而導致更加均勻和平滑的優(yōu)化結(jié)果。

這種全局約束學習過程對于以圖像生成為目標的二值模型來說是困難的，因為其有限的表示能力使得精細的蒸餾模仿難以直接調(diào)整模型以完全匹配全精度模型的方向。

與此同時，作者們注意到，使用U-Net作為骨干的擴散模型由于其基于卷積的結(jié)構和生成任務的要求，天然表現(xiàn)出空間局部性。

這與以往的判別模型不同，判別模型中的任務（如分類）只需整體特征提取，而不需要低層次的要求，這使得傳統(tǒng)的蒸餾方法不適用于具有空間局部性的生成任務。此外，大多數(shù)現(xiàn)有的擴散模型蒸餾方法專注于減少時間步數(shù)，并未解決圖像生成任務中對特征空間局部性的需求。

因此，鑒于現(xiàn)有損失函數(shù)在優(yōu)化二值擴散模型上的困難以及擴散模型的空間局部性，作者們提出了空間分塊蒸餾（SPD）。

具體來說，他們設計了一種新的損失函數(shù)，該損失函數(shù)在蒸餾之前將特征劃分為小塊，然后逐塊計算空間自注意力。雖然傳統(tǒng)的L2損失使二值擴散模型難以實現(xiàn)直接匹配，導致優(yōu)化挑戰(zhàn)，但注意力機制使得蒸餾優(yōu)化可以更側(cè)重于關鍵部分。

然而，對于完全二值化的擴散模型來說，這仍然是具有挑戰(zhàn)性的，因為高度離散的二值輸出信息有限，使得模型難以捕獲全局信息。

因此，作者們通過將中間特征劃分為多個小塊，并為每個小塊獨立計算空間自注意力，使得二值模型在優(yōu)化過程中能夠更好地利用局部信息。

SPD首先將二值擴散模型和全精度擴散模型的某個塊輸出的中間特征bi和fp?bxcxwxh劃分為p2個小塊：

然后，分別對每個小塊計算自注意力：

在正則化之后，計算并累加相應位置的損失：

其中|| · ||2表示L2函數(shù)。最后，總訓練損失被計算為：

其中，d表示在上采樣或下采樣過程中塊的數(shù)量，最終包括中間塊在內(nèi)，共有2d+1中間特征。是用于平衡損失項的超參數(shù)系數(shù)，默認為4。

如下圖所示，LDM模型在LSUN-bedroom數(shù)據(jù)集上的最后一個TimeStepBlock輸出的可視化。FP32表示全精度模型的輸出fp。Diff表示全精度模型輸出與二值化模型輸出之間的差異||fp-bi||。Ours表示自注意力機制。

作者們對上述提到的中間特征和自注意力進行了可視化。正如圖所示，作者們的空間分塊蒸餾（SPD）使得模型在每個小塊中更加關注局部信息。

實驗

作者們在各種數(shù)據(jù)集上進行了實驗，包括CIFAR-10 32×32、LSUN-Bedrooms 256×256、LSUN-Churches 256×256 和FFHQ 256×256，并在像素空間擴散模型和潛在空間擴散模型上進行了測試。

使用的評估指標包括Inception Score (IS)、Fréchet Inception Distance (FID)、Sliding Fréchet Inception Distance (sFID)、精度和召回率。截至目前，還沒有研究將擴散模型壓縮到如此極端的程度。

因此，作者們使用了經(jīng)典的二值化算法和最新的最先進的通用二值化算法作為基線。

他們提取了擴散模型中TimestepEmbedBlocks的輸出作為TBS和SPD的操作目標。且還在卷積層中采用了與ReActNet相同的快捷連接。詳細的實驗設置見附錄A。

主要結(jié)果

像素空間擴散模型：作者們首先在CIFAR-10 32×32數(shù)據(jù)集上進行了實驗。正如表1中所示，使用基線方法對擴散模型進行W1A1二值化導致了顯著的性能下降。然而，BiDM在所有指標上均表現(xiàn)出顯著改進，實現(xiàn)了前所未有的圖像質(zhì)量恢復。

具體而言，BiDM在IS指標上從4.23提升至5.18，F(xiàn)ID指標降低了27.9%。

以下為在CIFAR-10數(shù)據(jù)集上，使用100步的DDIM二值化結(jié)果。

潛在空間擴散模型：作者們的LDM實驗包括對LDM-4在LSUN-Bedrooms 256×256和FFHQ 256×256數(shù)據(jù)集上的評估，以及對LDM-8在LSUN-Churches 256×256數(shù)據(jù)集上的評估。

實驗使用了200步的DDIM采樣器，詳細結(jié)果見表2中。

在這三個數(shù)據(jù)集中，作者們的方法相對于最好的基線方法取得了顯著的改進。與其他二值化算法相比，BiDM在所有指標上都表現(xiàn)優(yōu)異。

在LSUN-Bedrooms、LSUN-Churches和FFHQ數(shù)據(jù)集中，BiDM的FID指標分別比基線方法降低了61.7%、30.7%和51.4%。與XNOR++相比，XNOR++在去噪過程中采用了固定的激活縮放因子，導致其激活的動態(tài)范圍非常有限，使得難以匹配擴散模型高度靈活的生成表示。

BiDM通過使微卷積k可學習，作用于動態(tài)計算的縮放因子，解決了這一挑戰(zhàn)。這種優(yōu)化使得各項指標取得了超過一個數(shù)量級的顯著改進。

在LSUN-Bedrooms和LSUN-Churches數(shù)據(jù)集中，BiDM的FID指標分別從319.66降低至22.74和從292.48降低至29.70。

此外，與最先進的二值化方法ReSTE相比，BiDM在多個指標上實現(xiàn)了顯著提升，特別是在LSUN-Bedrooms數(shù)據(jù)集中表現(xiàn)出顯著的改進。

消融實驗

作者們在LSUN-Bedrooms 256×256數(shù)據(jù)集上對LDM-4進行了全面的消融研究，以評估BiDM中每個提出的組件的有效性。

他們評估了提出的SPD和TBS方法的有效性，結(jié)果如表3所示。

當分別將SPD或TBS方法應用于LDM時，相較于原始性能，觀察到了顯著的改進。當引入TBS方法時，F(xiàn)ID和sFID分別從106.62和56.61大幅下降到35.23和25.13。

同樣，當添加SPD方法時，F(xiàn)ID和sFID分別顯著下降到40.62和31.61。其他指標也顯示出顯著的改善。這證明了作者們的方法在訓練過程中通過引入可學習的因子并加入跨時間步的連接，能夠持續(xù)將二值化模型的特征逼近全精度特征的有效性。

此外，當結(jié)合這兩種方法并將它們應用于LDM 時，較之單獨應用每種方法，觀察到額外的改進。這進一步證明了在補丁級別執(zhí)行全精度和二值模型之間的蒸餾能夠顯著提高二值模型的性能。

作者們還進行了額外的消融實驗，結(jié)果在附錄B中展示。

效率分析

作者們對模型在完全二值化下的推理效率進行了分析。

表4中的結(jié)果表明，BiDM在實現(xiàn)與XNOR基線相同的28.0倍內(nèi)存效率和52.7倍計算節(jié)省的同時，圖像生成能力顯著優(yōu)于基線模型，其FID從106.62降低至22.74。

總結(jié)

在本文中，作者們提出了BiDM，這是一種全新的完全二值化方法，將擴散模型的壓縮推向極限。

基于兩個觀察——不同時間步的激活特性和圖像生成任務的特點，他們分別從時間和空間角度提出了時間步友好的二值結(jié)構（TBS） 和 空間分塊蒸餾（SPD） 方法。

這些方法解決了完全二值化中表示能力的嚴重限制，以及高度離散的空間優(yōu)化挑戰(zhàn)。

作為首個完全二值化的擴散模型，BiDM在多個模型和數(shù)據(jù)集上表現(xiàn)出了顯著優(yōu)于現(xiàn)有最先進通用二值化方法的生成性能。

在LSUN-Bedrooms數(shù)據(jù)集上，BiDM的FID為22.74，遠超最先進方法的FID為59.44，成為唯一能夠生成視覺上可接受樣本的二值方法，同時實現(xiàn)了高達28.0倍的存儲節(jié)省和52.7倍的運算效率提升。

更多細節(jié)歡迎查閱原論文。

論文鏈接：
https://arxiv.org/abs/2412.05926
項目主頁：
https://github.com/Xingyu-Zheng/BiDM

— 完 —