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

Authors: Hui Yuan, Yifan Zeng, Yue Wu, Huazheng Wang, Mengdi Wang, Liu Leqi
Primary area: alignment, fairness, safety, privacy, and societal considerations
Keywords: Alignment, Preference Optimization, Large Language Model

Abstract:
Reinforcement Learning from Human Feedback (RLHF) has become the predominant approach for aligning language models (LMs) to be more helpful and less harmful. 
At its core, RLHF uses a margin-based loss for preference optimization, which specifies the ideal LM behavior only in terms of 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 results in 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 with 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 between the gradient of preferred log-probability and the gradient of dispreferred log-probability is large relative to the individual gradient norms. Furthermore, we theoretically investigate why such inner products can be large when aligning language models and empirically validate our findings. Empirical implications of our framework further extend to explaining important differences in the training dynamics of various preference optimization algorithms and suggesting future directions for improvement.

Decision: accept (poster)
Reviewer scores: 8, 6, 5, 5

Reviews and discussion:
<thread>
creation date:
1730130718394

thread_type:
official_review

<thread_object>
creator: reviewer_ai3q

<content>
summary:
This paper criticizes margin-based preference optimization for LM alignment, which is a core method in RLHF. While increasing preference margins can reduce unsafe responses, it also often reduces preferred response probabilities, due to synchronized probability shifts between chosen and rejected responses. The authors reveal when these shifts occur and suggest that the under-specification of behaviour in current margin-based methods limits model safety and effectiveness. The work proposes fine-grained, token-level approaches as potential improvements, though this is not really pursued very deeply in this work.

strengths:
- The paper provides a fairly nuanced and theoretically grounded critique of margin-based preference optimization, which is an increasingly relevant task (especially  in RL). Identifying conditions under which preference optimization can backfire is important for real-world use and may provide a relatively underdeveloped area of research
- The proposal to address preference optimization problems at the token level is interesting (though perhaps not surprising).

weaknesses:
- this is not meant to be a primarily empirical paper (apparently, it leans more on its theoretical aspects) but this work of course calls for broader empirical validation than it has. The proposed method may be evaluated across diverse tasks such as summarization, dialogue, and instruction-following. Other variants of the given LLMs (e.g., LLaMa) could have been explored, to demonstrate generalizability. 
- The analysis of Sec 4 focuses on synthetic, controlled scenarios. The kinds of complexities exist that in actual, diverse language data should be addressed.

questions:
- While the theoretical analysis of gradient inner products and synchronized probability shifts is fairly rigorous, what actionable insights for real-world RLHF applications are achievable?
- How would you evaluate the ‘safety-critical alignment tasks’ you mention in Sec 5 as being particularly ill-suited for margin-based preference optimization?
- Can you clarify more directly how token-level gradient inner products between chosen and rejected responses provide a reliable indicator for alignment behaviour?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730386051562

thread_type:
official_review

<thread_object>
creator: reviewer_uxg4

<content>
summary:
The paper identifies and analyzes a fundamental problem with margin-based preference optimization in language model alignment (like RLHF): these methods only specify how much better preferred responses should be compared to dispreferred ones, without specifying how the individual probabilities should behave. This leads to two unintended consequences: (1) sometimes the probability of dispreferred (e.g., unsafe) responses increases, and (2) when reducing dispreferred response probabilities, the probability of preferred responses often decreases too, even when they are ideal.

strengths:
The paper addresses a fundamental issue in language model alignment through RLHF by analyzing the under-specification problem in margin-based preference optimization - this is an important contribution given how widely RLHF is used.

weaknesses:
1  Presentation Issues:
The notation becomes dense and potentially confusing in Section 3.2, particularly around the general form analysis
Some figures (like Figure 4b) would benefit from more detailed explanations of their implications
The connection between the theoretical results and practical implications could be more clearly articulated

2  Limited Experimental Validation:  
Experiments are conducted on only one main dataset (TL;DR) and a simple sentiment classification task
Only tested on two main models (Mistral 7B and Llama-3 8B)
Lacks validation across diverse preference optimization scenarios and task types
More extensive empirical validation across different datasets and model architectures would strengthen the claims

3  Theoretical Assumptions:
The theoretical analysis in Section 4.1 relies on somewhat simplified model settings (learnable last linear layer, learnable logits)
Some assumptions about token probability relationships between chosen/rejected responses may not hold in practice
The extension of single-token results to multi-token cases could be more rigorously justified

questions:
N/A
</content>
</thread_object>

