🤖 AI

TetraJet-v2: Accurate NVFP4 Training for Large Language Models with Oscillation Suppression and Outlier Control

TetraJet-v2 は、すべての線形層で NVFP4 精度を活用し、重みの振動を効果的に抑制して外れ値を制御するために OsciReset と OutControl を導入する、大規模言語モデル向けのエンドツーエンドの 4 ビット完全量子化トレーニング手法であり、FP8 に比べて大幅な高速化を達成しつつ、ほぼ損失のない性能を実現します。

原著者： Yuxiang Chen, Yifan Liu, Xiaoming Xu, Pengle Zhang, Michael Beyer, Martin Rapp, Jun Zhu, Jianfei Chen

公開日 2026-05-06

📖 1 分で読めます☕ さくっと読める

CC BY 4.0

原著者： Yuxiang Chen, Yifan Liu, Xiaoming Xu, Pengle Zhang, Michael Beyer, Martin Rapp, Jun Zhu, Jianfei Chen

原論文は CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) でライセンスされています。 ✨ これは以下の論文のAI生成解説です。著者が執筆または承認したものではありません。技術的な正確性については原論文を参照してください。免責事項の全文を読む

巨大で優秀な生徒（大規模言語モデル）に膨大な量の情報を教えようとしている状況を想像してください。通常、これを効果的に行うには、完璧な照明と巨大なホワイトボードを備えた高品質で高価な教室（高精度コンピューティング）が必要です。これには莫大な電力とハードウェアコストがかかります。

この論文は、同じ生徒を小さな低コストの教室で、小さな黒板と薄暗い照明（4 ビット精度）を使って教えようとする新しい手法「TetraJet-v2」を紹介しています。その目標は、生徒が間違ったことを学んだり、直前に学んだことを忘れたりすることなく、トレーニングプロセスを 1.67 倍高速化し、大幅に低コスト化することです。

以下は、そのような「低品質な」教室で教えようとする際に通常発生する問題に対する解決策です。

1. 問題：「揺れる手」（重みの振動）

生徒が黒板に数字を書こうとしている状況を想像してください。完璧な教室では、「0.25」を正確に書けます。しかし、この低精度の教室では、整数または半整数（0 や 0.5 など）しか書けません。

もし生徒の真の答えがちょうど中間（0.25）にある場合、彼の手は揺れ始めます。ある瞬間は「0」を書き、次の瞬間は「0.5」を書き、また「0」に戻ります。彼はループに陥り、行ったり来たりしてしまいます。論文ではこれを**重みの振動（Weight Oscillation）**と呼んでいます。これはエネルギーを浪費し、生徒が正しい答えに落ち着くことを妨げます。

解決策：OsciReset（「中心化」のトリック）
著者たちはOsciResetと呼ばれる技法を開発しました。これは生徒の手を見守る教師を想像してください。教師が生徒が 2 つの選択肢の間で揺れ始めたと気づいた瞬間、教師は優しくチョークを掴み、安全域（「ビン中心」）の真ん中にピタリと合わせます。

なぜ機能するか： 揺れを即座に止めます。生徒を凍結させる（学習を停止させる）のではなく、新しいより滑らかな方向へ学習を続けるための安定した出発点を与えるだけです。これは、これらの巨大モデルにおいて、この特定の「揺れ」の問題が初めて解決されたものです。

2. 問題：「大声で叫ぶ者」（外れ値）

通常の教室では、ほとんどの生徒は普通の音量で話します。しかし、これらの巨大モデルでは、いくつかの特定の「チャネル」（マイクと考えるとよいでしょう）が突然最大音量で叫び始めます。これらは**外れ値（Outliers）**と呼ばれます。

黒板が小さすぎる（低精度である）ため、1 つのマイクが叫びすぎると、そのマイクが黒板を壊さないように、教師は他の全員の音量を下げなければなりません。これにより、他の生徒からの静かで重要なささやきがノイズの中に埋もれてしまいます。

解決策：OutControl（「VIP マイク」）
著者たちは、これらの叫んでいるマイクは予測可能であり、常に同じ少数のチャネルであることを発見しました。そこで、OutControlと呼ばれるシステムを作成しました。

仕組み： 彼らは早期に「VIP マイク」（外れ値）を特定し、それぞれに専用の高品質なマイクを与えます（FP8 などのより高い精度形式で保持します）。残りの生徒たちは、小さく安価な黒板（FP4）を使い続けます。
結果： 大声で叫ぶ者が他の全員の音量を下げさせることがなくなるため、静かなささやきもはっきりと聞こえるようになります。

3. 「ダブルブロック」黒板

この小さな黒板で数学が完璧に機能するようにするため、彼らは数字の書き方を再設計しました。巨大な数字を 1 つの小さな四角に収めようとする代わりに、ダブルブロックシステムを使用します。

