Tri-level method improves AI preference learning under shifting conditions

TPO: Tri-level Distributionally Robust Learning for OOD Direct Preference Optimization

Machine LearningArtificial Intelligence

Summary

Preference learning helps AI like language models understand what users like based on comparisons. However, user preferences can change in complex ways across different groups, contexts, and individual choices, making it hard for AI to stay accurate. The authors present a new method called TPO that carefully models these three layers of preference shifts to make AI's learning more robust to changes. Their approach proves to be reliable and improves performance compared to existing methods.

What this means in practice

  • For llm deployment teams: Improve the safety and reliability of AI language models by robustly aligning them to user preferences despite shifts in user groups and contexts.
  • For recommendation system builders: Enhance recommendation algorithms by accounting for multi-level changes in user preferences, ensuring better adaptation to new user groups and situations.

Authors

Chengtao Jian, Kai Yang

Abstract

Direct Preference Optimization (DPO) minimizes empirical risk over observed preference comparisons, implicitly assuming that the training preference distribution matches the deployment preference distribution. In practice, this i.i.d. assumption often fails, as preference distribution shift exhibits a natural hierarchical structure across three levels: group-level shift, prompt-level shift, and preference-level shift. Such multi-level distributional mismatch poses a critical challenge for safe LLM deployment, yet existing robust DPO methods typically either model uncertainty over the joint preference distribution or focus on individual sources of shift, leaving the hierarchical structure of OOD preference shifts underexplored. We propose TPO, a tri-level distributionally robust framework for OOD preference alignment, which jointly models all three levels of distributional uncertainty through a hierarchically structured rectangular ambiguity set. By instantiating it with Wasserstein distance, we derive a tractable and interpretable objective that combines a closed-form preference-level margin penalty, prompt-level adversarial embedding perturbation, and worst-case group reweighting, each governed by an independent radius. We establish finite-sample convergence guarantees and empirically validate the effectiveness of TPO over existing robust DPO methods.