Optimizing Language Models for Inference Time Objectives using Reinforcement Learning
This work addresses the efficiency and effectiveness of language models during inference, which is crucial for real-world applications, though it appears incremental as it builds on existing training methods with new objectives.
The researchers tackled the problem of optimizing language models for inference time performance by explicitly training them with inference objectives like pass@k and majority voting, showing that this approach significantly improves pass@k objectives on code generation tasks compared to baseline methods.
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.