Technical Summary: Explanation and Prevention of Alignment Collapse in Iterative RLHF

Problem Statement

Reinforcement Learning from Human Feedback (RLHF) is the prevailing paradigm for aligning Large Language Models (LLMs) with human intent. However, it is fundamentally limited by reward over-optimization, a manifestation of Goodhart's Law where strategies exploit imperfections in the Reward Model (RM) to achieve high scores while deviating from true human preferences.

While state-of-the-art pipelines have shifted toward iterative RLHF—where the RM is periodically retrained on data generated by the current strategy to create a robust, adaptive target—this work identifies a critical failure mode in this regime: alignment collapse. In iterative deployment, a feedback loop emerges where the strategy generates data that retrains the RM, which in turn updates the strategy. The authors argue that standard iterative RLHF, employing myopic optimizers (e.g., PPO, REINFORCE), fails to account for the strategy's influence on the future parameters of the RM. This myopia causes the strategy to systematically exploit the "blind spots" of the RM (areas with high estimation errors). Retraining the RM on this exploitative data reinforces exactly the errors being exploited, creating a pathological feedback process that diverges from the true human utility function.

Methodology and Theoretical Framework

1. Formulation as a Stackelberg Game

The authors formalize iterative RLHF as a Stackelberg game (a two-stage optimization problem).

• Leader: The policy ($\pi_\theta$) commits to parameters $\theta$ to maximize expected reward.
• Follower: The RM ($r_\phi$) observes the data distribution induced by the policy and updates parameters $\phi$ to minimize its loss $L(\theta, \phi)$.

2. Analytic Decomposition of the Gradient

By unrolling the two-stage optimization via implicit differentiation, the authors derive the true optimization gradient for the policy. They show that the total gradient $\nabla_\theta F(\theta)$ decomposes into two distinct terms:

1. Standard Policy Gradient: The term used by current algorithms, which treats the RM as a static entity ($\nabla_\theta J(\theta, \phi^*)$).
2. Parameter-Steering Gradient: A term capturing how the current policy outputs distort the future parameters of the RM ($\nabla_\phi J \cdot \frac{d\phi^*}{d\theta}$).

Using influence functions from robust statistics, the authors demonstrate that the parameter-steering gradient is equivalent to a policy gradient with respect to an implicit effective reward:
$$\tilde{r}_\theta(x, y) = r_{\phi^*}(x, y) + \langle \bar{g}_r(\theta), I_\theta(x, y) \rangle$$
Where:

• $r_{\phi^*}$ is the proxy reward.
• $I_\theta(x, y)$ is the influence function (sensitivity of RM parameters to the sample $(x,y)$).
• $\bar{g}_r(\theta)$ is the global reward gradient direction.

The second term, $\langle \bar{g}_r, I_\theta \rangle$, acts as a self-correcting mechanism. It penalizes samples that push the RM to overestimate its own reward (exploiting blind spots) and rewards samples that align with areas where the proxy is reliable.

3. Foresighted Policy Optimization (FPO)

To prevent alignment collapse, the authors propose Foresighted Policy Optimization (FPO), a mechanism design intervention that restores the missing steering term.

• Objective: Extend the standard RLHF objective with a penalty $R_{FPO}$ that regularizes the steering effect of the policy.
• Scalability: Exact computation of the steering term requires inverting the Hessian matrix of the RM loss, which is intractable for high-dimensional models. The authors implement FPO via a first-order approximation using the TracIn method.
  • They approximate the influence function via the loss gradient ($H^{-1} \approx I$).
  • They replace the global reward gradient with the local sample gradient.
  • This yields a tractable penalty proportional to the inner product of the reward gradient and the negative loss gradient: $\tilde{R}_{FPO}(x, y) \propto \langle \nabla_\phi r_\phi, -\nabla_\phi \ell \rangle$.
• Oracle-Free Variant: Since the true utility $U$ is unobservable, they propose a practical penalty $\bar{R}_{FPO}$ that penalizes the magnitude of the reward gradient ($-\|\nabla_\phi r_\phi\|^2$), effectively preventing the policy from generating data in highly volatile, poorly calibrated areas.

The framework is extended from pointwise rewards to standard pairwise comparisons (Bradley-Terry loss), demonstrating that the same structural failure and solution apply to modern iterative pipelines.

Main Contributions

1. Theoretical Unrolling: The work analytically unrolls the iterative Stackelberg game to prove that the policy's true optimization landscape is determined by an implicit effective reward comprising the proxy reward and a parameter-steering term.
2. Diagnosis of Alignment Collapse: It identifies that ignoring the parameter-steering term is the geometric cause of alignment collapse, where myopic policies drive the system into a regime of shortsighted reward manipulation.
3. FPO Mechanism: It introduces Foresighted Policy Optimization, a targeted penalty that explicitly regularizes the policy to account for its influence on future RM updates, restoring self-correcting dynamics.
4. Scalable Implementation: It bridges the gap to high-dimensional models by formulating FPO as a computationally feasible penalty using a first-order TracIn relaxation, requiring only one additional gradient evaluation per sample.

Experimental Results

1. Controlled Continuous Environments

• Linear RM: In a toy environment with a linear RM and a Gaussian utility peak, standard iterative RLHF converged to a suboptimal equilibrium where the policy maximized noise dimensions to exploit RM blind spots, leading to a drop in true utility. FPO successfully suppressed this hacking and converged to the human ideal.
• Neural Network RM: In a 10-dimensional continuous space with a 2-layer MLP RM, standard RLHF initially improved but eventually diverged from the true optimum due to over-optimization. FPO prevented this divergence and converged precisely to the human ideal.

2. LLM Alignment Pipeline

The authors simulated an iterative RLHF pipeline using Llama-3.2-1B (policy) and DeBERTa-v3-base (RM).

• Setup: The policy was optimized via Best-of-N (BoN) rejection sampling on the UltraFeedback dataset.
• Evaluation: Models were evaluated on TruthfulQA (817 prompts) using a blind Llama-3.3-70B judge.
• Results:
  • Relaxed FPO ($\tilde{R}_{FPO}$): Significantly outperformed standard RLHF (56.6% win rate, $p=0.014$), demonstrating that optimizing the implicit effective reward improves alignment.
  • Practical FPO ($\bar{R}_{FPO}$): Showed a non-significant advantage over the baseline (50.9% win rate, $p=0.41$) but maintained robustness.
  • Capability Trade-off: While the practical penalty improved performance on standard factual tasks, it was too conservative compared to the relaxed version on adversarial prompts. However, both FPO variants preserved the base model's general reasoning capabilities (MMLU, ARC-Challenge) within statistical margins, confirming that the regularization only constrains the policy in sensitive, exploitable areas.

Significance and Claims

The work claims to offer a rigorous geometric defense against reward hacking in iterative RLHF. By formulating the problem as a Stackelberg game, it reveals that the failure of current pipelines is not merely due to a lack of data or model capacity, but to a fundamental neglect of the policy's strategic influence on the reward model.

The authors claim that FPO provides a principled mechanism to restore the self-correcting dynamics of the true Stackelberg equilibrium. Unlike prior works that treat the RM as a static object or rely on static flatness bounds, FPO targets the sensitivity of RM parameters under retraining. The work suggests that preventing alignment collapse does not require correcting the RM's misspecification (which is often impossible), but rather preventing the policy from amplifying these errors through strategic exploitation.

The work concludes that while FPO offers a strong theoretical and empirical foundation, future work is required to extend the analysis to non-convex settings (since the current derivation assumes strong convexity of the RM loss) and to scale the experiments to larger model sizes.