清華大学とハルビン工業大学は大規模なモデルを 1 ビットに圧縮し、携帯電話で大規模なモデルを実行するという夢が実現しつつあります。

ビッグモデルが普及して以来、ビッグモデルを圧縮したいという人々の欲求は減ることはありません。大規模モデルは多くの面で優れた機能を備えているものの、導入コストが高いため、利用の敷居が高くなるためです。このコストは主に、スペースの占有と計算の複雑さから生じます。 「モデル量子化」は、大規模なモデルのパラメータを低ビット幅の表現に変換することでスペースを節約します。現在、主流の方法で、モデルのパフォーマンスをほとんど損なうことなく、既存のモデルを 4 ビットに圧縮できます。しかし、3 ビット未満の量子化は、研究者を怖がらせる乗り越えられない壁のようなものです。

図1: 量子化モデルの難解度は2ビットで急激に上昇する

最近、清華大学とハルビン工業大学が共同でarXivに発表した論文は、この障壁を突破する希望をもたらし、国内外の学界で大きな注目を集めました。この論文は 1 週間前に huggingface のホット ペーパーにも掲載され、有名な論文推薦者 AK によって推薦されました。研究チームは2ビット量子化レベルを直接飛ばし、1ビット量子化に大胆に挑戦したが、これはモデル量子化研究としては初めてのことだ。

論文タイトル: OneBit: 極低ビット大規模言語モデルに向けて

論文アドレス: https://arxiv.org/pdf/2402.11295.pdf

著者が提案した方法は「OneBit」と呼ばれ、事前トレーニング済みの大規模モデルを真の 1 ビットに圧縮するというこの研究の本質を的確に表しています。本論文では、モデルパラメータの 1 ビット表現の新しい方法と量子化モデルパラメータの初期化方法を提案し、量子化を考慮したトレーニング (QAT) を通じて高精度の事前トレーニング済みモデルの機能を 1 ビット量子化モデルに移行します。実験では、この方法は LLaMA モデルのパフォーマンスの少なくとも 83% を確保しながら、モデル パラメータを大幅に圧縮できることが示されています。

著者らは、モデルパラメータが 1 ビットに圧縮されると、行列乗算における「要素乗算」がなくなり、より高速な「ビット割り当て」演算に置き換えられ、計算効率が大幅に向上すると指摘しています。この研究の意義は、2ビット量子化のギャップを埋めるだけでなく、PCやスマートフォン上で大規模なモデルを展開できるようにすることにあります。

既存の研究の限界

モデル量子化は主に、モデルの nn.Linear レイヤー (Embedding レイヤーと Lm_head レイヤーを除く) を低精度表現に変換することで空間圧縮を実現します。これまでの研究[1,2]の基礎は、Round-To-Nearest（RTN）法を使用して、高精度の浮動小数点数を近くの整数グリッドに近似することです。これは と表現できます。

しかし、RTN ベースの方法では、非常に低いビット幅 (3 ビット未満) で精度が大幅に低下し、量子化後のモデル機能の低下も深刻になります。特に、量子化されたパラメータが 1 ビットで表現される場合、RTN のスケーリング係数 s とゼロ点 z は実用的な意味を失います。このため、RTN ベースの量子化方法は 1 ビットの量子化ではほとんど効果がなく、元のモデルのパフォーマンスを効果的に維持することが困難になります。

さらに、これまでの研究では、1 ビット モデルがどのような構造を採用する可能性があるかが調査されてきました。数か月前、BitNet [3]の研究は、モデルパラメータをSign(·)関数に渡して+ 1/-1に変換することで1ビット表現を実現しました。しかし、この方法ではパフォーマンスが大幅に低下し、トレーニング プロセスが不安定になるため、実際の適用が制限されます。

ワンビットフレームワーク

OneBit のメソッド フレームワークには、新しい 1 ビット レイヤー構造、SVID に基づくパラメーター初期化メソッド、量子化を考慮した知識蒸留に基づく知識転送が含まれています。

1. 新しい1ビット構造

OneBit の最終的な目標は、LLM の重み行列を 1 ビットに圧縮することです。真の 1 ビットでは、各重み値が 1 ビットでのみ表すことができること、つまり、可能な状態が 2 つだけあることが求められます。著者は、大規模モデルのパラメータにおいては、浮動小数点数の高い精度とパラメータ行列の高ランクという 2 つの重要な要素を考慮する必要があると考えています。

