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
Venue: ICLR 2026 Poster (2026)
Keywords: optimization
direct preference optimization
sharpness-aware minimization
learning dynamics
Primary area: optimization

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>.

Published as a conference paper at ICLR 2026
SHARPNESS-AWARE MINIMIZATION IN LOGIT SPACE EFFICIENTLY ENHANCES DIRECT PREFERENCE OPTI-
MIZATION
Haocheng Luo-
Zehang Deng"
Thanh- Toan Do
Mehrtash Harandi*
Dinh Phung! Trung Le Monash University, Australia "Swinburne University of Technology, Australia
fhaocheng.. 1uo, toan. mehrtash harandd, dinh.phung @rongsmmemonash.coo zehangdeng@ swi edu.au
ABSTRACT
Direct Preference Optimization (DPO) has emerged as a popular algorithm for aligning pretrained large language models with human preferences, owing its simplicity and training stability. However, DPO suffers from the recently iden- tified squeezing effect (also known as likelihood displacement), where the prob- ability of preferred responses decreases unintentionally during training. To un- derstand and mitigate this phenomenon, we develop a theoretical framework that models the coordinate-wise dynamics in logit space. Our analysis reveals that negative-gradien updates cause residuals to expand rapidly along high-curvature directions, which underlies the squeezing effect, whereas Sharpness-Awan Aware Min- imization (SAM) can suppress this behavior through its mewaatur-rreulariiznn effect. Building on this insight, investigate logits-SA.., computationally ef- ficient variant that perturbs only the output layer with negligible overhead. Exten- sive 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: //gith/gihhth com www.tiiiiinntooooo Lwc.ooooooooooteeoooo itggiits-- s- sam dpo
INTRODUCTION
Reinforcement learning from human feedback (RLHF) (Christiano et al. 2017 Stiennon et al. 2020 Ouyang et al., 2022) is a crucial technique for aligning pretrained large language models (LLMs) with human preferences ensure helpfulness, harmlessness, and safety (Bai et al. 2022 Dai al., 2023). Its pipeline typically comprises three stages: supervised fine-tuning (SFT), reward modeling, and policy optimization. Classical policy optimization methods such Proximal Policy Optimiza- tion (PPO) (Schulman et 2017). while widely used for their effectiveness, depend heavily on the quality of the learned reward model, rendering training complex and often unstable. Direct Pref- erence Optimization (DPO) (Rafailov et al.) 2024b) is a recently proposed and promising offline alternative that, by reparameterizing the implicit reward and optimizing closed-fome objective on preference data, trains the policy directly without explicitly fitting reward model, DPO has gained traction due to algorithmic simplicity and training stability.
Despite DPO and its many variants demonstrating state-of-the-ar performance across range of tasks, several potential issues remain. A particularly important one is the recently identified squeez-- ing effect (Ren & Sutherland] 2024) (also known as likelihood displacement (Razin et al.) (024 which describes an unintended decrease in the generation probability of preferred responses during DPO training, contrary to the intended goal of increasing it embodied in the DPO objective. This phenomenon can lead to performance degradation, reduced safety, and even alignment failure (Pal
al. 2024 Yuan et al., 2024 Rafailov et al. 2024a Tajwar et al. 2024 Pang et al.) 2024).

---

Published as a conference paper at ICLR 2026
To understand the mechanism behind the squeezing effect and to identify effective remedy, develop a theoretical framework that elucidates the learning dynamics in both the parameter space and the logit space. Our analysis shows that the negative-gradient updates induced by the negative objective associated with rejected answers DPO cause the residual vector to expand rapidly along high-curvature directions, namely along the eigenvectors associated with large eigenvalues of the Hessian, which is the source of the squeezing effect. This raises natural question: Can curvature- aware training mitigate this unintended drift?
We investigate Sharpness-Aware Minimization (SAM) (Foret et al. 2021). bilevel optimization method widely used in supervised learning, and establish its dynamics both the parameter and logit spaces. Our theory demonstrates that SAM effectively alleviates the squeezing effect through its intrinsic curvature regularization. Guided by these insights, we advocate using logits-SAM for DPO training, computationally efficient variant of SAM that perturbs only the output-layer pa- rameters. Although logits-SAM has been mentioned merely as a byproduct in prior work (Baek) al., 2024; Singh et al. 2025) and often overlooked, our study turns this neglected variant into practically useful and effective technique by integrating it into DPO, where it efficiently mitigates the squeezing effect and consistently improves performance. To the best of our knowledge, this is the first work to analyze and apply SAM in the context of DPO.
Our contributions are summarized follows:
We develop theoretical framework that connects the parameter space and the logit space through geometric properties, enabling a unified analysis of learning dynamics in both domains. This framework yields unified dynamical equations for gradient descent (GD) and SAM that precisely track coordinate-we wise evolution with controlled error terms.
Our analysis identifies the root cause of the squeezing effect: under negative- gradient updates, residuals expand rapidly along high-curvature directions. We rigorously show that SAM, through its intrinsic curvature regularization, effectively alleviates this phenomenon.
Bridging theory and practice, implement efficient variant, logits-SAM, which perturbs only the output-layer parameters. Unlike vanilla SAM, it incurs virtually no additional overhead. Ex- periments on Pythia-2.8B, Mistral-7B, and Gemma-2B-IT across multiple datasets and bench- marks validate its effectiveness, demonstrating consistent performance gains for DPO and its vari- ants.
2 PRELIMINARIES
2.1 PREFERENCE OPTIMIZATION
SFT-RLHF pipeline. Classical RLHF alignment proceeds in three phases: (i) supervised fine- tuning of base policy on instruction-folloviig data; (ii) reward modeling by fitting a scalar reward function pairwise human preferences; and (ii) policy optimization maximize the learned re- ward under KL regularizer toward a reference policy.
DPO reparameterization. DPO (Rafailov et al. 2024b) bypasses training an explicit reward model by expressing an implicit reward for policy To as log- likeliiiddood ratio to a fixed reference policy Tref (typically the SFT model):
To x)
re( y) log
Boo Z(x),
(1)
Fere (y x)
where 3 >> 0 is temperature and ((x)) a partition term independent of 0. Combining equation
with the Bradley-Terry preference model (Bradley & Terry 1952) p(y+ > y x)) <ro((x y
re (ac yr)) yields the standard DPO objective, optimized over a dataset D = {(x, y }} of prefereed/dispeffrred pairs:
LDPO(TO; Tree) = Bes log a log T x))
To To x
)
B log
(2)
Teet (y xx)
nee (y xr)
where a()) is the logistic function.
2

---

