Re-FORC: Adaptive Reward Prediction for Efficient Chain-of-Thought Reasoning

AI summary

Overview

Research area: Efficient inference for large language model reasoning — specifically metareasoning, reward forecasting, compute allocation, and test-time scaling.

Technical level: Intermediate. The core idea (predict how much a model will improve if it keeps thinking) is intuitive, but the paper formalizes it through Gittins index policies, Pandora's box decision theory, Markov chain selection, and Beta-distribution likelihood training.

Scope: The paper introduces Re-FORC, a lightweight adapter trained on frozen Qwen3 reasoning models (1.7B, 4B, 8B) that predicts the expected reward of generating a given number of additional thinking tokens, and uses those forecasts to make cost-aware stopping, model-selection, and test-time scaling decisions.

What This Paper Is About

Reasoning models can spend more inference-time computation to improve answers, but there is no principled way to decide how much computation a given query deserves. The paper's core claim is that current models lack the ability to predict the marginal benefit of additional reasoning tokens, so stopping, backtracking, and model-switching decisions are made heuristically. Re-FORC fills that gap by forecasting the expected reward as a function of how many future thinking tokens are generated, then using those forecasts to maximize a net-utility objective of the form J = E[R_best − λ T_total], where λ encodes the user's cost sensitivity.

Key Contributions

1. A forecasting functional for reasoning. Re-FORC learns ψ(t | x, z, π), the expected reward after generating exactly t additional thinking tokens from query x and partial trace z using model π. It is trained as a lightweight adapter on frozen base models, predicting Beta(α, β) reward distributions across a uniform token grid.

2. A Gittins-index, Pandora's-box decision rule. Building on a formulation of inference-time search as a Pandora's box problem, the paper derives a greedy policy that assigns each model-trace continuation a reservation value and advances the highest-index option or terminates when no continuation beats the current best reward.

3. Three compute-aware applications from one forecaster. Early stopping (Eq. 11), joint model and thinking-length selection (Eq. 13), and adaptive test-time scaling at both trace level (Eq. 14) and step level (Eq. 15).

4. User-controlled inference without retraining. A single λ parameter set at inference time controls the accuracy-compute trade-off, replacing token-count specifications, and the forecaster also yields upfront estimates of expected computation time.

Main Findings

• Early stopping saves compute while preserving accuracy. Compared with the S1 fixed-budget cutoff baseline, Re-FORC saves 26% compute on average for the Qwen3 4B reasoning model while maintaining accuracy; relative to the DeepConf confidence-based baseline, the saving is 21% for the same model. The abstract phrases this as reducing compute by up to 26%.

• Maximum accuracy advantage over S1 at equal compute is modest but consistent. +3.8 percentage points at 1.5k tokens for the 1.7B model, +4.2 pp at 2k tokens for the 4B, and +3.9 pp at 4k tokens for the 8B — larger models benefit most at higher token budgets.

• Model selection beats the largest single model. Across the three Re-FORC variants (Smallest, Highest Forecasted, and Pandora), peak accuracy is 1.7 percentage points above All-8B, while needing less compute to match the largest model's accuracy (12% in the main text; the Figure 3 caption states 10%).

• Adaptive test-time scaling improves over confidence-based baselines. At peak compute, Re-FORC outperforms the strongest DeepConf variant by +13.9 pp (+29.3% relative) on 1.7B, +8.4 pp (+13.9%) on 4B, and +7.5 pp (+12.3%) on 8B; the abstract reports this as a 9.9 percentage point average gain at maximum compute.

• Forecast quality improves as reasoning progresses. In Figure 1, Pearson correlation rises and MSE/MAE fall with reasoning progress across Qwen3 models, evaluated on MATH500, AMC 2024, AIME 2024, and AIME 2025.

• Re-FORC dominates baselines across token budgets. In Figure 2, Re-FORC outperforms S1, DeepConf, and unconstrained generation at nearly every token budget for all three model sizes, with the largest gap in the moderate-compute regime.

• Cost sensitivity changes where compute goes. In Figure 5, at high λ = 4.0 × 10⁻⁴ the models allocate a smaller proportion of compute to harder problems than at smaller λ values or under unconstrained generation; for the 8B model, on the easiest ~40% of problems, higher λ reduces the compute share spent on those problems.

• Step-level search has a low-compute cost. Re-FORC Pandora can beat Re-FORC-scaling at ≥50k tokens, but the simpler Re-FORC-scaling is preferable at ≤25k tokens because maintaining multiple partial traces incurs overhead.

• The method is training-agnostic. Re-FORC is independent of the base reasoning model's training procedure, unlike training-based length-control methods such as L1 and e1.

Methodology in Plain English

The researchers treat reasoning as a sequential decision problem: at each point the system can generate more thinking tokens, stop and answer, or switch to a different model. The right decision depends on the expected gain from continuing, which the paper argues is not directly available in current models.

