參數減少99.5%，媲美全精度FLUX！字節跳動等發布首個1.58-bit FLUX量化模型
文章鏈接：https://arxiv.org/pdf/2412.18653
項目鏈接：https://chenglin-yang.github.io/1.58bit.flux.github.io/
git主頁：https://github.com/Chenglin-Yang
亮點分析
- 1.58-bit FLUX，第一個將 FLUX 視覺 Transformer 的參數（共 119 億）減少 99.5% 至 1.58-bit 的量化模型，無需依賴圖像數據，大幅降低存儲需求。
- 開發了一個高效的線性內核，針對 1.58-bit 計算進行了優化，實現了顯著的內存減少和推理加速。
- 證明了1.58-bit FLUX在具有挑戰性的 T2I 基準測試中，性能與全精度 FLUX 模型相當。
總結速覽
解決的問題
- 當前文本生成圖像（T2I）模型，如 DALLE 3、Stable Diffusion 3 等，參數量巨大，推理時內存需求高，難以在資源有限的設備（如移動設備）上部署。
- 本文重點研究極低比特量化（1.58-bit）在 T2I 模型中的可行性，以減少存儲和內存需求，同時提升推理效率。
提出的方案
- 選用 FLUX.1-dev 模型作為量化目標，通過后訓練量化方法將其權重壓縮為 1.58-bit（值限制為 {-1, 0, +1}），無需訪問圖像數據。
- 開發專用的低比特操作優化內核，進一步提升推理效率。
應用的技術
- 1.58-bit 權重量化：使用類似 BitNet b1.58 的方法，將模型的線性層權重壓縮至 1.58-bit，并通過 2-bit 有符號整數存儲權重，從而實現極低比特化。
- 無監督量化方法：完全依賴 FLUX.1-dev 模型本身的自監督機制，無需依賴混合精度方案或額外的訓練數據。
- 定制化內核：針對低比特操作優化的推理內核，降低內存使用并縮短推理延遲。
達到的效果
- 存儲效率：模型存儲需求減少 7.7×，從 16-bit 壓縮到 2-bit。
- 推理效率：推理時的內存使用減少 5.1×，推理延遲顯著改善。
- 生成質量：在 GenEval 和 T2I Compbench 基準測試上，生成質量與全精度 FLUX 基本持平，驗證了方案的有效性和實用性。
實驗結果
設置
量化：使用一個校準數據集進行量化，數據集由 Parti-1k 數據集和 T2I CompBench 訓練集的提示語組成，共計 7,232 條提示語。整個過程完全不依賴圖像數據，不需要額外的數據集。量化將 FLUX 中 FluxTransformerBlock 和 FluxSingleTransformerBlock 的所有線性層權重壓縮至 1.58-bit，占模型總參數的 99.5%。
評估：在 GenEval 數據集 和 T2I CompBench 驗證集上評估 FLUX 和 1.58-bit FLUX，遵循官方的圖像生成流程。
- GenEval 數據集：包含 553 條提示語，每條提示語生成 4 張圖像。
- T2I CompBench 驗證集：包含 8 個類別，每個類別有 300 條提示語，每條提示語生成 10 張圖像，總計生成 24,000 張圖像進行評估。
- 所有圖像均以 1024 × 1024 的分辨率生成，適用于 FLUX 和 1.58-bit FLUX。
結果
性能：在 T2I Compbench 和 GenEval 基準測試中，1.58-bit FLUX 與全精度 FLUX 的性能表現相當，具體結果見表 1 和表 2。在應用自定義線性內核前后，性能變化微乎其微，進一步驗證了實現的準確性。
效率：如下圖 2 所示，1.58-bit FLUX 在模型存儲和推理內存上取得了顯著提升。在推理延遲方面，如下表 3 所示，特別是在低性能但易于部署的 GPU（如 L20 和 A10）上，改進更為顯著。
結論與討論
本文提出了 1.58-bit FLUX，將 99.5% 的 Transformer 參數量化至 1.58-bit，并通過自定義計算內核實現了以下改進：
- 存儲需求減少：模型存儲需求降低 7.7 倍。
- 推理內存減少：推理內存使用減少超過 5.1 倍。
盡管實現了這些壓縮效果，1.58-bit FLUX 在 T2I 基準測試中表現出與全精度模型相當的性能，同時保持了較高的視覺質量。希望 1.58-bit FLUX 能夠激勵社區開發更適合移動設備的模型。
當前局限性
關于速度改進的局限性
- 盡管 1.58-bit FLUX 降低了模型大小和內存消耗，但由于缺乏激活值量化和更高級的內核優化，其延遲改進有限。
- 鑒于目前取得的成果，希望激勵社區開發適用于 1.58-bit 模型的自定義內核實現。
關于視覺質量的局限性
- 如下圖 1、圖 3 和圖 4 所示，1.58-bit FLUX 能生成與文本提示高度一致的生動逼真的圖像，但在渲染超高分辨率細節時仍落后于原始 FLUX 模型。
- 計劃在未來研究中縮小這一差距。
本文轉自AI生成未來 ，作者：AI生成未來
原文鏈接:https://mp.weixin.qq.com/s/ZFJ1f7YZZEKoEs5Ie1YAbA
