TRACE、MoE強化学習におけるFP4学習とRolloutを整合

大規模言語モデルのポストトレーニングでは、強化学習が重要な役割を担っています。一方で、強化学習のコストはパラメータ更新だけで決まりません。モデルは何度もRolloutを生成する必要があり、生成時の計算量、メモリ使用量、KVキャッシュの容量が大きな負担になります。そのためRolloutをFP4のような低精度で実行する方法が注目されていますが、BF16で学習したモデルを後から量子化するだけでは、強化学習に必要な挙動を保てない場合があります。

課題は学習経路とRollout経路のずれ

既存のFP4強化学習手法は、学習経路とRollout経路の量子化精度を別々に改善することが一般的です。しかし、それぞれの局所的な量子化誤差が小さくても、両経路の出力やモデル挙動が一致するとは限りません。強化学習では、現在の方策が次の学習データを生成します。そのため、学習時に想定した計算と実際のRolloutで使われる計算が異なると、方策の更新とデータ生成の間にずれが生じる可能性があります。

TRACEは、この差を直接抑えることを目的とします。名称はTrain-Rollout Quantization Alignment via Compact Guidanceの略で、混合専門家（MoE）言語モデル向けに、二つの実行経路を整合させる量子化フレームワークとして提案されています。

TRACEの主な仕組み

• Rollout誘導型の量子化認識学習：Rollout側で得られた量子化結果を学習に戻し、学習側のFP4丸め判断を導きます。学習側が独立して量子化結果を選ぶのではなく、実際の生成で使われる低精度の挙動に近づけます。
• 生成パイプライン全体のFP4化：重みと活性値だけでなく、FP4 KVキャッシュによるRolloutにも対応します。長い生成でメモリを消費しやすいキャッシュも最適化対象に含める点が特徴です。
• 量子化情報の選択的キャッシュ：Rollout情報を学習へ渡すと、保存・通信コストが増加します。TRACEは深い層の仮数とスケール情報を選択的に保持し、誘導効果と追加オーバーヘッドのバランスを取ります。

結果と意味

論文では、4つの大規模MoE言語モデルを使い、推論、コード、長期的な強化学習タスクで評価しています。重み、活性値、KVキャッシュをまとめてFP4化したRolloutでも、BF16 Rolloutに近い強化学習性能を達成できたと報告されています。また、BF16で学習した方策を後からFP4化する方法と比べ、最終的なFP4性能も強い結果を示しました。報告された最大Rollout高速化は5.4倍です。

この研究の重要性は、低ビット化の目標を「高精度モデルをどこまで再現できるか」から、「実際に使う低精度経路に合わせて学習できるか」へ広げた点にあります。MoEモデルでは生成を繰り返すため、Rolloutの計算、キャッシュ、通信が運用上のボトルネックになりやすく、経路レベルの整合は実用的な意味を持ちます。

ただし、提示された資料は全モデル、タスク、ハードウェアごとの詳細な内訳を示していません。したがって5.4倍は論文で報告された最大値であり、すべての環境で得られる固定的な効果ではありません。それでもTRACEは、低精度強化学習では単独の量子化誤差だけでなく、学習とRolloutの実行経路そのものを揃えることが重要だと示しています。

出典：Hugging Face Daily Papers