To supply it, they attach a small adapter to a frozen reasoning model. The adapter reads the model's penultimate-layer hidden states for the query plus the partial trace, and outputs Beta distribution parameters for a set of token budgets on a uniform grid from 0 to 8192 in steps of 512. Values between grid points come from linear interpolation, and a softplus activation keeps the Beta parameters positive. Training maximizes the likelihood of observed rewards under the predicted Beta distributions.

Training data comes from problem instances sampled from DeepScaleR-Preview, with full unconstrained reasoning trajectories generated up to a maximum thinking length of 8192 tokens. Instead of regenerating continuations for every (partial trace, horizon) pair — which would cost O(|T| × N × L) trajectory samples — they reuse segments of existing trajectories, reducing the requirement to O(N) samples. Monte Carlo estimation uses N = 8 samples, except N = 4 for the 1.7B model on AIME 2025, with rewards clipped to (ε, 1 − ε) for ε = 10⁻⁶.

Decision-making follows the Gittins index policy from Pandora's box theory for the binary-reward case. The optimal compute budget minimizes λt / ψ(t | x, z, π) (Eq. 4), and the resulting Gittins index for a model at a state is 1 − λt*/ψ(t*) (Eq. 5). The agent continues reasoning only when this index exceeds the best reward found so far (Eq. 11); with multiple models it advances whichever model-trace pair has the highest index (Eq. 12) and stops when none beats the best reward.

Evaluation uses AMC 2024, MATH500, AIME 2024, and AIME 2025, with a maximum of 8192 thinking tokens, token costs measured in 512-token reasoning chunks, 32 samples per problem (4 for MATH500), and, for the test-time scaling figure, averaging over up to 100 randomly sampled size-k subsets per problem.

Why This Matters

The paper argues that no canonical λ exists, because the value of an agent depends on the user and the environment — a fruit-edibility checker that answers after the fruit spoils is useless, while a botanist categorizing species is unaffected by delay. Re-FORC makes that cost sensitivity an explicit, adjustable input rather than a fixed property of a trained model, which matters for any deployment where latency and accuracy trade off differently across queries and users.

Real-world applications:

• Latency-sensitive agents. Early stopping halts unpromising reasoning chains, with the paper reporting up to 26% compute reduction over fixed-budget cutoffs while maintaining accuracy.
• Tiered model fleets. Joint model and thinking-length selection chooses among 1.7B, 4B, and 8B options per query, reaching 1.7 pp higher peak accuracy than the largest model alone and matching its accuracy with less compute.
• User-configurable serving. A single λ set at inference time lets an operator shift along the accuracy-compute frontier without retraining, replacing manual token-count limits.
• Capacity planning for latency-sensitive services. The forecaster supplies upfront estimates of expected computation time, which the paper lists as a user-experience benefit.

Industry relevance: The work targets production inference economics directly — measured in thinking tokens and (for model selection) TFLOPs — and is designed to attach to existing models as a lightweight adapter and to work with external black-box systems, which lowers the barrier to adoption on top of already-deployed reasoning models.

Future Directions

• Reward assumptions. The closed-form Gittins index and the early-stopping rule rely on binary rewards. Extending the forecaster and decision rule to richer, non-binary reward structures is a natural open question, since the paper notes that computing Gittins indices generally requires knowledge of the reward distribution.

• Low-compute regimes. Step-level Pandora's interleaving across partial traces is less efficient than simple Re-FORC-scaling at ≤25k tokens, so reducing that overhead would broaden where multi-trajectory search helps.

• Connecting to training-based length control. The paper contrasts Re-FORC with training-based approaches such as L1 and e1 because it leaves the base model untouched; whether combining forecast-driven inference with trained length control yields further gains is not addressed.

• Independent forecasters. The paper notes concurrent work (Manvi et al., 2026) that also trains a forecaster for efficient adaptive reasoning. Comparing forecasters, and the question of whether a single adapter generalizes across base model families beyond the Qwen3 family tested here, remain open.

Target Audience

Researchers and engineers working on inference efficiency for large reasoning models: people who serve chain-of-thought models under cost or latency budgets, practitioners building model-routing or multi-model serving systems, and researchers interested in metareasoning, decision-theoretic compute allocation, or reward forecasting for LLMs. Readers with a background in reinforcement learning, bandit/Gittins index methods, or test-time scaling will get the most from the theory sections, while the application sections in Figures 2 through 4 are accessible to anyone who deploys reasoning models.

Authors’ abstract

We propose Re-FORC, an adaptive reward prediction method that, given a query, enables prediction of the expected future rewards as a function of the number of future thinking tokens. Re-FORC trains a lightweight adapter on reasoning models, demonstrating improved prediction with longer reasoning and larger models. Re-FORC enables: 1) early stopping of unpromising reasoning chains, reducing compute by up to 26\% compared to fixed-budget cutoffs, while maintaining accuracy, 2) optimized model and thinking length selection that outperforms the largest model alone--- reaching 1.7 percentage points higher peak accuracy while needing up to 12\% less compute to match the largest model's accuracy, 3) adaptive test-time scaling, which increases accuracy by 9.9 percentage points (on average at maximum compute) over confidence-based baselines. Re-FORC allows dynamic reasoning with length control via cost-per-token thresholds while estimating computation time upfront.