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

Thank you for acknowledging the importance of our work:

"Identifying conditions under which preference optimization can backfire is important for real-world use and may provide a relatively underdeveloped area of research."

and appreciating our work for its being theoretically grounded and interesting. We address your concerns individually as follows.

Weakness 1: This work calls for broader empirical validation than it has. (1) Other variants of the given LLMs (e.g., LLaMa) could have been explored. (2) The proposed method may be evaluated across diverse tasks such as summarization, dialogue, and instruction-following.

Response to Weakness 1: Thanks for your suggestions on experiments.

• To resolve your concern(1), we extended our experiments to cover datasets: TL;DR, UltraFeedback and models: Llama3.1-8b, Llama3.2-3b, Mistral-7b, Gamma2-9b. Among all combinations of datasets and models, empirical observations are consistently aligned with our theoretical framework, here is the link to plots of new results and more interpretations can be found in the "New Experiments"(link) section in our global rebuttal.
• For concern(2), we want to kindly remind you that our paper did not propose new methods thus requiring method evaluation may not be valid. The toy synthetic settings including the analysis and experiments we provided in Sec 4 is for building up our general intuition on why the gradient inner product can be large and thus causes the problematic behaviors in chosen and rejected logps during training, rather than for motivating new methods.

Weakness 2: The analysis of Sec 4 focuses on synthetic, controlled scenarios. The kinds of complexities exist that in actual, diverse language data should be addressed.

Response to Weakness 2: We emphasize that the purpose of Sec 4 is for building up a high-level intuition on why the gradient inner product can be large during alignment. The alignment process on real-world preference data is too complicated to have theoretical analysis, where the variety of language data is a main source of complications. Though being simple, the current data configuration captures fundamental problems in many preference datasets: between the chosen and rejected responses, many times, most of the tokens share similar meanings and only a few tokens carry distinguishing information to separate out chosen and rejected responses. Theoretically extend to some less constrained settings can be an interesting future direction. One may consider the setting where blocks of their tokens are paraphrase to each other that share similar meanings instead of the exact same tokens. In "Clarification on Theoretical Simplifications"(link) section in our common rebuttal, you can find more explanations of our theory.

Dear Reviewer ai3q,

We want to remind you that our rebuttals have been posted for a while, it's now approaching the end of the disscussion period. we are eager to hear your thoughts and any further feedback you might have.

Dear Reviewer ai3q,

As the discussion period is coming to a close, we would greatly appreciate it if you could let us know whether our rebuttal has fully addressed your concerns and questions. Regarding the two concerns and three questions listed in the review, we provided point-by-point responses, including:

• Added Extensive Experiments: the key phenomenon (synchronized shifts of chosen and rejected log-probabilities) investigated in this paper is widely observed for models: Llama3.1-8b, Mistral-7b, Gamma2-9b, Llama3.2-3b; and commonly used RLHF datasets: TL;DR, UltraFeedback. And the cause we identified (correlated gradient inner products) is verified in all settings.
• Clarifying and Justifying Theoretical Simplifications: We have clarified that the contributions of our theory are two-folded. The first half reveals the cause of the phenomenon, in which no assumptions are made. The second half is intended for gaining intuition, thus some simplifications are made. For this part, we explained how our theoretical settings capture the main characteristics of RLHF data, and discussed how our theoretical findings are generalized empirically in our paper.
• Actionable Items: We have proposed two actionable items (i.e., new learning objectives) in solving the simultaneous shifts problems using our theoretical framework.
• New Experiment to Show the Risk of the Identified Pitfall: We have added new experimental results to showcase how the synchronized shifts can potentially result in unreliable language model behavior in safety-critical alignment tasks.
• Token-Level Information in Alignment: We have provided additional discussion on how token-level information can be considered for language model alignment, and how our empirical results are pointing to this.

Thank you for your time and consideration.

Best regards,
Authors of Submission 13302

