🤖 AI

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

This paper formalizes reward model optimization under KL regularization as a Stackelberg game and proposes a simple reward shaping scheme that effectively mitigates base policy bias while preventing reward hacking, thereby significantly improving inference-time alignment performance.

Original authors: Haichuan Wang, Tao Lin, Lingkai Kong, Ce Li, Hezi Jiang, Milind Tambe

Published 2026-06-09

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Haichuan Wang, Tao Lin, Lingkai Kong, Ce Li, Hezi Jiang, Milind Tambe

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Core Problem: The Stubborn Intern

Imagine you have hired a brilliant but stubborn intern (the Base Large Language Model) to write emails for you. This intern has strong, ingrained habits. For example, if you ask for a neutral opinion on a political topic, the intern naturally leans heavily to the left because that’s how they were trained.

You want the intern to write exactly what you prefer. So, you give them a scorecard (the Reward Model) based on your personal tastes. You tell them: "If you write a neutral response, I’ll give you 2 points. If you write a left-leaning response, I’ll give you 1 point."

However, the intern is lazy and resistant to change. They don’t want to completely rewrite their personality. So, they follow a rule: "I will try to get your points, but I won’t stray too far from my natural instincts." In technical terms, this is called KL Regularization. It keeps the intern’s output close to their original, biased self.

The Problem: Because the intern is so stubborn, simply showing them your scorecard isn’t enough. They might still lean left because their "natural instinct" is stronger than your small reward difference. If you try to fix this by giving the neutral response infinite points, the intern might panic and start writing gibberish just to get those points (this is called Reward Hacking).

The Solution: The Boss and The Intern (A Stackelberg Game)

The authors propose treating this situation like a strategic game between a Boss (the Reward Model Provider) and an Intern (the LLM). This is known in game theory as a Stackelberg Game, where the Boss moves first, and the Intern reacts.

Instead of just handing the intern the raw scorecard, the Boss plays a smarter game. The Boss asks: "How should I tweak the scores so that, knowing how stubborn my intern is, they will actually do what I want?"

The paper proves mathematically that the best way to do this is through "Threshold Reward Shaping."

The Analogy: The "Pass/Fail" Gate

Think of the Boss’s strategy like setting a strict pass/fail gate for a job interview, rather than giving a nuanced grade.

The Threshold: The Boss picks a specific score line (the threshold).
The Shaping:
- If the intern’s natural output is below this line, the Boss gives it a score of 0 (Zero points).
- If the intern’s natural output is above this line, the Boss gives it a score of B (Maximum points).
- If it’s exactly on the line, it gets a middle score.

Why this works:
By turning the rewards into a sharp "All or Nothing" choice, the Boss forces the stubborn intern to make a bigger shift. The intern realizes, "If I don’t cross this threshold, I get nothing. If I do, I get the max reward." This exaggeration of rewards counteracts the intern’s stubbornness (bias) without causing them to break down and write nonsense (reward hacking).

How It Works in Practice (Inference-Time)

Most AI alignment happens during training (teaching the model from scratch). This paper focuses on Inference-Time Alignment, which is like coaching the intern while they are working, without retraining them.

Here is the step-by-step process the authors propose:

Sample: When you ask the AI a question, it first generates a few rough drafts using its natural, stubborn instincts.
Score: The system checks these drafts against your true preferences.
Calculate the Threshold: Using a quick calculation (Monte Carlo sampling), the system figures out the perfect "threshold" score for that specific question. It asks: "At what point does the average quality of the drafts meet my expectations?"
Reshape: The system then applies the "Pass/Fail" logic. It boosts the scores of the good drafts to the maximum and drops the bad ones to zero.
Decode: The AI uses these reshaped scores to pick the final words. Because the rewards are now sharply defined, the AI is pushed to generate the output you actually want, overcoming its natural bias.

The Results

The authors tested this method on popular AI models (like Llama and Qwen). They found that:

Better Alignment: The AI produced responses that matched user preferences much better than standard methods.
No Gibberish: The method avoided "reward hacking," meaning the AI didn’t start writing nonsense to cheat the system.
Efficient: This calculation is very fast. It adds almost no delay to the AI’s response time.
Winning: When judged by another AI (GPT-4), the responses generated with this method were rated as better or equal to the baselines 66% of the time.

Summary

In short, the paper argues that you shouldn’t just tell an AI what you like; you should strategically exaggerate the rewards to overcome the AI’s natural biases. By framing this as a game where the "Boss" (reward designer) anticipates the "Intern’s" (AI’s) stubbornness, they created a simple, effective method to make AI assistants more helpful and aligned with user intent.

