線性擴散模型LiT來了，用極簡線性注意力助力擴散模型AIPC時代端側部署

來自華人團隊，可實現快速訓練。

原標題：線性擴散模型LiT來了，用極簡線性注意力助力擴散模型AIPC時代端側部署
文章來源：機器之心
內容字數：13271字

LiT: 高效擴散模型助力AIPC時代

機器之心AIxiv專欄報道了香港大學與上海人工智能實驗室、華為諾亞方舟實驗室合作提出的高效擴散模型LiT。該模型探索了擴散模型中極簡線性注意力的架構設計和訓練策略，可在斷網狀態下于Windows筆記本電腦上離線部署，快速生成1K分辨率逼真圖片，顯著提升了文生圖應用的效率和便捷性。

1. 背景與挑戰

Diffusion Transformer模型在文生圖領域展現出巨大潛力，但其自注意力機制的二次計算復雜度限制了其在高分辨率場景和端側設備的應用。LiT團隊致力于解決這一問題，并探索了線性注意力機制在擴散模型中的應用。

2. 線性注意力與LiT的設計

線性注意力機制具有簡潔性和高并行化程度的優勢，非常適合大型模型。LiT采用簡化線性注意力機制，并通過實驗發現，使用更少的注意力頭可以在增加理論計算量的同時，不增加實際GPU延遲，實現了“免費午餐”效應。LiT的架構設計基于DiT，但將自注意力替換為線性注意力。

3. 高效訓練策略

LiT提出了五條高效訓練指導原則：1. 簡化線性注意力足以完成圖像生成任務；2. 使用更少的注意力頭；3. 從預訓練的DiT模型繼承權重，但不要繼承自注意力中的權重；4. 使用知識蒸餾加速訓練；5. 同時蒸餾噪聲預測結果和方差預測結果。

4. 實驗結果與驗證

在ImageNet 256×256和512×512基準測試中，LiT在訓練迭代次數遠少于DiT的情況下，實現了相當甚至更好的FID結果。在文生圖任務中，LiT-0.6B可在Windows筆記本電腦上離線生成1K分辨率逼真圖片。

5. 離線端側部署

LiT成功在Windows 11筆記本電腦上實現離線端側部署，無需網絡連接即可快速生成高分辨率圖像，展現了其在邊緣設備上的應用潛力，推動了AIPC時代的到來。

總而言之，LiT通過巧妙的線性注意力設計和高效的訓練策略，顯著提升了擴散模型的效率和可訪問性，為文生圖應用的商業化和普及化提供了新的可能性。

聯系作者

文章來源：機器之心
作者微信：
作者簡介：專業的人工智能媒體和產業服務平臺