ZIP-RC: Inference-Time Reward & Cost Prediction

• ZIP-RC is a framework that provides zero-overhead, introspective prediction of both expected reward and remaining compute cost during inference.
• It employs reserved output tokens to generate a joint probability distribution over terminal reward and predicted sequence length in a single forward pass.
• Empirical results on reasoning datasets demonstrate ZIP-RC’s ability to dynamically trade off accuracy, compute, and latency, outperforming traditional sampling methods.

ZIP-RC refers to Zero-overhead Inference-time Prediction of Reward and Cost, an adaptive inference framework equipping pre-trained LLMs with real-time introspection on both their expected terminal reward (e.g., task success) and the computation cost required to achieve it. The approach utilizes reserved output tokens to predict a joint distribution over future reward and sample length in the same forward pass as next-token generation, incurring no additional inference-time overhead. ZIP-RC is designed for efficient, meta-cognitively capable generation, allowing adaptive search strategies that dynamically trade off accuracy, compute, and latency as information unfolds during decoding (Manvi et al., 1 Dec 2025).

1. Motivation and Conceptual Framework

State-of-the-art LLMs conventionally lack self-assessment mechanisms during inference: they do not anticipate their own likelihood of success nor estimate the remaining effort required. Standard scaling methods, such as Best-of-N sampling, allocate compute and latency in a fixed manner, producing redundant samples even on easy tasks, while learned reward models or verifiers add substantial deployment cost and inference latency by requiring separate model invocations to estimate correctness or utility. Such deficiencies hinder models from making rational meta-cognitive decisions about when to branch, prune, or stop generation.

ZIP-RC addresses these limitations by embedding introspective reward-cost prediction into the base model itself, using reserved output logits to provide a categorical joint distribution over candidate reward and length outcomes, all computed in a single forward pass. This facilitates dynamic sampling policies driven by real-time estimations of utility.

2. Mathematical Formulation of Joint Reward–Cost Prediction

At each decoding step $t$, the LLM produces logits $z_t \in \mathbb{R}^{|\mathcal{V}|}$, partitioning the vocabulary $\mathcal{V}$ into generation tokens $\mathcal{V} \setminus \mathcal{R}$ and a reserved subset $\mathcal{R}$ supporting introspective prediction. For current prefix $s_t$, ZIP-RC defines two key random variables:

• $Z^{\pi}(s_t)$: Expected terminal reward (e.g., predicted correctness score at sequence end)
• $L^{\pi}(s_t)$: Predicted remaining length until termination (e.g., number of tokens to EOS)

These variables are discretized into $B_V$ and $B_T$ bins, yielding reserved tokens $r_{b,\ell} \in \mathcal{R}$ for bin pairs. The resulting auxiliary logits induce a categorical joint distribution

$$p_\theta(b, \ell\,|\,s_t) = \frac{\exp\left(z_t[r_{b,\ell}]\right)} {\sum_{b',\ell'}\exp\left(z_t[r_{b',\ell'}]\right)}$$

Marginal distributions, $q_\theta^V$ and $q_\theta^L$, and point estimates for expected reward and length are computed by summing over respective bins and averaging:

$$\mathbb{E}[Z^\pi(s_t)] \approx \sum_b \frac{v_b+v_{b+1}}{2}\,q_\theta^V(b|s_t)$$

$$\mathbb{E}[L^\pi(s_t)] \approx \sum_\ell \frac{t_\ell+t_{\ell+1}}{2}\,q_\theta^L(\ell|s_t)$$

This "joint introspection" allows the model to provide actionable reward-cost estimates for each sampled prefix with zero inference overhead.

3. Sampling Utility and Inference-time Decision Process

ZIP-RC operationalizes meta-cognitive search as a meta-MDP over candidate sets of prefixes $A_t$ at each step. Utility is formalized by a lower-bound "sampling utility" function

$$Q^{\text{Rollouts}}(S_t, A_t; \mathcal{H}_t^*) = \mathbb{E}\left[ \max_{s \in A_t} Z^\pi(s) \right] - \beta \left(\alpha \sum_{s} \mathbb{E}[L^\pi(s)] + (1-\alpha)\mathbb{E}[\max_{s \in A_t} L^\pi(s)] \right)$$

where $\alpha \in [0, 1]$ configures the trade-off between total compute and maximal latency, and $\beta > 0$ sets the cost scaling.

ZIP-RC sampling proceeds iteratively:

1. At each step, reserved logits predict the reward-cost joint for each candidate prefix.
2. Candidate multisets $A$ (up to $N_\text{max}$) are scored via the utility function.
3. The meta-action $A_t$ maximizing utility is selected, allocations are updated, and generation continues.
4. Prefixes with low predicted reward or excessive length are pruned dynamically.

Upon completion, the sampled sequence with highest predicted reward is selected as output.

4. Empirical Evaluation and Pareto Optimization

On diverse mathematical reasoning datasets (AIME 2024, AMC 2023, MATH-500, GSM8K), ZIP-RC delivers strict Pareto improvements over traditional baselines, including majority voting and reward-pruned Best-of-N. When tracing Pareto frontiers by sweeping cost parameters, ZIP-RC demonstrates superior accuracy at equal or lower compute/latency:

• On the Mixed set with Qwen3-1.7B, ZIP-RC achieves 92.2% accuracy vs. 91.0% for majority voting at GenCost $\approx$ 1.4 (normalized over three single-sample runs).
• On hardest subsets (AIME 2024), ZIP-RC boosts accuracy from 53.1% (majority vote) to 65.8%, a 12.7% absolute improvement at matched cost.
• Adaptive sample allocation: ZIP-RC branches more aggressively on hard prompts, minimizes resource allocation on easy ones, and avoids over-generation in compute-constrained regimes.

The utility function enables continuous interpolation between pure latency and pure compute-bound modes, unattainable via static heuristic scaling methods.

5. Interpretability and Practical Implications

ZIP-RC's introspective predictions expose real-time estimates of both "success likelihood" and "time-to-completion" at every token. Users and systems gain transparent signals enabling rational allocation of computation and branching decisions—contrasting with black-box majority voting or non-interpretable verifiers. Branch selection and pruning trace back directly to model-predicted reward and cost, facilitating post-hoc analysis and automated decision-making.

A plausible implication is that zero-overhead joint prediction could enable model-based resource arbitration in multi-agent or tool-augmented systems, transfer to other domains beyond mathematical reasoning, and serve as a foundation for enhanced diversity-promoting architectures.

6. Limitations and Future Developments

ZIP-RC's effectiveness depends on sample diversity—if candidate prefixes collapse to similar trajectories, improvement over static methods is limited. Current meta-action selection restricts to shared horizon and root multiplicity for tractability; extending to richer meta-action classes may provide further accuracy/computation trade-off gains. Prospective directions include cross-model resource allocation, generalization to non-mathematical tasks, and integration with external diversity-enhancing decoders.

This systematic framework marks a shift in reasoning models from heuristic, fixed-budget decoding toward principled, interpretable, and utility-maximized generation, enabled by zero-overhead introspective prediction of both reward and cost (Manvi et al., 1 Dec 2025).