🤖 machine learning

Distribution-Aware Reward: Reinforcement Learning over Predictive Distributions for LLM Regression

本論文は、連続順位確率スコアを用いて複数のデコードされたサンプルを評価することで、回帰タスクにおける校正された予測分布を生成するように大規模言語モデルを最適化する強化学習目的である分布認識報酬を導入し、従来の点推定トレーニング手法と比較して不確実性の推定、ランキング性能、および頑健性を向上させる。

原著者： Jungsoo Park, Hyungjoo Chae, Ethan Mendes, Jay DeYoung, Varsha Kishore, Wei Xu, Alan Ritter

公開日 2026-05-21

📖 1 分で読めます☕ さくっと読める

CC BY 4.0

原著者： Jungsoo Park, Hyungjoo Chae, Ethan Mendes, Jay DeYoung, Varsha Kishore, Wei Xu, Alan Ritter

原論文は CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) でライセンスされています。 ✨ これは以下の論文のAI生成解説です。著者が執筆または承認したものではありません。技術的な正確性については原論文を参照してください。免責事項の全文を読む

「分布認識型報酬：LLM 回帰のための予測分布上における強化学習」という論文を、平易な言葉と創造的な比喩を用いて解説します。

大きな問題：「一匹狼」対「チーム」

あなたが外気温の正確な値を推測しようとしている場面を想像してください。あなたは天気の説明を見て数値を提示できる、非常に賢い AI アシスタント（大規模言語モデル、または LLM）を持っています。

従来の方法（ポイント別報酬）：
現在、ほとんどの AI 訓練は、単一の生徒の答えを厳格に採点する教師のような仕組みで動いています。

AI が「72°F」と推測します。教師は確認します。「72 は実際の 70 に近いか？はい。よくやった。」
AI が「75°F」と推測します。教師は確認します。「75 は 70 に近いか？いいえ。よくない。」
欠点： AI は「一匹狼」になることを学びます。毎回的に的を外さないようにしようとします。しかし、間違えることを恐れるあまり、年間平均気温（例えば 60°F）を毎回推測し始めるかもしれません。それは安全ですが、多様性を失います。AI は「どの程度確信があるか」を伝えないのです。実際の気温が 70 で、AI が毎回 60 と言う場合、平均的には「近い」と言えますが、決して変動しないため、ひどい予測者です。

新しい方法（分布認識型報酬）：
この論文は、分布認識型報酬（DAR）と呼ばれる新しい手法を導入します。教師が AI に単一の推測を評価する代わりに、AI に12 種類の異なる推測を一度に行うよう求めます。

AI は言います。「68、69、70、71、72、73、74、75、76、77、78、あるいは 79 になる可能性があると思います。」
教師は単一の数値だけを見るのではなく、推測の全体を見ます。
目標： 教師は、そのグループが実際の答え（70）の中心に位置しつつ、AI が異なる可能性を考慮していることを示すために、十分に広がっていることを望みます。

中核となるアイデア：「留め置き除外」スコア

教師はどのようにしてどの推測が良いかを判断するのでしょうか？彼らは留め置き除外クレジット割り当てと呼ばれる巧妙なトリックを使用します。

AI の 12 個の推測を、的を狙う 12 人の弓使いのチームだと想像してください。

従来の方法： 教師は一人の弓使いを選び、その矢が的の中心に命中したかどうかを見てスコアを与えます。もしある弓使いが少し外れたが、チームがより広い範囲をカバーするのを助けたとしても、彼らは罰せられます。
新しい方法（DAR）： 教師は、「もしこの特定の弓使いをチームから取り除いたら、チーム全体の性能が悪化するか？」と問います。
- 弓使いを取り除くことでチームの広がりが狭すぎたり偏ったりするようになる場合、その弓使いは的の中心に最も近い矢を放っていなくても、高い報酬を得ます。
- 弓使いを取り除いても何も変わらない場合（つまり、別の弓使いがすぐ隣に立っている場合）、その弓使いは冗長であるため低い報酬を得ます。

これにより、AI は多様でありながら正確であることを学びます。「70 付近だと確信しているが、少し低いか高い可能性もある」と言うようになり、単一の数値を叫ぶだけにはなりません。

検証された場所

研究者たちは、この新しい「チーム採点」方法を 3 つの異なる種類の問題でテストしました。

