On the Generalization of Preference Learning with DPO

We sincerely appreciate your positive feedback and insightful comments! We address the questions and comments below in detail.

1. The paper assumes only the last layer of the transformer is trainable

We appreciate the concerns about this assumption. This choice serves to strike a good balance between theoretical tractability and practical relevance. This decomposition allows us to rigorously capture core dynamics driving generalization, and extract valuable insights into the factors that influence preference learning. The reward dynamics involved, even in last-layer, turned out to be quite non-trivial for theoretical analysis. Jointly analyzing all components of the model leads to intractable complexity.

While the analysis is linear with respect to $W$, it's important to note that the underlying feature representations $g(x)$ still encapsulate the rich non-linear behaviors of the LLM, which is captured by our theoretical analysis. Specifically, the characterization of non-linear features is detailed in Section 4.1, which is empirically validated on real-world LLMs (Lines 428-464). Additionally, this setting is relevant in many parameter-efficient fine-tuning methods, where the feature backbone is often frozen to prevent overfitting, and in black-box scenarios where the backbone is not exposed to the end-user. Therefore, our approach remains relevant to certain practical applications, and is not entirely divorced from practices.

We acknowledge that this linearization abstracts away some complexities, such as the interactions between $W$ and the evolving feature backbone. However, our empirical results validate that this approximation remains effective in capturing the behavior of real-world systems, such as LLaMA-2 fully fine-tuned with preference data, where the feature-map is allowed to change. As shown in Lines 466-474, the results are consistent with our theory, suggesting that the insights gained from our analysis extend to more general settings. Moreover, this decomposition provides a clear and interpretable starting point for understanding reward dynamics and generalization, which can be expanded in future work to include interactions with $g(x)$ or even nonlinear changes in the weight matrix.

For these reasons, we believe our framework provides a foundational step in analyzing preference learning, focusing on the most impactful and mathematically tractable components. We hope this clarification highlights why this decomposition is both a reasonable and productive choice for our theoretical contribution.

2. Multi-token case.

That's a great question. Our current analysis focuses on the single-token case as a foundational step, as it provides a more tractable framework to rigorously analyze the generalization error. Extending these results to the multi-token case is an important and challenging open problem.

That said, we believe similar meaningful results can be established for the multi-token case. As outlined in Section 4.3, we already analyze the reward dynamics for multi-token scenarios. Specifically, when the reward dynamics for a given example can be shown to lead to a consistently positive or negative reward, we can generalize this result to samples with similar embeddings, ensuring that preferences learned for one example apply reliably to similar cases. This suggests that the insights from our single-token framework are not only foundational but also extendable to capture the dynamics of multi-token generation, which we aim to address more comprehensively in future work.

3. Other fine-tuning settings

Thank you for the insightful question. We believe that the analysis can generalize to other fine-tuning settings, such as those involving low-rank updates like LoRA, as well as to other preference objectives, such as IPO. For example, if we wanted to consider any other objective function $f(r)$ where $r$ is the reward margin, the only part of the original dynamics, $\tau \dot{r_j} = \frac{1}{N}\sum_{i=1}^N \beta^2 \sigma(-r_i)C(x_i, x_j)$ that would change is that $\sigma(-r_i)$ would instead be $-f'(r_i)$. In the case of LoRA, the weight matrix updates are restricted to low-rank subspaces, which could introduce additional constraints on the reward margin dynamics. These constraints could be modeled by incorporating the rank limitation directly into the theoretical analysis. This would introduce non-linearities to the reward dynamics and allow us to study how the low-rank nature of the updates impacts generalization guarantees. We view this as an exciting opportunity for future work and are optimistic that our current analysis can provide a strong foundation for exploring these specialized settings.

We sincerely appreciate your feedback and insightful comments! We are encouraged that you recognize the importance and novelty of our research focus. We address the questions and comments below in detail.

1. Unclear assumption on LM, concern on the stochasticity

