---
title: 'Logits-SAM: Logit-Aware Optimization'
url: https://www.emergentmind.com/topics/logits-sam
type: topic
---

# Logits-SAM: Logit-Aware Optimization

Searching arXiv for the cited papers and related Logits-SAM work.
arXiv search query: "2405.09492 OR 2502.02407 OR 2603.18258 Logits-SAM"
Logits-SAM is a polysemous term in the recent sharpness-aware optimization literature. In one line of work, it denotes a continual-learning mechanism that couples soft-logit replay with Sharpness-Aware Minimization (SAM) so that perturbations are biased toward memory-preserving directions, as instantiated by MGSER-SAM [2405.09492]. In another, it names a failure mode of vanilla SAM in large-vocabulary language modeling, where apparent sharpness reduction is achieved predominantly by manipulating logit statistics rather than improving the geometry of the learned function [2502.02407]. In a third, more recent usage, it denotes an efficient SAM variant for Direct Preference Optimization (DPO) that perturbs only the output layer, thereby operating effectively in logit space with negligible overhead [2603.18258]. Across these usages, the common thread is that the interaction between SAM and logits is treated as the central object of analysis or intervention.

## 1. Terminological scope and historical placement

The literature does not use “Logits-SAM” in a single canonical sense. The 2024 continual-learning paper describes Logits-SAM as instantiated by MGSER-SAM, where soft-logit matching is integrated into replay-based continual learning and coupled to SAM in ER/DER++-style rehearsal [2405.09492]. The 2025 NLP study uses “Logits-SAM” to designate a pathological mechanism: SAM’s sharpness objective in language modeling is spuriously minimized through changes to logit statistics, especially confidence-related effects, rather than through improvements in function-space curvature [2502.02407]. The 2026 DPO work then reuses the term as the name of a constructive algorithm, “logits-SAM,” which perturbs only the output layer and is motivated by a logit-space curvature analysis of preference optimization [2603.18258].

This terminological divergence is consequential. It implies that “Logits-SAM” may refer to a replay regularizer, a diagnostic failure mode, or an efficient output-layer SAM approximation, depending on context. A common misconception is that these papers study the same algorithm. They do not. What unifies them is a focus on how SAM interacts with pre-softmax outputs, output-space curvature, or stored logits.

| Usage | Core meaning | Setting |
|---|---|---|
| MGSER-SAM | Soft-logit replay plus memory-guided SAM direction | Continual learning |
| Logits-SAM as diagnosis | Spurious sharpness minimization through logit statistics | NLP pretraining |
| logits-SAM for DPO | Output-layer-only SAM in logit space | Preference optimization |

This suggests that the term has evolved from a replay-oriented mechanism to a broader analytical lens on SAM’s behavior in output space.

## 2. Logits-SAM in continual learning: MGSER-SAM

In continual learning, the relevant problem is catastrophic forgetting under sequentially presented tasks. The continual-learning objective is stated as
\[
\theta^{*}  = \argmin_{\theta} (\mathcal{L}_{t} + \mathcal{L}_{s}),
\]
where
\[
\mathcal{L}_{t} = \mathbb{E}_{(\mathbf{x}_t, y_t)\sim \mathcal{D}_t} [l(f_\theta(\mathbf{x}_t), y_t)]
\]
and \(\mathcal{L}_{s}\) is the loss over previously learned tasks [2405.09492]. In ER, current-task minibatches are merged with memory minibatches and optimized through a combined replay loss. The paper evaluates task-incremental, class-incremental, and domain-incremental scenarios, with standard continual-learning metrics of average accuracy and forgetting.