Published as a conference paper at ICLR 2026
2.2 SHARPNESS AWARE MINIMIZATION
SAM regularizes training by explicitly penalizing parameter-space sharpmess: it chooses parameters that minimize the worst-case loss within an lg ball of radius p around 6. Concretely, for supervised learning with examples (xr, y) ~ D and per-example loss f(0; ar ar. y), the SAM objective is
min (x.)) max f(e t x. ||c|aaS
(3)
This formulation can be interpreted as a form of curvature regularization: by seeking minimizers whose neighborhoods exhibit consistently low loss, SAM favors flatter minima that often correlate with improved generalization. In practice, the inner maximization is approximated to first order by the perturbation E* E²(0)= = P Vef æ, y) / æ, y)l,, and one takes descent step using the gradient at the perturbed point, Vef xz
LEARNING DYNAMICS IN LOGIT SPACE
3.1 SETTING
We adopt the same theoretical setting as in Ren & Sutherland (2024), namely multiclass logistic classification, where the features of the samples (Malladi et al., 2023)). In this abstraction, we keep the per-example loss nonnegative, and model the DPO terms by allowing the esfective update direction take either sign: the preferred term (y+) corresponds to standard descent with > O, whereas the rejected term (y") induces negative- gradient updates. which we represent as an update with 7j O. For analytical convenience, we follow Ren & Sutherland (2024) and encode these greative-gradien updates as gradient descent with a negative effective learning rate, which equivalent using a positive learning rate on
fixed (also referred to the kernel regime
negated objective (see Appendix D for short derivation).
Prior work (Ren & Sutherland 2024) has shown that the resulting negative-gradient dynamics, in- cluding the characteristic squeezing effect, can be faithfully reproduced within this simplified set- ting. Several phenomena observed this multiclass logistic regression abstraction also emerge empirically during real LLM fine-tuning. These findings suggest that analyzing DPO through this framework offers a theoretically tractable and practically relevant perspective on the learning behav- ior of LLMs.
Let be a training example with one-hot label y E 10, 1T y = 1. In the fixed-feature (kernel) regime, ((x) Rd are fixed and
z' = W ('o(z)) ER,
p' softmax((2"
f(z', = Σ yk log pk.
where W E Vxd xd are trainable parameters, z' are the logits. For notational convenience, we write ((x) d. We use || to denote the l, norm for vectors and the Frobenius for matrices. We use & to denote the Kronecker product.
We denote the parameter Hessian by H V Vi f("", y) E RVavVd and μ == || We denote the logit gradient by g' V=f(=", 3) p' y E RRV. and denote the logit Hessian by H =
[2,(", y) E RVVV
3.2 THEORY
The theoretical results Ren & Sutherland (2024) demonstrate that the squeezing effect arises from rrgative-gradien updates. Specifically, they show that the probability of the ground-truth label necessarily decreases, while the probability of the model's most confident incorrect class increases. Building on this, we provide a finer-graine analysis the learning dynamics. We develop a unified modeling framework that tracks the residuals of all classes and characterizes their coordinate- wise evolution with linear convergence up to higher-order remainders. Using this framework, we further establish rigorous conditions under which SAM effectively mitigates the squeezing effect.
3

---

Published as a conference paper at ICLR 2026
For GD, first-order derivatives sufficient to characterize its dynamies. However, the intrinsic curvature regularization effect of SAM motivates us to further investigate the geometric structure of the parameter space through the Hessian matrix. To this end, develop theoretical framework that connects the geometry of the parameter space and the logit space, via the link between the parameter Hessian and the logit Hessian.
Proposition 3.1 (Geometry of the space; simplified version of PropositionnB.l) In coordinates,
Hw =(bØ") & Hz. Thus, if + O, then rank Hw) H2) Moreover, the second-order effect of any parameter perturbation depends only on the induced logits perturbation (') AW Φ.
This proposition establishes that all second-order effects in the parameter space, whose Hessian Hw lies in IR. VdvV, can be equivalently studied through the logit Hessian H₃ in RVV, thereby greatly simplifying the analysis of second-order dynamics. Next, building on this geometric corre- spondence, we derive unified dynamical equations for SAM in both the parameter space and the logit space. Unlike prior work, these dynamics simultaneously capture the evolution of all coordinates of the parameters, logits, and residuals, with precise control over the error terms.
Theorem 3.2 (SAM dynamics in parameter and logit space; informal version of Theorem B.2 Assume that we conduct the SAM update for W. Under mild assumptions, there exists a constant C> 0 such that the following expansions hold with O(1²²) remainders:
(parameters)
W W
TT
H2'g (T) )
R'w:
|R ll Cn,,
SAM's
(logits)
n ===nnppg' p' H'H'g ) z x'
+
|| S Snn
correction
(residuals)
g gt = pt11 - n H nup (H2) (p' y) +r |r/ll S CT2²
SAM's correction
where pf = P VT//lg'|| is the equivalent perturbation coefficient.
When p O, the dynamics reduce to standard GD. Expressed through the logit Hessian, both GD and SAM share a unified dynamical structure across parameter, logit, and residual spaces. In both parameter and logit space, GID amounts to scaling by the logit gradient, whereas SAM introduces an additional H₂ correction term that can be regarded as a preconditioning matrix, rescaling updates differently along the rigen-directions of Hz.
Moreover, the updates of the residual vector under GD and SAM are both preconditioned by H₂ (and, for SAM, by 22 This implies that if we choose the eigenvectors of the logit Hessian as a basis, the curvature coupling effects of both the first-order and second-orde -odder terms can be unified. To formalize this intuition, we show that g lies precisely in the column space of H thus we can select the nonzero eigenvectors of H- as a basis to obtain the coordinate representation of span [1} and
Proposition 3.3. H₂ symmetric positive semidefinite with ker (er(H,) rank(H, = V 1. Moreover. for the residual we have 1 "gg00 O, hence g E 1 rangel H2); in particular, given any eigenbusis of H2 restricted 1 9 admits a unique coordinate represen- tation in that basis
Corollary 3.4 (Modal dynamies in the eigenbasis of H²). Under the same assumptions as Theo- rem 3.2 For each I, let the spectral decomposition of the symmetric positive semidefiite matrix H be
V-1
H
k=1
where X/ > O, (v1) = dkg are the non-zero eigenvalues and eigenvectors. Define the modal coefficients of the residual p' y by
k= -,.,,,,,f* - Then there exists a constant C > 0 such that for all nonzero modes 1
et (vt) g'; 1+1 =
(v()
(4)
k (1 7 [Xg + ('(A)) 1)ck + r Iril S CTp.
(5)
SAM's correction
4

---

CD Published as a conference paper at ICLR 2026
Proofs deferred to Appendix B wise scalars in the eigenbasis of z. making SAM's effect transparent. We now characterize the additional SAM correction in two regimes.
The corollary diagonalizes the vector dynamics into coordinate-
Case I: Positive (corresponding to the yt objective in DPO). In this case, GD induces stronger contraction of the residual along high-curvature directions, i.c., those associated with large eigenvalues of Hz. The additional correction term introduced by SAM has the same sign as that of GD, thereby amplifying this contraction.
Case 2: Negative 7) (corresponding to the y objective in DPO). Here, GD causes the residual to expand more rapidly along high-curvature directions. Moreover, standard SAM with positive p exacerbates this phenomenon, leading to even faster expansion com- pared to GD. By contrast, choosing negative p counteracts this expansion.
Next, we extend our theoretical framework to the result of Ren & Sutherland (2024), which in- troduced the squeezing effect. For consistency with their notation, we let y E {sr y denote the ground-truth class index (with one-hot label y = ey), and y" arg maxj Pj be the most confident incorrect class.
To compare GD and SAM, we study the one-step confidence ratio
(GD)
P: (SAM)
SAM
P
which measures the multiplicative change of the class probability after update.
Lemma 3.5 (One-step confidence ratios under GD, Lemma I of & Sutherland (2024)). Con- sider the negative learning-ree rate case 1) < 0, for which the ground-truth label is y = y Then
GD 0 >> 1,
GD 1.
Lemma 3.5 formalizes the squeezing effect under GD with negative learning rate: the probability of the most confident incorrect class increases, while that of the ground-trer truth class decreases. Within framework, we next analyze the ratio of these two probabilities after one-step SAM update.
Corollary 3.6 (One-step confidence ratios under SAM, informal version of Corollary B.7 and Corol-
lary B.10). Under the same assumptions as Theorem 3.2 set qp > 0 and lp| = KV % Then, for
sufficiently small step size (y], the following inequality holds:
aty
SAM S a.y
(6)
and under additional technical conditions, the following inequality holds:
04 2 a
SAM
Here E fyt .y denotes the ground-trune truth label corresponding to the positive or negative learn- ing rate, respectively. Moreover, the inequalities strict whenever ply- (0, and ý +
The proof is deferred to Appendix B Corollary and Lemma 3.5 together imply that, when 1) < O, using SAM with a negative p moderates the growth of the most confident incorrect class and slows the decay of the ground-trui truth class, thereby preventing excessive expansion and premature collapse. Thus, for negative 7,, choosing negative p effectively alleviates the squeezing effect.
We empirically validate this prediction using 1000-dimensional toy example with three classes. We first train for 10 epochs using class 0 the label mimic SFT, and then switch to class while continuing training with negative learning rate. As shown in Figure [1 this setup reproduces the squeezing effect observed in prior work (Ren & Sutherland 2024): both modal coefficients expand rapidly, the probabilities of class and class 2 decrease. and only the probability of class 0 increases. Moreover, SAM with positive exacerbates this behavior, whereas SAM with negative p suppresses it, exactly predicted by theory.
Additionally, Corollary 3.6 indicates that for > O, using SAM with > 0 also has a beneficial effect alleviating the squeezing effect: the contraction of y" is accelerated, while the growth of the ground -round--ruthh y is enhanced. Taken together, these results suggest a simple rule: choosing with the same sign the learning rate alleviates the squeezing effect.
5

---

Published as a conference paper at ICLR 2026
1000 dinenssnninnlll boy exarggle
LOGD cite-dinntionaall ty osaple
GPT2 on WHOOPP Comparionn
Pythin 2 oo 7E.08
0.54
-300
Modglooghhing
120
25
Modgl
100
0.23
3350
40
-000
0.70
01.54
125
10
10
12
Pgech
Fporh
Epoch
Pgoch
(x)
(b)
(4)
Figure 1: Training dynamics under different settings. (a-b) 1000-dimensional toy example with
three classes, trained with negative learning rate under GD, SAM (p > 0), and SAM (p 0). Panel (a) shows the modal coefficients, and panel (b) shows the class residuals. (c) Real-data ex- periment 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_coo conf, which denotes the model's 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).
To verify this rule in real models, we track the probability dynamics of chosen, rejected, and most confident responses when fine-tuning GPT-2 (Radford et al., (019) on the WebGPT Comparisons dataset (Nakano et al. 2022) (Figure Tbb and Pythia-2.8B (Biderman et al., (023) on the TL:DR dataset (Stiennon et al. 2020) (Figure 1d). In both cases, SAM increases the probability of chosen
responses, slows the decrease rejected responses, and prevents the most confident responses from growing, consistent with our theoretical predictions.
3.3 FROM THEORY TO PRACTICE
In practice, an important challenge in applying SAM to DPO is that it requires an additional for- ward and backward pass, thereby nearly doubling the computational cost. However, our dynamical analysis shows that curvature regularization can still be achieved even when the perturbation is ap- plied solely in the logit space (with an appropriate choice of the sign of p). which also alleviates the squeezing effect. Motivated by this observation, we suggest using computationally efficient SAM variant that perturbs only the last layer, called logils-SAM, to improve the effectiveness and robustness of DPO. Its objective can be formulated as follows:
SAM (W, 0; = LDoo W +p
V LDPO W, 0: x. y
u. y 0;
w LDPo W, 0: x, yn,
where W denotes the parameters in the output layer, and 0 denotes the parameters except W.
Implementation. Unlike our theoretical setting, common DPO implementation I typically en- code the y objective negative while using a single positive learning rate, rather than assigning positive and negative rates to yt and y respectively. Accordingly, we adopt this convention in logits- SAM implementation. Our dynamical analysis further indicates that p should share the sign of the learning rate; hence we consistently use positive p. In Appendix we provide the full update formulas and derivation establishing the equivalence between the theoretical and practical settings, summarized in 5
This choice does not render our analysis of the negative learning rate redundant. For first- order methods such as GD, using negative objective with a positive learning rate is equivalent to using a positive objective with negative learning rate. Therefore, our analysis applies fully the case of negative objectives.
Remark.
The implementation pseudocode can be found in Algorithm D of Appendix We compute the perturbation manually using the hidden states from the penultimate layer and the parameters of the final layer, requiring only single full in standard SAM. Since the parameters of the final layer typically constitute only small fraction of all trainable parameters (e.g., 4.64% in Pyth and 1.81% in Mistral-7B), the additional
pass instead of the two full passes required
httpss//ggttbbbb.oomm//gggggggg

---

Published as a conference paper at ICLR 2026
training overhead introduced by logits-SAA Segits-SAMin is negligible. detailed comparison of wall-clock time and peak memory usage provided in Section 4.3
EXPERIMENTS
4.1 EXPERIMENTAL SETUP
Datasets. We conduct DPO training on three widely used datasets to evaluate our algorithm: Anthropic-HH (Bai al. 2022). the Reddit TL;DR summarization dataset (Stiennon et al. 2020).
and the UltrFeedback Binarized dataset (Cui et al. 2023).
Models. Following common practice, we adopt SFT models as our base models. We use Pythia- 2.8B (Biderman et al. (023) for experiments on Anthropic-HI and Reddit TL:DR, and Mistral- 7B-v0.1 (Jiang et 20.3)) for UltraFeedback. For Pythia-2.8B. we initialize from the Hugging
For the TL;DR experiments, we use the checkpoin which was SFT for one epoch on Reddit TL;DR. For Mistral- 7B-v0.1 1, we use the Alignment Handbook (Tunstall et al. 2023a) checkpoint Zephyr-76 (Tunstall ct al. 2023b), which was SFT for one epoch on UltraChat- 200k (Ding et al. 2023
Face open-source checkpoin which was SFT for one epoch on Anthropic-
Evaluation. For 2y8h we evaluate model performance on Anthropic-HH and Reddit TL;DR by measuring win rates (WR) against both the SFT baseline and the human-pre preeerreee re- sponses, using GPT-5-mini (version 2025-08-07
the automatic judge. Following the DPO paper, we set the decoding temperature to for HH and for TL:DR. For Mistral-7B-vO.1 we conduct
evaluation three popular open-ended instruction-following benchmarks: AlpacaEval 2 (Dubois jet al., 2024), Arena-Hard v0.1 (Li et al. 2024). and MT-Bench (Zheng et al., (023). Details of each benchmark be found in Appendix
We adopt the default generation parameters provided by each benchmark. Specifically, we report both length-controlled win rates (LC) and raw WR for
AlpacaEval 2, model WR for Arena Frena-Hard v0.1, I, and averaged judge scores (1-10) for MT-Bench, all following the standard evaluation protocols.
Baselines. We apply logits-SAM to DPO and two SOTA variants. SLiC-HF ZZhao al. 2023) and CPO (uu et al. 2024). We use AdamW optimizer (Lohcchllvv & Hutter 2019) in all experiments. For Pythia-2.8B, we set batch size 64 and learning rate 10 6 following the DPO paper; for Mistral-7B, we use batch size 128 and learning rate 5 10-7 following the Alignment Handbook's ' recommended settings.
Hyperparameters. For DPO, we adopt the recommended B values from the DPO paper and the Aligmment Handbook, which are widely used and well tuned. For SLiC-HF and CPO, we select hyperparameters following the tuning protocol from Meng et al. (2024b). For logits-SAM, we keep all hyperparameters identical to each corresponding baseline to ensure fairness; the only additional hyperparameter is p. which tune over {1 x 10 10-5, 10 1,1 10-33. Full hyperparameter settings are provided in Ta and Table 7 of Appendix
4.2 EXPERIMENTAL RESULTS
Performance of summarization and dialogue generation tasks. We present the results in Ta- ble [ We find that logits-SAM SAM consistently improves performance across both HH and TL;DR datasets. All three baselines (DPO, SLiC-HF, and CPO) achieve higher win rates against both SFT and chosen responses when augmented with logits-SAM. Notably, SLiC-HF shows the largest gains on HH (+6.60 pp vs SFT, +7.49 pp vs chosen), while CPO achieves strong improvements on TL;DR (+2.30 pp SFT, +6.03 pp vs chosen), demonstrating that logits-SAM provides stable and general- izable benefits across different optimization methods.

---

Published as a conference paper at ICLR 2026
Table 1: Evaluation results (WR %) HH and TL,DR datasets using Pythia-2.8B. The judge is
GPT-5-mini. The highest value within each method group (bascline vs. logits-SAM) is bolded.
HH
TL;DR
Method
vs SFT
vs chosen
vs SFT
vs chosen
DPO
70.52
56.35
84.21
34.78
DPO+Hogits-SAM
72.28
60.51
89.58
36.57
SLiC-HF
65.27
54.72
91.88
31.36
SLiC-HF+logits-ANN
71.87
62.21
94.40
32.80
CPO
66.60
58.19
90.99
39.38
CPO+logits-SAM
70.24
59.90
93.29
45.41
Table 2: Evaluation results on AlpacaEval (LC and WR), Arena-Hard v0.1 (WR), and MT--eench using Mistral-7B--O.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.
Method
AlpacaEval
Arena- Hard v0.1
MT-Bench (score)
LC (%)
WR (%)
WR (%)
DPO
13.08
10.96
19.0
5.49
DPO+logits-SAM
13.90
11.62
23.1
5.79
SLiC-HF
8.92
8.97
19.1
5.05
SLiC-HF+logir SAM
10.63
9.23
21.1
5.22
CPO
8.97
8.13
19.2
5.22
CPO+logits-SAd
13.32
11.78
21.4
5.49
Performance on open-ended instruction-following benchmarks. We present the results in Ta- ble The results demonstrate that combining logits-SAM with different DPO variants consistently yields performance gains across all benchmarks. On open-ended instruction -oollooionnfolooiingg evaluations, logits-SAM improves both length-controlled and original win rates on AlpacaEval 2 (e.g., with CPO: +4.35 pp LC, +3.65 pp WR), increases head-to-head win rate on Arena Hard v0.1 (e.g., with DPO: +4.1 pp WR), and provides steady gains on MT-Bench (e.g.. DPO: +0.30, SLiC-HF: +0.17, CPO:
+0.27), These findings indicate that logits-SAM is generally effective and robust enhancement
across diverse evaluation settings.
4.3 ADDITIONAL ANALYSIS
Compared to vanilla SAM, logits-SAM minimizes additional computational over- head. We report wall-clock training time and peak mem- ory on Pythia-2.8B trained on the Reddit TL:DR dataset (Figure 2) using data-parallel training (DDP) across two NVIDIA A100 GPUs with per-device batch size of 4.
Computational overhead.
AdamW
100
90-
Parcentage
80
72
70
59.39
69.36 GB
The results show that logits-SAM adds only 2-3% extra time, with negligible peak-memory overhead. By contrast, vanilla SAM practically infeasible for Pythia- 2.8B on 100s with DDP: it nearly doubles the step time (due extra full fowward-bacwaal pass) and requires a perturbation buffer comparable to the model size (for billion-parameter models, this entails more than GB of additional GPU memory), which leads out-of-memory even with batch size 1. These observations highlight the clear computational cost advantage of logits-SAM.
memory
Figure 2: Efficiency comparison.
Sensitivity analysis. We present a sensitivity analysis of the additional hyperparameter for logits-SAM in Table 3 The results indicate that, within reasonable range of p. performance is typically improved, whereas further enlarging P leads to marked degradation. Notably, unlike original SAM, logits-SAM perturbs only the output layer, so the appropriate scale of p is much smaller than the range (0.01-0.5) recommended in the SAM paper. We recommend starting the search for logits-SAM. p at 10 or 10 and, if resources permit, performing a finer sweep in this neighborhood.
8

---

Published as a conference paper at ICLR 2026
Table 3: Performance on HH and TL.DR datasets under different p values. Each entry reports win rate vs SFT (left) and vs chosen (right).
Dataset P = 0 (AdamW)
p= 10-3
p p= 10-4
p p= 10-3
= 10-2
HH
70.52 / 56.35
69.47 / 58.27
72.28 60.51
68.49 / 59.52
65.49 / 56.31
TL:DR
84.21 34.78
87.79 /3.97
89.58 / 36.57
84.25 / 29.93
81.56 / 29.31
Training Loss
Evaluation Loss
Evaluation Accuracy
0.70
AdamW Logits-SAM (p=1e:5)
0.76
0.56
Logits-SAM (p=1e-4)
0.75
Loos
Logits SAM (p=1e-3)
Loss
0.54
0.74
0.73
0.72
0.50
0.71
0.48
0.70
200 300
400
100
200
300
400
100
200
300
Training Steps
Training Steps
Training Steps
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 p.
Learning dynamics. In Figure 3. we compare the learning dynamics of AdamW and logits-SAM when training Mistral- on the UtraFeedback dataset. The figure reports training loss, evaluation loss, and evaluation accuracy across training steps, and includes curves for logits-SAN under multi- ple choices of p (1 x 10-5, 1 10-1, (0-1,11 1 10- We observe that for three values of p. logits-SAM achieves training loss that is fairly close to that of AdamW, yet consistently attains lower evaluation loss and higher evaluation accuracy. This indicates that logits-SAM provides better generalization than AdamW, and the fact that range of p values yields consistent improvements suggests that its benefits are robust to the choice of this hyperparameter.
Sharpness. To further probe the reasons underlying the generalization gains of logits- we measure the traces of the parameter Hessian and the logit Hessian at the final checkpoint of Mistral-7B. For AdamW, the traces are 1.337 x 10a 2.732 x 102 (parameter logit Hessian), while for logits-SAM they are reduced to 1.186 x 10 / 2.586 x 102. This reduction indicates that Iogits-SAM converges to a flatter solution, which is widely believed to be beneficial for generaliza- tion.
Extension to AI safety and on-policy setting, Razin et al. (2024) refer to the squeezing effect as likelihood displacement and find that, in AI safety scenarios, can reduce the model's harmful- response refusal rate, which leads to severe safety concerns. We evaluate the performance of logits- SAM in the same on-policy setting and Al safety scenario as in their work. The reference model is the instruction-tun tuned Gemma-2B-IT (Team et al.] 2024). and the evaluation is conducted using SorryBench (Xie et al., 2024). We train for one epoch with a learning rate x 10-6 6 and a batch size of 16. We compare the performance of the reference model, DPO, DPO with logits-SAM, CHES (Razin et al., 2024). and CHES with Dogits-SAM.. For CHES, we filter 50% of samples using the CHES score. Performance is measured by the harmful-respons. refusal rate and is reported in Table The results show that logits- SAM significantly improves performance in this setting. In particular, DPO with logits-SAM avoids the degradation in refusal rate and performs better than the reference model. Combining logits-SAM with the CHES method of Razin et al.) (2024) further increases the refusal rate, with an absolute improvement of approximately 9% on both the training and test sets. These findings indicate that logits-SAM can be effectively transferred to other settings and tasks.
Table 4: Train and test refusal rates for different methods on SorryBench (higher is better).
Ref model
DPO
DPO+logitS-SAM
CHES
CHES+logitS-SAM
Train Refusal
0.8054
0.7703
0.8135
0.8459
0.9324
Test Refusal
0.7231
0.7077
0.7538
0.7846
0.8769

---

Published as a conference paper at ICLR 2026
5 RELATED WORK
Reinforcement learning from human feedback. RLHF has emerged the de facto post-training
recipe for aligning large language models (Christiano et al. 2017 Ziegler et al. [011 Ouyang et al.
2022
Bai al., 2022), typically combining supervised fine-tuning (Zhou al. [203 Taori al.
2023
Conover et al. 2023] Wang et al. 2023b), reward modeling (Gao et al. 2023] Luo et al.
2023
Lambert et al. 2024). and policy optimization (Schulman et al. 2017] Anthony et al. 2017)
To reduce the complexity and instability of online preference optimization, offline methods such as
SLiC-HF (Zhao et al., 2023) and RRHF (Yuan et al. 2023) learn policies from comparisons using closed-form objectives. DPO (Rafailov et al. 2024b) is central example that maximizes the log- probability margin between preferred and rejected responses relative to reference policy. Thanks to its simplicity and training stability, DPO has rapidly gained popularity, spurring line of variants aimed at improving performance. For example, |Azar et al. (2024) propose IPO, more theoretically grounded variant; CPO (Xu et al. 2024) approximates the reference policy uniform to eliminate the reference term; f-DPO [Wang et al. (023a) generalizes DPO via a family of f-divergences; SimPO (Meng al. 2024a) uses nength-normalized scores that better reflect generation-time pref-
erences; and Cal-DPO Xiao et al. 2024) aligns the implicit reward scale with likelihoods.
Squeezing effect (likelihood displacement). The squeezing effect (Ren & Sutherland 2024 also known as likelihood displacement (Razin et al. 2024). refers to the recently identified phenomenon in which the probability of the ground-truth label is unintentionally reduced during DPO training.
This effect has been widely observed and can lead to performance degradation, reduced safety, and
even alignment failure (Pal et al. 2024 Yuan et al. 2024 Rafailov et al. 2024a Tajwar al. 2024]
Pang et al.] 2024). Several studies have attempted to mitigate this issue. |Asadi et al. (2025) ) con-
strain the shift of probability mass between preferred and rejected responses the reference and target policies. [iu et al. (2025) introduce a KL--divergence-baee basedd policy drift constraint to dynami- cally regularize policy updates. Razin et al. (2024) strengthen safety alignment by filtering samples that are likely to induce likelihood displacement based the CHES score between token embed- dings. Unlike existing approaches, which either focus on designing alternative objective functions or filter the training data, our method takes pure optimization-base perspective. It is therefore conceptually orthogonal to these techniques and can be used in combination with them.
Sharpness-aware minimization. A widely held belief in the deep learning community is that
flatter solutions typically generalize better (Hochreiter & Schmidhuber 1997, Keskar et al., 2016) 2017 Diang al. 2019: Xie al. Liu et al. (002)). Motivated by this view, SAM (Foret et al.
Dinh et al.
2022)) is a bilevel optimization method that explicitly seeks flatter minima, and it has
gained popularity for delivering consistent improvements across a wide range of supervised learning
tasks (Foret et al.
2021
Kwon et al., 2022 Kaddour et al. [2022] Liu et al. 2022 Kim et al. 2022
Li Giannakis]
2023
Nguyen et al., 2023a b] [Tuong et al. 2024] Luo et al.) 2025b Phan et al,
2025 Luo et al.
2025a). Most relevant to our work are its recent applications in LLMs. Singh et al.
(2025) propose Functional-SAM for LLM pretraining and demonstrate strong performance, while
Lee & Yoon (2025) apply SAM to Proximal Policy Optimization to improve robustness in both the reward and action spaces. Logits-SAN. is a byproduct mentioned in recent studies, yet it is often overlooked. Baak et al.] (2024) analyze the effect of label noise on SAM in linear regression and argue that Jacobian-SAM. the counterpart of Aigits-SAM.. plays the dominant role. Similarly, (iinhh let al. (2025) identify Jacobian-SAM,, also referred to Functional-SAM as more important and show that it can effectively improve the generalization performance of LLM pretraining.
CONCLUSION
We analyzed the squeezing effect in DPO via coordinate-wise dynamics in parameter and logit spaces. Our framework shows that GD with negative S drives residuals to expand along high- curvature directions, and that SAM suppresses this behavior via curvature regularization; in par- ticular, negative 7) calls for negative p. Motivated by this, we adopt logits-SAM, which perturbs only the output layer and adds negligible overhead, and demonstrate consistent gains in effective- ness and robustness across models and datasets. We expect these insights to inform curvature- aware preference optimization going forward.
10

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGEMENTS
Trung Le, Mchrtash Harandi, and Dinh Phung were supported by the ARC Discovery Project grant DP250100262. Trung Le and Mehrtash Harandi were also supported by the Air Force Office of Scientific Research under award number FA9550-23-S-0001.
REFERENCES
Zahra Rahimi Afzal, Tara Esmaeilbeig, Mojtaba Soltanalian, and Mesrob Ohannessian. Can the spectrum of the neural tangent kernel anticipate fine-tuning performance? In Adaptive Foundation Models: Evolving Al for Personalized and Efficient Learning, 2024.
Thomas Anthony, Zheng Tian, and David Barber. Thinking fast and slow with deep learning and tree search. Advances in neural information processing systems, 30, 2017.
Kavosh Asadi, Julien Han, Idan Pipano, Xingzi Xu, Dominique Perrault-Joncas, Shoham Sabach, Karim Bouyarmane, and Mohammad Ghavamzadeh. C2-dpo: Constrained controlled direct pref- erence optimization, 2025. URL
org wttpaar///.2..12200000
Mohammad Gheshlaghi Azar, Zhaohan Daniel Guo, Bilal Piot, Remi Munos, Mark Rowland, Michal Valko, and Daniele Calandriello. A general theoretical paradigm to understand learn- ing from human preferences. In International Conference on Artificial Intelligence and Statistics, pp. 4447 4455. PMLR, 2024.
Christina Baek, Zico Kolter, and Aditi Raghunathan. Why is sam robust to label noise? arXiv preprint arXiv:2405.03666 2024.
Yuntao Bai, Andy Jones, Kamal Ndousse, Amanda Askell, Anna Chen, Nova DasSarma, Dawn Drain, Stanislav Fort, Deep Ganguli, Tom Henighan, et al. Training helpful and harmless assistant with reinforcement learning from human feedback. arXiv preprint arXiv:2204 05862,, 2022.
Stella Biderman, Hailey Schoelkopf, Quentin Gregory Anthony, Herbie Bradley, Kyle Brien, Eric Hallahan, Mohammad Aflah Khan, Shivanshu Purohit, USVSN Sai Prashanth, Edward Raff, al. Pythia: A suite for analyzing large language models across training and scaling. In International Conference Machine Learning, pp. 2397-2430. PMLR, 2023.
Ralph Allan Bradley and Milton E Terry. Rank analysis incomplete block designs: I. the method of paired comparisons. Biometrika, 39(3/4):324 345, 1952.
Paul F Christiano, Jan Leike, Tom Brown, Miljan Martic, Shane Legg, and Dario Amodei. Deep reinforcement learning from human preferences. Advances in neural information processing sys- tems, 30, 2017.
Mike Conover, Matt Hayes, Ankit Mathur, Jianwei Xie, Jun Wan, Sam Shah, Ali Ghodsi, Patrick Wendell, Matei Zaharia, and Reynold Xin. Free dolly: Introducing the world's first truly open instructiontuned 1lm. 2023.
Ganqu Cui, Lifan Yuan, Ning Ding, Guanming Yao, Wei Zhu, Yuan Ni, Guotong Xie, Zhiyuan Liu, and Maosong Sun. Ultrafeedbcc:: Boosting language models with high-quality feedback, 2023.
Josef Dai, Xuehai Pan, Ruiyang Sun, Jiaming Ji, Xinbo Xu, Mickel Liu, Yizhou Wang, and Yaodong Yang. Safe rlhf: Safe reinforcement learning from human feedback. arXiv preprint 1rXiv:2310.1277: 2023.
Ning Ding, Yulin Chen, Bokai Xu, Yujia Qin, Zhi Zheng, Shengding Hu, Zhiyuan Liu, Maosong Sun, and Bowen Zhou. Enhancing chat language models by scaling high-quality instructional conversations, 2023.
Laurent Dinh, Razvan Pascanu, Samy Bengio, and Yoshua Bengio. Sharp minima can generalize for deep nets. In International Conference Machine Learning, pp. 1019-1028. PMLR, 2017.
11

---

Published as a conference paper at ICLR 2026
Yann Dubois, Balázs Galambosi, Percy Liang, and Tatsunori B Hashimoto. Length-controlled al- pacaeval: A simple way to debias automatic evaluators. arXiv preprint arXiv:2404.04475, 2024.
Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware minimiza- tion for efficiently improving generalization. In 9th International Conference on Learning Rep- resentations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021. URL
htttps:/cppenre//
net/ forum 6mmmmmmmmmmmmmmmm 1rM
Ieo Gao, John Schulman, and Jacob Hilton. Scaling laws for reward model overoptimization. In International Conference on Machine Learning, pp. 10835-10866. PMLR, 2023.
Sepp Hochreiter and Jürgen Schmidhuber. Flat minima. Neural computation, 9(1): 1997.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kerne:: Convergence and gen- eralization in neural networks. Advances in neural information processing systems, 31, 2018.
Uijeong Jang, Jason D Lee, and Emest K Ryu. Lora training in the ntk regime has no spurious local minima. arXiv preprint 1rXX:::4022.1186,, 2024.
Albert Q. Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chap- lot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, Lélio Renard Lavaud, Marie-Anne Lachaux, Pierre Stock, Teven Le Scao, Thibaut Lavril, Thomas Wang, Timothée Lacroix, and William El Sayed. Mistral 7b, 2023. URL https:
06825
Yiding Jiang, Behnam Neyshabur, Hossein Mobahi, Dilip Krishnan, and Samy Bengio. Fantastic generalization measures and where to find them. arXiv preprint arXiv: 1912.02/7, 2019.
Jean Kaddour, Linqing Liu, Ricardo Silva, and Matt Kusner. When do flat minima optimizers work? Advances in Neural Information Processing Systems, 35:16577-11695,, 2022.
Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Pe- ter Tang. On latch-batch training for deep learning: Generalization gap and sharp minima. arXiv preprint arXiv: 1009006609. 2016.
Minyoung Kim, Da Li, Shell X Hu, and Timothy Hospedales. Fisher sam: and sharpness aware minimisation. In International Conference on Machine Learning, pp. 11148- 11161. PMLR, 2022.
Information geometry
Jungmin Kwon, Jeongseop Kim, Hyunseo Park, and In Kwon Choi. Asam: Adaptive sharpness- aware minimization for scale-invariant learning of deep neural networks. In International Con- ference Machine Learning, pp. 5905-5914. PMLR, 2021.
Nathan Lambert, Valentina Pyatkin. Jacob Morrison, LJ Miranda, Bill Yuchen Lin, Khyathi Chandu, Nouha Dziri, Sachin Kumar, Tom Zick, Yejin Choi, et al. Rewardbench: Evaluating reward models for language modeling. arXiv preprint arXiv:2403. 13787, 2024.
Hyun Kyu Lee and Sung Whan Yoon. Flat reward in policy parameter space implies robust rein- forcement learning. In The Thirteenth International Conference on Learning Representattins, 2025.
Bingcong Li and Georgios B. Giannakis, Enhancngg http:/iiinnn.aaaanseeseeeereonii through variance suppression, 2023. URL httpps//aaaaai..ccc//aaa///////0
Tianle Li, Wei-Lin Chiang, Evan Frick, Lisa Dunlap, Tianhao Wu, Banghua Zhu, Joseph E. Gon- zalez, and Ion Stoica. From crowdsourced data to high-quality benchmarks: Arena-hard and benchbuilder pipeline, 2024. URL BtRp:ttps:/ arxivv org 220p://../2222..... 11939
Hong Liu, Sang Michael Xie, Zhiyuan Li, and Tengyu Ma. Same pre-training loss, better down- stream: Implicit bias matters for language models. In International Conference on Machine Learning, pp. 22188-22214, PMLR, 2023.
Yong Liu, Siqi Mai, Xiangning Chen, Cho-Jui Hsieh, and Yang You. Towards efficient and scalable sharpness-awee aware minimization. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 22360-12370, 2022.
12

---

Published as a conference paper at ICLR 2026
Zongkai Liu, Fanqing Meng, Lingxiao Du, Zhixiang Zhou, Chao Yu, Wenqi Shao, and Qiaosheng Zhang. Cpgd: Toward stable rule-based reinforcement learning for language models. arXiv preprint arXiv:2505.125500 2025.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization, 2019. URL https:
Haipeng Luo, Qingfeng Sun, Can Xu, Pu Zhao, Jianguang Lou, Chongyang Tao, Xiubo Geng, Qing- wei Lin, Shifeng Chen, and Dongmei Zhang. Wizardmath: Empowering mathematical reasoning for large language models via reinforced evol-instruct. arXiv preprint arXiv:2308. 2023.
Haocheng Luo, Mehrtash Harandi, Dinh Phung, and Trung Le. Unveiling m-sharpness through the structure of stochastic gradient noise. arXiv preprint arXiv:2509. /8001, 2025a.
Haocheng Luo, Tuan Truong, Tung Pham, Mehrtash Harandi, Dinh Phung, and Trung Le. Explicit eigenvalue regularization improves sharpness-aware minimization. In The Thirty-eighth Annual Conference Neural Information Processing Systems, 2025b.
Sadhika Malladi, Alexander Wettig, Dingli Yu, Danqi Chen, and Sanjeev Arora. A kernel-based view of language model fine-tuning. In International Conference on Machine Learning, pp- 23610-2366 23641. PMLR, 2023.
Yu Meng, Mengzhou Xia, and Danqi Chen. Simpo: Simple preference optimization with a reference-free reward. Advances in Neural Information Processing Systems, 37:124198-122235,, 2024a.
Yu Meng, Mengzhou Xia. and Dangi Chen. Simpo: Simple preference optimization with reference-free reward, 2024b.
14734
Reiichiro Nakano, Jacob Hilton, Suchir Balaji, Jeff Wu, Long Ouyang, Christina Kim, Christopher Hesse, Shantanu Jain, Vineet Kosaraju, William Saunders, Xu Jiang, Karl Cobbe, Tyna Eloundou, Gretchen Krueger, Kevin Button, Matthew Knight, Benjamin Chess, and John Schulman. Webgpt: Browser-assisted -answiin with human feedback, 2022. URL https://arxiv.
org/abs/ 2212. 09332
Van-Anh Nguyen, Trung Le, Anh Bui, Thanh-Toan Do, and Dinh Phung. Optimal transport model distributional robustness. Advances Neural Information Processing Systems, 36:24074-24087, 2023a.
Van-Anh Nguyen, Tung-Long Vuong, Hoang Phan, Thanh-Toan Do, Dinh Phung, and Trung Le. Flat seeking bayesian neural networks. Advances in Neural Information Processing Systems, 36: 30807-30820, 2023b.
Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al. Training language models to fol- low instructions with human feedback. Advances in neural information processing systems, 35:
27744-27744, 2022.
Arka Pal, Deep Karkhanis, Samuel Dooley, Manley Roberts, Siddartha Naidu, and Colin White. Smaug: Fixing failure modes of preference optimisation with dpo-positive. arXiv preprint arXiv:2402.132228 2024.
Richard Yuanzhe Pang. Weizhe Yuan, He He, Kyunghyun Cho, Sainbayar Sukhbaatar, and Jason Weston. Iterative reasoning preference optimization. Advances in Neural Information Processing Systems, 37:1 16617-16616677 2024.
Hoang Phan, Lam Tran, Quyen Tran, Ngoc Tran, Tuan Truong, Qi Lei. Nhat Ho, Dinh Phung, and Trung Le. Beyond losses reweighting: Empowering multi-task learning via the generalization perspective. In Proceedings of the IEEF/CVF International Conference Computer Vision, pp- 2440-2450, 2025.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, llya Sutskever, et al. Language models unsupervised multitask learners. OpenAl blog, 188):9 2019.
13

---

Published as a conference paper at ICLR 2026
Rafael Rafailov, Joey Hejna, Ryan Park, and Chelsea Finn. From to q": Your language model is secretly a q-function. arXiv preprint arXiv:2404.12358, 2024a.
Rafael Rafailov, Archit Sharma, Eric Mitchell, Stefano Ermon, Christopher D. Manning, and Chelsea Finn. Direct preference optimization: Your language model is secretly reward model,
2024b. URL htt.htpsss/acxxi.orr//aaaaaeee
10.90189000
Jeff Rasley, Samyam Rajbhandari, Olatunji Ruwase, and Yuxiong He. Deepspeed: System opti- mizations enable training deep learning models with over 100 billion parameters. In Proceedings of the 26th ACM SIGKDD international conference on knowledge discovery & data mining, pp. 3505 3506, 2020.
Noam Razin, Sadhika Malladi, Adithya Bhaskar, Danqi Chen, Sanjeev Arora, and Boris Hanin. Unintentional unalignment: Likelihood displacement in direct preference optimization. arXiv preprint (XXx:24400008447 2024.
Yi Ren and Danica Sutherland. Learning dynamics of Il finetuning. arXiv preprint arXiv:2407.10490, 2024.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347 2017.
Sidak Pal Singh, Hossein Mobahi, Atish Agarwala, and Yann Dauphin. Avoiding spurious sharpness minimization broadens applicability of sam. arXiv preprint #rXiv:2502.02400 2025.
Nisan Stiennon, Long Ouyang, Jeff Wu, Daniel M. Ziegler, Ryan Lowe, Chelsea Voss, Alec Radford, Dario Amodei, and Paul Christiano. Learning to summarize from human feedback. In NeurlPS, 2020.
Fahim Tajwar, Anikait Singh, Archit Sharma, Rafael Rafailov, Jeff Schneider, Tengyang Xie, Ste- fano Ermon, Chelsea Finn, and Aviral Kumar. Preference fine-tuning of 1lms should leverage suboptimal, on-policy data. arXiv preprint arXiv:2404. 14367, 2024.
Rohan Taori, Ishaan Gulrajani, Tianyi Zhang. Yann Dubois, Xuechen Li, Carlos Guestrin, Percy Liang, and Tatsunori B Hashimoto. Stanford alpaca: An instruction-followiin llama model, 2023.
Gemma Team, Thomas Mesnard, Cassidy Hardin, Robert Dadashi, Surya Bhupatiraju, Shreya Pathak, Laurent Sifre, Morgane Rivière, Mihir Sanjay Kale, Juliette Love, et al. Gemma: Open models based on gemini research and technology. arXiv preprint ar urXiv:2403.08295, 2024.
Tuan Truong, Quyen Tran, Quan Pham-Ngoc, Nhat Ho, Dinh Phung, and Trung Le. Improving generalization with flat hilbert bayesian inference. arXiv preprint arXiv:2410.04196, 2024.
Lewis Tunstall, Edward Beeching, Nathan Lambert, Nazneen Rajani, Shengyi Huang, Kashif Rasul, Alvaro Bartolome, Carlos M. Patiño, Alexander M. Rush, and Thomas Wolf. The Alignment Handbook, 2023a. URL
Lewis Tunstall, Edward Beeching, Nathan Lambert, Nazneen Rajani. Kashif Rasul, Younes Belkada, Shengyi Huang, Leandro von Werra, Clémentine Fourrier, Nathan Habib, Nathan Sarrazin, Omar Sanseviero, Alexander M. Rush, and Thomas Wolf. Zephyr: Direct distillation of Im alignment, 1694
2023b. URL htt.:ttpss//axiv..rrttt
Chaoqi Wang, Yibo Jiang, Chenghao Yang, Han Liu, and Yuxin Chen. Beyond reverse kl: Generalizing direct preference optimization with diverse divergence constraints, 2023a. URL.
16240
Guan Wang, Sijie Cheng, Xianyuan Zhan, Xiangang Li, Sen Song, and Yang Liu. Openchat: Ad- vancing open-source language models with mixed mixed-quality data, arXiv preprint arXiv:2309.1123: 2023b.
Teng Xiao, Yige Yuan, Huaisheng Zhu, Mingxiao Li, and Vasant G Honavar. Cal-dpo: Calibrated direct preference optimization for language model alignment. Advances in Neural Information Processing Systems, 37:1 1e:114289-111 114320, 2024.
14

---

Published as a conference paper at ICLR 2026
Tinghao Xie, Xiangyu Qi, Yi Zeng, Yangsibo Huang, Udari Madhushani Sehwag, Kaixuan Huang, Luxi He, Boyi Wei, Dacheng Li, Ying Sheng, et al. Sorry- Systematically evaluating large language model safety refusal. arXiv preprint arXiv:2406. 14598, 2024.
Zeke Xie, Issei Sato, and Masashi Sugiyama. A diffusion theory for deep learning dynamics: Stochastic gradient descent exponentially favors flat minima. arXiv preprint arXiv:2002.03444 03445, 2020.
Haoran Xu, Amr Sharaf, Yunmo Chen, Weiting Tan, Lingfeng Shen, Benjamin Van Durme, Kenton Murray, and Young Jin Kim. Contrastive preference optimization: Pushing the boundaries of 11m performance in machine translation, 2024. URL http:htttt///t htt.ctps:/arxi..og////o hwaggabs//20 w40.aa/....aaa..oeeeeeeeeeaae 08417
Lifan Yuan, Ganqu Cui, Hanbin Wang, Ning Ding, Xingyao Wang, Jia Deng. Boji Shan, Huimin Chen, Ruobing Xie, Yankai Lin, et al. Advancing Ilm reasoning generalists with preference trees. arXiv preprint ar 2Xivv2444.002078 2024.
Zheng Yuan, Hongyi Yuan, Chuanqi Tan, Wei Wang, Songlang Huang, and Fei Huang. Rrhf: Rank responses align language models with human feedback without tears. arXiv preprint arXiv:2304.055302, 2023.
Yao Zhao, Rishabh Joshi, Tianqi Liu, Misha Khalman, Mohammad Salch, and Peter J. Liu. Slic-hf: Sequence likelihood calibration with human feedback, 2023. URL https://arxiv..rr abss2305.2 10 10425
Liammin Zheng, Wei-Lin Chiang, Ying Sheng, Siyuan Zhuang, Zhanghao Wu, Yonghao Zhuang, Zi Lin, Zhuohan Li, Dacheng Li, Eric P. Xing, Hao Zhang, Joseph E. Gonzalez, and Ion Stoica. Judging llm-as-a-judge with mt-bench and chatbot arena, 2023. URL.httss:/aaxiiv.co
abs /206..306... 05685
Chunting Zhou, Pengfei Liu, Puxin Xu, Srinivasan Iyer, Jiao Sun, Yuning Mao, Xuezhe Ma, Avia Efrat, Ping Yu, Lili Yu, et al. Lima: Less is more for alignment. Advances in Neural Information Processing Systems, 36:55006-55021, 2023.
Daniel M Ziegler, Nisan Stiennon, Jeffrey Wu, Tom Brown, Alec Radford, Dario Amodei, Paul Christiano, and Geoffrey Irving. Fine-tuning language models from human preferences. arXiv preprint
2019.
15

---

Published as a conference paper at ICLR 2026
A ADDITIONAL RELATED WORK
Kernel and fixed feature regime of LLMs. In the context of LLMs, there a growing body of work that investigates model dynamics through the lens of kernels. A pioneering line of work by Malladi et al. (2023) uses Neural Tangent Kerne--based dynamics (Jacot et al.) 2018) to accurately characterize the behavior of LLM fine-tuning and achieves performance comparable to fine-tuning through kernel methods, under the fixed feature assumption. Affza et al. (2024) leverage the spec- of the NTK to predict the generalization performance of LLM fine-tuning, and Jang et al. (202) study the training dynamics of low-rank adaptation from an NTK perspective.
B FORMAL THEOREMS AND PROOFS
Notation. We use || to denote the & norm for vectors and the Frobenius norm for matrices. For linear and multilincar maps, denotes the corresponding operator norm induced by these norms. Proposition B.1 (Geometry of the logit space and the parameter-logit correspondence). Lel & RVV IR be C2, Fix an imput and feature map (xx) Rdd For W R Vdd set
2 = W WOERY,
F(W) = f(z).
Let H2 V=((z) and Hw == F(W be the Hessians. Equip I. Vxd with the Frobenius inner product (A, B)F A and Ry with the Euclidean inner product. Let
T : Vxd RRV
(a(WW) = AW é
be the differential of the map W We. and let T3 RV + RVxd be its adjoint with respect to these inner products, ie.. (T), v) = (AW, (v)) for all AW, v. Then (g = v ø The following statements hold.
(1) Pullback identity (operator form).
Hw T To
linear operators on R Vxd, Consequently, ó + O, then rank H rank(H, H
canonical Euclidean coordinates,
= (Vz((z)) ST,
Hw = (o0 @ H
(2) Pullback identity (bilinear form). For every AW, AW' I. Vxd the bilinear
HwW AW" = (AW), Ta(AW"))
Thus the second-orde effect of any parameter perturbation depends only on the induced logits perturbation (AW = AW 6.
(3) Surjectivity, kernel, and quotient-space view. If Φ + O, then T is surjective. For any preimage is
Az RVY
AW, = Az TT
with Ta(AW,) = =Az.
10
The kernel is
ker (Tx) = AW Rxdd AW2=01
dimension V(d 1). Consequently, Hw descends the quotient RV ker(Tss) _ R
Proof. A direct computation gives
(To v) (((AW))TT) u(WW vØJ) = (AW, vZ
hence
(b) v
16

---

Published as a conference paper at ICLR 2026
(1) Pullback identity (operator form). Let F(W) ( (Wo). The first differential of F is dF[AW] (V=((z), T((AW) (TX(V=((2)) AW).
so
VwF(W) = (z())(())) = (V+((z)) oJ. Differentiating once more and using d(V=()(=)AAz - H2 A= with A= = To (AW) yields, for all AW, AW'
F'F AW = (To) H₂ (AW'))
By the Riesz representation on (R Vxd, ()) this means
Hw T H To:
linear operators RVxd Using T% (v) = vo" and To (AW) = AW we obtain the explicit operator action
Hw (wW))) = T H2 (AW o)) = HH (AW Φ)) oJ = H AW () For the rank statement, assume Φ * Then Tş is surjective and T is injective. Hence
rank (ork(Hw)=a rTn H₂ TT) rank H TT) rank ran)
because range(Tp) = TR..
Under the canonical identification Xxdd R Vd obtained by stacking columns, denote by vec(-) the corresponding vectorization. Using the standard identity
Kronecker (matrix) form.
vec( A B) = B 6 A) vec (X),
we have
ved Hw (AW) = ved H2 AW (oo = ((ooT) @ H₂) vec AW)
Hence the matrix representation of Hw these coordinates
Hw (óo @ H
(2) Pullback identity (bilinear form). Recall that the inner products on Vxdd and RV induce
bilinear forms
(AW, Hw())'))r By the operator identity Hw T H2 To. for all AW, AW,AW'
Hw AW")
Hz (Az. H.A=)).
HwlAW, AW"l (AW, Hw(AW'))
= =AW, T3 H. Ta(A"))
= (To(AW), H. H.To("W") -AW),, Ta(A))))
Thus the second-order effect of any parameter perturbation depends only on the induced logits per- turbation To(AW) = AW d.
(3) Surjectivity, kernel and quotient view.
If é * O, then for any Az E V
Az Φ
AW.
satisfies TA(W..) Az,
|
so Ta is surjective. The same choice minimizes the Frobenius norm among all preimages (row- wise Mauchy-Schward The kernel ker(T,) (AW AW & 0). and rank-nullity gives dim (er(T,) = (id - 1). Finally, if AW, - AW, E ker(TT). then (AW,) = T (A())) and (AW)), H. TAWw))) Hw WwW, AW2),
H AW AW1] (A)))))) H2 Ho (AW1))
so the bilinear form descends to the quotient RV RVdd (er((rr))))) = R.V.
If then T 0 and Hw 0, the degenerate case.
17

---

Published as a conference paper at ICLR 2026
Theorem B.2 (Dynamics of SAM). Fix sample and set μ || > 0. Assume:
For each y, the loss f(; y) is C in and there exists a constant 3o such that for all
[z, w)ll S L,
1,
The step size [n] E (0, 1] and the SAM radius satisfies |pl kv/ll with constant 2 0.
Write F(W)= = f(Wó, o, y) and We.
Consider standard SAM:
V F(W!)
AW adv P VV (WWF(WW) F(W'III
W W AW
W W w P (W
adv:
We adopt the convention that when llg' O, the inner perturbation is 0, iee i 0.
Otherwise, define p' == P PV/F/llg')) Then, there exists constant C' > 0 (depending only on L, K) such that the following expansions hold with O(1²²) remainders:
(parameters)
W W W - n(9' o + H'g 67) + R'w.
|| R $ sCn,
(logits)
z 211 = n1 + p' H3gg) +r
|| S Cn²,
| S Cn? (logit gradient) g 1+1 = (1- n H - nup (H))²) g" + trg
In particular, for softmax cross-entrops where g' = p' - y and
np (residual) p'+1- - yy- = - nn - uu (H2)² (p' - y) + + | s sCt.
Proof. By Proposition B. (operator form),
w F() gg gΦ
(wW))) = H- AW (de)))
and To (AW = AW& with |T S || vT. Moreover, by the multilinear chain rule applied to F(W) f f(W o y)
F(W) [A AI,,,,,,,24. Asl (zf(=
Ix(A2), Is(As)).
(8)
hence the operator norm satisfies
sup W
W F(W)|| S sup
|| & L L
(9)
(1) Parameter update.
If |لاsll
> 0, then VwF(W') g and و| T |1 =
||g' Il VP, so
g'o
AWA e adv
AW wdvll = |p| - KV |n/-
(If ||g'|| = O, our convention sets AW! adv 0.) second-order Taylor expansion of V w F at W gives, for some 0 (0,
W FwW(W = Vw ((W)) + H A &dv
+ 1 v WWW WW +Wa A /WW ady Awt.) adv
By equation
and | W advll k//lll
V W F(.) AWW & AWA.
S 1 3/2
| 12 S Colnl,
for a constant
Co Co(L, pl, K). Using the operator identity from Proposition B.1
PVP |
H W AW: s H AW. adv (oo
Higg
ST 6² H TT.
Therefore
W W - (9g00 7 H 26 -17 -,
where || < C. ln). Setting R'y = -n R's yields || S CT² with C = C(L,H,
proving the parameter expansion.
18

---

Published as a conference paper at ICLR 2026
(2) Logit update.
Right-nultiplying by & and using μ ||
z-1 =W W = 17 (g' + 8D H + +
with |||||| < W || S Cr2² (absorbing v into C). This proves the logits expansion.
(3) logit gradient update. Since g Vzf(z, y) first-order Taylor expansion at 2' gives g g" (2 x²) f=' + , () Ax"], Az' =='=2'+1 z².
By assumption || & L and |Az'| || On) from the previous step, hence the remainder has norm & C rf. Substituting the logits expansion from step (ii) yields
g I n H 1 p² (H2)² + +
|| S S
after absorbing constants into C. This proves the logit gradient statement.
Combining (i)-iii) completes the proof, with a constant C' depending only on (L, M. k), and the bounds holding for all |v] E (0, 1] and |e| S kv/ul-
For softmax cross-entropy,
V=f(z,y) p(=) y,
H=(z) V2f(2)) = Diag(p(=)) p(z)p(z)T
Since p(=) A C [0. 1] for all z, every entry of the third derivative tensor f2,(2z, y) is
bounded polynomial in p(=) (hence in [0, 1)). Therefore there exists finite constant Lum (V)
depending only on such that
f y)
L (V),
h
1,2,3
In particular, f C and Assumption (1) of the theorem holds with L Lem
In the remainder of this section, we specialize the softmax cross-entropy loss.
Proposition B.3. H- is symmetric positive semidefinite with ker( (+2) = span and rank(H,) V L. Moreover. for the residual g we have 1 9= hence g E 1- range (H²);
in particular, given any eigenbasis of H2 restricted to 1 g admits a unique coordinate represen- tation in that basis.
Proof. Let p softmax(z) E 1)) that 1 1, and recall
H = Diag(p) - pp
For any v E RVV.
V Σ peus) poo))== Varp(u) 20,
H v
Σ Drut
hence H2 is symmetric positive semidefinite. Moreover, H iff Varp(u) 0. ie., E4 is constant across i. Since p; for all i, this means v = c1, thus
ker( H = span(1} =
rank (H)) = V - dim ker H V -
Then 1 g = 1T 1 T y so g E 1 For any symmetric matrix, range H₂) (ker ++;
using ker( H2) = span (1} yields 1 = range( H2), hence g E rrange( Hz).
Restrict H. the invariant subspace - Being symmetric, H₂ 11 admits an orthonormal eigen-
basis { V-1
associated with its positive eigenvalues. Since g 1- it has the unique expansion Us. with e =wx) g-
k= >>=1
Corollary B.4 (Modal dynamics in the cigenbasis H³). Under the same assumptions as The- orem B.2 For each t, let the spectral decomposition of the symmetric positive- desitive-semidefiie matrix
H be
H =
Σ X TT
19

---

Published as a conference paper at ICLR 2026
2i = dąg are the non-zero eigenvalues and eigenvectors. Define the modal coefficients of the residual g" by
where X/ > O, (o
c = (vt) T g',
k k=1,..., V
Then there exists constant C such that for all nonzero modes 1,
(of) g"t nu X + (22))) cle rke
1 s Cn²²
(10)
Proof. By Theorem
B.2 (residual expansion). we have
7) : 71 P- (H2)) + r
|| S Ctt²
(11)
Fix and let the eigendecomposition of H be H off N orthonormal. (The zero mode corresponding to = 0 is orthogonal to g" in the softmax- CE case and is therefore omitted.)
with X and
Project equation
11 onto the eigenvector ot:
(o)) = ( (I [I-Nμ nn (H2)² + (vt)
Using the eigen-relations H of and (H²) 2 = (21)) of and the definition et (xt) g', we obtain
( n - np
(1
- (AA)) + r r = (vf) r (T')
Finally, since || = 1 we have Iril S ||rس| || - Cn which is exactly equation 10 This completes the proof.
Following Ren & Sutherland 2024),
define the one-step confidence ratio as follows. Definition B.5 (One-step confidence ratio). For each class i E [V] and update rule a E (GD, define the one-step confidence ratio
(a)
P
As consequence of Theorem B.2 we obtain the following lemma.
Lemma B.6 (One-ster step ratio representation and factorization). Under the assumptions of Theo-
rem B.2 fix an iteration and set r/' = 77 μ. The ratio a; admits the representation
ΣV= E/= 85000 ez5
where
8 cxpl (1/) - リÿ) (p) v1)}}.
and the SAM correction appears multiplicatively
3SAM expl-ri
rf)} -'p'- explr - r
DJ.:- - (Hzgg))
gGD
curvalure jactor
remainder factor
and there exists constant C1 > that depends only on (I,, Ml vi) such that e 200n? S exp(r} r(] & 20,r for all i,j.
Proof. By Theorem B.2 (logits line).
Az 211 z' = - n/ (g' + "''H'g') + r
|| & 1r s Cin²,
where C depends only on (I,, 15) and the hypothesis |e| S ///nn/ is in force.
20

