【ITニュース解説】Gemma 4 E2B in Pure JAX on a Colab TPU: Google's 4-Bit Export Against an Exact Repack

大規模言語モデル（LLM）は、現代のIT分野において非常に重要な技術だが、その巨大さゆえに、動作させるためのリソースも膨大になるという課題がある。そこで、モデルの性能を維持しつつ、サイズを小さくして効率的に動作させるための技術が研究されており、その一つが「量子化」と呼ばれる手法だ。量子化とは、モデルが計算に使う数値の精度を落とし、より少ない情報量で表現できるようにすることで、モデルのファイルサイズを小さくし、計算速度を向上させることを目指す技術である。

今回、Googleが開発したLLMの一つである「Gemma 4 E2B」モデルの4ビット量子化版について、その品質を詳細に比較する検証が行われた。Gemma 4 E2Bは、「量子化対応学習（QAT: Quantization-Aware Training）」という特別な方法で訓練されている。これは、モデルが学習する時点から、計算に使われる重み（モデルが学習した知識を表す数値）を4ビットという限られた情報量で表現できるよう、段階的な数値の集合である「グリッド」に乗るように調整しながら学習を進めるものだ。このQATによって、最終的にモデルのサイズを小さくする準備が整っている。

GoogleはこのGemma 4 E2Bモデルの4ビット版を公開している。しかし、その公開されている4ビット版（以下、「Googleのexport」と呼ぶ）は、訓練時にモデルが学習した本来の重みのグリッドとは異なる方法で、重みを再度「丸め込む」、つまり再量子化して保存していることが判明した。具体的には、「各グループの最大重み ÷ 7.5」という固定のルールで再量子化されており、これにより重みの値が訓練時の本来の値からずれてしまう可能性が指摘された。

これに対し、筆者は、訓練時にモデルが学習した本来のグリッドと、そのグリッドの間隔を示す「ステップ」の値を正確に保持するように重みを再構築した4ビット版（以下、「repack」と呼ぶ）を作成した。このrepackとGoogleのexport、そして元の訓練済みモデル（QAT参照モデル）の三者を比較することで、重みの保存方法の違いがモデルの品質や効率にどのような影響を与えるかを検証したのである。検証は、Google ColabのTPU v5eというAIモデルの高速計算に特化したチップ上で、JAXという数値計算ライブラリを用いた独自のエンジンを使用して実施された。

検証の結果は、重みの保存方法の重要性を明確に示した。まず、「重みの正確性」という観点では、Googleのexportの重みは、QAT参照モデルの重みと比較して平均で6.7%もの相対誤差があった。これは、モデルが学習した本来の知識が失われていることを意味する。一方、repackの重みは、QAT参照モデルと比較してわずか0.19%という極めて小さな相対誤差に抑えられていた。さらに、repackでは、重みの値の約73.65%がビット単位でQAT参照モデルと完全に一致しており、訓練された重みを非常に高い精度で再現していることが確認された。Googleのexportで誤差が生じたのは、訓練時に使用されたグリッドのステップとは異なる固定ルールで重みが再量子化されたためであり、この不一致がモデルの知識の劣化を招いたのである。

次に、「モデルの出力品質」について比較した。AIモデルの真価は、与えられた入力に対してどれだけ適切で正確な出力を生成できるかにかかっている。この検証では、モデルが次に予測する単語（トークン）の確率分布や、実際に生成する文章の質が評価された。結果として、repackは、次のトークン予測の分布がQAT参照モデルに比べて342.6倍も近かった（KLダイバージェンスという指標で評価）。また、モデルが一番目と予測するトークンがQAT参照モデルと一致する割合も、repackが99.29%と非常に高かったのに対し、Googleのexportは85.96%にとどまった。さらに、「The history of the Roman Empire」というプロンプトを与えて128トークンの物語を生成させたところ、repackはQAT参照モデルと完全に同じ物語を生成したのに対し、Googleのexportは最初の数単語は同じだったものの、すぐにQAT参照モデルとは異なる単語を選び始め、内容が大きく乖離する物語となった。これらの結果は、重みのわずかな誤差が、モデルの出力に決定的な違いをもたらすことを示している。

「効率性」の面でも、repackはいくつかの点で優れていた。repackのダウンロードサイズは7.51GBで、Googleのexportの8.32GBよりも0.8GB小さかった。これは、Googleのexportに、モデルの動作に不要な重複する重みデータが含まれていたためだ。モデルを動作させるためにチップ上にロードされる重みの量は、repackとGoogleのexportともに6.56GBで同じだった。実行速度（1秒あたりに生成できるトークン数）については、repackが93.3トークン/秒、Googleのexportが93.7トークン/秒と、ほぼ同等だった。ただし、どちらの4ビット版も、元の16ビット精度のQAT参照モデル（135トークン/秒）よりは遅い結果となった。これは、今回の検証で使用されたエンジンが、4ビットの重みを毎回16ビットに展開してから計算を行っていたことが主な理由だ。重みを一度展開してメモリに保持するなどの最適化を行えば、速度は改善される可能性がある。

これらの検証結果を総合すると、repackはGoogleのexportと比較して、訓練されたモデルの知識をはるかに正確に保持しており、その結果として生成される文章の品質が劇的に向上することが明らかになった。さらに、ダウンロードサイズが小さく、実行速度やメモリ使用量は同等であるため、Gemma 4 E2Bモデルをシステムに導入する際には、repackを選択することが強く推奨される。repackはGoogleのexportと同じ「compressed-tensors W4A16」という圧縮形式を使用しているため、既存のシステムからの切り替えも、モデルのリポジトリ名を変更するだけで容易に行えるというメリットもある。

今回の検証は、AIモデルの重みの保存方法という、一見すると技術的な詳細に過ぎない部分が、モデル全体の性能にどれほど大きな影響を与えるかを示す好例となった。システムエンジニアがAIモデルを扱う際には、単に公開されているモデルを使用するだけでなく、そのモデルがどのように構築され、どのような特性を持っているかを深く理解することが、高品質で信頼性の高いシステムを構築する上で非常に重要となるだろう。