想像してください。大きなグループの生徒（128 人のブロック）が 1 つの大きな定規（グローバルスケール）を共有していますが、そのグループ内では 16 人ずつの小さなチームに分かれており、それぞれが独自の小さな定規を持っています。これにより、非常に小さな数から非常に大きな数まで、精度を失うことなく処理できます。

結果

これらすべてのトリック（中心化のトリック、VIP マイク、ダブルブロック黒板）を組み合わせると：

速度： 現在の標準（FP8）よりも1.67 倍高速にモデルをトレーニングしました。
精度： モデルは、高価で高精度な教室でトレーニングされた場合とほぼ同等に学習しました。安価な手法と高価な手法の間の性能ギャップを50% 以上削減しました。
規模： 最大 3 億 7000 万パラメータのモデルでテストし、2000 億語以上でトレーニングしました。

要約： TetraJet-v2 は、揺れを止め、大きな声を保護し、黒板をより効率的に整理することで、はるかに安価で高速なハードウェア上で巨大な AI モデルをトレーニングできるようにする新しいルールセットです。

技術概要：TetraJet-v2

1. 問題定義

大規模言語モデル（LLM）の前学習は、ますます高コスト化しており、低精度の完全量子化トレーニング（FQT）の必要性を駆り立てています。NVFP4 などの 4 ビット形式は、FP8 や BF16 に比べて大幅な効率向上をもたらしますが、この精度で安定した、損失の少ないトレーニングを実現することは依然として重要な課題です。これまでに試みられた完全 4 ビットトレーニングは、深刻な性能低下に悩まされてきました。

著者らは、単純な量子化誤差とは異なる、完全な FP4 トレーニングを阻害する 2 つの主要な最適化ボトルネックを特定しました：

重みの振動（Weight Oscillation）： 高精度のマスター重みに対する小さな更新が、その量子化表現を隣接するビン間（例：0 と 0.5 の間）で繰り返し反転させる、支配的な失敗モードです。この現象は、特に学習率が減少するにつれて最適化を妨げ、モデルが安定した局所最小値に収束するのを防ぎます。
アウトレイヤー特徴（Outlier Features）： 少数の活性化チャネルが巨大な大きさを持ち、FP4 が正確に表現できない動的範囲を作成します。これにより、同じ量子化グループ内のより小さな値の精度を圧迫する高いスケーリング係数が生じます。

NVIDIA の NVFP4 レシピなどの既存の方法は、しばしば混合精度（一部のブロックを BF16 に保持）に依存するか、ゼロから前学習を行う設定において、これらの最適化ダイナミクスを完全に解決できていません。

2. 手法：TetraJet-v2

TetraJet-v2 は、NVFP4 に特化したエンドツーエンドの 4 ビット FQT 手法であり、すべての線形層において活性化、重み、勾配にこの形式を適用します。この手法は、特定されたボトルネックを以下の 3 つの中核コンポーネントを通じて解決します：

A. NVFP4 線形層のための不偏ダブルブロック量子化

NVFP4 の数値的制約（8 ビットのスケーリングを用いた 16 要素グループを使用）内での精度を最大化するため、著者らはダブルブロック量子化スキームを提案します：

構造： 128 要素の外部ブロックはグローバルにスケーリングされ、16 要素の内部ブロックはローカルにスケーリングされます。これにより、テンソル全体にわたる絶対値の最大値を計算する必要がなくなり、テンソルごとのスケーリングと比較してハードウェアの親和性と精度が向上します。
不偏勾配推定： 線形層の設計は、以下の方法により不偏な勾配推定を確保します：
- 逆伝播においてテンソルを整合させ（ $X$ ではなく $\hat{X}$ を推定する）。
- 逆伝播において**確率的丸め（Stochastic Rounding）**を使用し、 $E[\text{round}(x)] = x$ となるようにして、勾配期待値におけるバイアスを防止する。
- 高精度を得るために $1 \times 16$ の重みグループ形状（非正方形）を採用し、転置重み行列に対して特定の整合を必要とする。

B. OsciReset：振動抑制

重みの振動に対処するため、著者らは LLM の前学習においてこのボトルネックを抑制するために設計された最初のアルゴリズムであるOsciResetを導入します。

検出： この手法は重みの最適化軌跡を追跡し、「振動リスク」指標を計算します。これは、量子化された重みが移動した距離（ $dist_Q$ ）とマスター重みが移動した距離（ $dist_M$ ）の比率として定義されます。高いリスクは、マスター重みの更新が小さいにもかかわらず、重みがビンを反転していることを示します。
メカニズム： 学習率が低く、振動リスクがしきい値を超えた場合、振動している要素のマスター重みを現在の量子化ビンの中心にリセットします。
利点： 重みを凍結したり、重い正則化を使用したりする（これらはグローバル最適化を害する）従来の手法とは異なり、OsciReset は重みが振動しにくい新しい軌道で最適化を継続することを可能にし、即座に性能を低下させることなく振動を抑制します。

