# Reinforcement Learning over Predictive Distributions for LLM Regression

arXiv: https://arxiv.org/abs/2605.20740 · HTML: https://mlsanity.com/paper/2605.20740/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 41 of its 57 references that are in the map.
- First posted on arXiv 20 May 2026.

## Builds on (most cited first, 30 of 41)

- [Training language models to follow instructions with human feedback](https://mlsanity.com/paper/2203.02155/) (2022)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Chain-of-Thought Prompting Elicits Reasoning in Large Language Models](https://mlsanity.com/paper/2201.11903/) (2022)
- [DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models](https://mlsanity.com/paper/2402.03300/) (2024)
- [Qwen3 Technical Report](https://mlsanity.com/paper/2505.09388/) (2025)
- [Simple and Scalable Predictive Uncertainty Estimation using Deep Ensembles](https://mlsanity.com/paper/1612.01474/) (2016)
- [What Uncertainties Do We Need in Bayesian Deep Learning for Computer Vision?](https://mlsanity.com/paper/1703.04977/) (2017)
- [Qwen2 Technical Report](https://mlsanity.com/paper/2407.10671/) (2024)
- [HybridFlow: A Flexible and Efficient RLHF Framework](https://mlsanity.com/paper/2409.19256/) (2024)
- [MoleculeNet: A Benchmark for Molecular Machine Learning](https://mlsanity.com/paper/1703.00564/) (2017)
- [The Entropy Mechanism of Reinforcement Learning for Reasoning Language Models](https://mlsanity.com/paper/2505.22617/) (2025)
- [Accurate Uncertainties for Deep Learning Using Calibrated Regression](https://mlsanity.com/paper/1807.00263/) (2018)
- [Understanding the Effects of RLHF on LLM Generalisation and Diversity](https://mlsanity.com/paper/2310.06452/) (2023)
- [Analyzing Learned Molecular Representations for Property Prediction](https://mlsanity.com/paper/1904.01561/) (2019)
- [Mol-Instructions: A Large-Scale Biomolecular Instruction Dataset for Large Language Models](https://mlsanity.com/paper/2306.08018/) (2023)
- [LlaSMol: Advancing Large Language Models for Chemistry with a Large-Scale, Comprehensive, High-Quality Instruction Tuning Dataset](https://mlsanity.com/paper/2402.09391/) (2024)
- [Beyond Binary Rewards: Training LMs to Reason About Their Uncertainty](https://mlsanity.com/paper/2507.16806/) (2025)
- [ChemLLM: A Chemical Large Language Model](https://mlsanity.com/paper/2402.06852/) (2024)
- [Evaluating and Calibrating Uncertainty Prediction in Regression Tasks](https://mlsanity.com/paper/1905.11659/) (2019)
- [Rewarding the Unlikely: Lifting GRPO Beyond Distribution Sharpening](https://mlsanity.com/paper/2506.02355/) (2025)
- [From Words to Numbers: Your Large Language Model Is Secretly A Capable Regressor When Given In-Context Examples](https://mlsanity.com/paper/2404.07544/) (2024)
- [LLM Processes: Numerical Predictive Distributions Conditioned on Natural Language](https://mlsanity.com/paper/2405.12856/) (2024)
- [Rewarding Doubt: A Reinforcement Learning Approach to Calibrated Confidence Expression of Large Language Models](https://mlsanity.com/paper/2503.02623/) (2025)
- [OmniPred: Language Models as Universal Regressors](https://mlsanity.com/paper/2402.14547/) (2024)
- [Restoring Calibration for Aligned Large Language Models: A Calibration-Aware Fine-Tuning Approach](https://mlsanity.com/paper/2505.01997/) (2025)
- [Outcome-based Reinforcement Learning to Predict the Future](https://mlsanity.com/paper/2505.17989/) (2025)
- [Scaling Open-Ended Reasoning to Predict the Future](https://mlsanity.com/paper/2512.25070/) (2025)
- [Understanding LLM Embeddings for Regression](https://mlsanity.com/paper/2411.14708/) (2024)
- [TRACT: Regression-Aware Fine-tuning Meets Chain-of-Thought Reasoning for LLM-as-a-Judge](https://mlsanity.com/paper/2503.04381/) (2025)
- [Regression-aware Inference with LLMs](https://mlsanity.com/paper/2403.04182/) (2024)

## Built on by (most cited first, 1 of 1)

- [Sample What You Say: Aligning Language Models to Sample the Distributions They State](https://mlsanity.com/paper/2609.34929/) (2026)

## Abstract (arXiv)

Large language models (LLMs) have emerged as flexible regressors capable of predicting real-valued quantities from heterogeneous inputs. Yet most LLM regression objectives optimize predictions independently, often yielding poor calibration. We introduce Distribution-Aware Reward (DAR), an on-policy reinforcement learning objective that instead jointly evaluates the empirical predictive distribution formed by multiple predictions for the same input. To translate this distribution-level objective into rollout-level rewards, we assign each prediction credit based on its leave-one-out contribution to the quality of the overall predictive distribution. This encourages predictions that are well-centered and appropriately dispersed around the target. We evaluate on three regression settings: a synthetic task probing interpolation and extrapolation, and two real-world scientific tasks involving code and molecular data. Across tasks, DAR produces better-calibrated uncertainty estimates while consistently reducing prediction error and improving ranking quality over supervised fine-tuning and pointwise reinforcement learning. Together, these results highlight the benefits of distribution-aware training for LLM regression.