We thank you for appreciating the importance of our work:

"This paper addresses a fundamental under-specification problem in margin-based preference optimization -- this is an important contribution given how widely RLHF is used."

We address your concerns individually as follows.

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

Our Response: Thanks for your suggestions on presentation. We have made revisions accordingly, where we (1) simplified the math presentation in Sec 3.2 for general objectives; (2) revised the interpretation of experiment plots including Fig 2, Fig 3 and Fig 4; and (3) strengthened our explanation on the connection between experiments and theory in Sec 3.3 and Sec 4.2. For a More detailed summary of our paper revision, please refer to the "Paper Revisions"(link) section in our common rebuttal and feel free to check out our revised paper.

Limited Experimental Validation: Experiments are conducted only on TL;DR dataset in addition to a simple sentiment classification task, and only on two main models (Mistral 7B and Llama-3 8B).

Our Response: Thanks for your suggestions on experiments. We extended our experiments to cover datasets: TL;DR, UltraFeedback and models: Llama3.1-8b, Llama3.2-3b, Mistral-7b, Gamma2-9b. Among all combinations of datasets and models, empirical observations are consistently aligned with our theoretical framework, here is the link to plots of new results and more interpretations can be found in the "New Experiments"(link) section in our common rebuttal.

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.

Our Response:

• We want to clarify that the synthetic settings in Section 4.1 is for building up our general intuition on why the gradient inner product can be large and thus causes the problematic behaviors in logps. We acknowledge that this setting is toy but the insight obtained from our analysis is non-trivial. In the "Clarification on Theoretical Simplifications" (link) section in our common rebuttal, you can find more explanations on why we made some theoretical simplifications and a detailed list of the practical implications of our theory. Furthermore, in Section 4, our theoretical findings are verified when fine-tuning full parameters of GPT-2, extending our findings beyond simplified model settings (learnable last layer/logits).
• For extension of single-token results to multi-token case: we have revised Corollary 2 and provided a rigorous proof for it in Appendix, feel free to check it in our revised version.

Dear Reviewer UXg4,

As the discussion period is coming to a close, we would greatly appreciate it if you could let us know whether our rebuttal has fully addressed your concerns and questions. Regarding the two concerns and three questions listed in the review, we provided point-by-point responses, including:

• Paper Revision: We have revised the paper to provide easy-to-follw theoretical results, and a clear discussion on the connection between theoretical results and empirical investigation.
• Added Extensive Experiments: the key phenomenon (synchronized shifts of chosen and rejected log-probabilities) investigated in this paper is widely observed for models: Llama3.1-8b, Mistral-7b, Gamma2-9b, Llama3.2-3b; and common RLHF datasets: TL;DR, UltraFeedback. And the cause we identified (correlated gradient inner products) is verified in all settings.
• Clarifying and Justifying Theoretical Simplifications: We have clarified that the contributions of our theory are two-folded. The first half reveals the cause of the phenomenon, in which no assumptions are made. The second half is intended for gaining intuition, thus some simplifications are made. For this part, we explained how our theoretical settings capture the main characteristics of RLHF data, and discussed how our theoretical findings are generalized empirically in our paper.

Thank you for your time and consideration.

Best regards,
Authors of Submission 13302

Response to Other Questions:

• Some papers focus on distributional framework, which insist the model to encounter a broader range of sample types to cover the distribution's full support. Could the authors discuss the potential impacts of their idealized setup compared to a distributional approach?

Response: Thanks for pointing us to the line of research on distributional RLHF. In our understanding, margin-based objectives are aligning models to a target distribution controlled by some scalar preference, where the scalar preference is unknown but to learn from preference data. Further constraining that the logp of positive sample to increase and the logp of negative sample to decrease (the idealized setup) implies a constraint that the positive samples should have a positive score and the negative samples score negative under the scalar preference function.

• 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?

