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

TL;DR

This work reframes LLM alignment as a Stackelberg game in which the reward-model designer (leader) chooses a reward function to maximize user utility under a KL-regularized follower (the LLM). It proves the optimal reward strategy is a threshold-based model that can be efficiently computed via Monte Carlo estimation, and introduces a soft relaxation (SRS) to enhance robustness. The approach is integrated into inference-time alignment methods (CD and ARGS) and validated on established benchmarks, showing consistent gains in average reward with minimal inference overhead and strong GPT-4 judging results. By bounding and shaping rewards, the method mitigates reward hacking and adapts to diverse user preferences without retraining, offering practical benefits for scalable, per-request alignment.

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.

Figures (12)

• Figure 1: We illustrate the Stackelberg game formulation of LLM alignment. In this framework, the reward model provider acts as the leader by selecting a reward model, while the LLM policy responds as the follower by solving the resulting alignment problem. The reward model provider’s goal is to choose an optimal reward model that maximizes user utility.
• Figure 2: Reward and GPT-4 win-tie rate as a function of the inference-time reward strength $\frac{1}{\beta}$. The Win-Tie rate is compared with base model with no alignment. Solid lines denote the reward given by the reward model ,and dashed lines denote the Win-Tie rate.
• Figure 3: Solid lines correspond to the left axis (Reward) and dashed line corresponds to the right axis (Diversity/Coherence).
• Figure 4: Solid lines correspond to the left axis (Reward) and dashed line corresponds to the right axis (Diversity/Coherence).
• Figure 5: Controlled Decoding (CD) performance as a function of reward strength. Performance degrades under the Skywork reward model and shows negligible improvement under UltraRM, indicating that CD is sensitive to the choice of reward model.

Theorems & Definitions (15)

• Definition 1: Threshold reward
• Theorem 1: Optimality of threshold reward
• Definition 2
• Theorem 2
• Theorem 3
• Corollary 1
• Lemma 1
• proof
• Lemma 2
• proof