---

Published as a conference paper at ICLR 2026
For any increment Az,
E,e
Exet)
Pi(z' + Az) Pi(z*)
Σ, exp(p22 - Azie
Σ, Byets
With the above Az,
B5AM expl (9) gi) expl - - (
explr) - r remainder factor
A fJ, - (Hig'")
BgD
curvature factor
From || S Ciri, we have e 201n S explrf - S e for all i,j.
We begin by proving the first part of Corollary 3.6 which relies on the same assumptions as Theo- rem B.2. except that we set p = K KY instead of requiring p S ky/n
Corollary B.7 (One-step confidence ratio of y" under SAM vs. GD). Under the assumptions of
Theorem B.2 fix an iteration t and set ní' = n and [p] kv/nl. Let be the ground-truth label
and y" arg max P3 the most confident incorrect class. Assume the sign condition T/p" O.
Then there exists 7ou = no(p', H |lggg I, > such that, for all step sizes < |17] S 1/o the following one-step inequalities hold:
SAM I ag
GD
a y
Moreover, the inequality is strict whenever SAM), equality holds.
ply. (0, and 5 + O. In particular, when p (no
Proof. Define
D.DD = Σ 50 c
D. E30 e expl-
DSAM SAAM E35
with
@SAM- = GGD exp( - n'p" p"
A, j exp(r) - -
دJ. (Hig") ( -H''g')
By the remainder bounds of Lemma B.6 e 20'n D S DSAM < 2020n D..
With H diag(p') pp') and g" =p' Ey.
(H'g"): = pt (p: - y1 C*),
C²:= := Σ(pt) pyy
Then C2 S S''- and one checks: for = y", A S for all j. and D/100 < for some j whenever
pzz E (0, 1).
Next, by ex 2 1 + x and the sign structure of j,
Dy DGD y"
Σ (v)) e'''"' 2 1 1'p' Σ w (0000'' (-A),w)) 2 1 1+ cy. n'p",
for some Cg- > 0 whenever ply- 1); here w = BCD / are positive weights.
Now use the scaling |p| n then r/ p' O(733/2), whereas =0 =1 O(1²). Hence there exists no > 0 (depending only on (p'. H |1 H1, N, L)) such that for 0 < inl S nu.
SAM 2 e 20n2² Dy 2 D (1+ 1 2'
Since a; = (E, e we obtain for 0 < |n| S nuu
SAM 0y'
GD < agy
with strict inequality under the stated nondegeneracy conditions (because then cy- > 0). If || = (so 30 by convention), equality holds. This completes the proof.
21

---

Published as a conference paper at ICLR 2026
The proof of Corollary [B.7 relies on the sign of A,.s. which holds automatically for i y² How- ever, this need not be true for an arbitrary ground-truth label y, except in certain special cases such as binary classification. Therefore, sumption. Intuitively, we require that most of the probability mass concentrates on the two classes of interest, y and y"
further introduce the following notation and technical as-
Notation.
We define
Py By p= E (0,
S : + Pi
7: 7:=1-8 S,
(Py) = 40y -
s
and let B : :2
Assumption B.8 (Top-2 dominance and feasible curvature gap). there existy 7/0 > 0 such that
Under the notation above, assume'
4B e2 p/y-
S 7o S (Py) 6r.
Remark. For any fixed Py (0, Abin (Py) 40, D,y) 2 is a positive constant, and hence Abin (Py) 67 remains strictly positive for all sufficiently small while the left-hand side scales as (T) as - 0. Hence, whenever the tail mass T exists non-empty interval of feasible "o satisfying the assumption. Intuitively, this corresponds to the regime where most probability mass is concentrated on the two dominant classes y and y".
(p' + pfj-) ) is sufficiently small, there
Lemma B.9 (Top-2 reduction of the curvature gap). Under the notation above, for softmax cross- entropy.
(H'g)) "Ha==""gg);; By :
(H2g'),
Avia (Py) + 5
15"1 S 6T.
Proof.
Let = p'y. b. b= p S
T. u: By := a/S, and
PT=
(
j (1.r))
For softmax cross-entropy, we have for any i,
(Hig'): pí (p) Vs-
-
A straightforward expansion then yields the following explicit form:
A = Atiu (u) + (S - 1)+ -u (1=)== -
-1)+ -1) 1)+ (4w² Ga² 1--41)($ (100-1)(s-1+(((0 1) Sg.
where Apin (u) = 4u(1 - u
Using S T, we have
-7 5² (1 - 7)² -2 +
- 1= -1
-37 37 -33
Substituting these into the expansion yields
Ay v Aun(u) = -((ur + -Cu)) + +(2u 12u-1))8g
where
A(u) == u 2(1 - u 2
+ 3(4u² .u + 4u - 1) = 12u3 - 22n² + 1lu - 1,
B(u) = (1 2n²) 3(4u² Gu 4u 1) 12 20u² + 11u 2 (2u 1)² 2),
Cuu) 4n 6u + 4u (2u 12 2u +1+1).
Hence, for all u E [0, 1], the coefficient bounds satisfy
AA(w)
BBu)I < 2,
[C(u)] 1
Finally, since r E [0, implies 72 < and 3 < T, and Sq $ r, we obtain
A Apin S [A(u) + 2 + (C()) + 22u 1|5g S 6r.
Setting Ç¹ = A Ag Abin (b,) completes the proof.
22

---

Published as a conference paper at ICLR 2026
Corollary B.10 (One-step confidence ratio of y under SAM vs. GD). Under the assumptions of
Theorem B.2 fix iteration and set n/ n and ppl /. ///n.. Let y be the ground-truti label
and y arg maxjzy P'-- Assume the sign condition rf pt > 0. Assume further that Assumption B.8 holds iteration f.
Then there exists no = no(p', H Ilgg 14, (o)) > such that, for all step sizes 0 < [n] S nj. the following inequality holds:
4y
SAM
2 By
GD
Moreover, the inequality is strict whenever 5e # 0 and p'y- E (0, 1). In particular, when P 0 (no SAM), equality holds.
Proof. Let E := n'p' > 0, and let R [V] fu, v² Define
D = Σ 30 es Dy: = Σ b 50 のt expf -
DSAM = Σ
BSAM e"),
with AJ.i == (H'g)) - (H'g'): ). By Lemma B.6 SAM S 200rn² D
Introduce weights w.y (2) = -GG e+ DCD so that E, w) and
Dy Σ w - 1 + m. (2) (e (0- 1) Σ jER
/ (e (e-EAisv 1).
DCD
We prove the claim in five steps.
(1) uniform bound on ا For softmax-CE, |H S 1/2 and llg' S v2, hence
| S 22/2 and thus for all i.j.
A (H'gg) (Hig')ll 2 S B,
B B V2.
(2) tail weight concentration under top-2 dominance. Since each component of g" lies in Therefore,
[-1, 1], we have for any i.j, e -2/0'1 s B S ell'',
Σ Σjen e Σx GGD
2w'' Σ Ej = =11'' Σ = edll'T
Σ JER
v 2/71| Σ eth
jER
Moreover, for in'l sufficiently small, w 2 b/ (since wy. and 440'l 2 1/2).
= BGD py- Ex BGD pt 2 e 1
(3) control the tail term via centering.
Using lea 1| S |xr|el=] and Step (1),
Σ E (e - 1)
(v) 3
E Je - BB
Σ w
JER
JER
JER
Combining with Step (2) and choosing sufficiently small so that {B < 1 and 4/77 S (which implies eBB < and e Aln'l < e), we obtain the explicit bound
(b) (e (5 11 Σ u JER
S e S SEBBT
S
p/
By Assumption B.8 (the feasibility side), T < 4862 No. Hence
E
{ = 4 11/070 70 I w 70.
2
where the last inequality uses v0(4) > ,.
23

---

Published as a conference paper at ICLR 2026
(4) top-2 main term and curvature gap. By Lemma Assumption B.8 implies " "o > O. Since (Air... 0, the second-order Taylor bound gives cm
B.9 Ayr 2 Abin (Py) 6r. Hence
2² 2/2 for u 2 0, hence
5² (v)
5² y
uy (8) (5 - fu
(2)
A 6.. "
(4 1/2 -fwlg
(v) To
(v)
B²
Choose further so that & ; since B² = 2, this yields _ " B² S
(1)
L m 70-
(v)
(5) combine bounds. Putting Steps (3) together,
Dy DCD
5 w To + (x)()) gr 70 =
E w 70-
S
Ew 7o +
4
Therefore D, & D (1 { y0 (1)) 70). 70) Finally, using D I 22nnn D, and the scaling
n'p' ( while 220 + (17²), we may further reduce yuu so that the O(7²) factor is (1) dominated by Lug Yoo yielding DSAM - DDD
Since a; = (E, / this implies aaAM 2 dy GD. Strictness holds when & > 0, wy
(v)
> 0 and
To 0 (in particular when 5t * 0 and ply. (0, 1)); if p 0 then & = 0 and equality holds. O
C IMPLEMENTATION
Algorithm Logits- SAM (one training step)
Require: model parameters 0, batch (ar, (ar,tt n.t.t,""" y radius p
1: W + --mmhedd. weight
2: Forward pass:
3: Compute hidden states H and pre-perturbation loss free
4: g g+ &we
5: - --99/199
6: Perturbed forward pass:
7: logits <- linear(H,, W + e)
8: Compute perturbed loss Eosst
9: Update:
10: Backpropagate poon
11: Optimizer step on 0
D EQUIVALENCE BETWEEN THEORETICAL AND PRACTICAL SETTINGS
Table 5: Comparison between theoretical and practical settings of DPO with SAM. Although the signs differ yo the resulting dynamics are equivalent. For yt, the settings coincide.
Class
Objective
Learning rate
Setting
Positive objective log p
Positive (17
Positive
Theory Practice
(Theory)
Positive objective f == log P
Negative (1) < 0)
Negative
Theory
y (Practice)
Negative objective log p
Positive (17 > 0)
Positive
Practice
Equivalence of sign conventions for yr Theoretical setting (J log 17 < 0. p- < 0):
0 sheory
O. Vf(0.).
Practical setting uses f -f n -n 0:
0 0, v ((.)) = O: --- Otheory
24

---

Published as a conference paper at ICLR 2026
Method
Objective
Hyperparameter
SLiC-HF
max(0, 00, 8 - log (B À log no (ywr x)
2)) + log Te( x))
X
[0.1, 0.5, 1.0, 10.0);
e
(0.5, 1.0, 2.0, 10.0)
CPO
log (B log (V x) - 8 log no (yi x)) À Algg my (ywr z)
X = 1.0: B
E
(0.01, 0.05, 0.13
Table 6: Objectives and hyperparameters for SLiC-HF and CPO.
Method
Pythia-2.8B
Mistral-7B
DPO
10-1
10-5
SLiC-HF
x 10-3
1 x 10
CPO
x 10
x 10-5
Table 7: Choice of p for logits-SAM.
For SAM, theoretical perturbation and update:
(0()
Et
theory
0 gtheary
0: - Vf(0,
theory.
[V(0))]l
Practical setting uses f = f. p=
-p- > 0, -1] > 0:
Vf(0.)
vf(O,)
gucc
theory
E
[VS(0.)]
[V))
grac 0: - VV e @theory
E ADDITIONAL EXPERIMENTAL DETAILS
Benchmark details. AlpacaFval Dubois et al. 2024) is large-scale preference benchmark for open-ended instruction following that uses L..M-as-2 a---a----uugee calibrated to human preferences; its evaluation set contains 805 single-turn instructions, and models are typically compared pairwise settings against bascline. Arena-Hard v0.1 (Li et al. (024) is a challenging subset of difficult user instructions mined from Chatbot Arena; enables fine-grained, head-to-head comparisons between models via pairwise judging and comprises 500 hard prompts. MT-Bench (Zheng et al. 2023) is a multi-turn dialogue benchmark that tests a model's ability to handle diverse conversational tasks across several categories; the standard evaluation set consists of 80 multi-turn questions.
Training details. For experiments on Pythia-2.8B, we use two NVIDIA A100 GPUs with data- parallel training under DDP; for Mistral-7B. we use four NVIDIA A100 GPUs with parallel training via DeepSpeed ZeRO-3 (Rasley et al. 2020).
F LLM USAGE STATEMENT
In preparing this manuscript, we employed an LLM as an auxiliary tool. Specifically, the LLM was used assist with proofreading, formatting, grammar checking, and verification of mathematical proofs. All content and conclusions remain the responsibility of the authors.
25


Reviews:
[{"review_id": "lfivKTitTh", "reviewer": "Reviewer_DsoC", "summary": "This paper aims to address the squeezing effect of DPO and proposes a new DPO method with leveraging sharpness-aware minimization (SAM). Comprehensive theoretical and experimental analyses have been conducted to demonstrate the effectiveness of the SAM. As such, while this work has some limitations. I still recommend an acceptance.", "questions": "Please refer to weaknesses.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "This work has the following strengths:\n\n1.\tThis work studies on an important problem.\n2.\tThis work proposes to leverage SAM to mitigate the squeezing effect of DPO, with providing comprehensive theoretical evidences.\n3.\tExtensive experiments on real-world datasets have been conducted to verify the efficacy of the proposed method.", "weaknesses": "I also have some concerns: \n\n1.\tOn the description of “gradient descent with a negative learning rate”: The phrase “gradient descent with a negative learning rate” is unconventional, as learning rates in DPO are typically positive. This may cause unnecessary confusion for readers. From the theoretical analysis, I understand that DPO may, in certain cases, apply a reversed gradient direction. However, it would be more precise to present this as theoretically equivalent to using a negative learning rate, rather than stating that the learning rate is negative outright. This clarification is important to prevent misinterpretation.\n\n2.\tGap between theoretical setting and practical application: For ease of theoretical derivation, the authors formulate DPO as a multi-class logistic classification problem. While this abstraction facilitates analysis, it may not fully reflect practical usage scenarios of DPO. Additional explanation should be provided to justify the relevance and applicability of the theoretical setting to real-world DPO implementations.\n\n3.\tMany techniques in the theoretical section appear to follow the framework introduced by Ren & Sutherland. This is not necessarily a critical limitation, as research naturally builds upon prior work. However, the following questions should be addressed: (1) What are the specific theoretical challenges in integrating SAM into DPO? (2) What are the novel theoretical contributions of this work compared with Ren & Sutherland’s approach?\n\n\n\n4.\tI also some concerns on the experiments: 1) The “squeezing effect” is not a new concept and has been studied in recent work (e.g., [a1], [a2]). How does the proposed method perform compared with these recent baselines? 2) It would be better to include more experiments on diverse backbones and benchmark. \n\n[a1] C2-DPO: Constrained Controlled Direct Preference Optimization (arxiv'25)\n[a2] Unintentional unalignment: Likelihood displacement in direct preference optimization (ICLR’25)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ozYoVXn3J1", "reviewer": "Reviewer_5zL9", "summary": "The paper analyzes the ``squeezing'' effect in DPO—where preferred responses lose likelihood—by deriving a unified logit-space dynamic that links GD and SAM and reveals high-curvature modes as the culprit. It proves that SAM, applied with a coefficient sharing the sign of the effective learning rate, damps these unstable modes. Leveraging this, **logits-SAM** perturbs only the output layer, preserving curvature-aware regularization with ~2–3% overhead. Experiments on Pythia-2.8B and Mistral-7B across HH, TL;DR, UltraFeedback, AlpacaEval 2, Arena-Hard, and MT-Bench show consistent gains and reduced sharpness, with robust performance for small ρ (≈1e-5–1e-4).", "questions": "see weakness", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. **Clear Theoretical Contribution:** Provides a unified logit-space dynamical analysis that links GD and SAM, pinpointing high-curvature mode amplification as the mechanism behind DPO ``squeezing,'' and proving sign-aligned SAM mitigates it.\n\n2. **Practical, Efficient Method:** Introduces logits-SAM, an output-layer perturbation that retains curvature-aware regularization with minimal overhead (~2–3%) and seamless integration into existing DPO/SLiC-HF/CPO pipelines.", "weaknesses": "1. **Theory–practice gap:** Core analysis relies on first/second-order approximations in logit space (fixed features, softmax CE), which may not fully capture nonlinearity and parameter coupling in deep, decoder-only LMs.\n\n2. **Final-layer perturbation bias:** Restricting SAM to the output layer improves efficiency but may miss sharp directions arising in earlier blocks/attention layers, potentially undercutting robustness on harder distributions.\n\n3. **DPO-specific framing:** The mitigation is analyzed through the ``negative learning rate / squeezing'' view of DPO; it is unclear how strongly the guarantees transfer to alternative preference-learning objectives or online RLHF settings.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "2n7WAFn7EG", "reviewer": "Reviewer_6UP6", "summary": "This paper seeks to address the \"squeezing effect\" (or \"likelihood displacement\") issue in DPO fine-tuning. This is a potentially harmful phenomenon, where the model performance degrades during DPO when it should be improving. The authors aim to theoretically analyze this behavior, and propose the root cause is due to interaction between the DPO objective and gradient descent. It is stated that rapid model updates expanding along high-curvature directions in the logit space explains the \"squeezing effect\". To remedy this, the authors propose using Sharpness-Aware Minimization (SAM), which high level seeks to minimize to \"flat\" surfaces. Since full SAM is too slow, the authors introduce a more practical variant, logits-SAM, which perturbs the outer-layer parameters only.", "questions": "- Section 3.2 explicitly argues that for the negative learning rate in DPO, \"one should choose a negative $\\rho$\". Section 3.3 says that the implementation \"consistently use a positive $\\rho$\". The claim that this is equivalent hold for first-order methods, but is there a leap to assume it holds for curvature-aware methods like SAM? \n\n- Have the authors conducted experiments on larger more currently relevant models beside GPT-2?\n\n- DPO (Eq. 2) already has an implicit KL regularization term. AdamW has weight decay. How does SAM's curvature regularization interact with these forms of regularization? \n\n- Why is perturbing only the last layer sufficient? Proposition 3.1 links the full parameter Hessian to the logit Hessian via a features map. Yet this is the output of all other layers. Can the authors give any intuition on why this is sufficient to regularize the entire system? Or is GPT-2 the limitation?", "limitations": "", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- The authors present rigorous theoretical diagnosis for a known DPO failure mode. Correlating the \"squeezing effect\" to \"high-curvature directions\" in the logit space  is a specific, actionable, and clear. The theoretical contributions and discussions are also clearly written, and presented as an elegant explanation for an existing harmful phenomenon. The clear theory to practical pipeline is also well motivated. Section 3.2 predicts how SAM should behave, Fig1-a,b confirm this prediction on a toy problem, and Fig1-c shows it on a real GPT-2 model.\n\n- The message being delivered is also succinct: choose the SAM hyperparameter with the same sign as the learning rate to alleviate the squeezing effect. \n\n- The paper's primary practical contribution is the logits-SAM algorithm. This identifies that only the last layer needs perturbation, and the authors propose a simple practical \"fix\". Beyond just merely stating the instability of DPO and analyzing its failure modes, this additional practical solution is a strong addition.", "weaknesses": "- This method introduces a new highly sensitive hyperparameter ($\\rho$). Table 3 demonstrates that while $\\rho=10^{-4}$ might yield good results, the slightly larger $\\rho=10^{-3}$ actually performs worse than the baseline. The optimal range seems extremely small. The majority of preference fine-tuning methods already encompass highly sensitive hyperparameters, and discussions about reference-based versus reference-free preference fine-tuning, in order to shift towards more robust solutions that are cheaper to tune. \n\n- Logits-SAM is not a new method. There are several works in this area recently, making this paper's contribution mildly incremental. The novelty of this work is in its application of logits-SAM in this particular setting. [1], [2], [3]\n\n- Efficiency is highly model-dependent. The claim that overhead is negligible requires the final output later to be a tiny fraction of the total model parameters. This is true for many decoder-only models, but it is not a universal guarantee. Hence the experiments conducted are all on small models like GPT-2 with toy datasets, and are largely outdated. \n\n[1] Wen, Kaiyue, Tengyu Ma, and Zhiyuan Li. \"How does sharpness-aware minimization minimize sharpness?.\" arXiv preprint arXiv:2211.05729 (2022).\n\n[2] Li, Dongqi, et al. \"Sharpness-aware minimization for out-of-distribution generalization.\" International Conference on Neural Information Processing. Singapore: Springer Nature Singapore, 2023.\n\n[3] Foret, Pierre, et al. \"Sharpness-aware minimization for efficiently improving generalization.\" arXiv preprint arXiv:2010.01412 (2020).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "8g6GbqFXu0", "reviewer": "Reviewer_vz7v", "summary": "This paper investigates the \"squeezing effect\" in DPO, where preferred response probabilities unintentionally decrease. Through a theoretical analysis of logit space dynamics, the authors identify rapid expansion along high-curvature directions under standard optimization as the cause. They propose that Sharpness-Aware Minimization (SAM) mitigates this via curvature regularization.\n\nTo apply this efficiently, they advocate for logits-SAM, perturbing only the final layer. Empirical results on Pythia-2.8B and Mistral-7B across various alignment tasks show that logits-SAM consistently improves DPO performance with negligible computational overhead.", "questions": "1. The theoretical framework links the mitigation of the squeezing effect (when modeling the $y^-$ dynamics with a conceptual $\\eta < 0$) to using a negative perturbation radius ($\\rho < 0$). Since $\\rho$ in standard SAM represents the magnitude of perturbation within a neighborhood, could the authors provide further intuition or clarification on the interpretation and role of a negative $\\rho$ in this theoretical context, especially regarding how it achieves the desired curvature regularization?\n2. Could the authors provide empirical results from the real-world DPO experiments (similar to Figure 1c) that explicitly track the probability dynamics of the most confident incorrect prediction ($y^*$)?", "limitations": "", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. The paper tackles the important and recently identified \"squeezing effect\" issue in DPO, which is a practical limitation of the algorithm.\n2. It provides a novel theoretical framework connecting parameter-space and logit-space dynamics, specifically analyzing the role of curvature in the squeezing effect under GD and how SAM counteracts it. The insight about matching the sign of the perturbation radius $\\rho$ with the effective learning rate $\\eta$ is particularly interesting.\n3. Based on their theoretical insights into curvature and the squeezing effect, the paper identifies the computationally efficient logits-SAM variant as a suitable mitigation strategy and successfully applies it to the DPO context, demonstrating its practical value for improving alignment.", "weaknesses": "1. The core theoretical analysis simplifies DPO by modeling the dynamics associated with preferred ($y^+$, using positive $\\eta$) and dispreferred ($y^-$, using negative $\\eta$) responses separately within a multiclass classification framework. However, the actual DPO loss function couples the gradients for $y^+$ and $y^-$ through the sigmoid of their log-likelihood ratio difference. This theoretical separation doesn't fully capture the coupled nature of the true DPO gradient dynamics, potentially limiting the direct applicability of conclusions drawn from the simplified, independent analyses (e.g., regarding negative $\\eta/\\rho$).\n2. A key aspect of the squeezing effect is the unintended probability increase of the most confident incorrect prediction ($y^\\star$). While the theory (Corollary 3.6) and toy example (Fig 1b) suggest SAM suppresses this increase, the main real-world experiment (Fig 1c) only tracks the chosen ($y^+$) and rejected ($y^-$) probabilities. It lacks direct empirical evidence showing SAM also prevents the rise of $y^*$ probabilities during actual DPO training.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)