SAM is first inserted into replay, producing ER-SAM through the min–max objective
\[
\min_{\bm{\theta} \max_{\|\bdelta \|_2 \leq \rho} L_{total}(\bm{\theta}+\bdelta),
\]
with first-order ascent perturbation
\[
\bdelta^{\star} \approx \rho \frac{\nabla_{\bm{\theta} \mathcal{L}_{total} (\bm{\theta})}{\|\nabla_{\bm{\theta} \mathcal{L}_{total} (\bm{\theta})\|_2}.
\]
The difficulty identified in continual learning is that the single combined loss \(\mathcal{L}_{total}=\mathcal{L}_t+\mathcal{L}_s\) can induce conflicting gradient directions between the current task and memory. The paper explicitly notes that when the angle between \(\nabla \mathcal{L}_t\) and \(\nabla \mathcal{L}_s\) exceeds \(90^\circ\), the averaged direction can be suboptimal, expressing the stability-plasticity dilemma.

MGSER-SAM addresses this through a soft-logits-augmented memory term
\[
\mathcal{\hat{L}_{s} = \mathbb{E}_{(\mathbf{x}, y) \sim \mathcal{B} [l(f_\theta(\mathbf{x}), y)] + \mathbb{E}_{(\mathbf{x'}, \mathbf{z'}) \sim \mathcal{B} [\|h_{\bm{\theta} (\mathbf{x'})} - \mathbf{z'} \|_2],
\]
where \(h_{\bm{\theta}(\mathbf{x'})}\) are pre-softmax logits and \(\mathbf{z'}\) are stored soft logits saved in memory when \(\mathbf{x'}\) was learned [2405.09492]. The formulation uses MSE on logits rather than KL-based distillation, and no temperature \(T\) is introduced in the paper’s equation. The combined soft-total loss is \(\hat{\mathcal{L}}_{total}=\mathcal{L}_t+\hat{\mathcal{L}}_s\).

The distinctive “Logits-SAM” component is the perturbation used in descent evaluation:
\[
\bm{g}^{MGSER-SAM}  = \nabla_{\bm{\theta} \mathcal{\hat{L}_{total} (\bm{\theta}) |_{\bm{\theta}+ \bdelta^{\star} + \nabla_{\bm{\theta} \mathcal{\hat{L}_{s} (\bm{\theta}) |_{\bm{\theta} ,
\]
with
\[
\bm{\theta}_{t+1}   = \bm{\theta}_{t} - \eta \bm{g}^{MGSER-SAM}.
\]
In words, the descent gradient is evaluated not only at \(\bm{\theta}+\bdelta^\star\) but at a further shifted point incorporating \(\nabla \hat{\mathcal{L}}_s\), thereby biasing the update toward memory-consistent gradients. The paper does not introduce an explicit cosine-similarity penalty; the alignment is implemented by construction through the additive perturbation.

Algorithmically, the method uses reservoir sampling, a current-task batch \(\mathcal{B}_t\), a memory batch \(\mathcal{B}_1=\{(\mathbf{x},y)\}\) for empirical replay loss, and a memory-logits batch \(\mathcal{B}_2=\{(\mathbf{x'},\mathbf{z'})\}\) for soft-logit matching. The paper sets the same batch size for \(\mathcal{B}_t\), \(\mathcal{B}_1\), and \(\mathcal{B}_2\), and uses equal weights for the two terms in \(\hat{\mathcal{L}}_s\) for simplicity.

## 3. Theoretical interpretation: logit paths, functional paths, and spurious sharpness minimization

The 2025 NLP paper reframes the interaction between SAM and logits at a more general level. Its starting point is the standard SAM objective
\[
\min_w \max_{\|\epsilon\|_p \le \rho} L(w + \epsilon),
\]
with practical perturbation
\[
\epsilon^* = \rho \frac{\nabla_w L(w)}{\|\nabla_w L(w)\|_p},
\]
and update
\[
w_{t+1} = w_t - \eta \nabla_w L(w_t + \epsilon^*).
\]
For cross-entropy with logits \(z \in \mathbb{R}^K\), the paper states
\[
\nabla_z L = p - y,
\qquad
\nabla^2_z L = \operatorname{diag}(p) - p p^\top,
\]
where the logit Hessian is positive semidefinite and shrinks as the distribution becomes more one-hot [2502.02407].

The key decomposition uses \(\nabla_w L(w) = J(w)^\top \nabla_F L\), where \(F(w)\) denotes the stacked network outputs over the dataset and \(J(w)=\partial F/\partial w\). The Hessian is split as
\[
H = H_{GGN} + H_{func},
\]
with
\[
H_{GGN} = J^\top (\nabla^2_F L) J
\]
and \(H_{func}\) the indefinite functional Hessian arising from second derivatives of the model outputs with respect to parameters. This yields two additive routes by which SAM can reduce its penalty term: a **logit path**, mediated by changes in loss sensitivity to logits, and a **functional path**, mediated by changes in the function Jacobian.

The paper defines normalized contributions
\[
\alpha_{logit} = \frac{\|H_{GGN} v^*\|^2}{\|H v^*\|^2},
\quad
\alpha_{func} = \frac{\|H_{func} v^*\|^2}{\|H v^*\|^2},
\quad
\alpha_{cross} = \frac{2 \langle H_{GGN} v^*, H_{func} v^* \rangle}{\|H v^*\|^2},
\]
so that \(\alpha_{logit}+\alpha_{func}+\alpha_{cross}=1\) [2502.02407]. In vision, \(\alpha_{logit}\) starts near \(0\) and then becomes comparable to \(\alpha_{func}\); in NLP on C4, \(\alpha_{logit} \gg \alpha_{func}\) for most of training, and \(\alpha_{cross}\) tends to be negative. The paper interprets this as SAM following the easy logit-statistic shortcut rather than improving the function’s geometry.

This is the sense in which the paper calls the phenomenon “Logits-SAM.” Increasing confidence through temperature-like rescaling can shrink \(\nabla^2_z L\) and hence \(H_{GGN}\) without necessarily improving perplexity or the Jacobian statistics of the learned function. A plausible implication is that low measured sharpness is not, by itself, evidence of function-space regularization in large-vocabulary language modeling.

## 4. Corrective methods: Functional-SAM and preconditioned SAM

Having diagnosed Logits-SAM as a failure mode, the 2025 paper proposes two remedies. The first is Functional-SAM, which keeps SAM’s finite-difference structure but suppresses the logit route. Let
\[
g = \nabla_w L(w), \qquad v^* = \frac{g}{\|g\|_p}, \qquad w' = w + \rho v^*.
\]
Functional-SAM then uses
\[
g_{FnSAM} = -[J(w')]^\top \nabla_F L(w)
\]
and updates \(w\) through an outer optimizer such as AdamW [2502.02407]. The paper’s interpretation is that the usual SAM gradient can be decomposed into a logit path term, a functional path term, and the base gradient; Functional-SAM discards the logit path and keeps the functional path.

The second remedy is preconditioned SAM. If \(P \approx V_t^{-1}\) is the elementwise inverse of Adam’s second-moment accumulator, the perturbation becomes
\[
\epsilon = \rho \frac{P v^*}{\|P v^*\|_p},
\]
and the perturbed gradient is then passed to AdamW. The proposed motivation is that SAM’s perturbation is otherwise taken in unpreconditioned Euclidean geometry, whereas Adam-like optimizers use a different diagonal geometry. This mismatch, according to the paper, makes SAM align with dominant \(H_{GGN}\) directions and exacerbates Logits-SAM.

The combined method, preconditioned Functional-SAM, uses the same preconditioned perturbation inside the Jacobian-scoped update:
\[
g_{pFnSAM} = -[J(w + \epsilon)]^\top \nabla_F L(w).
\]
The paper states that the combination increases robustness and consistently ameliorates Logits-SAM, especially at scale [2502.02407].

The reported empirical pattern supports the diagnosis. On a 2M model in the fixed-length setting, the evaluation losses are 3.90 for AdamW, 3.91 for SAM, 3.89 for Functional-SAM, 3.88 for preconditioned SAM, and 3.86 for preconditioned Functional-SAM. At 23.9M in Chinchilla-like training, AdamW yields eval 3.69, SAM 3.71, preconditioned Functional-SAM 3.66, and preconditioned SAM 3.63, while curvature metrics such as \(\lambda_{\max}(H)\), \(\operatorname{tr}(H)\), and \(\operatorname{tr}(H_{GGN})\) are all reduced by SAM-family methods [2502.02407]. The paper emphasizes that vanilla SAM can lower \(\lambda_{\max}(H)\) while worsening validation loss, which it takes as evidence of spurious minimization through logit curvature alone.

## 5. Logits-SAM for Direct Preference Optimization

The 2026 DPO paper returns to a constructive meaning of logits-SAM, but now in preference optimization rather than continual learning. DPO is formulated over preference pairs \((x,y^+,y^-)\), with target policy \(\pi_\theta(y|x)\) and fixed reference policy \(\pi_{ref}(y|x)\). The paper analyzes DPO through a coordinate-wise dynamical system in logit space and identifies the “squeezing effect” or likelihood displacement: the probability assigned to preferred responses can decrease unintentionally during training [2603.18258].

In the fixed-feature regime, the model uses last-layer parameters \(W^t \in \mathbb{R}^{V \times d}\), fixed features \(\phi(x)\), logits \(z^t = W^t \phi\), probabilities \(p^t = \operatorname{softmax}(z^t)\), logit gradient \(g^t = p^t - y\), and logit Hessian
\[
H_{z,t} = \operatorname{Diag}(p^t) - p^t (p^t)^\top.
\]
A key proposition states
\[
H_W = (\phi \phi^\top)\otimes H_z,
\]
and \(\operatorname{rank}(H_W)=\operatorname{rank}(H_z)\) if \(\phi \neq 0\), so parameter-space second-order effects can be studied entirely in logit space [2603.18258].

The SAM-corrected logit dynamics are given by
\[
z^{t+1} = z^{t} - \eta\,\mu\Big(g^{t} + \tilde\rho^{\,t}\,H_{z,t}\,g^{t}\Big) +R_z^{t},
\]
and
\[
g^{t+1} = \Big(-\eta\,\mu\,H_{z,t}-\eta\,\mu\,\tilde\rho^{\,t}\,H_{z,t}^{2}\Big)\,g^{t} +R_g^{t},
\]
with \(\tilde\rho^{\,t}=\rho \sqrt{\mu}/\|g^t\|\) and \(\|R_z^t\|,\|R_g^t\| \le C\eta^2\) [2603.18258]. In the eigenbasis of \(H_{z,t}\), residual coordinates satisfy
\[
e_{k}^{t+1} = \Big( 1 - \eta\,\mu\,\big[\lambda_{k}^{t}  + \tilde\rho^{\,t}(\lambda_{k}^{t})^{2}\big] \Big)\, e_{k}^{t} + r_{k}^{t},
\]
which the paper interprets as showing that negative-gradient updates expand residuals most strongly along high-curvature modes, thereby producing squeezing.

The practical algorithm perturbs only the output layer \(W\), typically the lm_head:
\[
\mathcal{L}^{\text{logits-SAM}_{\mathrm{DPO}(\theta;\,W;\,x,y^+,y^-) = \mathcal{L}_{\mathrm{DPO}\!\left( \theta;\,W +\rho \frac{\nabla_{W}\,\mathcal{L}_{\mathrm{DPO}(\theta;\,W;\,x,y^+,y^-)} {\big\|\nabla_{W}\,\mathcal{L}_{\mathrm{DPO}(\theta;\,W;\,x,y^+,y^-)\big\|}, \;x,y^+,y^- \right).
\]
Equivalently, with logits \(z=W\phi\), one may view the method as applying a normalized perturbation directly in logit space. The implementation computes a pre-perturbation DPO loss, forms \(e=\rho g/\|g\|\) on \(W\), caches the penultimate hidden states \(H\), recomputes only the final linear projection with \(W+e\), and backpropagates the perturbed loss once. Because only the output layer is recomputed, the paper reports negligible overhead.

## 6. Empirical profile, computational trade-offs, and points of contention

The empirical picture depends strongly on domain and on which “Logits-SAM” meaning is intended.

In continual learning, MGSER-SAM is reported to outperform replay baselines across task-IL, class-IL, and domain-IL benchmarks. The paper highlights that on S-CIFAR10, class-IL, \(M=1000\), “MGSER-SAM obtains \(78.51\%\) testing accuracy in class-IL, which is \(24.4\%\) higher than the ER method,” and that on S-MNIST, class-IL, \(M=2000\), “MGSER-SAM achieves \(93.29\%\) testing accuracy in class-IL, which is \(4.2\%\) and \(17.6\%\) higher than the ER and DER++ methods respectively” [2405.09492]. It also reports \(70.08\%\) task-IL accuracy on S-CIFAR100 with \(M=400\), which is \(6.3\%\) higher than DER++, and on S-TinyImageNet with \(M=1000\) it reports best \(21.22\%\) class-IL and \(54.01\%\) task-IL performance. The paper states that MGSER-SAM achieves the lowest forgetting on each benchmark. The cost is approximately \(2\times\) backpropagation relative to ER because SAM requires one pass to compute \(\bdelta^\star\) and another for the descent gradient.

In NLP pretraining, the 2025 paper reaches nearly the opposite operational conclusion for vanilla SAM: it “degraded validation loss vs AdamW and was twice as expensive per step,” and “for SAM, increasing \(\rho\) worsened performance; optimal \(\rho\) was effectively 0 (no SAM)” [2502.02407]. The corresponding claim is not that sharpness-aware methods fail universally, but that naive SAM fails because it is dominated by logit-statistics regularization. Functional-SAM and preconditioned variants are introduced specifically to avoid this spurious route.

In DPO, output-layer logits-SAM is presented as a computationally efficient favorable approximation. The paper reports that on Pythia-2.8B, DPO yields win rates of 70.52 vs SFT and 56.35 vs chosen on HH, whereas DPO+logits-SAM yields 72.28 and 60.51; on TL;DR, DPO yields 84.21 and 34.78, whereas DPO+logits-SAM yields 89.58 and 36.57 [2603.18258]. It also reports gains when combined with SLiC-HF and CPO, improvements on Mistral-7B-v0.1 across AlpacaEval 2, Arena-Hard, and MT-Bench, and safety improvements for Gemma-2B-IT on SorryBench. On Pythia-2.8B trained on Reddit TL;DR with two A100 GPUs, the additional wall-clock time is reported as only about 2–3%, with negligible peak memory, whereas vanilla SAM nearly doubles step time and requires a perturbation buffer comparable to model size.

These results clarify an important controversy. “Logits-SAM” is not intrinsically beneficial or intrinsically harmful. In one paper, it names the problematic route by which SAM spuriously lowers measured sharpness in language modeling [2502.02407]. In another, it names the deliberate restriction of SAM to logit-space curvature in DPO, where output-layer control is the intended mechanism and empirically helps [2603.18258]. The continual-learning use is again different: there, soft logits are used to steer the perturbation toward rehearsal-consistent directions [2405.09492].

## 7. Conceptual synthesis and open directions

Taken together, the three lines of work establish that logits are not a peripheral detail in SAM-family methods. They can serve as replay targets in continual learning, as a confounding shortcut in large-vocabulary language modeling, or as the explicit domain in which curvature control is applied for preference optimization.

A unifying interpretation is that these papers all separate, implicitly or explicitly, the geometry of the network’s function from the geometry induced by the loss at the output layer. In MGSER-SAM, soft logits preserve past predictive structure and inject memory gradients into the SAM perturbation, thereby reconciling current-task and replay directions [2405.09492]. In the NLP diagnosis, the Gauss–Newton term
\[
H_{GGN} = J^\top (\nabla_F^2 L) J
\]
captures a logit-side curvature component that can dominate optimization and anti-align with the functional route [2502.02407]. In DPO, the identity
\[
H_W = (\phi\phi^\top)\otimes H_z
\]
justifies restricting SAM to the output layer, because the curvature structure relevant to squeezing factors through the logit Hessian \(H_z\) [2603.18258].

Several limitations recur across the literature. All SAM-family methods still introduce extra computation relative to plain first-order training, though the amount varies sharply by implementation. MGSER-SAM depends on the quality and representativeness of replay memory and stored logits, and it uses equal weights for cross-entropy and logit-MSE without tuning [2405.09492]. Functional-SAM and preconditioned SAM preserve roughly the \(2\times\) per-step structure of SAM, even if they redirect the regularization path [2502.02407]. DPO logits-SAM is efficient, but it perturbs only the output layer, so deeper-layer sources of squeezing are not directly regularized, and performance degrades when \(\rho\) is too large [2603.18258].

The open questions identified in the source literature point in related directions: sharper definitions of function-space versus parameter-space or output-space sharpness; more expressive preconditioners than diagonal Adam inverse; broader curvature control in intermediate layers; adaptive \(\rho\) schedules; and combinations with objective-level or data-centric interventions such as CHES, IPO, SimPO, or Cal-DPO [2502.02407; 2603.18258]. This suggests that “Logits-SAM” is likely to remain less a single fixed algorithm than a family of perspectives on how sharpness-aware optimization should interact with output-space structure.

Source: https://www.emergentmind.com/topics/logits-sam