Sharpness-Aware Minimization in Logit Space Efficiently Enhances Direct Preference Optimization
This addresses a key limitation in aligning large language models with human preferences, offering an efficient solution for practitioners, though it is incremental as it builds on existing DPO and SAM methods.
The paper tackled the squeezing effect in Direct Preference Optimization (DPO), where preferred response probabilities decrease during training, by proposing logits-SAM, a Sharpness-Aware Minimization variant that perturbs only the output layer, and experiments on models like Pythia-2.8B and Mistral-7B showed consistent improvements in DPO effectiveness.
Direct Preference Optimization (DPO) has emerged as a popular algorithm for aligning pretrained large language models with human preferences, owing to its simplicity and training stability. However, DPO suffers from the recently identified squeezing effect (also known as likelihood displacement), where the probability of preferred responses decreases unintentionally during training. To understand and mitigate this phenomenon, we develop a theoretical framework that models the coordinate-wise dynamics in logit space. Our analysis reveals that negative-gradient updates cause residuals to expand rapidly along high-curvature directions, which underlies the squeezing effect, whereas Sharpness-Aware Minimization (SAM) can suppress this behavior through its curvature-regularization effect. Building on this insight, we investigate logits-SAM, a computationally efficient variant that perturbs only the output layer with negligible overhead. Extensive experiments on Pythia-2.8B, Mistral-7B, and Gemma-2B-IT across multiple datasets and benchmarks demonstrate that logits-SAM consistently improves the effectiveness of DPO and integrates seamlessly with other DPO variants. Code is available at https://github.com/RitianLuo/logits-sam-dpo.