Inference-time Alignment with Rewards in Besov Spaces: Provable Advantages of Feature Learning and Multi-Step Policy Updates

Inference-time alignment, the approach of adapting pre-trained models to reward feedback during inference, has proven empirically effective at improving language-model performance. Despite its success, theoretical foundations remain underdeveloped, especially in practical settings where neural networks are employed as reward models. In this paper, we explore the advantages of neural networks and how to effectively train them for inference-time alignment. Assuming that the true reward function lies in Besov spaces to capture the non-uniform smoothness, we compare neural networks to linear estimators and show that feature learning capability of neural networks is crucial for improving performance. We further analyze algorithms for training neural-network reward estimators. Specifically, we consider a multi-step algorithm that alternates between sampling from the current policy and refitting the reward estimator, and prove that it improves the regret, especially when the true reward exhibits local structure.

Large language models can often produce better answers when they try several possible responses and choose the most promising one. The difficult part is deciding which response is best: this depends on a reward model, a separate system that scores responses, and score errors can push the language model toward bad answers. We study how to build and use such reward models when only a limited number of scored examples are available. Our analysis covers cases where answer quality changes smoothly in some regions but has sharper local patterns in others, making it hard for simple scoring methods to find the best responses. We prove that neural-network reward models can learn these local patterns and outperform simpler methods based on fixed features. We also analyze an iterative approach that alternates between generating responses and retraining the reward model on new scored examples. Our results show that repeated updating can move the language model closer to high-quality responses, especially when the best responses are concentrated in structured regions. This gives a theoretical explanation for why neural networks and iterative updating can improve language-model outputs at inference time, and states the assumptions behind these guarantees.