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

• URL: http://arxiv.org/abs/2605.20740v1
• Date: Wed, 20 May 2026 05:43:40 GMT
• Title: Distribution-Aware Reward: Reinforcement Learning over Predictive Distributions for LLM Regression
• Abstract summary: Large language models can predict real-valued quantities from heterogeneous inputs such as text, code, and molecular strings.<n>Most training objectives score each decoded floating-point number independently, improving point estimates without ensuring calibrated predictive distributions.<n>We introduce Distribution-Aware Reward, an on-policy reinforcement learning objective whose main contribution is to train language models to produce better predictive distributions.
• Score: 31.507053461689523
• License: http://creativecommons.org/licenses/by/4.0/
• 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.

Related papers

• Unsupervised Confidence Calibration for Reasoning LLMs from a Single Generation [2.526814143603023]
  Reasoning language models can solve increasingly complex tasks, but struggle to produce the calibrated confidence estimates necessary for reliable deployment.<n>We introduce a method for unsupervised confidence calibration of reasoning LLMs when only a single generation is available at inference time.<n>Our approach uses offline sampling on unlabeled data to derive a self-consistency-based proxy target, then distills this signal into a lightweight deployment-time confidence predictor.
  arXiv  Detail & Related papers  (2026-04-21T13:25:25Z)
• Evidence-based Distributional Alignment for Large Language Models [58.65469623911573]
  We propose Evi-DA, an evidence-based alignment technique that improves the fidelity and robustness of LLM-based distribution estimation.<n>Given a target country, Evi-DA retrieves related World Values Survey items and their answer distributions, predicts a coarse Welzel value signature for each option, and infers the country-conditioned answer distribution in a structured format.
  arXiv  Detail & Related papers  (2026-03-03T03:34:06Z)
• Inference-Time Scaling of Diffusion Language Models with Particle Gibbs Sampling [70.8832906871441]
  We study how to steer generation toward desired rewards without retraining the models.<n>Prior methods typically resample or filter within a single denoising trajectory, optimizing rewards step-by-step without trajectory-level refinement.<n>We introduce particle Gibbs sampling for diffusion language models (PG-DLM), a novel inference-time algorithm enabling trajectory-level refinement while preserving generation perplexity.
  arXiv  Detail & Related papers  (2025-07-11T08:00:47Z)
• Quantile Regression with Large Language Models for Price Prediction [15.277244542405345]
  Large Language Models (LLMs) have shown promise in structured prediction tasks, including regression.<n>We propose a novel quantile regression approach that enables LLMs to produce full predictive distributions.<n>A Mistral-7B model fine-tuned with quantile heads significantly outperforms traditional approaches for both point and distributional estimations.
  arXiv  Detail & Related papers  (2025-06-07T04:19:28Z)
• Self-Boost via Optimal Retraining: An Analysis via Approximate Message Passing [58.52119063742121]
  Retraining a model using its own predictions together with the original, potentially noisy labels is a well-known strategy for improving the model performance.<n>This paper addresses the question of how to optimally combine the model's predictions and the provided labels.<n>Our main contribution is the derivation of the Bayes optimal aggregator function to combine the current model's predictions and the given labels.
  arXiv  Detail & Related papers  (2025-05-21T07:16:44Z)
• Towards Robust and Interpretable EMG-based Hand Gesture Recognition using Deep Metric Meta Learning [37.21211404608413]
  We propose a shift to deep metric-based meta-learning in EMG PR to supervise the creation of meaningful and interpretable representations. We derive a robust class proximity-based confidence estimator that leads to a better rejection of incorrect decisions.
  arXiv  Detail & Related papers  (2024-04-17T23:37:50Z)
• Distributionally Robust Post-hoc Classifiers under Prior Shifts [31.237674771958165]
  We investigate the problem of training models that are robust to shifts caused by changes in the distribution of class-priors or group-priors. We present an extremely lightweight post-hoc approach that performs scaling adjustments to predictions from a pre-trained model.
  arXiv  Detail & Related papers  (2023-09-16T00:54:57Z)
• Distributionally Robust Learning for Multi-source Unsupervised Domain Adaptation [9.359714425373616]
  Empirical risk often performs poorly when the distribution of the target domain differs from those of source domains.<n>We develop an unsupervised domain adaptation approach that leverages labeled data from multiple source domains and unlabeled data from the target domain.
  arXiv  Detail & Related papers  (2023-09-05T13:19:40Z)
• Value-Distributional Model-Based Reinforcement Learning [59.758009422067]
  Quantifying uncertainty about a policy's long-term performance is important to solve sequential decision-making tasks. We study the problem from a model-based Bayesian reinforcement learning perspective. We propose Epistemic Quantile-Regression (EQR), a model-based algorithm that learns a value distribution function.
  arXiv  Detail & Related papers  (2023-08-12T14:59:19Z)
• Improving Adaptive Conformal Prediction Using Self-Supervised Learning [72.2614468437919]
  We train an auxiliary model with a self-supervised pretext task on top of an existing predictive model and use the self-supervised error as an additional feature to estimate nonconformity scores. We empirically demonstrate the benefit of the additional information using both synthetic and real data on the efficiency (width), deficit, and excess of conformal prediction intervals.
  arXiv  Detail & Related papers  (2023-02-23T18:57:14Z)
• Leveraging Unlabeled Data to Predict Out-of-Distribution Performance [63.740181251997306]
  Real-world machine learning deployments are characterized by mismatches between the source (training) and target (test) distributions. In this work, we investigate methods for predicting the target domain accuracy using only labeled source data and unlabeled target data. We propose Average Thresholded Confidence (ATC), a practical method that learns a threshold on the model's confidence, predicting accuracy as the fraction of unlabeled examples.
  arXiv  Detail & Related papers  (2022-01-11T23:01:12Z)
• Estimating and Evaluating Regression Predictive Uncertainty in Deep Object Detectors [9.273998041238224]
  We show that training variance networks with negative log likelihood (NLL) can lead to high entropy predictive distributions. We propose to use the energy score as a non-local proper scoring rule and find that when used for training, the energy score leads to better calibrated and lower entropy predictive distributions.
  arXiv  Detail & Related papers  (2021-01-13T12:53:54Z)
• Unlabelled Data Improves Bayesian Uncertainty Calibration under Covariate Shift [100.52588638477862]
  We develop an approximate Bayesian inference scheme based on posterior regularisation. We demonstrate the utility of our method in the context of transferring prognostic models of prostate cancer across globally diverse populations.
  arXiv  Detail & Related papers  (2020-06-26T13:50:19Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.