そのため、著者らは量子化による精度の低下を補うために FP16 形式の 2 つの値ベクトルを導入しています。この設計は、元の重み行列の高ランクを維持するだけでなく、値ベクトルを通じて必要な浮動小数点精度を提供し、モデルのトレーニングと知識の転送を容易にします。 1 ビット線形層の構造を FP16 高精度線形層の構造と比較すると、次のようになります。

図3: FP16線形層とOneBit線形層の比較

(a) 左側は FP16 精度モデル構造、(b) 右側は OneBit フレームワークの線形レイヤーです。 OneBit フレームワークでは、値ベクトル g と h のみが FP16 形式を維持し、重み行列は完全に ±1 で構成されていることがわかります。このような構造は、精度とランクの両方を考慮に入れており、安定した高品質の学習プロセスを保証するために非常に有意義です。

OneBit はモデルをどの程度圧縮しますか?著者は論文の中で計算を行っている。 4096*4096 の線形レイヤーが圧縮されていると仮定すると、OneBit には 4096*4096 の 1 ビット マトリックスと 2 つの 4096*1 の 16 ビット値ベクトルが必要です。ビットの総数は 16,908,288 で、パラメータの総数は 16,785,408 です。平均すると、各パラメータは約 1.0073 ビットしか占有しません。この圧縮範囲はこれまでになく、まさに1ビットラージモデルと言えます。

2. SVIDに基づいて量子化モデルを初期化する

完全にトレーニングされた元のモデルを使用して量子化モデルをより適切に初期化し、より優れた知識移転効果を促進するために、著者らは「値と符号に依存しない行列分解 (SVID)」と呼ばれる新しいパラメータ行列分解法を提案しました。この行列分解法は、符号と絶対値を分離し、絶対値に対してランク 1 近似を実行します。元の行列パラメータを近似する方法は、次のように表すことができます。

ここでのランク 1 近似は、特異値分解 (SVD) や非負値行列分解 (NMF) などの一般的な行列分解法によって実現できます。次に、著者らは、この SVID 法が演算順序を交換することで 1 ビット モデル フレームワークと一致し、パラメータの初期化を実現できることを数学的に示しています。さらに、この論文では、記号行列が分解プロセス中に元の行列を近似する役割を果たすことも証明されています。

3. 知識の蒸留による元のモデルの機能の移転

著者らは、大規模モデルの超低ビット幅量子化を解決する効果的な方法は、量子化を考慮したトレーニング QAT である可能性があると指摘しています。 OneBit モデル構造では、知識蒸留を使用して量子化されていないモデルから学習し、機能ベクトル化モデルの移行を実現します。具体的には、生徒モデルは主に教師モデルのロジットと隠れ状態によって誘導されます。

トレーニング中は、値ベクトルと行列の値が更新されます。モデルが量子化された後、Sign（・）以降のパラメータが保存され、推論および展開時に直接使用されます。

実験と結果

OneBit は、FP16 Transformer、従来のトレーニング後の量子化の強力なベースライン GPTQ、量子化を考慮したトレーニングの強力なベースライン LLM-QAT、最新の 2 ビット重み量子化の強力なベースライン OmniQuant と比較されます。また、1ビット重み量子化に関する研究は存在しないため、著者は自身のOneBitフレームワークにのみ1ビット重み量子化を使用し、他の方式には2ビット量子化設定を採用しており、これは典型的な「弱者が強者より優位」である。

モデル選択に関しては、著者らは OneBit の有効性を実証するために、1.3B から 13B までのさまざまなサイズのモデルと、OPT および LLaMA-1/2 のさまざまなシリーズも選択しました。評価指標に関しては、著者は、以前のモデル定量化の 2 つの主要な評価次元、すなわち検証セットの困惑度と常識推論のゼロショット精度を採用しました。

表1: OneBitとベースライン手法のパフォーマンスの比較（OPTモデルとLLaMA-1モデル）

表2: OneBitとベースライン手法のパフォーマンスの比較（LLaMA-2モデル）

