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

📄 arXiv: 2605.20740v1 📥 PDF

作者: Jungsoo Park, Hyungjoo Chae, Ethan Mendes, Jay DeYoung, Varsha Kishore, Wei Xu, Alan Ritter

分类: cs.LG, cs.AI, cs.CL

发布日期: 2026-05-20

备注: 21 pages, 5 figures

💡 一句话要点

提出Distribution-Aware Reward，用于提升LLM回归任务中预测分布的质量。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 强化学习 语言模型 回归任务 预测分布 不确定性估计

📋 核心要点

  1. 现有语言模型回归方法侧重于点估计优化，忽略了预测分布的校准，限制了其在需要不确定性估计场景的应用。
  2. Distribution-Aware Reward通过强化学习直接优化预测分布，鼓励模型生成既准确又具有适当分散性的预测。
  3. 实验表明，该方法在多个任务上优于传统方法，显著提升了秩相关性，并改善了不确定性诊断。

📝 摘要（中文）

大型语言模型可以预测来自文本、代码和分子字符串等异构输入的实值数量，但大多数训练目标独立地对每个解码的浮点数进行评分，从而改进点估计，而不能确保校准的预测分布。这限制了需要候选排序或不确定性估计的应用。我们引入了Distribution-Aware Reward，这是一种在线强化学习目标，其主要贡献是训练语言模型，使其为回归任务生成更好的预测分布，而不仅仅是针对标量目标优化单个解码输出。我们的方法将多个解码样本视为经验预测分布，使用连续排序概率分数对其进行评估，并根据每个rollout对分布质量的边际贡献分配leave-one-out信用，奖励既准确又适当分散的预测。我们在受控高斯混合任务、代码性能预测和来自SMILES字符串的分子性质预测中评估了我们的方法。在所有任务中，我们的方法都优于监督微调和逐点强化学习基线，具有很强的秩相关性增益，包括在KBSS上提高了6个点的Spearman相关系数。在MoleculeNet上，它仅使用SMILES字符串，但仍与强大的基于图和3D分子的模型竞争。进一步的分析表明，我们的方法减轻了rollout多样性崩溃并改善了不确定性诊断，表明直接优化预测分布使语言模型回归更稳健且更好地校准。

🔬 方法详解

问题定义：现有的大型语言模型在回归任务中，通常采用点估计的方式进行训练，即针对每个目标值单独优化预测结果。这种方法忽略了预测分布的整体质量，导致模型无法提供可靠的不确定性估计，限制了其在需要评估预测置信度的场景下的应用。现有方法的痛点在于缺乏对预测分布形状的有效建模和优化。

核心思路：论文的核心思路是通过强化学习，直接优化语言模型生成的预测分布。具体来说，将语言模型生成多个样本视为一个经验预测分布，并使用连续排序概率分数（CRPS）来评估该分布的质量。通过奖励那些能够提高分布质量的预测，鼓励模型生成更准确、更分散的预测结果。这样设计的目的是让模型不仅关注单个预测的准确性，更关注整体预测分布的合理性。

技术框架：整体框架包含以下几个主要步骤：1) 使用语言模型生成多个样本，构成经验预测分布；2) 使用连续排序概率分数（CRPS）评估该分布的质量；3) 基于每个样本对CRPS的贡献，计算强化学习的奖励；4) 使用强化学习算法（如PPO）更新语言模型的参数，使其能够生成更好的预测分布。关键模块包括语言模型、CRPS计算模块和强化学习优化器。

关键创新：最重要的创新点在于提出了Distribution-Aware Reward，即一种基于预测分布质量的强化学习奖励函数。与传统的逐点奖励函数不同，该奖励函数考虑了整个预测分布的形状，能够更有效地引导模型学习生成校准良好的预测分布。与现有方法的本质区别在于，它直接优化预测分布，而不是仅仅优化单个预测值。

关键设计：关键设计包括：1) 使用连续排序概率分数（CRPS）作为评估预测分布质量的指标；2) 使用leave-one-out的方式计算每个样本对CRPS的贡献，作为强化学习的奖励；3) 使用PPO算法进行强化学习优化。具体参数设置包括PPO的学习率、折扣因子、以及CRPS的计算方式等。论文中还提到使用SMILES字符串作为分子性质预测的输入。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，Distribution-Aware Reward在多个任务上显著优于监督微调和逐点强化学习基线。在KBSS任务上，Spearman相关系数提高了6个点。在MoleculeNet分子性质预测任务上，该方法仅使用SMILES字符串作为输入，就达到了与基于图和3D分子模型的竞争水平。此外，该方法还能够有效缓解rollout多样性崩溃，并改善不确定性诊断。

🎯 应用场景

该研究成果可应用于多种需要不确定性估计的场景，例如药物发现、材料设计、金融风险评估等。通过提升语言模型预测分布的质量，可以更准确地评估预测结果的置信度，从而辅助决策，降低风险。未来，该方法有望推广到其他类型的回归任务，并与其他不确定性量化技术相结合，进一步提升预测的可靠性。

📄 摘要（原文）

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.