Research map / Large language models / Reasoning and chain of thought

Compute Aligned Training: Optimizing for Test Time Inference

NeurIPS 2026 · Large language models › Reasoning and chain of thought · approach: Test-time compute scaling · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes Compute Aligned Training (CAT), which treats test-time strategies (Pass@N, majority vote, best-of-N) as operators on the base policy and reweights SFT/RL gradients by a diagonal scaling factor so training matches test-time inference.

Problem

Standard SFT and RL optimize single-sample likelihood or expected reward, which misaligns with test-time strategies that aggregate or filter multiple outputs (Pass@N, majority vote, best-of-N), degrading test-time scaling.

Approach

Model inference strategies as differentiable operators T on the base policy and minimize the negative log-likelihood under the induced test-time distribution; under a diagonal approximation this reduces to re-weighting standard CE or RL gradients by a closed-form marginal-utility factor w(p, N) (e.g., N p(1-p)^{N-1}/(1-(1-p)^N) for Pass@N SFT).

Key results

• On MATH with Mistral-7B, the Pass@64-trained SFT model reaches 67.6% Pass@64 vs 59.8% for the SFT baseline, at a cost of 2.28 points in Pass@1.
• RL with Pass@16 reward-weighting reaches 40.0% Pass@32 vs 35.8% for standard RL on MATH.
• In protein best-of-N RL (complementary hydrophobicity), standard RL reaches 8.89 expected max reward at N=32 while BoN-2 reaches 9.96.

Surprising findings

• Pass@1 accuracy drops for strategy-aligned models (a 'search tax'), while Pass@k at larger budgets improves.
• Majority-vote SFT gradient weight approaches k ≈ N/2 for hard samples, implying variance inflation that grows with N.

Limitations

• Diagonal approximation ignores off-diagonal gradient terms and may underestimate learning signal.
• Majority vote SFT suffers variance inflation for hard examples as N grows.
• Experiments are limited to relatively simple strategies on a single math dataset plus a small protein task.

Experimental setup

Datasets: MATH

Benchmarks: MATH

Metrics: Pass@k Maj@k expected max reward

Baselines: SFT (cross-entropy) Standard RL (GRPO) Direct Coverage Optimization

Builds on: Direct Coverage Optimization GRPO

Closest papers by meaning

• Is Best-of-N the Best of Them? Coverage, Scaling, and Optimality in Inference-Time Alignment (ICML 2025 · similarity 0.91)
• On the Limits of Test-Time Compute: Sequential Reward Filtering for Better Inference (ICML 2026 · similarity 0.90)
• Optimizing Test-Time Compute via Meta Reinforcement Finetuning (ICML 2025 · similarity 0.90)
• Optimizing Language Models for Inference Time Objectives using Reinforcement Learning (ICML 2025 · similarity 0.90)
• Inference Scaling Laws: An Empirical Analysis of Compute-Optimal Inference for LLM Problem-Solving (ICLR 2025 · similarity 0.90)
• TTRL: Test-Time Reinforcement Learning (NeurIPS 2025 · similarity 0.89)