技術概要：ASTOR – コード用 LLM 向けのユーティリティ誘導型マルチタスク強化学習

問題定義

検証可能な報酬を用いた強化学習（RL）は、コーディングタスクにおける大規模言語モデル（LLM）の事後学習において有効であることが実証されている。しかし、現在のパラダイムは、コード生成用、単体テスト生成用など、タスク固有の RL 専門モデルを個別に訓練する方式に依存している。このアプローチは、タスク数に比例してメモリおよび計算コストが線形に増加するという課題を伴う。一方、単一の専門モデルに依存すると、そのモデルが特定の訓練ドメイン外で適用された際に、性能が急激に低下する。

マルチタスク強化学習（MTRL）は、単一のモデルをすべてのタスクにわたって訓練することで統合的な解決策を提供するが、既存の手法は異質なコーディングタスクを均一に扱うことにより、以下の 2 つの重大な制限に直面している：

1. データスケジューリング（「何を学ぶか」）： 現在の手法は、混合タスクプールからのサンプリングまたは固定されたカリキュラムに従っており、タスク間の相互影響を無視している。あるタスク（例：コード生成）の訓練が、別のタスク（例：単体テスト生成）にどのように正の影響または負の影響を与えるかを考慮できていない。
2. 方策最適化（「どのように学ぶか」）： 既存のアプローチは、すべてのタスクに対して均一な更新戦略と単一の KL 正則化係数を適用する。これは最適ではない。なぜなら、コーディングタスクには多様な最適化要件が存在するためである。構文が厳格なタスク（例：単体テスト生成）は、実行可能性を維持するために保守的な更新を必要とする一方、意味的に柔軟なタスク（例：コミットメッセージ生成）は、より広範な探索の恩恵を受ける。

手法：ASTOR

著者らは、ASTOR（ユーティリティ駆動型協調によるマルチタスクコード強化学習フレームワーク）を提案する。これは、タスクユーティリティという統合された信号を中心に据えた 2 つの結合モジュールを通じてこれらの課題を解決するフレームワークである。タスクユーティリティは、タスクの学習ポテンシャルとタスク間シナジーの両方を捉える。

1. 階層的ユーティリティルーティングデータスケジューリング

このモジュールは、推定されたタスクユーティリティに基づいて訓練予算を動的に割り当てることで、「何を学ぶか」に対処する。

• 二重レベルのユーティリティ推定：
  • タスクレベル： ユーティリティは 2 つのコンポーネントで定義される：
    • 学習ポテンシャル： ロールアウト報酬の分散によって測定される。高い分散は、成功と失敗を含む多様な結果のセットを示しており、実行可能な最適化信号を示唆する。
    • タスク間シナジー： 現在のタスクと他のタスク間の勾配コサイン類似度によって測定される。正の類似度は建設的な整合性を示し、負の類似度は競合を示唆する。
  • プロンプトレベル： タスク内では、ユーティリティは現在の学習ポテンシャル（報酬分散）と歴史的進捗（平均報酬の改善）によって決定される。
• 階層的割り当て：
  • タスク割当枠の配分： グローバルな訓練バッチサイズは、集約されたタスクユーティリティに対する温度スケーリング付きソフトマックスを用いてタスク間で分割される。これにより、高いポテンシャルとシナジーを持つタスクがより多くの訓練データを受け取ることを保証する。
  • プロンプトの優先順位付け： 各タスクの割り当てられた枠内で、プロンプトはプロンプトレベルのユーティリティに基づいた重み付きサンプリングによって選択され、モデルが積極的に改善しているサンプルが優先される。

2. 適応的ユーティリティ較正方策最適化

このモジュールは、各タスクの現在のユーティリティに基づいて KL 正則化係数（$\beta_k$）を動的に較正することで、「どのように学ぶか」に対処する。

• メカニズム： タスク $k$ の KL 係数は、タスク固有のベース（$\beta_{base,k}$）とタスクユーティリティに条件付けられた動的乗数に分解される：
  $$\beta_k(t) = \beta_{base,k} \cdot (1 + \lambda \cdot U(t)(T_k))$$
