Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Robust LLM Alignment via Distributionally Robust Direct Preference Optimization

Zaiyan Xu, Sushil Vemuri, Kishan Panaganti, Dileep Kalathil, Rahul Jain, Deepak Ramachandran

TL;DR

The paper tackles distribution shift in aligning LLMs to human preferences by formulating distributionally robust direct preference optimization (DPO) with two variants: Wasserstein DPO (WDPO) and KL-divergence DPO (KLDPO). It provides finite-sample, strong-convexity-based guarantees for robust policy parameter convergence at rate $O(n^{-1/4})$ under a log-linear policy, and offers tractable gradient-based algorithms that integrate into existing LLM alignment pipelines. Empirically, WDPO and KLDPO outperform standard DPO under various preference shifts, demonstrated across Emotion Alignment, ArmoRM multi-objective alignment, and OpenLLM Leaderboard tasks, using multiple model scales. The work delivers a principled approach to maintaining alignment under real-world distributional changes, with theoretical and practical insights into robustness, scalability, and sample efficiency.

Abstract

A major challenge in aligning large language models (LLMs) with human preferences is the issue of distribution shift. LLM alignment algorithms rely on static preference datasets, assuming that they accurately represent real-world user preferences. However, user preferences vary significantly across geographical regions, demographics, linguistic patterns, and evolving cultural trends. This preference distribution shift leads to catastrophic alignment failures in many real-world applications. We address this problem using the principled framework of distributionally robust optimization, and develop two novel distributionally robust direct preference optimization (DPO) algorithms, namely, Wasserstein DPO (WDPO) and Kullback-Leibler DPO (KLDPO). We characterize the sample complexity of learning the optimal policy parameters for WDPO and KLDPO. Moreover, we propose scalable gradient descent-style learning algorithms by developing suitable approximations for the challenging minimax loss functions of WDPO and KLDPO. Our empirical experiments using benchmark data sets and LLMs demonstrate the superior performance of WDPO and KLDPO in substantially improving the alignment when there is a preference distribution shift.

Robust LLM Alignment via Distributionally Robust Direct Preference Optimization

TL;DR

The paper tackles distribution shift in aligning LLMs to human preferences by formulating distributionally robust direct preference optimization (DPO) with two variants: Wasserstein DPO (WDPO) and KL-divergence DPO (KLDPO). It provides finite-sample, strong-convexity-based guarantees for robust policy parameter convergence at rate under a log-linear policy, and offers tractable gradient-based algorithms that integrate into existing LLM alignment pipelines. Empirically, WDPO and KLDPO outperform standard DPO under various preference shifts, demonstrated across Emotion Alignment, ArmoRM multi-objective alignment, and OpenLLM Leaderboard tasks, using multiple model scales. The work delivers a principled approach to maintaining alignment under real-world distributional changes, with theoretical and practical insights into robustness, scalability, and sample efficiency.

Abstract

A major challenge in aligning large language models (LLMs) with human preferences is the issue of distribution shift. LLM alignment algorithms rely on static preference datasets, assuming that they accurately represent real-world user preferences. However, user preferences vary significantly across geographical regions, demographics, linguistic patterns, and evolving cultural trends. This preference distribution shift leads to catastrophic alignment failures in many real-world applications. We address this problem using the principled framework of distributionally robust optimization, and develop two novel distributionally robust direct preference optimization (DPO) algorithms, namely, Wasserstein DPO (WDPO) and Kullback-Leibler DPO (KLDPO). We characterize the sample complexity of learning the optimal policy parameters for WDPO and KLDPO. Moreover, we propose scalable gradient descent-style learning algorithms by developing suitable approximations for the challenging minimax loss functions of WDPO and KLDPO. Our empirical experiments using benchmark data sets and LLMs demonstrate the superior performance of WDPO and KLDPO in substantially improving the alignment when there is a preference distribution shift.

Paper Structure

This paper contains 35 sections, 22 theorems, 107 equations, 5 figures, 4 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Related Work
  3. Preliminaries
  4. Distributionally Robust DPO
  5. Distributionally Robust DPO
  6. Theoretical Analysis
  7. Estimation Error for WDPO
  8. Estimation Error for KLDPO
  9. Tractable (Approximate) Algorithms
  10. Experiments
  11. Experimental Setup
  12. Experiment Results
  13. Conclusions
  14. Acknowledgments
  15. Useful Technical Results
  16. ...and 20 more sections

Key Result

Theorem 1

Let $\delta\in(0,1)$. With probability at least $1-\delta$, we have where $\gamma=\frac{\beta^2e^{4\beta B}}{(1+e^{4\beta B})^2}$ and $K=\lvert \log\sigma(-4\beta B) \rvert$, $\lambda$ is the regularity number defined in assum:uniform-data-cov-assumption.

Figures (5)

  • Figure 1: If the training population predominantly uses preference model 1 (P1), a non-robust RLHF/DPO model will favor Completion 1 (C1). However, deploying this model to a test population that prefers model 2 (P2), which favors Completion 2, leads to poor performance. Our distributionally robust DPO (WDPO/KLDPO) addresses this by optimizing across an uncertainty set of preference models, ensuring robust performance under preference shifts.
  • Figure 2: DPO, WDPO, and KLDPO in Emotion Alignment. Models are trained on preferences derived from convex (left two plots) and geometric (right two plots) mixtures of anger and fear objectives from the Emotion dataset saravia-etal-2018-carer. To simulate preference shift, evaluation is performed at mixing coefficients $\alpha \neq \alpha_o$, where $\alpha_o = 0.1$ is used during training. See \ref{['sec:experiment-setup']}.
  • Figure 3: DPO, WDPO, and KLDPO in ArmoRM Multi-objective Alignment. LLaMA-3.2-1B-Instruct models are trained on preferences derived from three equally weighted objective pairs: (1) Ultrafeedback-Truthfulness and Helpsteer-Complexity, (2) Ultrafeedback-Helpfulness and Helpsteer-Coherence, and (3) Helpsteer-Correctness and Helpsteer-Helpfulness (left to right plots). We train all models for 4 epochs. To simulate preference shift, models are evaluated on five individual objectives, Helpsteer-Helpfulness, Helpsteer-Correctness, Helpsteer-Coherence, Ultrafeedback-Honesty, and the overall ArmoRM score, three of which were not used during training.
  • Figure 4: Evaluation of WDPO, KLDPO and DPO on $r^*_{\mathrm{convex}}(\alpha)$ in ArmoRM Multi-objective Alignment.
  • Figure : WDPO Algorithm

Theorems & Definitions (37)

  • Definition 1: Joint data-generating distribution
  • Remark 1
  • Remark 2
  • Theorem 1: Estimation error of $\theta^{\mathrm{W}}_n$
  • proof : Proof sketch
  • Proposition 1: Estimation error of (non-robust) DPO
  • Remark 3
  • Theorem 2: Estimation error of $\theta^{\mathrm{KL}}_n$
  • Remark 4
  • Proposition 2: Worst-case distribution (informal)
  • ...and 27 more