</thread>
<thread>
creation date:
1730519275077

thread_type:
official_review

<thread_object>
creator: reviewer_5b6v

<content>
summary:
The paper investigates issues in margin-based preference optimization for language model alignment, with a focus on the Direct Preference Optimization (DPO) algorithm. It highlights that margin-based loss specifies only the ideal margin behavior, overlooking individual term behaviors, which often results in synchronized log probability changes between chosen and rejected responses. The authors theoretically derive conditions for this phenomenon starting with DPO, and extend their insights to other margin-based algorithms. They examine cases where the derived condition holds, noting that higher token similarity between chosen and rejected responses can intensify synchronized movement, and validate these findings empirically.

strengths:
This paper stands out in its originality by addressing the under-explored limitations of margin-based preference optimization in RLHF, offering fresh insights into its impact on language model alignment. The theoretical contributions, specifically the gradient inner product conditions, demonstrate high quality, providing a clear, mathematically grounded framework that identifies when alignment failures may occur under existing loss structures. 
In addition to an in-depth focus on the DPO algorithm, the paper comprehensively reviews existing methods in Table 2 and extends its findings in Section 3.2, which gives additional quality to the paper.
The suggested condition in DPO is further substantiated with its practical implication, with experimental validation in Section 4.

weaknesses:
The paper centers primarily on DPO in Sections 3.1 and 4, with limited coverage of other margin-based optimization algorithms. To better align with reader expectations, it would be beneficial to introduce their main focus of DPO explicitly in the introduction and outline how subsequent sections will compare DPO against alternative approaches. Alternatively, devoting more space to a detailed examination of other algorithms could create a more balanced narrative.

The experimental results in Section 3.3 primarily observe phenomena regarding preferred and dispreferred response probabilities increasing or decreasing together—findings that have already been discussed in previous work, such as Pal et al. (2024), and does not fully investigate the conditions under which these probability trends occur. Furthermore, the experiment relies on a single dataset, limiting the depth of empirical validation. 

The empirical observations are concise but lack important methodological details. For instance, it would be valuable to clarify how log probabilities are computed in Figure 1 (e.g., whether they are averaged over the training samples), and to include additional measures such as standard deviations or confidence intervals to improve the figure's reliability. Since Figure 1 is intended as a motivating visualization, providing a additional figure for dedicated to empirical results (Section 3.3) would allow for a more thorough presentation of the experimental findings without cluttering the introductory material.

For minor clarification, it would be better to change the interval of the grid to clearly show the trend of margin in Figure2.

questions:
1. The paper presents a clear delineation between preferred and dispreferred samples. 
However, some papers focus on distributional framework. This approaches insists the model to encounter a broader range of sample types to cover the distribution's full support, supporting the original RLHF sampling method.
- Go, Dongyoung, et al. "Aligning language models with preferences through f-divergence minimization." arXiv preprint arXiv:2302.08215 (2023).
- Korbak, Tomasz, et al. "On reinforcement learning and distribution matching for fine-tuning language models with no catastrophic forgetting." Advances in Neural Information Processing Systems 35 (2022): 16203-16220.

 Could the authors discuss the potential impacts of their idealized setup compared to a distributional approach? This clarification would help assess the implications of choosing one framework over the other in alignment training.

2. Training with both high- and moderate-quality responses, where one is better but another may serve as a hard negative example, can push the model to better distinguish between fine-grained preference levels. The ideal condition that assume clear difference of the quality  between preferred and dispreferred samples, seems to take a different approach. Does this work incorporate hard negative samples, or is it constrained to cases with explicit distinctions?

3. The findings on token-level gradients provide an intriguing direction for refining preference optimization. Could the authors elaborate on any potential methods to leverage token-level distinctions more effectively within their model, perhaps as a future research direction or via fine-grained adjustments?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730666582472

thread_type:
official_review

<thread_object>
creator: reviewer_ftet

<content>
summary:
This paper considers a common failure mode of RLHF algorithms like DPO, in which, although the difference between preferred and dispreferred outputs is guaranteed to increase, there is no guarantee that the log likelihood of the preferred option will increase, or that the log likelihood of the dispreferred option will decrease.  This paper derives a general criterion that can be used to identify pairs in which either both will increase, or both will decrease.

strengths:
The theoretical analysis was enjoyable to read. The examples demonstrate the proposed criterion quite clearly.

weaknesses:
The provided examples are tremendously simplified, which is useful to demonstrate the applicability of the theoretical analysis, but which reduces confidence that the analysis will have much impact on real LLM practice.

