Distribution-Aware Reward Estimation for Test-Time Reinforcement Learning

TL;DR

We address the fragility of majority-vote rewards in test-time reinforcement learning by showing MV discards information and becomes biased under correlated rollouts. We introduce Distribution-Aware Reward Estimation (DARE), which leverages an uncertainty-aware empirical rollout distribution $\hat{p}(\hat{y})$, augments it with an exploration bonus $b(y)$, and applies distribution pruning with threshold $\tau$ to stabilize learning. The approach yields a distribution-based proxy reward $R_{\text{dist}}$ and improved policy updates via GRPO, achieving substantial gains on challenging reasoning benchmarks, including a $25.3\%$ relative improvement on AIME 2024 and a $5.3\%$ improvement on AMC. DARE also enhances out-of-distribution generalization and training efficiency, demonstrating the value of preserving rollout diversity and uncertainty in reward signals for test-time adaptation. Examples of key formulas include $R_{\text{dist}}(\hat{y}) = g(\hat{p}(\hat{y}))$, $r(y_i) = r_{\text{dis}}(y_i) + \alpha b(y_i)$, and $\tilde{p}(y_i) = \frac{\hat{p}(y_i) \mathbf{1}[\hat{p}(y_i) \ge \tau]}{\sum_k \hat{p}(y_k) \mathbf{1}[\hat{p}(y_k) \ge \tau]}$. These design choices lead to robust, information-rich rewards for test-time learning.

Abstract

Test-time reinforcement learning (TTRL) enables large language models (LLMs) to self-improve on unlabeled inputs, but its effectiveness critically depends on how reward signals are estimated without ground-truth supervision. Most existing TTRL methods rely on majority voting (MV) over rollouts to produce deterministic rewards, implicitly assuming that the majority rollout provides a reliable learning signal. We show that this assumption is fragile: MV reduces the rollout distribution into a single outcome, discarding information about non-majority but correct actions candidates, and yields systematically biased reward estimates. To address this, we propose Distribution-AwareReward Estimation (DARE), which shifts reward estimation from a single majority outcome to the full empirical rollout distribution. DARE further augments this distribution-based reward with an exploration bonus and a distribution pruning mechanism for non-majority rollout exploration and reward denoise, yielding a more informative and robust reward estimation. Extensive experiments on challenging reasoning benchmarks show that DARE improves optimization stability and final performance over recent baselines, achieving relative improvements of 25.3% on challenging AIME 2024 and 5.3% on AMC.

Figures (6)

• Figure 1: (a): MV vs. Our Method. The distribution-based reward with an exploration bonus encourages the model to explore low-uncertainty rollouts and mitigates the confirmation bias of MV. (b): Distribution Pruning denoise distribution information, reduces reward variance and stabilizes optimization.
• Figure 2: Overview of the proposed DARE framework. Given a test query, multiple rollouts are sampled from the policy model. Rollout-level probabilities are computed based on empirical frequency and uncertainty, followed by exploration bonus and distribution pruning to calculate the final reward used to update policy.
• Figure 3: OOD generalization of Qwen2.5-Math-1.5B. Each subfigure shows evaluation on OOD benchmarks after adaptation on a training set. Bars indicate pass@1 accuracy for the original model, TTRL, and DARE, with DARE consistently improving performance.
• Figure 4: OOD generalization of Qwen3-1.7B. Each subfigure shows evaluation on OOD benchmarks after adaptation on a training set. Bars indicate pass@1 accuracy for the original model, TTRL, and DARE, with DARE consistently improving performance.
• Figure 5: Impact of rollout correlation on Qwen3-1.7B. The x-axis represents Rollout Correlation, and the y-axis shows Pass@1. TTRL performance drops sharply with increasing correlation, while DARE degrades smoothly, demonstrating robustness.

Theorems & Definitions (3)

• Theorem 2.1: Information Collapse under Majority Voting
• Theorem 2.2: Latent-Conditioned Bias of MV
• Proposition 2.3: Marginal Consistency under Exchangeable Rollouts