從時間和空間優化擴散模型
原標題:突破瓶頸!北航ETH等首次將擴散模型完全量化至1bit,28倍存儲節省+52.7倍效率提升
文章來源:量子位
內容字數:13960字
BiDM: 將擴散模型量化到1比特極限,實現SOTA
來自北航、ETH等機構的研究人員提出了一種名為BiDM的新方法,首次將擴散模型(DMs)的權重和激活完全二值化,并在圖像生成任務中取得了顯著的性能提升,創造了新的SOTA。
1. 擴散模型量化的挑戰
擴散模型因其高質量的生成能力而備受關注,但其高計算成本限制了其在資源受限場景下的應用。模型壓縮,特別是量化,成為關鍵。現有研究主要集中在權重量化,而對權重和激活同時進行1比特量化(完全二值化)則極具挑戰性:擴散模型的中間表示與時間步密切相關,高度動態的激活范圍在二值化后難以保持;同時,生成模型需要輸出完整的圖像,高度離散的參數和特征空間使得訓練過程難以收斂。
2. BiDM方法的創新
BiDM旨在克服完全二值化帶來的挑戰,包含兩項主要創新:
- 時間步友好二值結構 (TBS): 針對擴散模型激活特征與時間步高度相關的特點,TBS利用可學習的激活二值化器匹配動態激活范圍,并設計跨時間步的特征連接,利用相鄰時間步特征相似性增強表示能力。
- 空間分塊蒸餾 (SPD): 針對圖像生成任務的空間局部性,SPD將特征劃分為小塊,并對每個小塊計算空間自注意力,引導二值擴散模型更好地學習局部特征,從而更有效地進行蒸餾。
TBS通過可學習的縮放因子和跨時間步連接,自適應地調整激活范圍并利用時間步之間的相似性。SPD則通過空間分塊和自注意力機制,更有效地進行蒸餾,從而提升模型性能。
3. 實驗結果與分析
實驗結果表明,BiDM在多個數據集(CIFAR-10,LSUN-Bedrooms,LSUN-Churches,FFHQ)和模型(像素空間擴散模型和潛在空間擴散模型)上均取得了顯著的性能提升,遠超現有SOTA方法。例如,在LSUN-Bedrooms 256×256數據集上,BiDM的FID分數為22.74,遠低于現有最佳基線方法的59.44,同時實現了28倍的存儲節省和52.7倍的操作效率提升。
消融實驗驗證了TBS和SPD的有效性,兩者結合能取得最佳效果。效率分析也表明BiDM在極高的壓縮率下仍能保持良好的生成質量。
4. 結論
BiDM是首個實現擴散模型完全二值化的成功方法,它在極端壓縮率下取得了令人印象深刻的生成性能,為擴散模型在資源受限場景下的應用提供了新的可能性。
聯系作者
文章來源:量子位
作者微信:
作者簡介:追蹤人工智能新趨勢,關注科技行業新突破