文章鏈接：https://arxiv.org/pdf/2412.18653
項(xiàng)目鏈接：https://chenglin-yang.github.io/1.58bit.flux.github.io/
git主頁：https://github.com/Chenglin-Yang

亮點(diǎn)分析

• 1.58-bit FLUX，第一個(gè)將 FLUX 視覺 Transformer 的參數(shù)（共 119 億）減少 99.5% 至 1.58-bit 的量化模型，無需依賴圖像數(shù)據(jù)，大幅降低存儲(chǔ)需求。
• 開發(fā)了一個(gè)高效的線性內(nèi)核，針對(duì) 1.58-bit 計(jì)算進(jìn)行了優(yōu)化，實(shí)現(xiàn)了顯著的內(nèi)存減少和推理加速。
• 證明了1.58-bit FLUX在具有挑戰(zhàn)性的 T2I 基準(zhǔn)測(cè)試中，性能與全精度 FLUX 模型相當(dāng)。

總結(jié)速覽

解決的問題

• 當(dāng)前文本生成圖像（T2I）模型，如 DALLE 3、Stable Diffusion 3 等，參數(shù)量巨大，推理時(shí)內(nèi)存需求高，難以在資源有限的設(shè)備（如移動(dòng)設(shè)備）上部署。
• 本文重點(diǎn)研究極低比特量化（1.58-bit）在 T2I 模型中的可行性，以減少存儲(chǔ)和內(nèi)存需求，同時(shí)提升推理效率。

提出的方案

• 選用 FLUX.1-dev 模型作為量化目標(biāo)，通過后訓(xùn)練量化方法將其權(quán)重壓縮為 1.58-bit（值限制為 {-1, 0, +1}），無需訪問圖像數(shù)據(jù)。
• 開發(fā)專用的低比特操作優(yōu)化內(nèi)核，進(jìn)一步提升推理效率。

應(yīng)用的技術(shù)

• 1.58-bit 權(quán)重量化：使用類似 BitNet b1.58 的方法，將模型的線性層權(quán)重壓縮至 1.58-bit，并通過 2-bit 有符號(hào)整數(shù)存儲(chǔ)權(quán)重，從而實(shí)現(xiàn)極低比特化。
• 無監(jiān)督量化方法：完全依賴 FLUX.1-dev 模型本身的自監(jiān)督機(jī)制，無需依賴混合精度方案或額外的訓(xùn)練數(shù)據(jù)。
• 定制化內(nèi)核：針對(duì)低比特操作優(yōu)化的推理內(nèi)核，降低內(nèi)存使用并縮短推理延遲。

達(dá)到的效果

• 存儲(chǔ)效率：模型存儲(chǔ)需求減少 7.7×，從 16-bit 壓縮到 2-bit。
• 推理效率：推理時(shí)的內(nèi)存使用減少 5.1×，推理延遲顯著改善。
• 生成質(zhì)量：在 GenEval 和 T2I Compbench 基準(zhǔn)測(cè)試上，生成質(zhì)量與全精度 FLUX 基本持平，驗(yàn)證了方案的有效性和實(shí)用性。

實(shí)驗(yàn)結(jié)果

設(shè)置

量化：使用一個(gè)校準(zhǔn)數(shù)據(jù)集進(jìn)行量化，數(shù)據(jù)集由 Parti-1k 數(shù)據(jù)集和 T2I CompBench 訓(xùn)練集的提示語組成，共計(jì) 7,232 條提示語。整個(gè)過程完全不依賴圖像數(shù)據(jù)，不需要額外的數(shù)據(jù)集。量化將 FLUX 中 FluxTransformerBlock 和 FluxSingleTransformerBlock 的所有線性層權(quán)重壓縮至 1.58-bit，占模型總參數(shù)的 99.5%。

評(píng)估：在 GenEval 數(shù)據(jù)集 和 T2I CompBench 驗(yàn)證集上評(píng)估 FLUX 和 1.58-bit FLUX，遵循官方的圖像生成流程。

• GenEval 數(shù)據(jù)集：包含 553 條提示語，每條提示語生成 4 張圖像。
• T2I CompBench 驗(yàn)證集：包含 8 個(gè)類別，每個(gè)類別有 300 條提示語，每條提示語生成 10 張圖像，總計(jì)生成 24,000 張圖像進(jìn)行評(píng)估。
• 所有圖像均以 1024 × 1024 的分辨率生成，適用于 FLUX 和 1.58-bit FLUX。

結(jié)果

性能：在 T2I Compbench 和 GenEval 基準(zhǔn)測(cè)試中，1.58-bit FLUX 與全精度 FLUX 的性能表現(xiàn)相當(dāng)，具體結(jié)果見表 1 和表 2。在應(yīng)用自定義線性內(nèi)核前后，性能變化微乎其微，進(jìn)一步驗(yàn)證了實(shí)現(xiàn)的準(zhǔn)確性。

效率：如下圖 2 所示，1.58-bit FLUX 在模型存儲(chǔ)和推理內(nèi)存上取得了顯著提升。在推理延遲方面，如下表 3 所示，特別是在低性能但易于部署的 GPU（如 L20 和 A10）上，改進(jìn)更為顯著。

結(jié)論與討論

本文提出了 1.58-bit FLUX，將 99.5% 的 Transformer 參數(shù)量化至 1.58-bit，并通過自定義計(jì)算內(nèi)核實(shí)現(xiàn)了以下改進(jìn)：

• 存儲(chǔ)需求減少：模型存儲(chǔ)需求降低 7.7 倍。
• 推理內(nèi)存減少：推理內(nèi)存使用減少超過 5.1 倍。

盡管實(shí)現(xiàn)了這些壓縮效果，1.58-bit FLUX 在 T2I 基準(zhǔn)測(cè)試中表現(xiàn)出與全精度模型相當(dāng)?shù)男阅?，同時(shí)保持了較高的視覺質(zhì)量。希望 1.58-bit FLUX 能夠激勵(lì)社區(qū)開發(fā)更適合移動(dòng)設(shè)備的模型。

當(dāng)前局限性

關(guān)于速度改進(jìn)的局限性

• 盡管 1.58-bit FLUX 降低了模型大小和內(nèi)存消耗，但由于缺乏激活值量化和更高級(jí)的內(nèi)核優(yōu)化，其延遲改進(jìn)有限。
• 鑒于目前取得的成果，希望激勵(lì)社區(qū)開發(fā)適用于 1.58-bit 模型的自定義內(nèi)核實(shí)現(xiàn)。

關(guān)于視覺質(zhì)量的局限性

• 如下圖 1、圖 3 和圖 4 所示，1.58-bit FLUX 能生成與文本提示高度一致的生動(dòng)逼真的圖像，但在渲染超高分辨率細(xì)節(jié)時(shí)仍落后于原始 FLUX 模型。
• 計(jì)劃在未來研究中縮小這一差距。

本文轉(zhuǎn)自AI生成未來 ，作者：AI生成未來

原文鏈接:https://mp.weixin.qq.com/s/ZFJ1f7YZZEKoEs5Ie1YAbA