Research map / Large language models / Reasoning and chain of thought
Optimizing Language Models for Inference Time Objectives using Reinforcement Learning
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 RL gradient estimators that directly optimize k-sample inference-time objectives (pass@k, majority voting) using a leave-one-out control variate and a biased variant, showing gains on MATH and CodeContests.
Problem
Standard RL fine-tuning optimizes mean single-sample reward, ignoring that models are often deployed with multi-sample inference procedures such as pass@k or majority voting.
Approach
Formulates general k-sample objectives f(y_1..y_k) and derives an unbiased REINFORCE-style gradient with a leave-one-out control variate; further introduces a biased variant that subtracts the mean advantage, trading bias for lower variance and yielding sparser, more KL-efficient updates.
Key results
- On CodeContests test with Llama 3.1 70B Instruct, PPO with biased pass@k reaches pass@100 of 56.3 vs 41.1 for PPO mean (TACO easy pass@100 53.8 vs 39.5).
- Policy gradient with pass@k objective (k=8) improves CodeContests test pass@10 to 41.7 vs 35.1 for the mean policy gradient, while mean gives the best pass@1 (26.2 vs 20.6).
- Majority-voting variants show only slight gains over the policy gradient baseline on MATH; the 70B model shows little trade-off because MATH is near saturated (~90% pass@4).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Train with pass@k objective (k=8) instead of mean objective | Policy Gradient, Llama 3.1 70B Instruct · CodeContests valid | pass@100: 38.1 -> 54.9 | Helped |
| Train with pass@k objective (k=8) instead of mean objective | Policy Gradient, Llama 3.1 70B Instruct · CodeContests valid | pass@1: 24.6 -> 21.7 | Hurt |
| Train with biased pass@k objective (k=8) instead of mean objective | PPO, Llama 3.1 70B Instruct · CodeContests test | pass@100: 41.1 -> 56.3 | Helped |
| Train with pass@k objective (k=8) instead of mean objective | PPO, Llama 3.1 70B Instruct · CodeContests test | pass@1: 24.9 -> 21.2 | Hurt |
| Train with pass@k objective (k=8) instead of mean objective | Policy Gradient, Llama 3.1 70B Instruct · CodeContests valid | pass@10: 32.9 -> 42.2 | Helped |
| Increase the number of gradient samples from k=2 to k=8 (ablation) | Pass@k gradient estimates on MATH · MATH training set | pass@k | Helped |
| Increase the number of gradient samples k (ablation) | Policy Gradient pass@k estimates on MATH · MATH training set | pass@k | No clear effect |
| Use k-sample pass@k gradient estimates instead of mean policy gradient (ablation) | Llama 3 3B and 8B · MATH training set | pass@k | Helped |
| Use k-sample gradient estimates instead of mean policy gradient (ablation) | Llama 3 3B and 8B · MATH training set | mean performance | Hurt |
| Use k-sample gradient estimates instead of baseline (ablation) | Pass@k training on HARP, Llama 3 70B · HARP | pass@4 | Helped |
| Use majority-voting k-sample policy gradient instead of baseline (ablation) | Llama 3 8B · MATH training set | majority voting at k=4 | Helped |
| Use unbiased k-sample policy gradient instead of policy gradient baseline (ablation) | Llama 3 8B · MATH test set | majority voting | Hurt |
Surprising findings
- Mean policy gradient often retains an advantage on pass@k at evaluation on MATH despite lower training pass@k, which the authors attribute to a training-evaluation generalization gap.
- On the 70B MATH setting the mean policy gradient remains competitive even on k-sample objectives, since the task is nearly saturated.
Limitations
- Applying the method naively to multi-step inference algorithms (self-correction, reflection) requires careful credit assignment and large compute.
- Trade-off: optimizing pass@k sacrifices pass@1 performance.
- Majority-voting gains are modest and evaluation results are less consistent with training results.
Experimental setup
Datasets: MATH HARP CodeContests TACO
Benchmarks: MATH CodeContests TACO
Metrics: pass@1 pass@k majority voting accuracy pass@10 pass@100
Baselines: mean policy gradient PPO mean GRPO Dr. GRPO
Builds on: REINFORCE PPO GRPO leave-one-out control variate
Closest papers by meaning
- T1: Advancing Language Model Reasoning through Reinforcement Learning and Inference Scaling (ICML 2025 · similarity 0.90)
- TTRL: Test-Time Reinforcement Learning (NeurIPS 2025 · similarity 0.90)
- Compute Aligned Training: Optimizing for Test Time Inference (NeurIPS 2026 · similarity 0.90)
- Accelerating RL for LLM Reasoning with Optimal Advantage Regression (NeurIPS 2025 · similarity 0.90)
- Is Best-of-N the Best of Them? Coverage, Scaling, and Optimality in Inference-Time Alignment (ICML 2025 · similarity 0.89)
- On the Limits of Test-Time Compute: Sequential Reward Filtering for Better Inference (ICML 2026 · similarity 0.89)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.