ByteDance社など、全精度FLUXに匹敵する99.51 TP3Tパラメータ削減を実現した初の1.58ビットFLUX量子化モデルを発表！

ハイライト分析

• 1.58ビット FLUX最初は フラックス Vision Transformerのパラメータ（合計1,190億）は、99.5%で1.58ビットの量子化モデルに縮小され、画像データに依存する必要がなくなり、ストレージ要件が劇的に削減される。
• 1.58ビット計算に最適化された効率的な線形カーネルを開発し、大幅なメモリ削減と推論の高速化を実現。
• 証明する 1.58ビット FLUX 難易度の高いT2Iベンチマークでは、全精度のFLUXモデルに匹敵する性能を示した。

簡単なまとめ

問題解決

• DALLE 3やStable Diffusion 3のような現在のT2I（text-to-generate-image）モデルは、膨大な数のパラメータを持ち、推論に必要なメモリ容量が大きいため、リソースが限られたデバイス（モバイルデバイスなど）への導入が困難である。
• 本論文では、推論効率を向上させながら、ストレージとメモリ要件を削減するために、T2Iモデルにおける超低ビット量子化（1.58ビット）の実現可能性に焦点を当てる。

提案されたプログラム

• 定量化の対象としてFLUX.1-devモデルが選ばれ、その重みは画像データにアクセスすることなく、ポストトレーニング定量化法によって1.58ビット（値は{-1, 0, +1}に制限）に圧縮された。
• 推論効率をさらに向上させるための低ビット演算最適化専用カーネルの開発。

応用技術

• 1.58ビットの重量定量化BitNet b1.58と同様の方法でモデルの線形レイヤーの重みを1.58ビットに圧縮し、2ビットの符号付き整数で重みを格納することで、極めて低いビット化を実現しています。
• 教師なし定量的手法これは、FLUX.1-devモデル自体の自己監視メカニズムにのみ依存しており、混合精度スキームや追加の学習データに依存する必要はありません。
• カスタマイズ・カーネル低ビット演算に最適化された推論カーネルにより、メモリ使用量と推論レイテンシを削減。

パンチ

• 貯蔵効率モデルのストレージ要件は7.7倍に削減され、16ビットから2ビットに圧縮された。
• 推論効率推論時のメモリ使用量を5.1倍削減し、推論レイテンシを大幅に改善。
• 品質を生み出すGenEvalベンチマークとT2I Compbenchベンチマークでは、生成品質は全精密FLUXとほぼ同等であり、このスキームの有効性と有用性が実証された。

結果

セットアップ

定量的定量化は、Parti-1kデータセットとT2I CompBenchトレーニングセットからのキュー、合計7,232キューからなる較正済みデータセットを用いて行われた。全プロセスは画像データから完全に独立しており、追加のデータセットは必要ない。定量化は、FLUXのFluxTransformerBlockとFluxSingleTransformerBlockのすべての線形層の重みを1.58ビットに圧縮し、モデルパラメータ全体の99.5%を占める。

評価GenEvalデータセットとT2I CompBench検証セットにおけるFLUXと1.58ビットFLUXの評価。

• GenEvalデータセット553のプロンプトが含まれ、各プロンプトは4つの画像を生成します。
• T2I CompBench検証セット評価は8つのカテゴリーで行われ、各カテゴリーにつき300のプロンプトが用意され、各プロンプトは10枚の画像を生成し、合計24,000枚の画像が評価対象となった。
• すべての画像は、FLUXおよび1.58ビットFLUXでは1024×1024の解像度で生成されている。

結局

パフォーマンス1.58ビットFLUXの性能は、表1および表2に示すように、T2I CompbenchおよびGenEvalベンチマークにおいて、全精度FLUXの性能に匹敵する。 カスタム線形カーネルを適用する前後での性能変化はわずかであり、実装の精度がさらに検証された。

効率以下の図 2 に示すように、1.58 ビット FLUX は、モデルストレージと推論メモリにおいて大幅な改善を達成しています。推論レイテンシに関しては、以下の表3に示すように、L20やA10のような低性能ながら導入が容易なGPUにおいて、特に大きな改善が見られます。

結論と考察

本稿で紹介するのは 1.58ビット FLUX99.5%の場合、Transformerパラメータは1.58ビットに定量化され、計算カーネルをカスタマイズすることで以下の改善が達成された：

• ストレージ要件の削減モデルのストレージ要件を7.7分の1に削減。
• 推論記憶の減少推論メモリの使用量を5.1倍以上削減。

このような圧縮効果が得られているにもかかわらず、1.58ビットFLUXは、T2Iベンチマークにおいて、高いビジュアル品質を維持しながら、全精密モデルに匹敵する性能を示した。以下のことが期待される。 1.58ビット FLUX モバイル機器により適したモデルを開発するよう、コミュニティを動機付けることができる。

現在の制限

速度向上の制限

• 1.58ビットFLUXは、モデルサイズとメモリ消費量を削減するものの、活性化値の定量化やより高度なカーネル最適化ができないため、レイテンシの改善には限界がある。
• これまでの成果を考えると、コミュニティが1.58ビットモデル用のカスタムカーネル実装を開発する気になることが期待される。

ビジュアル・クオリティの限界について

• 以下の図1、図3、図4に示すように、1.58ビットFLUXは、テキストキューと非常に一致した鮮やかでリアルな画像を生成しますが、超高解像度の詳細をレンダリングする場合には、まだオリジナルのFLUXモデルに遅れをとっています。
• 今後の研究でこのギャップを埋める予定である。

1.58ビットFLUXと1.58ビットFLUXの視覚的比較。1.58ビットFLUXは1.58ビット量子化したFLUXと同等の生成品質を示す。一貫性を保つため、各比較のすべての画像は同じ潜在ノイズ入力を使って生成されている。1.58ビットFLUXはカスタム1.58ビットカーネルを使用。一貫性を保つため、各比較のすべての画像は、同じ潜在ノイズ入力を使用して生成されています。

推論コードとウェイトが公開される： https://github.com/Chenglin-Yang/1.58bit.flux