Technical Summary: Reward Shaping for (Inference-Time) Alignment: A Stackelberg Game Perspective

Problem Statement
The paper addresses the suboptimality of standard alignment pipelines that directly optimize a Large Language Model (LLM) policy using a reward model learned from user preference data, subject to KL-divergence regularization against a base policy. The authors argue that this direct optimization is flawed because the KL constraint forces the aligned policy to remain close to the base policy. If the base policy contains biases that conflict with user preferences, the KL regularization prevents the model from fully aligning with those preferences. While amplifying rewards can counteract this bias, it risks "reward hacking" (producing incoherent or degenerate outputs to maximize reward). Consequently, the paper poses the fundamental question: How should the reward model be optimally shaped to maximize user utility while mitigating reward hacking under KL regularization?

Methodology
The authors formalize the alignment process as a Stackelberg game involving two players:

The Leader (Reward Model Provider): Aims to maximize the user’s expected utility. The leader selects a reward model $r$ to guide the LLM, anticipating the LLM’s response. The leader is not constrained to use the true user utility function $r_U$ but can choose any reward model $r$ within a bounded range $[0, B]$ to mitigate reward hacking.
The Follower (LLM): Best-responds to the leader’s chosen reward model by optimizing an alignment objective that balances expected reward against the KL divergence from the base policy $\rho_{base}$ .

The authors derive a closed-form characterization of the optimal reward model $r^*$ . They prove that the optimal solution admits a threshold structure: for a given prompt, responses with true utility above a specific threshold $m^*(x)$ are assigned the maximum reward $B$ , while those below are assigned 0. The optimal threshold $m^*(x)$ is defined as the self-consistent expected utility the LLM will deliver under this shaped reward.

To compute this threshold efficiently, the authors propose Stackelberg Reward Shaping (SRS). Since the exact expectation is intractable, they use Monte Carlo samples from the base policy to estimate the threshold via a bisection algorithm on a helper function. To improve robustness against the discontinuity of the hard threshold, they introduce a soft threshold variant using a sigmoid transformation, controlled by a shaping strength hyperparameter $\alpha$ .

The framework is integrated into inference-time alignment methods, specifically Controlled Decoding (CD) and Alignment as Reward-Guided Search (ARGS). In this setting, SRS is applied offline to shape the rewards used to train Q-value estimators, which then guide the decoding process without retraining the base model weights.

Key Contributions

Theoretical Insight: Demonstrates that directly using a reward model learned from user preferences is suboptimal under KL-regularized objectives due to the distortion caused by base policy bias.
Game-Theoretic Formulation: Formulates reward model optimization as a Stackelberg game, proving that the optimal reward model has a threshold structure. They provide an efficient algorithm to approximate this optimal model using Monte Carlo sampling from the base policy.
Empirical Validation: Integrates SRS into existing inference-time alignment methods (CD and ARGS). The method requires minimal inference-time overhead and consistently improves performance across multiple evaluation settings.

Results
The authors evaluate SRS on the HH-RLHF and SHP datasets using Qwen3-8B and Llama3-8B-Instruct as backbones. Key findings include:

Reward Improvement: SRS consistently achieves the highest average reward across all evaluation settings compared to baselines (Vanilla, Minmax, Meanstd).
Quality Preservation: SRS maintains diversity and coherence metrics comparable to or better than baselines.
Human Evaluation: In head-to-head comparisons judged by GPT-4, SRS achieves average win–tie rates exceeding 66% against all baselines (Vanilla, Minmax, and Meanstd), indicating that improvements are not driven by reward hacking.
Robustness: SRS demonstrates robustness to input distribution shifts (generalizing from SHP to HH-RLHF and HarmfulQA) and is less sensitive to outliers than heuristic methods like Minmax.
Efficiency: For SRS-CD, reward shaping is offline with no inference overhead. For SRS-ARGS, the overhead is negligible (~0.03 seconds per 128-token response).

Significance
The paper claims significance in identifying a critical flaw in standard alignment practices: the assumption that direct reward maximization aligns with user utility under KL constraints. By framing reward design as an incentive mechanism problem via Stackelberg games, the work provides a principled, theoretically grounded approach to reward shaping. It offers a practical, lightweight solution for inference-time alignment that improves user utility without requiring model retraining, thereby facilitating more flexible and effective adaptation to diverse user preferences.