Technical Summary: Distribution-Aware Reward for LLM Regression

Problem Statement

Large Language Models (LLMs) have demonstrated the ability to predict real-valued quantities from heterogeneous inputs (text, code, molecular strings). However, standard training objectives for LLM regression—such as Supervised Fine-Tuning (SFT) or Reinforcement Learning (RL) with pointwise rewards (e.g., negative Mean Squared Error)—treat each decoded floating-point number as an independent prediction.

This approach optimizes individual point estimates without ensuring that the set of sampled predictions forms a calibrated predictive distribution. Consequently, models often produce predictions that are individually close to the target but collectively poorly calibrated, overly narrow, or one-sided. This "mean-seeking" behavior collapses the predictive distribution, failing to capture necessary uncertainty or relative ordering. This limitation hinders applications requiring candidate ranking, uncertainty estimation, or robust decision-making under uncertainty, where the distribution of possible values matters as much as the single best estimate.

Methodology: Distribution-Aware Reward (DAR)

The authors propose Distribution-Aware Reward (DAR), an on-policy reinforcement learning objective designed to optimize the quality of the entire predictive distribution formed by multiple decoded samples, rather than optimizing individual rollouts in isolation.

Core Mechanism

1. Empirical Predictive Distribution: For a given input $x$, the model samples $K$ rollouts, producing a set of predictions $\{p_{ik}\}_{k=1}^K$. This set is treated as an empirical predictive distribution $F_i$.
2. Scoring via CRPS: The quality of this distribution is evaluated against the ground-truth target $y_i$ using the Continuous Ranked Probability Score (CRPS). CRPS is a proper scoring rule that rewards distributions that are both accurate (concentrated near the target) and appropriately dispersed (avoiding degeneracy).
   • The reward for the full set is defined as $S_i = -\text{CRPS}(F_i, y_i)$.
3. Leave-One-Out Credit Assignment: Since RL requires a scalar reward for each individual rollout to update the policy, DAR converts the distribution-level score into rollout-level feedback. It calculates the marginal contribution of each rollout $p_{ik}$ by comparing the CRPS of the full set ($S_i$) against the CRPS of the set with that rollout removed ($S_i^{(-k)}$).
   • The rollout reward is: $R_{ik}^{\text{DAR}} = S_i - S_i^{(-k)}$.
   • If removing a rollout degrades the distribution (lowers the negated CRPS), that rollout receives a positive reward, even if its individual pointwise error is high. Conversely, redundant rollouts that do not improve the distribution's spread or centering receive lower or negative rewards.

Theoretical Distinction

The paper contrasts this with pointwise squared loss (MSE), which is theoretically "mean-seeking" and encourages the predictive distribution to collapse to a single point (the conditional mean). In contrast, CRPS is "distribution-seeking," optimizing for the full conditional distribution. DAR leverages this by using leave-one-out credit assignment to encourage a set of rollouts that are well-centered and usefully dispersed.

Key Contributions

1. Novel RL Objective: Introduction of DAR, which shifts the optimization goal from minimizing individual point errors to maximizing the quality of the empirical predictive distribution.
2. Leave-One-Out Credit Assignment: A specific mechanism to attribute rewards to individual rollouts based on their marginal contribution to the collective distribution quality, allowing the model to reward "useful spread" even when individual samples are less accurate.
3. Mitigation of Diversity Collapse: The method addresses the tendency of RL fine-tuning to reduce output diversity (policy entropy) and collapse predictions toward the mean, ensuring that sampled rollouts remain diverse enough to cover the ground truth.

Experimental Results

The authors evaluated DAR on three distinct tasks: a synthetic Gaussian-mixture regression task, code performance prediction, and molecular property prediction.

1. Synthetic Distributional Regression

• Task: Predicting targets from a heteroscedastic two-component Gaussian mixture.
• Findings: DAR achieved the highest Spearman correlation among LLM-based methods, significantly outperforming SFT and pointwise MSE rewards. It better preserved the global ordering of the target function and adapted well to extrapolation regions where other neural baselines struggled.

2. Code Performance Prediction

• Task: Predicting kernel latency (KBSS) and peak memory usage (APPS) from source code.
• Findings: DAR achieved the strongest rank correlation (Spearman $\rho$) and competitive or lower Mean Absolute Error (MAE) compared to SFT and pointwise RL baselines.
  • On KBSS, DAR improved Spearman correlation by 6 points over the strongest baseline.
  • On APPS, it improved Spearman correlation by 4 points.
  • This indicates superior sensitivity to input-dependent variation and more reliable candidate ranking.

3. Molecular Property Prediction (MoleculeNet)

• Task: Predicting physicochemical properties (FreeSolv, ESOL, Lipophilicity) from SMILES strings.
• Findings: Despite using only SMILES strings (lacking explicit graph or 3D inductive biases used by strong GNN baselines), DAR remained competitive.
  • It achieved statistically significant improvements in RMSE on FreeSolv and Lipophilicity compared to other LLM training strategies (SFT, MSE, Value Head).
  • It demonstrated that distribution-level optimization can yield robust regression behavior even without specialized molecular representations.

Uncertainty Diagnostics

• Calibration: Analysis showed a strong log-Pearson correlation (0.91) between rollout standard deviation and prediction error for DAR, indicating that the model's uncertainty estimates were well-calibrated to actual errors.
• Diversity: During training, DAR maintained higher rollout entropy and a higher "bracket rate" (the fraction of examples where the ground truth lies between the minimum and maximum rollout predictions) compared to MSE-based RL, confirming the mitigation of diversity collapse.

Significance and Claims

The paper claims that directly optimizing predictive distributions makes LLM-based regression more robust and better calibrated. By treating multiple decoded samples as a distribution and rewarding their collective quality, DAR improves rank-based performance (crucial for candidate selection) and uncertainty estimation without sacrificing pointwise accuracy.

The authors position this work as a step toward using LLMs not just as text generators or point estimators, but as flexible, probabilistic regressors for scientific domains. They emphasize that while LLM regression offers flexibility in handling heterogeneous inputs (text, code, strings), the proposed distribution-aware approach is necessary to unlock the full potential of these models for tasks requiring uncertainty quantification and reliable ordering, areas where traditional pointwise training fails. The work does not claim to replace specialized regression models (like GPs or GNNs) in all domains but demonstrates that distribution-aware RL can make general-purpose LLMs competitive in scientific regression settings.