We would like to clarify that our abstraction of the model output does not alter the autoregressive nature of LLMs nor prevent stochasticity in learning. We are happy to explain this in detail below.

First, modeling $Wg(x)$ is fully compatible with the autoregressive nature of LLMs and is consistent with both single-token prediction and multi-token predictions (Section 4.3). Here $g(x)$ encodes the non-linear, autoregressive features extracted by the LLM, and $W$ maps these features to the output logits (with the same dimensionality as the number of tokens). For example, in multi-token predictions, as described in Section 4.3 (Lines 340-367), the autoregressive nature is preserved by conditioning the prediction of each token $y_t$ on the preceding tokens $y_1, \dots, y_{t-1}$:

where $p(y_t | y_1, \dots, y_{t-1})$ is based on $\text{softmax}(Wg(x))$, with $g(x)$ corresponding to the embedding at the current step, and the tokens lying in the vocabulary space $\mathcal{V}$. This abstraction aligns fully with the operation of LLMs in practice.

Second, $W$ is updated stochastically, allowing adaptation to the preference data distribution. By treating $W$ as a trainable component optimized via stochastic gradient descent, our abstraction captures the inherent stochasticity in learning while focusing on key factors influencing preference learning. This allows us to isolate the dynamics of preference learning while maintaining analytical tractability. Importantly, the responses in the dataset are fixed, and we update the output distribution directly with respect to these fixed preferred/not-preferred responses, rather than modeling a sampled generation.

Lastly, it's important to note that the underlying feature representations $g(x)$ still encapsulate the rich non-linear behaviors of the LLM, which is captured by our theoretical analysis. Specifically, the characterization of non-linear features is detailed in Section 4.1, which is empirically validated on real-world LLMs (Lines 428-464). Additionally, this setting is relevant in many parameter-efficient fine-tuning methods, where the feature backbone is often frozen to prevent overfitting, and in black-box scenarios where the backbone is not exposed to the end-user. Therefore, our approach remains relevant to certain practical applications, and is not entirely divorced from practices.

Our empirical results further validate that our theory remains effective in capturing the behavior of real-world systems, such as LLaMA-2 fully fine-tuned with preference data, where the feature map is allowed to change. As shown in Lines 466-474, the results are consistent with our theory, suggesting that the insights gained from our analysis extend to more general settings. Moreover, this decomposition provides a clear and interpretable starting point for understanding reward dynamics and generalization, which can be expanded in future work.

For these reasons, we believe our framework provides a foundational step in analyzing preference learning, focusing on the most impactful and mathematically tractable components. In the revised manuscript, we will expand on these points and ensure the connection between our abstraction and the autoregressive nature of LLMs is explicitly clarified. Thank you for your thoughtful feedback.

2. Presentation of Section 3.2

Thank you for the careful read and helpful feedback. While we dedicated significant effort to polishing our paper, we agree that Section 3.2 could benefit from additional work. We will provide a more detailed illustration of Equation (6) as a continuous approximation to the SGD in Equation (4). As suggested, we will carefully revise this section to better articulate the underlying assumptions, conditions, and the model class being analyzed, and ensure the presentation is precise and accessible. Your comments are greatly appreciated!

I would like to thank the authors for their detailed responses. Despite this, my concerns are not addressed.

assumptions on LM in Sc3.2

I am still puzzled about $Wg(x)$. For a theorectically-oritended paper, a clear explanation of theorectical assumptions and why these assumptions should be of crucial importance, and again Sc3.2 is hard to follow with many elements not explained. Can you update Sc 3.2 with more explanations on what are the dimensions, what can this assumption represent?

Although the authors do admit that Equation (6) are some continuous approximations, but necessary details are still not provided.

Theorem 4.2's bound looks to be vacuous

For the updated bound, I am also still confused so far. When the $\epsilon$ is small, the exponential term is close to 1, this bound is bounded by like $K\cdot Q^2$, which grows quadractically with the sample size $Q$. I do not fully understand why this provides a sharper bound, and could provide meaningful pratical guidance.

