Multi-Task GRPO: Reliable LLM Reasoning Across Tasks

あなたは、数学の問題から論理パズルまで、あらゆる種類のパズルを解くことができる超スマートなロボットアシスタントを訓練していると想像してください。あなたはこのロボットが、難しい問題に直面しても混乱したり失敗したりすることなく、あらゆることに対応できる真のジェネラリスト（汎用的な能力を持つ者）になることを望んでいます。人工知能の世界では、これを「推論（reasoning）」と呼び、ロボットは「大規模言語モデル（LLM）」と呼ばれます。これらのモデルをより良くするために、科学者たちは「強化学習（Reinforcement Learning）」という手法を用います。これは、ロボットが正解したときにはハイタッチ（報酬）を与え、間違えたときには「もう一度やってみて」と優しく促すようなものです。この手法の一つとして、GRPOと呼ばれる非常に巧妙な方法があります。これは、ロボットがそれぞれの試行からどれだけ学ぶべきかを正確に判断するための方法です。しかし、ここには落とし穴があります。もし、一度に10種類の異なるパズルを学ばせようとすると、ロボットはしばしば行き詰まってしまいます。簡単なパズルには非常に詳しくなる一方で、難しいパズルを完全に無視してしまったり、あるいは特定の種類の謎解きに固執するあまり、他のスキルの解き方を忘れてしまったりすることがあるのです。大きな疑問は、一つの脳をあらゆる分野で等しく優秀にし、一つのスキルさえも取り残されることなく、どうすれば全てのスキルを向上させることができるのか、ということです。

これこそが、Multi-Task GRPO (MT-GRPO) と呼ばれる新しい手法が取り組んでいる課題です。標準的なトレーニングプロセスを、学習速度が大きく異なる生徒たちのクラスを教える教師だと考えてみてください。もし教師が単に「平均的な」生徒だけに焦点を当ててしまうと、賢い子供たちは退屈し、苦戦している子供たちはさらに取り残されてしまうかもしれません。標準的なアプローチでは、簡単なタスクがトレーニングを「乗っ取って」しまい、平均的なスコアは高く見えるものの、実は裏で難しい課題には密かに失敗している、という状況を許してしまいます。著者たちは、単にスコアを平均化するだけでは信頼できるロボットは作れない、つまり、クラスの中で最も成績が悪い生徒も向上していることを確認する必要があるのだと気づきました。

これを解決するために、彼らは二部構成の戦略を考案しました。第一に、彼らは動的な「重み付けシステム」を作成しました。これは、常にスコアボードを見守るコーチのようなものです。もしロボットが「カウントダウン（数学ゲーム）」では素晴らしい成績を収めているのに、「ゼブラパズル（論理パズル）」ではひどい成績である場合、コーチは即座に焦点を切り替えます。コーチは、ゼブラパズルの練習問題を増やし、数学のカテゴリーの練習問題を減らし始めます。しかし、ここがトリッキーな点です。時として、ロボットがゼブラパズルに挑戦しても、適切に解こうとさえしなかったためにスコアが「ゼロ」になり、その試行から何も学べないことがあります。もしコーチが単に問題の数を数えるだけなら、ゼブラパズルの指導を行っているつもりでも、実際には数学の課題からしか学習していないことになります。なぜなら、ゼブラパズルの試行は「ゼロ・グラディエント（学習が空の状態）」だからです。

この問題を解決するために、彼らの発明の第二の部分である「比率保持サンプラー（Ratio-Preserving Sampler）」があります。これは、非常に厳格な司書のようなものです。たとえロボットが無用な練習の試行を捨てたとしても、最終的な「有用な」練習問題の山が、依然として正しいミックス（組み合わせ）を保っていることを保証します。もしコーチが、練習の30%を論理パズルにしたいと考えた場合、司書は、不要なものを取り除いた後でも、残った「良い」問題の30%が依然として論理パズルであることを確実にします。これにより、難しいタスクにおける「失敗した試行」が多いという理由だけで、簡単なタスクが誤ってトレーニングセッションを支配してしまうことを防ぎます。

この新しい手法の結果は非常に有望です。チームが3つの異なる推論タスクを組み合わせてテストした際、この新しいMT-GRPO法は、標準的な手法と比較して、最も成績が悪かったタスクの精度を16〜28%向上させ、さらにDAPOと呼ばれる別の高度な手法と比較して6%向上させました。さらに驚くべきことに、最も難しいタスクにおいて、競合他社よりも50%少ないトレーニングステップで、精度**50%**の閾値に到達しました。9つの異なるタスクを用いたより大規模なテストにおいても、この手法は他を圧倒し続け、より複雑で多様な課題にも対応できるスケールアップが可能であることを示しました。著者たちは、「つまみ（$\lambda$ と呼ばれるパラメータ）」を調整することで、モデルが「平均的に良くなること」よりも「最も弱いスキルを修正すること」にどれだけ優先順位を置くかを制御できることを見出しました。

要約すると、この論文は、ロボットがどの問題を練習するかをより賢く選択し、そしてそれらの問題が実際にトレーニングセッションに到達することを確実にすることで、AIを単に簡単なことだけでなく、あらゆる分野で確実に有能にすることができると示唆しています。これは、AIが特定の狭い分野に特化するのではなく、自身のスキルを一つも取り残すことなく、多種多様な現実世界の課題を真に推論できるようにするための一歩です。