Response: Our results in Sec 3 apply to the setting with hard negative samples, for example, in our new experiments, our results are validated on UltraFeedback where positive and negative samples may not have clear distinctions. Results in Sec 4 may not apply to hard negative samples.

• Could the authors elaborate on any potential methods to leverage token-level distinctions?

Response: Sure. The following token-wise objective could potentially leverage token-level distinctions to improve alignment:

$-\log\sigma \left( \sum_{i=1}^L \mathbb{I}\set{u_w^i \geq r }\log \frac{\pi_{\theta}(y^i_w|x, y_w^{<i})}{\pi_\text{ref}(y^i_w|x, y_w^{<i})}- \mathbb{I}\set{u_l^i \geq r } \log \frac{\pi_{\theta}(y^i_l|x, y_w^{<i})}{\pi_\text{ref}(y^i_l|x, y_l^{<i})} \right) + \eta \left(\|\mathbb{I}\set{u_w \geq r}\|_1 + \|\mathbb{I}\set{u_l \geq r}\|_1\right),$

where $\eta\in \mathbb{R}_+, r \in \mathbb{R}$ are hyper-parameters and $u_w \in \mathbb{R}^L, u_l\in \mathbb{R}^L$ are learnable weights depending on $(x,y_w^{<i}), (x,y_l^{<i})$ respectively, interpreted as the confidence in considering token $i$ significant. The loss is inspired by sparsity-related ideas (e.g., LASSO), where the learnable masks $\mathbb{I}\set{u^i \geq r }$ ideally pick out the significant tokens in each response that enlarge the margin. The $\ell_1$ regularizer on the token-wise mask imposes sparsity on it.

Thank you for recognizing the originality and fresh insights presented in our paper. We appreciate your praise for our framework as a well-qualified, mathematically grounded, and comprehensive approach to understanding alignment failures.

Weakness 1: The paper centers primarily on DPO in Sections 3.1 and 4, with limited coverage of other margin-based optimization algorithms.

Response to Weakness 1:

• The purpose of Sec 3.1 (deriving the condition for DPO) is serving as a warm-up case to smoothly ramp up to the general class of margin-based objectives. Our main focus of Sec 3 is surely general margin-based objectives, rather than DPO only, as we further use the proposed condition to explain the different training dynamics of different objectives in Sec 3.3. In our revision, we improved the presentation of these sections to better balancing DPO and other objectives.
• As for Sec 4, we intentionally focused on DPO in order to establish fundamental understanding of the gradient inner product in language data. This is because the gradient inner product between chosen and rejected responses, is more essentially related to the configuration of data and model, rather than specific alignment objective. Focusing on the DPO algorithm gives theoretical simplicity to isolate the problem essentials from subtle algorithmic designs.

Weakness 2: The experimental results in Section 3.3 primarily observe phenomena regarding preferred and dispreferred response probabilities increasing or decreasing together. (1) This findings 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. (2) Furthermore, the experiment relies on a single dataset, limiting the depth of empirical validation.

Response to Weakness 2:

• For your concern (1), we want to remind you that we have plots to verify our proposed cause (the condition you are referring to) for the logps behavior, Fig 5 in Appendix. From Fig 5, we can see the inner product between the gradients of preferred logp and the gradients of preferred dispreferred logp is positive with significant correlation. This result validates our proposed cause and its link to the correlated gradients, which is not covered in any previous literature.
• For your concern (2), we extended our experiments to cover datasets: TL;DR, UltraFeedback and models: Llama3.1-8b, Llama3.2-3b, Mistral-7b, Gamma2-9b. Here is the link to our new experiments. More details can be found in the "New Experiments"(link) section in our common rebuttal.

Weakness 3: Presentation of empirical observations and important methodological details: (1) clarify how log probabilities are computed in Figure 1 (e.g., whether they are averaged over the training samples). (2) include additional measures such as standard deviations or confidence intervals to improve the figure's reliability. (3) providing a additional figure for dedicated to empirical results (Section 3.3). (4) It would be better to change the interval of the grid to clearly show the trend of margin in Figure2.

