Reproduce “Robust LLM Alignment via Distributionally Robust Direct Preference Optimization” on OpenResearch

About this paper

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

This research introduces Distributionally Robust Direct Preference Optimization (DRDPO) with Wasserstein DPO (WDPO) and Kullback–Leibler DPO (KLDPO) to enhance LLM alignment robustness against preference distribution shifts. These methods consistently outperformed standard DPO in diverse tasks, with KLDPO notably improving IFEval scores to 0.75 (from 0.62) and MATH scores to 0.22 (from 0.03) for LLaMA-3.1-8B.

Key insights

• The population-level robust DPO loss functions (for both Wasserstein and KL formulations) are shown to be strongly convex under specific assumptions, ensuring unique optimal policy parameters.
• Theoretical analysis indicates that robust policy parameters for WDPO and KLDPO converge at a rate of O(n^-1/4), which is slower than non-robust DPO but provides guarantees under distribution shift.
• The minimax objective of robust DPO can be efficiently approximated for practical LLM training via adding a gradient regularization term (for WDPO) or implementing a loss-sensitive re-weighting scheme (for KLDPO).

Summary by alphaXiv.