# Optimizing Language Models for Inference Time Objectives using Reinforcement Learning

arXiv: https://arxiv.org/abs/2503.19595 · HTML: https://mlsanity.com/paper/2503.19595/

## Key facts

- Cited by 53 papers in the map (data through 14 Oct 2026).
- Builds on 32 of its 50 references that are in the map.
- First posted on arXiv 25 Mar 2025.

## Builds on (most cited first, 30 of 32)

- [Training language models to follow instructions with human feedback](https://mlsanity.com/paper/2203.02155/) (2022)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models](https://mlsanity.com/paper/2402.03300/) (2024)
- [ReAct: Synergizing Reasoning and Acting in Language Models](https://mlsanity.com/paper/2210.03629/) (2022)
- [Self-Consistency Improves Chain of Thought Reasoning in Language Models](https://mlsanity.com/paper/2203.11171/) (2022)
- [Measuring Mathematical Problem Solving With the MATH Dataset](https://mlsanity.com/paper/2103.03874/) (2021)
- [Model-Agnostic Meta-Learning for Fast Adaptation of Deep Networks](https://mlsanity.com/paper/1703.03400/) (2017)
- [Deep reinforcement learning from human preferences](https://mlsanity.com/paper/1706.03741/) (2017)
- [Let's Verify Step by Step](https://mlsanity.com/paper/2305.20050/) (2023)
- [Constitutional AI: Harmlessness from AI Feedback](https://mlsanity.com/paper/2212.08073/) (2022)
- [Asynchronous Methods for Deep Reinforcement Learning](https://mlsanity.com/paper/1602.01783/) (2016)
- [Fine-Tuning Language Models from Human Preferences](https://mlsanity.com/paper/1909.08593/) (2019)
- [Self-RAG: Learning to Retrieve, Generate, and Critique through Self-Reflection](https://mlsanity.com/paper/2310.11511/) (2023)
- [Competition-Level Code Generation with AlphaCode](https://mlsanity.com/paper/2203.07814/) (2022)
- [Scaling Laws for Reward Model Overoptimization](https://mlsanity.com/paper/2210.10760/) (2022)
- [Large Language Models Cannot Self-Correct Reasoning Yet](https://mlsanity.com/paper/2310.01798/) (2023)
- [Solving math word problems with process- and outcome-based feedback](https://mlsanity.com/paper/2211.14275/) (2022)
- [Importance Weighted Autoencoders](https://mlsanity.com/paper/1509.00519/) (2015)
- [Is DPO Superior to PPO for LLM Alignment? A Comprehensive Study](https://mlsanity.com/paper/2404.10719/) (2024)
- [RLEF: Grounding Code LLMs in Execution Feedback with Reinforcement Learning](https://mlsanity.com/paper/2410.02089/) (2024)
- [On the Convergence Theory of Gradient-Based Model-Agnostic Meta-Learning Algorithms](https://mlsanity.com/paper/1908.10400/) (2019)
- [Inference-Aware Fine-Tuning for Best-of-N Sampling in Large Language Models](https://mlsanity.com/paper/2412.15287/) (2024)
- [Variational inference for Monte Carlo objectives](https://mlsanity.com/paper/1602.06725/) (2016)
- [Variational Best-of-N Alignment](https://mlsanity.com/paper/2407.06057/) (2024)
- [InfAlign: Inference-aware language model alignment](https://mlsanity.com/paper/2412.19792/) (2024)
- [Improving Policies via Search in Cooperative Partially Observable Games](https://mlsanity.com/paper/1912.02318/) (2019)
- [What Makes Large Language Models Reason in (Multi-Turn) Code Generation?](https://mlsanity.com/paper/2410.08105/) (2024)
- [HARP: A challenging human-annotated math reasoning benchmark](https://mlsanity.com/paper/2412.08819/) (2024)
- [Techniques for Learning Binary Stochastic Feedforward Neural Networks](https://mlsanity.com/paper/1406.2989/) (2014)

## Built on by (most cited first, 30 of 53)

- [Pass@k Training for Adaptively Balancing Exploration and Exploitation of Large Reasoning Models](https://mlsanity.com/paper/2508.10751/) (2025)
- [CWM: An Open-Weights LLM for Research on Code Generation with World Models](https://mlsanity.com/paper/2510.02387/) (2025)
- [Pass@K Policy Optimization: Solving Harder Reinforcement Learning Problems](https://mlsanity.com/paper/2505.15201/) (2025)
- [Outcome-based Exploration for LLM Reasoning](https://mlsanity.com/paper/2509.06941/) (2025)
- [Rewarding the Unlikely: Lifting GRPO Beyond Distribution Sharpening](https://mlsanity.com/paper/2506.02355/) (2025)
- [Self-Improving Language Models for Evolutionary Program Synthesis: A Case Study on ARC-AGI](https://mlsanity.com/paper/2507.14172/) (2025)
- [Maximum Likelihood Reinforcement Learning](https://mlsanity.com/paper/2602.02710/) (2026)
- [Differential Smoothing Mitigates Sharpening and Improves LLM Reasoning](https://mlsanity.com/paper/2511.19942/) (2025)
- [Representation-Based Exploration for Language Models: From Test-Time to Post-Training](https://mlsanity.com/paper/2510.11686/) (2025)
- [The Best of N Worlds: Aligning Reinforcement Learning with Best-of-N Sampling via max@k Optimisation](https://mlsanity.com/paper/2510.23393/) (2025)
- [Scalpel vs. Hammer: GRPO Amplifies Existing Capabilities, SFT Replaces Them](https://mlsanity.com/paper/2507.10616/) (2025)
- [From Code Foundation Models to Agents and Applications: A Comprehensive Survey and Practical Guide to Code Intelligence](https://mlsanity.com/paper/2511.18538/) (2025)
- [Risk-Sensitive RL for Alleviating Exploration Dilemmas in Large Language Models](https://mlsanity.com/paper/2509.24261/) (2025)
- [Don't Waste Mistakes: Leveraging Negative RL-Groups via Confidence Reweighting](https://mlsanity.com/paper/2510.08696/) (2025)
- [Certified Self-Consistency: Statistical Guarantees and Test-Time Training for Reliable Reasoning in LLMs](https://mlsanity.com/paper/2510.17472/) (2025)
- [SetPO: Set-Level Policy Optimization for Diversity-Preserving LLM Reasoning](https://mlsanity.com/paper/2602.01062/) (2026)
- [Large Reasoning Models are not thinking straight: on the unreliability of thinking trajectories](https://mlsanity.com/paper/2507.00711/) (2025)
- [Understanding Diversity Collapse in RLVR via the Lens of Overtraining](https://mlsanity.com/paper/2606.15455/) (2026)
- [Best-of-Majority: Minimax-Optimal Strategy for Pass@k Inference Scaling](https://mlsanity.com/paper/2510.03199/) (2025)
- [Thinking with Map: Reinforced Parallel Map-Augmented Agent for Geolocalization](https://mlsanity.com/paper/2601.05432/) (2026)
- [Expanding LLM Agent Boundaries with Strategy-Guided Exploration](https://mlsanity.com/paper/2603.02045/) (2026)
- [Why Pass@k Optimization Can Degrade Pass@1: Prompt Interference in LLM Post-training](https://mlsanity.com/paper/2602.21189/) (2026)
- [Polychromic Objectives for Reinforcement Learning](https://mlsanity.com/paper/2509.25424/) (2025)
- [LaDi-RL: Latent Diffusion Reasoning Prevents Entropy Collapse in Reinforcement Learning](https://mlsanity.com/paper/2602.01705/) (2026)
- [Mining Hidden Thoughts from Texts: Evaluating Continual Pretraining with Synthetic Data for LLM Reasoning](https://mlsanity.com/paper/2505.10182/) (2025)
- [Beyond Mode Collapse: Distribution Matching for Diverse Reasoning](https://mlsanity.com/paper/2605.19461/) (2026)
- [GuidedSampling: Steering LLMs Towards Diverse Candidate Solutions at Inference-Time](https://mlsanity.com/paper/2510.03777/) (2025)
- [Poly-EPO: Training Exploratory Reasoning Models](https://mlsanity.com/paper/2604.17654/) (2026)
- [Emergence of Exploration in Policy Gradient Reinforcement Learning via Retrying](https://mlsanity.com/paper/2606.00151/) (2026)
- [Regret Analysis of Retry-Based Bandits](https://mlsanity.com/paper/2605.20854/) (2026)

## Abstract (arXiv)

In this work, we investigate the merits of explicitly optimizing for inference time algorithmic performance during model training. We show how optimizing for inference time performance can improve overall model efficacy. We consider generic inference time objectives with k samples, with a focus on pass@k and majority voting as two main applications. With language model training on reasoning datasets, we showcase the performance trade-off enabled by training with such objectives. When training on code generation tasks, we show that the approach significantly improves pass@k objectives compared to the baseline method.