arXiv:2605.20740cs.LGcs.AI2026-05

让大模型回归预测更准且有可信度,直接优化概率分布。

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

论文配图:Distribution-Aware Reward: Reinforcement Learning over Predictive Distributions for LLM Regression
图 1 · 摘自论文原文
  • 用强化学习优化预测分布而非单个数值
  • 在多个任务上提升排名相关性,最高增6点斯皮尔曼系数
  • 适合需要置信度和排序的回归场景

大语言模型可从文本、代码、分子字符串等异构输入中预测实数值,但现有训练目标独立评分每个解码浮点数,仅优化点估计而无法保证预测分布校准,限制了需候选排序或不确定性估计的应用。本文提出分布感知奖励(Distribution-Aware Reward),一种基于策略梯度的强化学习目标,核心是训练语言模型生成更优的回归预测分布,而非仅优化单个输出与标量目标的匹配。该方法将多个解码样本视为经验预测分布,使用连续排名概率评分(CRPS)评估,并基于留一法信用分配机制,根据每次回溯对整体分布质量的边际贡献进行奖励,激励模型输出既准确又适当分散的预测。我们在高斯混合任务、代码性能预测和基于SMILES字符串的分子属性预测上评估该方法。结果表明,在所有任务上均优于监督微调和逐点强化学习基线,显著提升排名相关性,其中在KBSS任务上斯皮尔曼相关性提高6点;在MoleculeNet数据集上仅使用SMILES字符串,仍可媲美强图结构和三维分子模型。进一步分析显示,该方法缓解了回溯多样性崩溃问题,改善了不确定性诊断能力,说明直接优化预测分布可使语言模型回归更鲁棒、更校准。

原文摘要 · Abstract (English)

Large language models can predict real-valued quantities from heterogeneous inputs such as text, code, and molecular strings, but most training objectives score each decoded floating-point number independently, improving point estimates without ensuring calibrated predictive distributions. This limits applications requiring candidate ranking or uncertainty estimation. We introduce Distribution-Aware Reward, an on-policy reinforcement learning objective whose main contribution is to train language models to produce better predictive distributions for regression tasks, rather than only optimizing individual decoded outputs against scalar targets. Our method treats multiple decoded samples as an empirical predictive distribution, evaluates it with the Continuous Ranked Probability Score, and assigns leave-one-out credit based on each rollout's marginal contribution to distribution quality, rewarding predictions that are both accurate and appropriately dispersed. We evaluate our method on a controlled Gaussian-mixture task, code performance prediction, and molecular property prediction from SMILES strings. Across tasks, our method improves over supervised fine-tuning and pointwise reinforcement learning baselines, with strong rank-correlation gains, including a 6-point Spearman improvement on KBSS. On MoleculeNet, it uses only SMILES strings yet remains competitive with strong graph-based and 3D molecular models. Further analyses show that our method mitigates rollout diversity collapse and improves uncertainty diagnostics, suggesting that directly optimizing predictive distributions makes language model regression more robust and better calibrated.

回归预测强化学习分布校准语言模型

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。