本文にスキップ
AI News HubLIVE
サイト内リライト4 分で読了

個別に学習し、統合する:混合エキスパートを用いたモジュラー型ポストトレーニング

記事の要約

BAR(Branch-Adapt-Route)は、言語モデルの能力を段階的に学習する手法です。ドメインエキスパートを独立に訓練し、混合エキスパートモデルに統合し、他のエキスパートに影響を与えずに任意のエキスパートをアップグレードできます。

ソースAi2 Blog
個別に学習し、統合する:混合エキスパートを用いたモジュラー型ポストトレーニング
誤りを報告

訂正窓口はまだ利用できません。記事情報をコピーして保存できます。

訂正案内
本文へ

事前学習の後、言語モデルは実用的になるために一連の中間およびポストトレーニング段階を経ます。しかし、これらの段階後にモデルを更新または拡張することはしばしば困難です。最も信頼性の高い選択肢である最初からの再トレーニングは高コストであり、元のトレーニング設定への完全なアクセスが必要です。新しいデータでの追加トレーニングは安価ですが、モデルが既に持っている能力を失う可能性があります。また、ポストトレーニングは通常複数の段階(それぞれ独自のデータと目的を持つ)を含むため、新しいスキルを追加するには、既存の能力を壊さないように各段階を再実行または調整する必要があります。

これらの問題に対処するため、Allen Institute for AI(Ai2)の研究チームはBAR(Branch-Adapt-Route)を提案しました。これはモジュラー型ポストトレーニングの手法です。単一のモデルをすべてのデータで一度に訓練する代わりに、独立したドメインエキスパートを訓練し、混合エキスパート(MoE)アーキテクチャでそれらを統合します。各エキスパートは、他のエキスパートに影響を与えることなく、開発、アップグレード、または交換が可能です。チームはレシピ、テクニカルレポート、および手法の検証に使用したチェックポイントを公開しています。

BARの動機は、以前のFlexOlmoの研究にあります。FlexOlmoは、MoEベースのモジュラートレーニングが事前学習に有効であることを示しました。しかし、このレシピはポストトレーニングには移行できないことが判明しました。その理由は、事前学習が主に知識表現を更新するのに対し、ポストトレーニングは新しい出力形式、推論パターン、安全性制約などの行動変化を導入し、注意層、埋め込み、言語モデルヘッドなどの共有パラメータの変更を必要とするからです。例えば、FlexOlmoのアプローチを強化学習と検証報酬(RLVR)に直接適用すると、報酬曲線は完全に平坦でした。すべての共有パラメータが凍結されているため、モデルは学習できませんでした。これにより、ポストトレーニング専用の新しいレシピの開発が動機づけられました。

BARは3つの段階からなります。第1段階:独立エキスパート訓練。各ドメインエキスパートは、2つのエキスパートからなるMoEとしてインスタンス化されます。1つは凍結された「アンカー」エキスパート(ベースモデルのFFN重みを保持)、もう1つは訓練可能なエキスパートです。エキスパートは、そのドメインが必要とする訓練段階を経ます。実験では、数学とコードは中間トレーニング、教師あり微調整(SFT)、RLVRを経由し、ツール使用と安全性はSFTのみを使用しました。重要な技術的貢献は、段階ごとの共有パラメータの漸進的な凍結解除スケジュールです。中間トレーニングではすべての共有層が凍結され、SFTでは埋め込み層と言語モデルヘッドが凍結解除され(新しい特殊トークンを導入するドメインに必要)、RLVRでは注意層を含むすべての共有パラメータが凍結解除されます。また、各エキスパートはドメイン固有データと汎用SFTデータの混合で訓練されます。ドメインのみのSFTは、ドメイン内性能は強いが、汎用能力を著しく低下させるため、この混合は重要です。

第2段階:エキスパート統合。訓練後、すべてのエキスパートを単一のMoEモデルに統合します。エキスパート間で分岐した共有パラメータ(SFTまたはRLVR中に凍結解除されたため)は単純に平均化されます。この平均化は、ドメイン固有評価において個々のエキスパートと比較して測定可能な性能低下をほとんど引き起こさないことがわかりました。第3段階:ルーター訓練。最後に、他のすべてのエキスパートと共有重みを凍結した状態で、MoE内部のルーターを訓練します。SFTデータの層化5%サンプルで効果的なルーティングが可能であり、この段階は高速で低コストです。

実験では、BARモデルは19のベンチマークで、初期からの再トレーニングを必要としないすべてのベースラインを上回りました。平均して、BARはポストトレーニングのみの再トレーニングモデルを上回り(49.1 vs 47.8)、特に数学(+7.8)とコード(+4.7)で大きな改善を示しました。これは、モジュラートレーニングの構造的利点によるものです。単一パイプラインでは、後期の数学とコードのRLが初期のSFT段階で学習された安全性能力を低下させる可能性がありますが、モジュラートレーニングはこれを完全に回避します。注目すべき点として、中間トレーニング後の高密度モデル統合は壊滅的に失敗し(全体で6.5)、一方BARは中間トレーニングなしでも大きくリードしました。

BARの最も実用的な特性の1つは、エキスパートを独立してアップグレードできることです。研究では2種類のアップグレードを示しました。より高品質なデータとRLで訓練されたコードエキスパートに置き換えると、コード性能が16.5ポイント向上し、他のドメインはほぼ変化しませんでした。既存の数学エキスパートにRLを追加すると、数学性能が13ポイント向上しました。どちらの場合も、影響を受けるエキスパートと軽量ルーターのみを再訓練する必要があり、単一パイプラインではモデル全体の再訓練が必要です。これにより、BARのドメイン更新コストは線形にスケールし、単一モデルの更新コストはほぼ二次的にスケールします。

研究から得られた実用的な教訓:ポストトレーニングは事前学習よりも柔軟性を必要とし、漸進的な凍結解除が不可欠です。ドメインのみのSFTは不十分で、汎用SFTデータとの混合が重要です。凍結解除後の重み平均化は驚くほど効果的です。また、推論時にすべてのエキスパートをアクティブにする必要はなく、4/5のエキスパートでほぼ同等の性能が得られます。今後の課題として、高密度モデルを変換する代わりにネイティブな疎アーキテクチャから開始することで、モジュラーアプローチの効率とスケーラビリティが向上する可能性があります。

要点と分析を開く

記事インテリジェンス

エンジニア上級

要点

  • BARは独立したドメインエキスパートを訓練し、MoEアーキテクチャで統合することでモジュラー型のポストトレーニングを実現。
  • 共有パラメータの段階的な凍結解除が重要:SFTでは埋め込みと言語モデルヘッド、RLではアテンション層。
  • BARは再トレーニングを必要としないベースラインを上回り、線形コストでエキスパートの独立アップグレードが可能。
  • 汎用SFTデータとドメインデータの混合が汎用能力維持に不可欠。

要点と分析は自動生成され、誤りを含む場合があります。原典をご確認ください。