Train separately, merge together: Modular post-training with mixture-of-experts｜個別訓練、統合結合：Mixture-of-Expertsを用いたモジュール化ポストトレーニング | AIニュース最前線

2026年4月20日 Jacob Morrison, Sanjay Adhikesaven, Akshita Bhagia, Matei Zaharia, Noah A. Smith, Sewon Min - Ai2 [モデル](https://huggingface.co/collections/allenai/branch-adapt-route) [技術レポート](https://allenai.org/papers/bar) [コード](https://github.com/allenai/FlexOlmo/blob/jacobm-flex-post-train/scripts/BAR/bar_scripts.md) 事前学習（pretraining）終了後、言語モデルは実用的な有用性を獲得するために一連の中・後の学習段階を経て、指示のフォロー、問題解決のための推論、ツールの確実な呼び出しなどの能力を習得します。しかし、これらの段階を経てモデルを更新または拡張することはしばしば困難です。最も確実なオプションである、新しい能力を最初から組み込んでゼロから再学習（retraining from scratch）する方法は高コストであり、元の学習環境への完全なアクセスを必要とします。新しいデータでさらに学習させる方法はコストが低いですが、既存の能力を失う可能性があります。また、後の学習（post-training）は通常複数の段階からなり、それぞれが独自のデータと目的を持つため、新しいスキルを追加するには、以前の成果を損なうことなくそれらに対応するために各段階を再実行または調整する必要があります。 私たちは、これらの課題を回避するモジュラー型ポストトレーニングのためのレシピである**BAR（Branch-Adapt-Route）**を紹介します。単一のモデルをすべてのデータに対して一度に訓練するのではなく、BARは独立したドメイン固有のエキスパート（各々が独自の完全なトレーニングパイプラインを通じて訓練される）を構築し、Mixture-of-Experts（MoE：エキスパートの混合）アーキテクチャを通じてそれらを統合モデルに構成します。各エキスパートは、他のエキスパートに触れることなく、開発、アップグレード、または置き換えが可能です。 私たちは、このアプローチを検証するために使用された[レシピ](https://github.com/allenai/FlexOlmo/blob/jacobm-flex-post-train/scripts/BAR/bar_scripts.md)、[技術レポート](http://allenai.org/papers/bar)、および[チェックポイント](https://huggingface.co/collections/allenai/branch-adapt-route)を公開しています。 ## 背景と動機 私たちの以前の[FlexOlmo](https://arxiv.org/abs/2502.13613)に関する研究は、モジュラー型MoEベースのトレーニングが事前学習（pretraining）において効果的であることを示しました。共有されたベースから分岐し、すべての共有レイヤーを凍結した状態でドメイン固有のフィードフォワードネットワーク（FFN：Feed-Forward Network）エキスパートを訓練し、それらを再度マージすることができます。しかし、私たちはこのレシピがポストトレーニングには適用されないことを発見しました。その理由は後から考えれば直感的です。事前学習は主にFFNレイヤーに存在する知識表現を更新しますが、ポストトレーニングでは、新しい出力形式、推論パターン、安全制約といった行動面の変化を導入する必要があり、これにはアテンションレイヤー、埋め込み（embeddings）、言語モデリングヘッドなどの共有パラメータへの変更が必要です。 例えば、検証済み報酬を用いた強化学習（RLVR）の段階でFlexOlmoのアプローチを直接試した際、報酬曲線は完全に平坦な状態となり、すべての共有パラメータを凍結したままではモデルが学習できませんでした。これが、ポストトレーニング専用の新しいレシピを開発する動機となりました。 ## BARの動作原理 BARには3つのステージがあります： **ステージ1：独立したエキスパートのトレーニング。** 各ドメインのエキスパートは、2つのエキスパートを持つMoE（Mixture of Experts）として実装されます。1つは基盤モデルのFFN（Feed-Forward Network：前層ネットワーク）重みを保持する凍結された「アンカー」エキスパート、もう1つは学習可能なエキスパートです。各エキスパートは、そのドメインが必要とするトレーニングステージを通過します。私たちの実験では、数学とコードの分野はミッドトレーニング（Mid-training）、教師ありファインチューニング（SFT：Supervised Fine-Tuning）、そしてRLVRを経験します。一方、ツール使用と安全性の分野はSFTのみを行います。 この技術的な主要な貢献は、ステージ全体にわたる共有パラメータに対する**段階的な凍結解除スケジュール**です。 - 中間トレーニング：すべての共有レイヤーを凍結（事前トレーニングと同じ。知識の習得はFFNの更新のみで十分に捉えられるため）。 - SFT：埋め込み層と言語モデル化ヘッドを凍結解除。これは、新しい特殊トークン（例：ツール使用のための関数呼び出しフォーマット）を導入するドメインにおいて必要不可欠です。これを解除しないと、ツール呼び出し性能評価に使用したツール使用ベンチマークであるBerkeley Function Calling Leaderboard (BFCL)において、私たちのツール使用エキスパートは20.3というスコアに留まりました。凍結解除により、スコアは46.4まで向上しました。 - RLVR：注意機構を含むすべての共有パラメータを凍結解除。強化学習（RL）は、エキスパートのFFNが対応できる範囲を超えた分布シフトを引き起こします。各エキスパートは、**ドメイン固有のデータと一般的なSFTデータの混合**でもトレーニングされます。私たちはこれが重要であることを発見しました：ドメイン固有のSFTのみでは、ドメイン内でのパフォーマンスは高いものの、指示のフォローや知識といった一般的な能力が著しく低下します。 **ステージ2：エキスパートのマージ。** トレーニング後、すべてのエキスパートを単一のMoEモデルにマージします。エキスパートの実行間で分岐した共有パラメータ（SFTやRLVR中に凍結解除されたため）は、単に平均化されます。私たちは、この平均化により、個々のエキスパートと比較してドメイン固有の評価において測定可能なパフォーマンスの損失はほとんどないか、全くないことを確認しました。 **ステージ 3: ルーターのトレーニング。** 最後に、MoE 内のルーターを、他のすべてのエキスパートと共有重みを凍結した状態でトレーニングします。SFT データの層別抽出 5% サンプルで十分なルーティング効果が得られることがわかっており、このステージは高速かつ低コストで実行できます。 ## 評価全般での高いパフォーマンス 当社のモデルはすべて 7B スケール以上で、完全にポストトレーニングされた Olmo 2 ベースモデルの上に、数学、コード、ツール使用、安全性に関するエキスパートをトレーニングしています。（FlexOlmo アーキテクチャはこのモデルを中心に構築されており、新しいデータセットやポストトレーニングの改善が、元のリリース構成を超えてモデルをどのように強化できるかを探索するための有用なテストベッドを提供するため、Olmo 2 を使用しています。）19 のベンチマークにわたる 7 つの評価カテゴリをカバーする 6 つのベースラインと比較します。以下に報告されているすべてのスコアはカテゴリレベルの平均値です（100 点満点、高いほど良好）。ベンチマークごとの詳細な内訳については、技術レポートを参照してください。 いくつかの点が際立っています： **平均的に、BAR は中途からの再トレーニングを必要としないすべてのベースラインを上回ります。** BAR はポストトレーニングのみの再トレーニングよりも全体的に優れており（49.1 対 47.8）、特に数学（+7.8）とコード（+4.7）で大きな改善が見られます。これはモジュラートレーニングの構造的優位性によるものです。モノリシックなパイプラインでは、数学やコードに関する後期の RL が、早期の SFT ステージで学習された安全性の能力を低下させる可能性があります。モジュラートレーニングでは、各ドメインのパイプラインが独立しているため、この問題が完全に回避されます。 **中間学習後の密モデルマージは壊滅的に失敗する**。中間学習によりモデルの出力が十分に乖離し、単純な重み平均では機能しないほぼ動作しないモデル（ベンチマーク全体で6.5のスコア）が生成される。中間学習を行わない場合でも、マージ結果はBARを大幅に下回る（全体スコアで36.9対49.1）。 **BTX（各エキスパートを完全に独立した密モデルとして訓練する手法）は、同じドメインごとのデータと学習段階を使用しているにもかかわらずBARに劣る**（全体スコアで46.7対49.1）。共有パラメータなしでの訓練はより大きな乖離を生み、ルーティングによる組み合わせを困難にする。 **中間学習付きの完全再訓練は依然として性能の上限**（50.5）を示すが、元の事前学習チェックポイントへの完全なアクセスと、ゼロからの全データ再処理が必要であり、大半のオープンウェイトモデルにとって実用的ではなく、完全なアクセスが可能でも高コストである。 ## モジュール型アップグレード BARのもっとも実用的な有用性の一つは、エキスパートを独立してアップグレードできる点である。私たちは2種類のアップグレードを示す： - 新データへのアップグレード：高品質なデータと強化学習（RL）で訓練されたコード専門モデルに置き換えると、統合モデルにおけるコード性能は+16.5ポイント向上し、他のすべてのドメインはほぼ変化しません。 - 訓練ステージの追加：既存の数学専門モデルに、その監督教師あり学習（SFT）の上に強化学習（RL）を追加すると、統合モデルにおける数学性能は+13ポイント向上し、他のドメインへの影響は最小限に留まります。 どちらの場合も、影響を受けた専門モデルと軽量なルーティングネットワークのみを再訓練する必要があります。単一モデルのパイプラインでは、これらのアップグレードのいずれかを行うために、すべてのドメインにわたってフルモデルを再訓練する必要があります。これにより、BARはドメイン更新に対して線形の費用スケーリングを実現し、単一モデルの再訓練が事実上2次関数的なコスト（各ドメインの更新には全ドメインの再処理が必要）と比較して、大幅な効率化をもたらします。 ## 私たちが学んだこと いくつかの実践的な教訓： - 後期学習（Post-training）には、事前学習（Pretraining）よりも高い柔軟性が求められます。すべての共有レイヤーを凍結する FlexOlmo のレシピは事前学習には有効ですが、後期学習の段階では機能しなくなります。段階的な凍結解除が不可欠であり、特に新しいトークンを持つドメインにおいては、強化学習（RL）中の注意機構の凍結解除や、埋め込み層/言語モデルヘッドの凍結解除が重要です。 - ドメイン固有の教師あり学習（SFT）だけでは不十分です。専門家が自身のドメインデータのみで訓練されると、そのドメイン内のパフォーマンスは向上しますが、汎用的な能力が損なわれます。一般的な SFT データと混合することが極めて重要です。 - 凍結解除後の重み平均化は、驚くほど良好な結果をもたらします。各専門家が SFT および RLVR（Reinforcement Learning from Verifiable Rewards）中に共有パラメータを独立して変更しているにもかかわらず、分岐したパラメータを単純に平均化しても、測定可能な劣化はほとんど、あるいは全く生じません。 - すべての専門家が常にアクティブである必要はありません。推論時に 5 つの専門家のうち 4 つを有効にするだけで、全 5 つを使用した場合とほぼ同等のパフォーマンスが達成され、より効率的なルーティング戦略の余地があることが示唆されています。 ## 今後の展望 実際の大規模モデル開発はすでにモジュール化されており、異なるチームが異なる機能に取り組んでおり、新しいデータセットが異なるタイムラインで登場しており、単一のドメインの改善のためにパイプライン全体を再実行するコストは正当化しにくいものです。BAR は、この現実とトレーニングプロセスを整合させるレシピを提供します。 完全な再学習が依然として性能の上限を決定します。しかし、個々の能力について反復的に改善するチームにとって、BARはモデルの一部を独立してアップグレードし、劣化なく個別にトレーニングされたエキスパートを組み合わせ、単一の学習シーケンスですべてのドメインを実行することに伴う破滅的忘却を回避する方法を提供します。自然な次のステップは、密なモデルを引き上げる（アップサイクル）のではなく、元々疎なアーキテクチャから始めることであり、これによりモジュラーアプローチの効率性とスケーラビリティの両方が向上する可能性があります。 最新のAi2ニュースに関する月次アップデートを受け取るために登録してください。

個別訓練、統合結合：Mixture-of-Expertsを用いたモジュール化ポストトレーニング

背景や根拠まで確認しますか？

関連記事

調べる

選ぶ

サイト