Response to Weakness 3 Thanks for suggestion, we have polished our figures in our revision, including: (1) clarifying how the logp plots are computed and averaged on the evaluation data split in the caption of Fig 1. and (4) zooming in the grid range so that the trend in logps are more clearly shown. We are still working on (2), (3) and considering what's the best way to present, will update in our draft after we figure it out.

Thank you for your positive evaluation on our paper and your supportive comment that "the theoretical analysis was enjoyable to read, the examples demonstrate the proposed criterion quite clearly".

Weakness 1: 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.

Response to Weakness 1 Could you please clarify on what are the examples you are referring to here, the examples of margin-based objectives in Sec 3, or, the sentiment task we used in Sec 4 to validate our analysis on gradient inner product? Either way, we want to resolve your concerns by clarifying our theoretical analysis. Our analysis has two folds:

• The starting point of our theory is to answer question (1): What causes the both-go-up/both-go-down behavior of chosen and rejected logps. The analysis in this part applies to and is supported by non-simplified realistic alignment practice (Sec 3.2).
• In Sec 4, the purpose is to build up our intuition on why the gradient inner product can be large, thus we analyze toy synthetic setting. Though being simple, the current data configuration captures a fundamental characteristic of many preference datasets: between the chosen and rejected responses, many times, most of the tokens or share similar meanings and only a few tokens carry distinguishing information to separate out chosen and rejected responses. In "Clarification on Theoretical Simplifications"(link) section in our common rebuttal, you can find more explanations of our theory and assumptions.

Question 1: 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?

Response to Question 1: Yes. When the both-go-up/both-go-down logps occurs during alignment, the model may end up with having learned unwanted/wrong behaviors. Though there isn't much previous literature addressing this failure and linking it to the synchronized logps changes, we add a new experiment on the sentiment alignment task to reproduce the failure and show that the accuracy of sentiment prediction is decreasing as the chosen and rejected logps are both decreasing during DPO (Rebuttal_Fig_3). More details for this experiment can be found in the "New Experiments"(link) section in our common rebuttal.

Question 2: The claim that the log likelihood of a dispreferred output can go up is intriguing, but this seems to not be a problem that has concerned any previous authors. The paper mainly address the alternative: the log probability of the preferred option might go down in theory.

Response to Question 2:

• Based on our gradient condition, when the inner product $\langle \nabla \log \pi_w, \nabla \log \pi_l \rangle$ is large, whether the logps of preferred and dispreferred outputs are both going up or both going down depends on which one in $\|\nabla \log \pi_w\|^2$ and $\|\nabla \log \pi_l\|^2$ is greater than the other ($w$:preferred, $l$:dispreferred}). When $\|\nabla \log \pi_w\|^2$ is larger, then both logps will go up. Actually, in our Theorem 3 where both logps are proven to go down, the assumption on logits plays a key role in making sure $\|\nabla \log \pi_l\|^2$ is greater. If we make assumption on logits the other way around, then both logps will go up. The reason we adopted the current assumption is because after SFT, it is more likely to happen that: $s_{w,i}[j_i^*] \geq s_{l,i}[j_i^*]$ and $s_{w,i}[j] \leq s_{l,i}[j]$ for $j \neq j^*_i$ where $j^*_i$ is the $i$-th token in chosen response, which is also adopted by [1].
  [1] Pal et al. Smaug: Fixing failure modes of preference optimisation with dpo-positive.

Thanks for your comment and coming up with the simple vs complex examples. We understand you have concern on: when the pair of responses are getting more different from each other, will the pair be less correlated so that the correlation is not probelmetic?

We have new experiments showing that the gradient correlation is still significantly large (Rebuttal Fig2) and thus causes the problematic behavior in log-probabilities (Rebuttal Fig1), on the popular RLHF dataset Ultrafeedback where the responses are more different from each other in each pair, to give you an example:

• Prompt: "how can I develop a habit of drawing daily"

• Chosen: "Developing a daily habit of drawing can be challenging but with consistent practice and a few tips, it can become an enjoyable and rewarding part of your daily routine. Here are some strategies to help you develop the habit of drawing daily: 1. Set a specific time:... 2. Set a specific duration:... 3. Start small and simple:... 4. Use a variety of tools and mediums:...5. Take breaks and rest:...6. Challenge yourself:... 7. Track your progress:..."

• Rejected: "As an AI language model, I cannot personally develop habits for you. But, here are some tips for developing a habit of drawing daily: 1. Start small:... 2. Set a schedule:...3. Make it fun:...4. Use resources:... 5. Surround yourself with inspiration:..."

In the example above, the rejected sample appears to have a negative attitude and less helpful compared to the chosen one, and the tips given by the two responses diverge. We want to make two comments regarding your example and conjecture:

• We agree with your conjecture that the gradient can be less correlated when the pair of responses getting extremely and contrastively different, like in your example with toxic response. But it may not be the case of how RLHF data are colllected in reality: chatgpt collects preference data bey deploying close model variants, so the actual pair being collected is more like the case you tested under the prompt "did people really land on the moon?" then the toxic example. This procedure is documented in [1] and we believe that this practice has become a convention for developing other LLMs.
• Empirically, we show that for common RLHF datasets, even when the data pairs have less lexical overlaps, we still observe positively correlated gradients and margin-based objectives have synchronized likelihood shifts. Our analysis in Sec 4 based on toy settings should not be considered limiting the scenarios to which the pitfall of margin-based RLHF we identified in this paper applies.

[1] Stiennon et al.. "Learning to summarize with human feedback."

Once again, the pitfall (synchronized likelihood shifts, high gradient correlation) wildly applies to the real-world practice, the purpose of our analysis in Sec 4 is seeking for intuition on why the high gradient correlation appears in RLHF data and to generalize the analysis is an intersting future direction. Hope it can resolve your concern and let us know if you have further questions.

We appreciate your positive comments on our paper and have corrected the typo accordingly. For presentation clarity, we have moved part of the content to the Appendix.

2. Paper Revisions for Better Presentation

We want to first thank reviewer UXg4 and 5B6v for their suggestion on our paper presentation. During rebuttal, we revise our paper to improve presentation by taking in those suggestions. Our revision is uploaded, and here we give a summary of our major revisions:

• Simplified math presentation. We optimized the presentation of math derivations by refactoring Sec 3.1 and Sec 3.2. Specifically, (1) we highlight and term our key finding "gradient entanglement": Recall our key finding characterized by equations $\Delta \log \pi_w \approx C \cdot \left( \|\nabla \log \pi_w \|^2 - \langle \nabla \log \pi_w , \nabla \log \pi_l \rangle \right),$ $\Delta \log \pi_l \approx C \cdot \left( \langle \nabla \log \pi_w , \nabla \log \pi_l \rangle-\|\nabla \log \pi_l \|^2 \right), (*)$ which shows that the change in the chosen probability is coupled with the gradient of the rejected one, and vice versa. So we name the inherent flaw of margin-based losses as "gradient entanglement." (2) We move some less important derivations to the Appendix to improve clarity and conciseness.
• Practical implications with more clarity. In Sec 3.3 and Sec 4.2, we revised the interpretation of empirical results in Fig 2, Fig 3 and Fig 4, as well as strengthened our explanation on the connection between experiments and theory.
• Polished plots to include more details. Fig 1 now has an update caption to include implementation details on how the logps are tracked; Fig 2 has been zoomed in to better show the trend of curves.
• Rigorous extension of single-token results to multi-token. We revised Corollary 2 and provided rigorous proof for it.