aiqbench analysis
This research introduces Distribution-Aware Reward (DAR), a reinforcement learning framework that improves Large Language Model regression by jointly optimizing predictive distributions rather than isolated point estimates. By assigning credit based on leave-one-out contributions, DAR achieves better-calibrated uncertainty and lower prediction errors across synthetic and scientific tasks like code generation and molecular data analysis. The study demonstrates that distribution-aware training significantly outperforms standard supervised fine-tuning and pointwise reinforcement learning methods for real-valued prediction tasks.
Reinforcement Learning over Predictive Distributions for LLM Regression
arXiv:2605.20740v2 Announce Type: replace-cross Abstract: 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.
Abstract: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.
Machine Learning (cs.LG); Artificial Intelligence (cs.AI); Computation and Language (cs.CL)
arXiv:2605.20740 [cs.LG]
(or arXiv:2605.20740v2 [cs.LG] for this version)
https://doi.org/10.48550/arXiv.2605.20740 Focus to learn more arXiv-issued DOI via DataCite
arXivLabs is a framework that allows collaborators to develop and share new arXiv features directly on our website.
Both individuals and organizations that work with arXivLabs have embraced and accepted our values of openness, community, excellence, and user data privacy. arXiv is committed to these values and only works with partners that adhere to them.
Have an idea for a project that will add value for arXiv's community? Learn more about arXivLabs.
Full text under the publisher's open licence · 250 words
Related coverage
Exploding variance of means of exponentials: least-squares to the rescue – Machine Learning Research Blog
francisbach.com · Sep 27, 2026
WCM: A World Critic Model for Vision-Language-Action Reinforcement Learning
arXiv · Aug 3, 2026
Reinforcement Learning Techniques for the Optimization of Target Polarization in Nuclear Physics Scattering Experiments
arXiv · Oct 6, 2026
The Alignment Problem: Computer Science Meets Analytic Philosophy
LessWrong · Oct 18, 2026