合成数学（トレーニングジム）： 答えが複雑な波状のパターンで変化する架空の数学問題を作成しました。
- 結果： 従来の方法（単一推測）は混乱し、波を平坦化してしまいました。新しい方法（チーム推測）は、未見の領域であっても、波状のパターンを完璧に追従しました。
コードパフォーマンス（プログラマー）： AI にコンピュータコードの実行速度やメモリ使用量を推測させるよう求めました。
- 結果： 新しい方法はランキングにおいて非常に優れていました。100 個のコードがあり、どれが最も遅いかを知りたい場合、新しい方法は従来の方法よりもはるかに正確にそれらをソートできました。平均的な速度を推測するだけでなく、高速なコードと低速なコードの違いを理解していました。
分子特性（化学者）： AI に化学式（テキスト文字列として記述）を与え、水への溶解度などの特性を予測させました。
- 結果： AI はテキスト（3 次元化学モデルではない）しか見ていませんでしたが、専門的な化学コンピュータと同等かそれ以上の性能を発揮しました。科学者にとって重要なのは、可能な値の範囲を予測することであり、その点で優れていました。

なぜこれが重要なのか

この論文は、AI が単一の推測の精度だけでなく、推測の**形状（分布）**を気にするように訓練することで、AI は以下のように変化すると主張しています。

ランキング能力の向上： どの項目が「より高い」または「より低い」可能性が高いかを判断できます。
不確実性の理解： 無茶な推測をしているときと、確信を持っているときを区別できます。
堅牢性の向上： 毎回同じ退屈な答えを出すように崩壊しません。

結論

従来の方法は、テストの正確な答えを暗記するように生徒を訓練するものだと考えてください。新しい方法（DAR）は、生徒が概念を十分に理解し、答えが変動する可能性がある理由を説明しながら、妥当な答えの範囲を与えるように訓練します。これにより、AI は科学や工学にとって、数値そのものと同じくらい「誤差範囲」を知ることが重要な、はるかに信頼性の高いツールとなります。

技術的サマリー：LLM 回帰のための分布認識型報酬

問題定義

大規模言語モデル（LLM）は、テキスト、コード、分子文字列などの異種入力から実数値を予測する能力を実証している。しかし、LLM 回帰の標準的な訓練目的（教師あり微調整（SFT）やポイント単位の報酬（例：負の平均二乗誤差）を用いた強化学習（RL）など）は、各復号化された浮動小数点数を独立した予測として扱う。

このアプローチは、個々の点推定値を最適化するが、サンプリングされた予測の集合が較正された予測分布を形成することを保証しない。その結果、モデルは個々のターゲットに近い予測を生成する一方で、集合的には較正が不十分、過度に狭小、または片寄ったものになることが多い。この「平均指向」の振る舞いは予測分布を崩壊させ、必要な不確実性や相対的な順序付けを捉え損なう。この限界は、候補のランキング、不確実性の推定、不確実性下での堅牢な意思決定を必要とする応用を阻害する。これらの応用では、単一の最良推定値と同様に、可能な値の分布が重要である。

手法：分布認識型報酬（DAR）

著者らは、個々のロールアウトを孤立して最適化するのではなく、複数の復号化サンプルによって形成される予測分布全体の質を最適化するように設計されたオンポリシー強化学習目的である**分布認識型報酬（Distribution-Aware Reward: DAR）**を提案する。

中核メカニズム

経験的予測分布: 入力 $x$ に対して、モデルは $K$ 個のロールアウトをサンプリングし、予測の集合 $\{p_{ik}\}_{k=1}^K$ を生成する。この集合は経験的予測分布 $F_i$ として扱われる。
CRPS によるスコアリング: この分布の質は、連続ランク確率スコア（CRPS）を用いて、真のターゲット $y_i$ $y_{i}$ に対して評価される。CRPS は、正確性（ターゲットの近くに集中すること）と適切な分散性（退化を避けること）の両方を備えた分布に報酬を与える適切なスコアリング則である。
- 完全な集合に対する報酬は $S_i = -\text{CRPS}(F_i, y_i)$ と定義される。
留め置きのクレジット割り当て: RL は方策を更新するために個々のロールアウトごとのスカラー報酬を必要とするため、DAR は分布レベルのスコアをロールアウトレベルのフィードバックに変換する。各ロールアウト $p_{ik}$ $p_{ik}$ の限界貢献度を、完全な集合の CRPS（ $S_i$ $S_{i}$ ）と、そのロールアウトを除去した集合の CRPS（ $S_i^{(-k)}$ $S_{i}^{(- k)}$ ）を比較することで計算する。
- ロールアウト報酬は $R_{ik}^{\text{DAR}} = S_i - S_i^{(-k)}$ である。
- あるロールアウトを除去することで分布が悪化（負の CRPS が低下）する場合、そのロールアウトは個々のポイント単位の誤差が高くても正の報酬を受ける。逆に、分布の広がりや中心化を改善しない冗長なロールアウトは、低い、あるいは負の報酬を受ける。