C. OutControl：アウトレイヤー制御

アウトレイヤーを管理するため、TetraJet-v2 はハイブリッド戦略を採用します：

ランダムアダマール変換（RHT）： 逆伝播計算（$dX $と$ dW$）に適用し、チャネル間でアウトレイヤーのエネルギーをより均等に分散させ、量子化分散を軽減します。
静的アウトレイヤー保持： アウトレイヤーチャネルは構造的に一貫している（特定のチャネルが常に高いノルムを持つ）ことを認識し、短い較正ウィンドウ中に上位 $p\%$ （例：10%）のアウトレイヤーチャネルを特定します。これらのチャネルは、順伝播と逆伝播の両方でより高い精度（例：FP8）で保持され、ネットワークの残りは FP4 のままです。この静的な選択により、バッチごとの動的調整に伴う計算オーバーヘッドを回避します。

3. 主な貢献

不偏ダブルブロック量子化： 不偏勾配推定と LLM 向けの最適なブロック形状（ $1 \times 16$ ）を備えた新規の NVFP4 線形層設計。既存の NVFP4 レシピを上回る性能。
OsciReset： 前学習中の LLM における重み振動を抑制する最初の効果的なアルゴリズム。低ビットトレーニングにおける不安定性の主要な源を解決。
OutControl： 逆伝播には RHT を、順伝播と逆伝播の両方には静的アウトレイヤー保持を組み合わせる混合精度アルゴリズム。精度を維持する。
包括的な評価： OLMo-2 モデル（最大 370M パラメータ、212B トークン）における広範な前学習実験により、フルスタックの有効性を実証。

4. 実験結果

著者らは、ゼロからトレーニングされた OLMo-2 モデル（70M、150M、370M パラメータ）上で TetraJet-v2 を評価しました。

性能ギャップの縮小： 提案されたすべての手法を組み込んだ TetraJet-v2-full は、モデルサイズ全体で完全精度（BF16）トレーニングとの性能ギャップを**平均 51.3%**削減しました。
パープレキシティ（PPL）： 212B トークンでトレーニングされた 370M モデルにおいて、TetraJet-v2-full は検証用 PPL で25.11を達成しました。これは BF16 の 23.70 や NVIDIA NVFP4 ベースラインの 26.20 と比較されます。
下流タスク： 一連のベンチマーク（ARC、BoolQ、MMLU など）において、すべての FP4 手法の中で平均して最高の性能を達成しました。370M モデルは平均 43.60% の精度に達しました。
効率性： CUDA カーネル実装は、RTX 5090 ハードウェア上で FP8（TransformerEngine）に対してエンドツーエンドで 1.67 倍の高速化を達成しました。振動抑制とアウトレイヤー制御の追加された複雑さにもかかわらず、実用的な性能向上が確認されました。
アブレーション研究：
- OsciReset を除去すると、深刻な性能低下と損失の反発が生じました。
- OutControl は順伝播と逆伝播の両方で精度を向上させました。
- この手法は、ハイパーパラメータの選択（振動しきい値、開始時間など）に対して堅牢であることが示されました。

5. 意義と主張

本論文は、TetraJet-v2 が LLM の完全な低ビット（4 ビット）前学習を実用的かつ高精度トレーニングと競争可能なものにするための重要な一歩であると主張しています。

最適化対表現： 著者らは、低ビットトレーニングにおける性能ギャップは量子化誤差（表現）のみによるものではなく、以前の手法が体系的に解決できなかった最適化の問題（振動とアウトレイヤー）によって決定的に駆動されていることを強調しています。
一般化可能性： 視覚トランスフォーマー向けに設計された以前の振動緩和技術（多くの場合微調整や特定のハイパーパラメータ調整に依存）とは異なり、OsciReset は、最小限の調整でゼロから LLM 前学習に適用可能な、シンプルで効果的な手法として提示されています。
実用性： FP8 に対して 1.67 倍の高速化を達成しながら、BF16 との精度ギャップを 50% 以上縮小することで、この研究は完全量子化トレーニングが計算効率とモデル品質の両方を提供し、基盤モデル研究のハードウェアおよびエネルギーコストを削減する可能性を実証しています。

著者らは、実験の規模（最大 370M パラメータ）に関する限界と、FP6 $\times$ FP4 のような混合精度戦略に対する現在の主流ハードウェアのサポートの欠如を認めています。これらは、さらなる精度向上のための将来の方向性として検討されました。