Overall, given the current manuscript (which authors haven't updated the main contents yet as far as I can see), although I like the interesting topic and this research direction the paper is focusing on, I think it's hard for readers who may come from practitioners of alignment fields to understand the whole logic of this paper and to seek relevant insights provided.

In the rebuttal, the authors have mentioned the updates to be done in the revised manuscript. I have been waiting to see it, but haven't seen that yet, for which I would like to kindly remind that ICLR allows you to update your manuscript towards a better version. Since the current manuscript still does not address my concerns, I will keep my score.

We sincerely appreciate your constructive feedback and insightful comments! We are encouraged that you recognize our contribution of offering rigorous theoretical foundation for preference learning and bridging theory and practice. We address the questions and comments below in detail.

1. The theoretical analysis decomposes the model output into the product of W g(x), whereas we only consider the dynamic of W. It is reduced to a linear analysis.

We appreciate the concerns about this assumption. This choice serves to strike a good balance between theoretical tractability and practical relevance. This decomposition allows us to rigorously capture core dynamics driving generalization, and extract valuable insights into the factors that influence preference learning. The reward dynamics involved, even in the last-layer, turned out to be quite non-trivial for theoretical analysis. Jointly analyzing all components of the model leads to intractable complexity.

While the analysis is linear with respect to $W$, it's important to note that the underlying feature representations $g(x)$ still encapsulate the rich non-linear behaviors of the LLM, which is captured by our theoretical analysis. Specifically, the characterization of non-linear features is detailed in Section 4.1, which is empirically validated on real-world LLMs (Lines 428-464). Additionally, this setting is relevant in many parameter-efficient fine-tuning methods, where the feature backbone is often frozen to prevent overfitting, and in black-box scenarios where the backbone is not exposed to the end-user. Therefore, our approach remains relevant to certain practical applications and is not entirely divorced from practices.

Lastly, our empirical results validate that this modeling remains effective in capturing the behavior of real-world systems, such as LLaMA-2 fully fine-tuned with preference data, where the feature map is allowed to change. As shown in Lines 466-474, the results are consistent with our theory, suggesting that the insights gained from our analysis extend to more general settings. Moreover, this decomposition provides a clear and interpretable starting point for understanding reward dynamics and generalization, which can be expanded in future work to include interactions with $g(x)$ or even nonlinear changes in the weight matrix.

For these reasons, we believe our framework provides a foundational step in analyzing preference learning, focusing on the most impactful and mathematically tractable components. We hope this clarification highlights why this decomposition is both a reasonable and productive choice for our theoretical contribution.

2. Clarification on the Gaussian distribution

Thank you for raising this important question. The assumption is motivated by the fact that Gaussian distributions are sufficient to capture the variance of embeddings while enabling a more straightforward and tractable analysis. This choice allows us to focus on deriving meaningful generalization bounds and reward dynamics without introducing unnecessary complexity.

To support this assumption, we have conducted empirical verification based on the Llama-2 model. Specifically, we observe that across different personas, the average norm of the covariance matrix is 0.058, and the average squared distance from the mean is 0.227. These empirical findings suggest that the Gaussian assumption provides a reasonable approximation of the underlying structure of preference data. While the covariance structure could be further refined, we opted for a single parameter $v$ to enhance interpretability.

We also recognize that real-world preference distributions might deviate from strict Gaussianity. Future work could explore more flexible models that better capture the nuances of preference data, such as non-parametric clustering methods or other distributional assumptions. For this work, the Gaussian mixture model serves as a practical and reasonable starting point, enabling us to make progress on the challenging problem of understanding generalization in preference learning.

We sincerely appreciate your feedback and insightful comments! We address the questions and comments below in detail.

1. Clarification on pre-trained LLM

We fully agree with your viewpoint on considering the fact that LLM is pre-trained. We indeed considered pre-trained LLM in our theoretical analysis, making our results practically relevant. In particular, we characterize our data distribution based on pre-trained LLMs and how DPO optimizes with respect to this data distribution. This characterization is detailed in Section 4.1, which is empirically validated on real-world, pre-trained LLMs, such as LLaMA-2, and alignment datasets designed to capture the challenges posed by diverse human preferences (Lines 428-464). Specifically, we characterize the structure of preference data, comprising of a shared component and concept-specific directions, to derive stronger guarantees compared to scenarios with less structured preference data. Our empirical validation demonstrates that our theory applies to pre-trained models undergoing fine-tuning for alignment. We do observe the empirical trends align with the generalization bounds, which validates the correctness of our theory.

It is worth noting that pretraining enables data efficiency by structuring embeddings in ways that facilitate subsequent fine-tuning. Our results align with this intuition, as demonstrated in our generalization guarantees. Pretraining contributes to these factors by embedding high-quality priors, implicitly enhancing alignment performance with less preference data. While our framework does not directly model the pretraining process, the theoretical results are consistent with and supported by the benefits of pretraining.

Finally, we acknowledge that certain aspects of pretraining, such as its specific inductive biases, could further refine our framework. However, this does not diminish the utility or relevance of our analysis. Instead, it underscores the generality and extensibility of our theoretical framework. For instance, future work could incorporate additional assumptions about pretrained representations, such as their cluster structure or alignment with preference datasets, to derive more specialized results. In this work, our focus has been on building a theoretical foundation that can be applied broadly, allowing capturing the dynamics of alignment for both well-pretrained and less-pretrained models.

2. For LLMs, I think the traditional generalization bound analysis doesn't seem particularly interesting (i.e., in-distribution)

You raise an insightful point. This is something we have also thought about deeply.

Although LLMs exhibit impressive generalization capabilities, enabling them to handle many cases they have not explicitly seen, they also frequently fail in situations where the data is similar to what they have encountered. For preference learning, this inconsistency underscores the importance of understanding when and why LLMs behave as expected.

A critical first step is determining whether the model can reliably distinguish what is preferred from what is not. This is the central question we address in our work, and it is foundational to ensuring alignment in real-world applications. For example, even in settings that should theoretically be straightforward, such as in-distribution preference prediction for HH-RLHF-trained models, test accuracy can be only 75%. This highlights that even within the same distribution, generalization is far from perfect or well-understood. Scaling to more complex preference distributions makes generalization even harder, as our theoretical framework indicates. These observations reinforce the importance of focusing on this classical setting and ensuring it is thoroughly understood before extending the analysis to more complex out-of-distribution scenarios.

Lastly, our work provides a deeper understanding of multi-token generation by analyzing reward dynamics, which further contributes to unraveling why and when LLMs align with expected behaviors. Future work could leverage our theoretical framework as a foundation to address the broader and impactful question of OOD generalization in preference learning. We recognize that they represent expansive and complex problems that cannot be addressed comprehensively in the scope of a single work. This makes it all the more important to take measured steps, starting with a thorough understanding of in-distribution generalization as a foundation.

3. Gap between DPO and on-policy: This may limit the real-world impact of follow-up works, but I think it's a relatively minor issue here.

Thank you for pointing this out, and we will keep this in mind!

Dear reviewers,

We understand it’s been an overwhelming season, and everyone is juggling many responsibilities. This makes us especially grateful for the time and attention you’ve given to our submission despite these challenges.

We hope our responses have addressed your concerns and clarified the key points of our paper. As a theory-focused work, we have made every effort to carefully validate our assumptions and theoretical results based on real-world LLMs. Closing the gap between theory and practice is a non-trivial moonshot, and our work represents a significant stride toward narrowing it. We believe this contribution will inspire a line of future works to deepen the theoretical understanding of alignment, which is urgently needed for the community.

If there are any remaining questions or suggestions, we welcome your feedback and will address them promptly. Thank you again for your thoughtful insights and for supporting this important area of research.

Sincerely,

Authors