理論的区別

本論文は、これを理論的に「平均指向」であり、予測分布を単一点（条件付き平均）に崩壊させることを促すポイント単位の二乗誤差（MSE）と比較する。対照的に、CRPS は「分布指向」であり、完全な条件付き分布を最適化する。DAR はこれを利用し、留め置きのクレジット割り当てを用いて、適切に中心化され、有益に分散したロールアウトの集合を促す。

主要な貢献

新規 RL 目的: 個々の点誤差の最小化から、経験的予測分布の質の最大化へと最適化目標をシフトさせる DAR の導入。
留め置きのクレジット割り当て: 集合的な分布の質に対する個々のロールアウトの限界貢献度に基づいて報酬を割り当てる特定のメカニズム。これにより、個々のサンプルの精度が低くても「有益な分散」に対して報酬を与えることが可能になる。
多様性崩壊の緩和: RL 微調整における出力多様性の低下（方策エントロピーの減少）と予測の平均への崩壊の傾向に対処し、サンプリングされたロールアウトが真の値をカバーするのに十分な多様性を維持することを保証する。

実験結果

著者らは、合成ガウス混合回帰タスク、コード性能予測、分子特性予測の 3 つの異なるタスクにおいて DAR を評価した。

1. 合成分布回帰

タスク: 異方性の 2 成分ガウス混合からのターゲット予測。
知見: DAR は LLM ベースの手法の中で最も高いスピアマン相関を達成し、SFT やポイント単位の MSE 報酬を大幅に上回った。これはターゲット関数のグローバルな順序付けをよりよく維持し、他のニューラルベースラインが苦労する外挿領域にも適応した。

2. コード性能予測

タスク: ソースコードからのカーネルレイテンシ（KBSS）とピークメモリ使用量（APPS）の予測。
知見: DAR は、SFT やポイント単位の RL ベースラインと比較して、最も強いランク相関（スピアマン $\rho$ $ρ$ ）を達成し、競合する、あるいはより低い平均絶対誤差（MAE）を示した。
- KBSS において、DAR は最強のベースラインに対してスピアマン相関を6 ポイント改善した。
- APPS において、スピアマン相関を4 ポイント改善した。
- これは、入力依存の変動に対する優れた感度と、より信頼性の高い候補ランキングを示している。

3. 分子特性予測（MoleculeNet）

タスク: SMILES 文字列からの物理化学的特性（FreeSolv、ESOL、親油性）の予測。
知見: 強力な GNN ベースラインが使用する明示的なグラフや 3D 的帰納バイアスを持たない SMILES 文字列のみを使用しているにもかかわらず、DAR は競争力があった。
- FreeSolv と親油性において、他の LLM 訓練戦略（SFT、MSE、Value Head）と比較して、RMSE に統計的に有意な改善を達成した。
- 専門的な分子表現がなくても、分布レベルの最適化が堅牢な回帰行動をもたらすことを示した。

不確実性の診断

較正: 分析により、DAR においてロールアウトの標準偏差と予測誤差の間に強い対数ピアソン相関（0.91）が確認された。これはモデルの不確実性推定が実際の誤差に対してよく較正されていることを示している。
多様性: 訓練中、DAR は MSE ベースの RL と比較して、より高いロールアウトエントロピーとより高い「ブケット率」（真の値が最小および最大ロールアウト予測の間にある事例の割合）を維持し、多様性崩壊の緩和を確認した。

意義と主張

本論文は、予測分布を直接最適化することで、LLM ベースの回帰をより堅牢で、よりよく較正されたものにするとしている。複数の復号化サンプルを分布として扱い、その集合的な質に対して報酬を与えることで、DAR は点誤差を犠牲にすることなく、ランクベースの性能（候補選択に不可欠）と不確実性の推定を向上させる。

著者らは、この研究を、LLM を単なるテキスト生成機や点推定機としてではなく、科学分野向けの柔軟な確率的回帰器として使用するための一歩と位置づけている。LLM 回帰は異種入力（テキスト、コード、文字列）の処理において柔軟性を提供すると強調しつつ、不確実性の定量化と信頼性の高い順序付けを必要とするタスクにおいて、従来のポイント単位の訓練が失敗する領域でこれらのモデルの潜在能力を最大限に引き出すためには、提案された分布認識型アプローチが必要であると述べている。この研究は、すべての分野で専門的な回帰モデル（ガウス過程や GNN など）を置き換えることを主張するものではないが、分布認識型 RL が汎用 LLM を科学的回帰の文脈において競争力のあるものにし得ることを実証している。