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

TL;DR

This paper introduces Distribution-Aware Reward, a reinforcement learning method that trains large language models to produce well-calibrated predictive distributions for regression tasks, enhancing uncertainty estimation and ranking.

Contribution

It proposes a novel RL objective that optimizes predictive distributions directly, improving calibration and robustness over traditional pointwise methods.

Findings

Improves rank-correlation metrics, including a 6-point Spearman gain on KBSS.

Achieves competitive molecular property prediction using only SMILES strings.

Mitigates rollout diversity collapse and enhances uncertainty diagnostics.

Abstract

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.