questions:
Is there any strong argument that the problematic behavior addressed by this paper (both go up, or both go down) is important in real-world settings?

I was intrigued by the claim that the log likelihood of a dispreferred output can go up, but this seems to not be a problem that has concerned any previous authors.  The alternative approaches discussed in the paper all address the problem that the log probability of the preferred option might go down, and that outcome is the outcome suggested by the last example analyzed in the paper.
</content>
</thread_object>

</thread>
<thread>
creation date:
1731519194466

thread_type:
public_comment

<thread_object>
creator: ~huaiguang_cai1

<content>
title:
typos

comment:
Really nice paper. The theoretical results are appealing. I want to point out the typos (they do not influence the correctness of the paper but may improve readability）:

1. **Lines 197 to 198 and 262 to 263**: The first `η` in each line should include a negative sign. This adjustment aligns with the negative sign in the loss function. The specific reason is:

   ```
   f(θ - η x) ≈ f(θ) + f'(θ)(-η x) then  f(θ - η x) - f(θ) ≈  f'(θ)(-η x)
   ```

2. **Lines 211, Case 3**: The expression should be `log π_w increases more`.

Thank you for the hard work on this paper. It has truly inspired me.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732327657790

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Author's Rebuttal to Common Questions #1: New Experiments

comment:
### 1. **New Experiments**
We want to first thank all reviewers for their suggestions on our experiments.
- **More Extensive Experimental Validation**
    During rebuttal, we extend our experiments to broader datasets and models, including:
    
    - **Datasets:** TL;DR, UltraFeedback.
    - **Models:** Llama3.1-8b, Mistral-7b, Gamma2-9b, and Llama3.2-3b. 

    Among all combinations of datasets and models, empirical observations are consistently aligned with our theoretical framework. A full suite of plots can be founded here: [link to plots](https://anonymous.4open.science/r/Rebuttal_ICLR-12D0), and we summarize our results as follows:
    - **Summary of Results:** 
        - In all experiments, a synchronized increase or decrease in chosen and rejected logps is observed ([Rebuttal_Fig_1](https://anonymous.4open.science/r/Rebuttal_ICLR-12D0/Rebuttal_Fig_1:%20All_Logps.png)).
        - In all experiments, positive inner product between the gradient of the chosen logps and the gradient of the rejected logps are always observed ([Rebuttal_Fig_2](https://anonymous.4open.science/r/Rebuttal_ICLR-12D0/Rebuttal_Fig_2:%20Grad_InnerProduct.png)), and the correlation is significant: the cosine similarity is ranging from 0.3 to 0.9 across all models and datasets, validating our theoretical framework.

- **New Empirical Evidence: Why is the synchronized logp behavior harmful to alignment?**
In addition to extending our existing empirical results to more datasets and models, we provide new empirical evidence to showcase the harm that the problematic logp behaviors can exhibit to alignment.
    - **Experimental Set-up:** We run the DPO algorithm for the long suffix task: `It has a positive/negative sentiment based on my judgement.` We track the log-probabilities of chosen and rejected responses (whole sentence) as well as the log-probability on the exact token `positive` or `negative`.
    - **Results:** Plots of sentence-wise and token-wise logps can be found here [Rebuttal_Fig_3](https://anonymous.4open.science/r/Rebuttal_ICLR-12D0/Rebuttal_Fig_3:%20Logp_Behavior_Harms.png). The interpretation is: (1) chosen and rejected logps on the whole sentence are both decreasing (*Rebuttal_Fig_3(a)*), which is already observed in Section 4.2, verifying the synchronized logp behavior. (2) *Rebuttal_Fig_3(b)* shows evidence on the harmfulness of the synchronized logp behavior: the probability of the true sentiment label token `positive` or `negative` in the chosen response is indeed decreasing during DPO. This result suggests that DPO will hurt the accuracy in sentiment classification, leading to an unwanted outcome. More generally, this suggests that the synchronized updates may cause the model to learn unwanted behaviors in broader alignment tasks.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper studies a significant problem with certain preference tuning algorithms, namely, that their objective can be improved even when the probability of the preferred response is reduced or when the probability of the dispreferred response increases, when the preferred and dispreferred responses move in the same direction. The paper performs a theoretical analysis of when this happens, and points to directions for better algorithms. While there were some concerns about simplifying assumptions as well as the scope of the experiments, I believe these are reasonable for a theoretical paper; indeed, simpler settings can lead to larger insights. I believe the results in this paper may be of significant interest to the community.