教師ありファインチューニング(SFT)において、トレーニングデータに高品質な思考連鎖(CoT)推論トレースを作成することは、しばしば非現実的で非常に高コストです。その結果、SFT中に推論をスキップし、入力と出力のみでトレーニングすることを選択するかもしれません。しかし、推論はAmazon Nova 2ファミリのモデルの中核的な能力であり、Novaや他のフロンティアモデルにおいて、特にコーディングや数学のような難しい問題で性能向上に大きく寄与することが示されています。Amazon Nova 2のカスタマイズでは、SFTや強化学習ファインチューニング(RFT)などの手法を通じて、特定のドメインで思考モデルの力を活用できます。SFTでは、これらの利点を引き出すための重要な要件は、SFTデータセットに黄金のCoTトレース(つまり、強力な教師モデルからの思考トレースをさらに検証・洗浄したもの)が存在することです。
本稿では、SFTカスタマイズにおいて推論トレースを欠くデータセットに対して思考トークンを生成するアイデアを探求します。まず推論抑制問題を検討し、次に自己蒸留推論(SDR)を導入し、3つのベンチマークで検証し、実用的な推奨事項を提供します。SDRはベースのAmazon Nova 2 Liteモデルの思考連鎖を非推論データセットの代用として再利用します。実験を通じて、SDR導入の影響(目標性能の向上から壊滅的忘却の緩和まで)をいくつか観察しました。これは、トレーニング中にモデルからモデル自身への情報蒸留(自己蒸留として知られる)の重要性を発見する研究の増加と一致しています。
以下の図は、目標性能と数学保持に関して、SDRを通常のSFTおよびモデルマージと比較したものです。モデルマージは、以前に獲得したスキルを保持する簡単なアプローチで、SFTチェックポイントを(SFT前の)ベースモデルにマージします。しかし、マージはしばしばSFTチェックポイントで得られた利得を犠牲にします。これは、目標性能と汎用性能のトレードオフとして現れる壊滅的忘却によるものです。一方、SDRは、汎用性能のほとんどを保持しつつ、純粋なSFT(マージなし)と同等またはそれ以上の目標性能を達成することで、この問題を軽減します。
壊滅的忘却のため、データセットへのSFTが以前の知識の完全な喪失につながることが観察されます。例えば、通常のSFTでは、数学性能がベースの70%から平均6%に低下します。SDRを使用すると、ほぼ完全に回復できます。自己蒸留推論により、目標性能を犠牲にすることなくベースモデルの能力を保持でき、多くの場合目標性能も向上します。モデルマージは壊滅的忘却を緩和する一般的な解決策ですが、SDRはこの保持をより効果的に実現します。
SDRは、モデル自身の推論トレースでデータセットを拡張することで、トレーニング中の正則化を提供します。別個の教師モデルや事後補間を必要としません。このアプローチは人間によるアノテーションを必要とせず、ドメインに関係なく既存のSFTデータセットに適用可能であり、目標性能と汎用性能の両方を保持する点でモデルマージよりも効果的です。
モデルマージと比較して、自己蒸留は数学性能でほぼ同等(70%対68%)でありながら、目標性能を平均6.5%以上同時に向上させます。この発見は、壊滅的忘却を緩和する強力なメカニズムとして自己蒸留を提唱する最近の研究と一致しています。
背景:推論トークンのないSFTは推論能力の喪失につながる可能性 アプローチを紹介する前に、なぜ推論トレースのないデータセットでのSFTがモデル性能を低下させるのかを理解することが重要です。このセクションでは、推論抑制問題を検討し、推論が目標性能に与える影響を定量化し、既存の回復メカニズムとしてのモデルマージについて議論します。
推論抑制問題 推論モードを有効にした状態で非推論データセットに対してSFTトレーニングを行うと、重大な問題が発生します。つまり、明示的に推論モードをオンにしても、モデルは推論中に推論能力を失います。この現象は、トレーニング目標が推論トークンと出力トークンの両方の損失を計算するために発生すると仮定します。トレーニングデータに中間推論ステップのない入力-出力ペアのみが含まれている場合、モデルは一貫した推論トレースを生成するための教師信号を受け取りません。損失関数は本質的に最終出力に直接寄与しないトークンを罰し、モデルに推論メカニズムをバイパスするよう訓練します。この動作は近道学習を示しており、モデルは堅牢な推論パターンを学習する代わりに擬似相関に依存します。
推論の影響の定量化 抑制問題にもかかわらず、トレーニングと推論の両方で推論をオンにすると、目標性能に大きな利点がもたらされます。公平な比較のために、トレーニングと推論の設定を一貫させます。つまり、トレーニングと推論の両方で推論オンと両方で推論オフを比較します。実験では、トレーニングと推論の両方で推論をオンにすると、大幅な改善が見られることが明らかになりました。次の表は、LLaVA CoTデータセットで異なるLoRAマージ重みを使用した場合の効果を示しています。
マージ重み | 目標性能(推論オン) | 目標性能(推論オフ) | 差分 0.0(ベース) | 12.30% | 12.30% | 0% 0.3 | 34.38% | 35.28% | -0.9% 0.5 | 60.51% | 48.80% | +11.7% 0.7 | 66.67% | 54.05% | +12.6% 1.0(完全マージ) | 65.17% | 47.90% | +17.3% 表1:推論とマージ重みが目標性能に与える影響。トレーニングと推論のモードは一貫しています。
回復メカニズムとしてのモデルマージ 推論トークンのないSFTユースケースにおける推論抑制問題に対処するための現在の推奨アプローチは、モデルマージを推論スキルを回復するメカニズムとして使用することです。モデルマージがドメイン固有のファインチューニング後に数学やコーディングの能力を回復できるのと同様に、ファインチューニングされたモデルとベースモデルとの加重補間によって推論能力を回復できることがわかります。この回復は、汎用性能指標(数学やコーディングのベンチマークなど)の回復と強く相関しており、推論が複数の下流スキルを支える基本的な能力であることを示唆しています。
これをよりよく理解するために、この実験では次のベンチマークを使用します:
- ToolACE:ツール呼び出しベンチマーク(10Kサンプル)。ユーザークエリと利用可能なツール定義に基づいて、正しい関数呼び出しを生成するモデルの能力を評価します。
- CoCoHD:長文脈構造化抽出(730サンプル)。米国議会公聴会記録からメタデータを抽出します。
- GovReport:長文脈抽象的要約(17Kサンプル)。米国政府報告書の要約を生成します。
- Invoice-OCR:マルチモーダル文書理解(約500サンプル)。スキャンされた請求書画像から構造化フィールドを抽出します。
- CaptionGen:ビデオキャプション(1800サンプル)。ビデオクリップの説明キャプションを生成します。
モデルマージによって推論が出現するかどうかを調査します。命令追従ベンチマークでは、推論なしでトレーニングされたモデル上の推論トークンの数をカウントします。x軸はカスタムモデルとベースモデル間のマージ重みを示し、0.0はベースモデル、1.0はマージなしのSFTモデルに対応します。
カスタマイズへの影響 これらの発見は、カスタマイズされたモデルを展開する場合に重要な意味を持ちます。第一に、トレーニングと推論で同じ推論モードを維持することが最適なパフォーマンスのために重要です。第二に、介入なしでは、ドメイン固有のSFTが汎用推論能力を著しく低下させる可能性があります。第三に、モデルマージは事後解決策を提供しますが、競合する目標のバランスを取るためにマージ重みの慎重な調整が必要です。これらのトレードオフにより、事後介入を必要とせずに推論能力を保持するトレーニングアプローチの開発が促進されました。
SFTにおける推論の役割 SFTにおける推論トレースは、強化学習ファインチューニング(RFT)、正則化、自己蒸留パラダイムを結びつける3つの重要な機能を果たします。第一に、推論トレースは暗黙的なポリシー正則化として機能し、ファインチューニングされたモデルをベースモデルのポリシーに近づけることで、壊滅的忘却と能力劣化を防ぎます。これはRLHFにおけるKL正則化と同様です。第二に、推論はプロセス監視を提供し、結果監視だけでなく、RFTにおけるステップバイステップの報酬形成と類似しており、より堅牢な汎化につながります。第三に、自己蒸留(モデルが自身の予測から学習する)は、Born-Again NetworksからConstitutional AIに至るまで、さまざまなドメインで効果が証明されています。推論トレースは、最終出力だけでなくモデルの問題解決プロセスを捉えることで、特に豊かな自己監視を提供します。これは最近の継続学習研究で示されています。
最近の研究では、効果的なファインチューニングモデルはパラメータ空間でベースモデルに近接していることが示されており、私たちのアプローチはこの洞察を利用して、ベースモデル自身の推論トレースをトレーニングターゲットとして適用し、モデルの内部表現と一致する学習信号を生成すると同時に、推論プロセス全体にわたって密なトークンレベルの監視を提供します。
自己蒸留推論 推論トレースのないデータセットに対するSFTカスタマイズのために、自己蒸留推論(SDR)を使用する効果的なアプローチを提案します。プロセスは3段階で動作します。まず、ベース推論モデル(Amazon Nova 2 Liteなど)を使用して、SFTデータセットの各サンプルの推論トレースを生成します。次に、生成された推論トレースを元の出力の前に追加してトレーニングデータを拡張します。最後に、推論モードをオンにしてモデルをファインチューニングし、推論トークンと出力トークンの両方に監視を提供します。
SDRはSFTワークフローにいくつかの利点をもたらします。推論トレースを作成するための人間の労力が不要であるため、アノテーションコストがゼロです。推論トレースはベースモデル自身の問題解決アプローチと一致しています。この手法はドメインに関係なく既存のSFTデータセットに拡張可能であり、異なる教師モデルを使用する柔軟性も提供します。
SFTのための推論トレースの構築 Amazon Bedrockにクエリを実行して、特定のデータセットの推論トレースを取得します。思考連鎖を構築するには2つのアプローチがあります(以下、AIコスト制御のため省略)。