從時(shí)間和空間優(yōu)化擴(kuò)散模型

原標(biāo)題：突破瓶頸！北航ETH等首次將擴(kuò)散模型完全量化至1bit，28倍存儲節(jié)省+52.7倍效率提升
文章來源：量子位
內(nèi)容字?jǐn)?shù)：13960字

BiDM: 將擴(kuò)散模型量化到1比特極限，實(shí)現(xiàn)SOTA

來自北航、ETH等機(jī)構(gòu)的研究人員提出了一種名為BiDM的新方法，首次將擴(kuò)散模型（DMs）的權(quán)重和激活完全二值化，并在圖像生成任務(wù)中取得了顯著的性能提升，創(chuàng)造了新的SOTA。

1. 擴(kuò)散模型量化的挑戰(zhàn)

擴(kuò)散模型因其高質(zhì)量的生成能力而備受關(guān)注，但其高計(jì)算成本限制了其在資源受限場景下的應(yīng)用。模型壓縮，特別是量化，成為關(guān)鍵。現(xiàn)有研究主要集中在權(quán)重量化，而對權(quán)重和激活同時(shí)進(jìn)行1比特量化（完全二值化）則極具挑戰(zhàn)性：擴(kuò)散模型的中間表示與時(shí)間步密切相關(guān)，高度動態(tài)的激活范圍在二值化后難以保持；同時(shí)，生成模型需要輸出完整的圖像，高度離散的參數(shù)和特征空間使得訓(xùn)練過程難以收斂。

2. BiDM方法的創(chuàng)新

BiDM旨在克服完全二值化帶來的挑戰(zhàn)，包含兩項(xiàng)主要?jiǎng)?chuàng)新：

1. 時(shí)間步友好二值結(jié)構(gòu) (TBS): 針對擴(kuò)散模型激活特征與時(shí)間步高度相關(guān)的特點(diǎn)，TBS利用可學(xué)習(xí)的激活二值化器匹配動態(tài)激活范圍，并設(shè)計(jì)跨時(shí)間步的特征連接，利用相鄰時(shí)間步特征相似性增強(qiáng)表示能力。
2. 空間分塊蒸餾 (SPD): 針對圖像生成任務(wù)的空間局部性，SPD將特征劃分為小塊，并對每個(gè)小塊計(jì)算空間自注意力，引導(dǎo)二值擴(kuò)散模型更好地學(xué)習(xí)局部特征，從而更有效地進(jìn)行蒸餾。

TBS通過可學(xué)習(xí)的縮放因子和跨時(shí)間步連接，自適應(yīng)地調(diào)整激活范圍并利用時(shí)間步之間的相似性。SPD則通過空間分塊和自注意力機(jī)制，更有效地進(jìn)行蒸餾，從而提升模型性能。

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

實(shí)驗(yàn)結(jié)果表明，BiDM在多個(gè)數(shù)據(jù)集（CIFAR-10,LSUN-Bedrooms,LSUN-Churches,FFHQ）和模型（像素空間擴(kuò)散模型和潛在空間擴(kuò)散模型）上均取得了顯著的性能提升，遠(yuǎn)超現(xiàn)有SOTA方法。例如，在LSUN-Bedrooms 256×256數(shù)據(jù)集上，BiDM的FID分?jǐn)?shù)為22.74，遠(yuǎn)低于現(xiàn)有最佳基線方法的59.44，同時(shí)實(shí)現(xiàn)了28倍的存儲節(jié)省和52.7倍的操作效率提升。

消融實(shí)驗(yàn)驗(yàn)證了TBS和SPD的有效性，兩者結(jié)合能取得最佳效果。效率分析也表明BiDM在極高的壓縮率下仍能保持良好的生成質(zhì)量。

4. 結(jié)論

BiDM是首個(gè)實(shí)現(xiàn)擴(kuò)散模型完全二值化的成功方法，它在極端壓縮率下取得了令人印象深刻的生成性能，為擴(kuò)散模型在資源受限場景下的應(yīng)用提供了新的可能性。

聯(lián)系作者

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