• 適応ロジック：
  • 高いユーティリティ（強い学習信号）を持つタスクは、より大きな $\beta_k$ を受け取り、勾配信号が集中している際のオーバーシュートを防ぐために KL 制約を厳格化する。
  • 低いユーティリティを持つタスクは、より小さな係数を受け取り、方策に改善の余地がある場合に、より攻撃的な更新を可能にする。
• 報酬システム： このフレームワークは、ソフトウェア開発ワークフロー全体にわたる 4 つの異質なタスクをサポートする：コード I/O 予測（実行の正しさ）、コード生成（通過率）、単体テスト生成（カバレッジ）、コミットメッセージ生成（テキスト類似度）。それぞれが、検証可能で実行ベースの報酬信号を利用する。

主要な貢献

1. コード用 LLM における MTRL の最初の研究： この論文は、コードに関するマルチタスク RL を「何を学ぶか」と「どのように学ぶか」という 2 つの結合された意思決定を中心に定式化し、両方を支配する統合信号としてタスクユーティリティを導入する。
2. ASTOR フレームワーク： 以下の機能を備えた協調フレームワーク：
   • 階層的ユーティリティルーティングデータスケジューリング： 学習ポテンシャルとシナジーに基づいてタスク間で予算を割り当て、プロンプトに優先順位を付ける。
   • 適応的ユーティリティ較正方策最適化： 異質なタスクの訓練状態に一致するように、タスクごとの KL 正則化を動的に調整する。
3. 実証的検証： 単一の ASTOR モデルが、タスク固有の専門モデルおよび既存の MTRL ベースラインの両方を凌駕することを示す広範な実験。

実験結果

著者らは、ASTOR を 2 つのベースモデル（Qwen2.5-Coder-7B および Qwen3-8B）で、4 つの代表的なコーディングタスクにわたって評価した。

• 単一タスク学習（STL）との比較： ASTOR は強力なマルチタスク性能を達成し、タスク固有の専門モデルと同等かそれ以上の性能を発揮する。Qwen2.5-Coder-7B において、全体の平均性能で最良の単一タスクモデルを**9.0%上回る。Qwen3-8B においては、改善幅は9.5%**である。
• MTRL ベースラインとの比較： ASTOR は、Joint Learning（均一混合）、Curriculum Learning（固定された逐次順序）、Model Merging（事後の専門モデル統合）の 3 つのベースラインを大幅に上回る。
  • 最も強力な MTRL ベースライン（Joint Learning）に対する改善幅は、モデルおよび指標によって**7.5% から 19.2%**の範囲にある。
  • 具体的には、Qwen2.5-Coder-7B において、ASTOR は Joint Learning を**12.8%**上回る。
• アブレーション研究： データスケジューリングモジュール（タスク割り当てまたはプロンプト優先順位付け）のいずれか、または方策最適化モジュール（動的またはタスク固有の正則化）のいずれかを除去すると、性能が著しく低下する。これは、異質なタスクを調整するために両方のコンポーネントが不可欠であることを確認する。

意義と主張

この論文は、ASTOR がマルチタスクコード RL におけるユーティリティ駆動型協調の有効性を検証することを主張する。タスクを均一に扱わず、学習ポテンシャルとシナジーに基づいてリソースを割り当て、タスク固有のニーズに基づいて更新を制約することで、単一のモデルが複数の専門モデルを展開する際の線形スケーリングコストなしに多様な能力を達成できる。

著者らは、この研究を、単一のモデル内でソフトウェア開発ワークフロー全体（理解、実装、テスト、ドキュメント作成）を処理でき、展開の障壁を低減しつつ、すべての段階における性能を維持または向上させる、効率的で統合された AI 支援コーディングシステムへの一歩として位置づけている。この研究は、タスク・言語のカバレッジおよびベース KL 係数の手動調整に関する限界を認めており、これらを将来の課題として提案している。