表 1 と 2 は、1 ビット量子化における他の方法と比較した OneBit の利点を示しています。検証セット上の量子化モデルの難解さの点では、OneBit は FP16 モデルに最も近いです。ゼロショット精度に関しては、OneBit 量子化モデルは、OPT モデルの個々のデータセットを除いて、ほぼ最高のパフォーマンスを実現します。残りの 2 ビット量子化方法では、両方の評価メトリックで大きな損失が見られます。

OneBit は、より大きなモデルでパフォーマンスが向上する傾向があることに留意してください。つまり、モデルのサイズが大きくなるにつれて、FP16 精度モデルはパープレキシティの低減効果がほとんどなくなりますが、OneBit ではパープレキシティの低減効果がさらに高まります。さらに、著者らは、超低ビット幅の量子化には量子化を考慮したトレーニングが必要になる可能性があることも指摘しました。

図4: 常識推論課題の比較

図5: 世界知識比較

図6: いくつかのモデルのスペース占有率と平均ビット幅

図 4 ～ 6 では、異なるアプローチで取得されたいくつかの種類の小さなモデルのスペース使用量とパフォーマンス損失も比較しています。これには、完全にトレーニングされた 2 つのモデル Pythia-1.0B と TinyLLaMA-1.1B、および低ランク分解で取得された LowRank Llama と OneBit-7B が含まれます。 OneBit-7B は平均ビット幅が最も小さく、占有スペースも最も小さいにもかかわらず、常識推論能力においては他のモデルよりも優れていることがわかります。著者はまた、このモデルが社会科学の分野で深刻な知識の忘却に直面すると指摘している。全体として、OneBit-7B は実用的なアプリケーション価値を実証しています。図 7 に示すように、OneBit 量子化 LLaMA-7B モデルは、命令の微調整後に流暢なテキスト生成機能を発揮します。

図7: OneBitフレームワークによって量子化されたLLaMA-7Bモデルの機能

議論と分析

1. 効率性

表3: 異なるLLaMAモデルにおけるOneBitの圧縮率

表 3 は、さまざまなサイズの LLaMA モデルに対する OneBit の圧縮率を示しています。 OneBit のモデルの圧縮率は 90% を超えており、これは前例のない数値であることがわかります。モデルのサイズが大きくなるにつれて、OneBit の圧縮率が上がることは注目に値します。これは、量子化に関与しない埋め込みレイヤーのパラメータの割合がどんどん小さくなるためです。前述のように、モデルが大きくなるほど、OneBit によってもたらされるパフォーマンスの向上が大きくなり、これは、より大きなモデルでの OneBit の利点を示しています。

図8: モデルサイズとパフォーマンスのトレードオフ

超低ビット量子化ではパフォーマンスが多少低下する可能性がありますが、図 8 に示すように、サイズとパフォーマンスのバランスが良好になります。著者らは、特にモデルをモバイル デバイスに展開する場合、モデルのサイズを圧縮することが非常に重要であると考えています。

さらに、著者らは 1 ビット量子化モデルの計算上の利点も指摘しました。パラメータは純粋なバイナリなので、1 ビットで 0/1 で表すことができ、間違いなく多くのスペースを節約できます。高精度モデルにおける行列乗算の要素ごとの乗算は、効率的なビット演算に変換することができ、ビット割り当てと加算のみで行列乗算を完了できるため、大きな応用の見通しがあります。

2. 堅牢性

バイナリ ネットワークは一般に、トレーニングが不安定だったり、収束が困難だったりするなどの問題に直面します。著者が導入した高精度の値ベクトルのおかげで、モデルトレーニングの順方向計算と逆方向計算の両方が非常に安定しています。 BitNet は以前 1 ビットのモデル構造を提案しましたが、この構造では完全にトレーニングされた高精度モデルから機能を転送することが困難です。図 9 に示すように、著者はさまざまな学習率を試して BitNet の転移学習能力をテストし、教師の指導下では収束がより困難であることがわかりました。これは、OneBit の安定したトレーニング値も間接的に証明しました。

図9: さまざまな学習率でトレーニングした後のBitNetの量子化能力

論文の最後で、著者は超低ビット幅の将来の研究の方向性についても示唆した。たとえば、より優れたパラメータ初期化方法、より少ないトレーニング コスト、またはアクティベーション値の量子化のさらなる検討などです。

より技術的な詳細については、元の論文を参照してください。