Table of Contents
Fetching ...

Why Is RLHF Alignment Shallow? A Gradient Analysis

Robin Young

TL;DR

It is proved that gradient-based alignment inherently concentrates on positions where harm is decided and vanishes beyond, and an objective based on recovery penalties that creates gradient signal at all positions is derived, providing theoretical grounding for empirically successful data augmentation techniques.

Abstract

Why is safety alignment in LLMs shallow? We prove that gradient-based alignment inherently concentrates on positions where harm is decided and vanishes beyond. Using a martingale decomposition of sequence-level harm, we derive an exact characterization of alignment gradients. The gradient at position $t$ equals the covariance between the conditional expected harm and the score function. This implies that positions beyond the harm horizon where the output's harmfulness is already determined receive zero gradient signal during training. This explains empirical observations that KL divergence between aligned and base models concentrates on early tokens. Consequently, standard alignment objectives cannot produce deep alignment, regardless of optimization quality. We introduce the concept of harm information $I_t$, which quantifies each position's influence on harm, and prove that equilibrium KL divergence tracks this quantity. Finally, we derive an objective based on recovery penalties that creates gradient signal at all positions, providing theoretical grounding for empirically successful data augmentation techniques.

Why Is RLHF Alignment Shallow? A Gradient Analysis

TL;DR

It is proved that gradient-based alignment inherently concentrates on positions where harm is decided and vanishes beyond, and an objective based on recovery penalties that creates gradient signal at all positions is derived, providing theoretical grounding for empirically successful data augmentation techniques.

Abstract

Why is safety alignment in LLMs shallow? We prove that gradient-based alignment inherently concentrates on positions where harm is decided and vanishes beyond. Using a martingale decomposition of sequence-level harm, we derive an exact characterization of alignment gradients. The gradient at position equals the covariance between the conditional expected harm and the score function. This implies that positions beyond the harm horizon where the output's harmfulness is already determined receive zero gradient signal during training. This explains empirical observations that KL divergence between aligned and base models concentrates on early tokens. Consequently, standard alignment objectives cannot produce deep alignment, regardless of optimization quality. We introduce the concept of harm information , which quantifies each position's influence on harm, and prove that equilibrium KL divergence tracks this quantity. Finally, we derive an objective based on recovery penalties that creates gradient signal at all positions, providing theoretical grounding for empirically successful data augmentation techniques.
Paper Structure (37 sections, 21 theorems, 81 equations)

This paper contains 37 sections, 21 theorems, 81 equations.

Key Result

Proposition 2

The sequence $(h_t)_{t=0}^T$ is a martingale with respect to the filtration generated by $(y_t)$:

Theorems & Definitions (61)

  • Definition 1: Conditional Expected Harm
  • Proposition 2: Martingale Property
  • proof
  • Definition 3: Harm Innovation
  • Proposition 4: Doob Decomposition
  • proof
  • Definition 5: Harm Information
  • Proposition 6: Variance Reduction Characterization
  • proof
  • Corollary 7: Variance Decomposition
  • ...and 51 more