Reward Shaping for Inference-Time Alignment: A Stackelberg Game Perspective

TL;DR

This paper introduces a Stackelberg game framework for optimizing reward models in inference-time alignment of large language models, effectively balancing reward amplification and bias mitigation to improve user utility.

Contribution

It formalizes reward model optimization as a Stackelberg game and proposes a simple reward shaping scheme that approximates the optimal solution with minimal overhead.

Findings

Consistently improves average reward in alignment tasks.

Achieves win-tie rates over 66% against baselines.

Seamlessly integrates into existing methods.

Abstract

Existing alignment methods directly use the reward model learned from user preference data to optimize an LLM policy, subject to KL regularization with respect to the base policy. This practice is suboptimal for maximizing user's utility because the KL regularization may cause the LLM to inherit the bias in the base policy that conflicts with user preferences. While amplifying rewards for preferred outputs can mitigate this bias, it also increases the risk of reward hacking. This tradeoff motivates the problem of optimally designing reward models under KL regularization. We formalize this reward model optimization problem as a Stackelberg game, and show that a simple reward shaping scheme can effectively approximate the optimal reward model. We empirically evaluate our method in inference-time alignment settings and demonstrate that it integrates seamlessly into existing alignment methods with minimal overhead. Our method consistently improves average reward and achieves win-tie rates exceeding 66% against all baselines, averaged across evaluation settings.