cs.LGMar 18, 2026

Sharpness-Aware Minimization in Logit Space Efficiently Enhances Direct Preference Optimization

Authors: Haocheng Luo, Zehang Deng, Thanh-Toan Do, Mehrtash Harandi, Dinh Phung, Trung Le

Organizations: Monash University, Australia · Swinburne University of Technology, Australia

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

Abstract

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.

Figures & tables

Figure 1: Training dynamics under different settings. (a–b) 1000-dimensional toy example with three classes, trained with a negative learning rate under GD, SAM ( ρ>0\rho>0 ), and SAM ( ρ<0\rho<0 ). Panel (a) shows the modal coefficients, and panel (b) shows the class residuals. (c) Real-data experiment on WebGPT Comparisons with GPT-2, comparing GD and SAM: the panel reports the log-probabilities of the chosen responses, the rejected responses, and max_conf , which denotes the model’s most confident response. (d) Real-data experiment on the TL;DR dataset with Pythia-2.8B, showing the same three curves (chosen, rejected, and max_conf ).
MethodHHTL;DR
vs SFTvs chosenvs SFTvs chosen
DPO70.5256.3584.2134.78
DPO+logits-SAM72.2860.5189.5836.57
SLiC-HF65.2754.7291.8831.36
SLiC-HF+logits-SAM71.8762.2194.4032.80
CPO66.6058.1990.9939.38
Table 1: Evaluation results (WR %) on HH and TL;DR datasets using Pythia-2.8B. The judge is GPT-5-mini. The highest value within each method group (baseline vs. logits-SAM) is bolded .
MethodAlpacaEval 2Arena-Hard v0.1MT-Bench
 LC (%)WR (%)WR (%)(score)
DPO13.0810.9619.05.49
DPO+logits-SAM13.9011.6223.15.79
SLiC-HF8.928.9719.15.05
SLiC-HF+logits-SAM10.639.2321.15.22
CPO8.978.1319.25.22
Table 2: Evaluation results on AlpacaEval 2 (LC and WR), Arena-Hard v0.1 (WR), and MT-Bench using Mistral-7B-v0.1. Judges are GPT-4 Turbo for AlpacaEval 2, and GPT-4.1 for Arena-Hard v0.1 and MT-Bench. The highest value within each method group (baseline vs. logits-SAM) is bolded .
Figure 2: Efficiency comparison.
Datasetρ=0 (AdamW)\rho=0\ (\text{AdamW})ρ=10−5\rho=10^{-5}ρ=10−4\rho=10^{-4}ρ=10−3\rho=10^{-3}ρ=10−2\rho=10^{-2}
HH70.52 / 56.3569.47 / 58.2772.28 / 60.5168.49 / 59.5265.49 / 56.31
TL;DR84.21 / 34.7887.79 / 33.9789.58 / 36.5784.25 / 29.9381.56 / 29.31
Table 3: Performance on HH and TL;DR datasets under different ρ\rho values. Each entry reports win rate vs SFT (left) and vs chosen (right).
Figure 3: Learning dynamics of Mistral-7B on UltraFeedback. We compare AdamW and logits-SAM in terms of training loss, evaluation loss, and evaluation accuracy, and report curves for logits-SAM under different values of ρ\rho .
 Ref modelDPODPO+logits-SAMCHESCHES+logits-SAM
Train Refusal0.80540.77030.81350.84590.9324
Test Refusal0.72310.70770.75380.78460.8769
Table 4: Train and test refusal rates for different methods on SorryBench (higher is better).
Appendix figures & tables3 assets

Supplementary material from the paper’s appendix.

Appendix

ClassObjectiveLearning rateρ\rhoSetting
y+y^{+}Positive objective f=−log⁡pf=-\log pPositive ( η>0\eta>0 )PositiveTheory = Practice
y−y^{-} (Theory)Positive objective f=−log⁡pf=-\log pNegative ( η<0\eta<0 )NegativeTheory
y−y^{-} (Practice)Negative objective f=log⁡pf=\log pPositive ( η>0\eta>0 )PositivePractice
Appendix
Table 5: Comparison between theoretical and practical settings of DPO with SAM. Although the signs differ for y−y^{-} , the resulting dynamics are equivalent. For y+y^{+} , the settings coincide.
MethodObjectiveHyperparameter
SLiC - HFmax⁡ ⁣(0,  δ−log⁡πθ(yw∣x)+log⁡πθ(yl∣x))  −  λlog⁡πθ(yw∣x)\displaystyle\max\!\bigl(0,\;\delta-\log\pi_{\theta}(y_{w}\mid x)+\log\pi_{\theta}(y_{l}\mid x)\bigr)\;-\;\lambda\log\pi_{\theta}(y_{w}\mid x)λ∈{0.1, 0.5, 1.0, 10.0}\lambda\in\{0.1,\,0.5,\,1.0,\,10.0\} ; δ∈{0.5, 1.0, 2.0, 10.0}\delta\in\{0.5,\,1.0,\,2.0,\,10.0\}
CPO−log⁡σ ⁣(βlog⁡πθ(yw∣x)−βlog⁡πθ(yl∣x))  −  λlog⁡πθ(yw∣x)\displaystyle-\log\sigma\!\bigl(\beta\log\pi_{\theta}(y_{w}\mid x)-\beta\log\pi_{\theta}(y_{l}\mid x)\bigr)\;-\;\lambda\log\pi_{\theta}(y_{w}\mid x)λ=1.0\lambda=1.0 ; β∈{0.01, 0.05, 0.1}\beta\in\{0.01,\,0.05,\,0.1\}
Appendix
Table 6: Objectives and hyperparameters for SLiC-HF and CPO.
MethodPythia-2.8BMistral-7B
DPO1×10−41\times 10^{-4}1×10−51\times 10^{-5}
SLiC-HF1×10−31\times 10^{-3}1×10−41\times 10^{-4}
CPO1×10−41\times 10^{-4}1×10−51\times 10^{-5}
Appendix
Table 7: Choice of ρ\rho for logits-SAM.
More in this topic

