A Common Pitfall of Margin-based Language Model Alignment: Gradient Entanglement

TL;DR

The paper identifies a fundamental pitfall in margin-based language model alignment: margin-based losses only constrain the difference between chosen and dispreferred outputs, under-specifying their behavior individually and enabling gradient entanglement. It formalizes this through a gradient-inner-product condition that governs when the chosen output can increase while the dispreferred output decreases; empirical and theoretical analyses show that entanglement often causes synchronized movement of both log-probabilities, impacting safety and fidelity. The authors derive a general gradient-condition framework for margin-based losses, explain differences among DPO and other losses, and validate findings with synthetic and sentiment-task experiments. To mitigate entanglement, they propose two designs—Normalized Preference Optimization and Sparse Preference Optimization—that reweight gradients or emphasize token-level contrasts. Collectively, the work provides a principled lens to understand RLHF dynamics and suggests practical algorithmic directions to improve alignment outcomes.

Abstract

Reinforcement Learning from Human Feedback (RLHF) has become the predominant approach for language model (LM) alignment. At its core, RLHF uses a margin-based loss for preference optimization, specifying ideal LM behavior only by the difference between preferred and dispreferred responses. In this paper, we identify a common pitfall of margin-based methods -- the under-specification of ideal LM behavior on preferred and dispreferred responses individually, which leads to two unintended consequences as the margin increases: (1) The probability of dispreferred (e.g., unsafe) responses may increase, resulting in potential safety alignment failures. (2) The probability of preferred responses may decrease, even when those responses are ideal. We demystify the reasons behind these problematic behaviors: margin-based losses couple the change in the preferred probability to the gradient of the dispreferred one, and vice versa, often preventing the preferred probability from increasing while the dispreferred one decreases, and thus causing a synchronized increase or decrease in both probabilities. We term this effect, inherent in margin-based objectives, gradient entanglement. Formally, we derive conditions for general margin-based alignment objectives under which gradient entanglement becomes concerning: the inner product of the gradients of preferred and dispreferred log-probabilities is large relative to the individual gradient norms. We theoretically investigate why such inner products can be large when aligning language models and empirically validate our findings. Empirical implications of our framework extend to explaining important differences in the training dynamics of various preference optimization algorithms, and suggesting potential algorithm designs to mitigate the under-specification issue of margin-based methods and thereby improving language model alignment.

Figures (7)

• Figure 1: Training dynamics of the chosen and rejected log probabilities on the TL;DR dataset stiennon2020learning, with log probabilities reported on the evaluation set. As the margin between the two increases, the chosen and rejected log-probabilities exhibit synchronized increases and decreases per step. In Figure \ref{['fig:mistra-7b']}, both chosen and rejected log-probabilities have an overall trend of increasing, especially towards the end of training, whereas in Figure \ref{['fig:llama3-8b']}, both have a trend of decreasing.
• Figure 2: Training dynamics of the chosen and rejected log-probabilities on the TL;DR dataset for different algorithms trained on Mistral 7B. The corresponding plot for Llama3 8B is in Figure \ref{['fig:llama3-8b-all-logp']} (Appendix \ref{['app:logp-plots']}). For SimPO and IPO, the log-probabilities are normalized by the response length, while in the other plots, the log-probabilities are of entire responses. All algorithms exhibit synchronized increases and decreases in the chosen and rejected log-probabilities. We also provide the cosine similarity plots between $\nabla_\theta \log \pi_w$ and $\nabla_\theta \log \pi_l$ in Appendix \ref{['app:logp-plots']} (Figure \ref{['fig:tldr-grad-ip']}).
• Figure 3: Training dynamics of the chosen and rejected log probabilities for sentiment tasks.
• Figure 4: Gradient cosine similarity behaviors on the sentence-level and token-level for sentiment tasks. Figure \ref{['fig:sentiment-grad-ip']} gives the cosine similarity between $\nabla_\theta \log \pi_w$ and $\nabla_\theta \log \pi_l$ for DPO on single token, short suffix and long suffix datasets, defined as: $\frac{\langle \nabla_\theta \log \pi_w, \nabla_\theta \log \pi_l \rangle}{\|\nabla_\theta \log \pi_w\| \|\nabla_\theta \log \pi_l\|}$. Figure \ref{['fig:sentiment-token-heat']} shows the token-wise gradient similarity for an instance in the prefix+suffix task.
• Figure 5: Training dynamics of the chosen and rejected log probabilities during DPO, observed across models: Llama3.1-8B dubey2024llama, Llama3.2-3B, Mistral-7B jiang2023mistral7b and Gemma2-9B team2024gemma on TL;DR stiennon2020learning and UltraFeedback cui2024ultrafeedback datasets. Log probabilities are averaged on the evaluation splits.

Theorems & Definitions (4)

• Theorem 1
• Corollary 2
• Theorem 3
• Remark