LLM Alignment

Large language model alignment studies methods for making models follow human intentions, values, and safety constraints across contexts.

See all 607 papers

Explore similar work

CardsList
Oct 6, 2026

LSC-DPO: Learning-Signal-Controlled Direct Preference Optimization

Direct Preference Optimization (DPO) has become a standard reward-model-free approach for aligning language models with preference data. However, as the scaled preference margin grows during training, the logistic DPO loss becomes progressively less sensitive to further changes. We study DPO from a loss-level geometric perspective and identify the sigmoid factor as a learning signal that characterizes the local sensitivity of the objective. Based on this view, we propose Learning-Signal-Controlled Direct Preference Optimization (LSC-DPO), which dynamically regulates the learning signal near a target regime. A log-space analysis establishes conditions for stable tracking of the target learning-signal regime. Experiments on AlpacaEval 2, MT-Bench, and Anthropic-HH show that LSC-DPO consistently improves over DPO and strong preference-optimization baselines. We further find that different coefficient initializations induce distinct transient learning-signal trajectories even when their later signal levels become similar. Based on this observation, we derive a signal-budget compensation rule that adjusts the target learning signal to compensate for these transient differences. The resulting compensation substantially reduces performance variation across coefficient initializations.

Yang Qu, Yusheng Han, Chengjia Feng +1

Northeastern University

Direct Preference OptimizationLLM AlignmentPreference Learning
AskPDF
Jun 10, 2026

Boosting Direct Preference Optimization with Penalization

Offline preference optimization has become a practical substitute for reinforcement learning from human feedback, but pairwise objectives such as Direct Preference Optimization (DPO) and its variants use only the chosen and rejected responses stored in a static dataset. This leaves a useful signal unused: the response that the reference model itself would generate for the same prompt. We propose Direct Preference Optimization with Penalization (DPOP), a simple extension of DPO that augments the base preference loss with a gated penalty on reference-greedy responses. DPOP activates this penalty only when the current policy still assigns a lower likelihood to the preferred response than to the rejected response. On AlpacaEval 2.0, DPOP improves length-controlled win rate over DPO, SimPO, and AlphaDPO on both Llama-3-8b-it and Gemma-2-9b-it, achieving relative gains of 5.3% and 4.4% over baselines on the two models, respectively. Ablations further show that a SimNPO-style length-normalized penalty is stronger than NPO and token-level unlikelihood in this setting.

Pengwei Sun

Stanford University.

Pairwise Preference LearningLLM AlignmentDirect Preference OptimizationOffline Preference Optimization
AskPDF
May 4, 2026

Gradient-Gated DPO: Stabilizing Preference Optimization in Language Models

Direct Preference Optimization (DPO) improves relative preference by increasing the margin between chosen and rejected responses, but this objective does not specify how probability mass should be redistributed to achieve that margin. Consequently, preference separation can improve even as probability mass collapses away from both preferred and unrelated responses, a phenomenon associated with the squeezing effect. We show that this pathology is localized: continued negative updates become destructive when rejected responses enter low-probability valleys, where further suppression induces harmful redistribution through the coupled softmax geometry. We introduce Gradient-Gated Preference Optimization (Gate-DPO), which uses the current probability geometry of each rejected response to selectively attenuate its contribution, limiting excessive rejected suppression while delaying preference-loss saturation. Gate-DPO is objective-agnostic and composes with existing preference methods. Across four architectures, two preference datasets, and multiple objectives, gating consistently reduces squeezing and improves chosen-response likelihood, while exhibiting robust reductions in rejected-response collapse across architectures and objectives. Mass-dynamics analysis further shows that large preference margins can conceal substantial suppression of unrelated responses, while gating achieves healthier redistribution. Together, our results show that preference margin alone is insufficient to characterize optimization health and that selectively controlling probability dynamics provides a simple, general mechanism for stabilizing preference optimization.

Inoussa Mouiche, Abderaouf Bahi

WASP Lab, School of Computer Science University of Windsor, ON N9B 3P4

LLM AlignmentDirect Preference OptimizationRL Post-Training
AskPDF