Logits-SAM: Logit-Aware Optimization
- Logits-SAM is a term that defines approaches leveraging pre-softmax logit manipulations to guide sharpness-aware minimization across diverse optimization settings.
- It integrates techniques like soft-logit replay in continual learning, identifies diagnostic failure in NLP pretraining, and enables efficient output-layer perturbations in DPO.
- Empirical evaluations indicate that logit-based adaptations can improve task accuracy and training stability while managing computational trade-offs in optimization.
Searching arXiv for the cited papers and related Logits-SAM work. arXiv search query: "(Li et al., 2024) OR (Singh et al., 4 Feb 2025) OR (Luo et al., 18 Mar 2026) 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 (Li et al., 2024). 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 (Singh et al., 4 Feb 2025). 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 (Luo et al., 18 Mar 2026). 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 (Li et al., 2024). 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 (Singh et al., 4 Feb 2025). 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 (Luo et al., 18 Mar 2026).
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
where
and is the loss over previously learned tasks (Li et al., 2024). 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 can induce conflicting gradient directions between the current task and memory. The paper explicitly notes that when the angle between and exceeds , 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 0 are pre-softmax logits and 1 are stored soft logits saved in memory when 2 was learned (Li et al., 2024). The formulation uses MSE on logits rather than KL-based distillation, and no temperature 3 is introduced in the paper’s equation. The combined soft-total loss is 4.
The distinctive “Logits-SAM” component is the perturbation used in descent evaluation: 5 with
6
In words, the descent gradient is evaluated not only at 7 but at a further shifted point incorporating 8, 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 9, a memory batch 0 for empirical replay loss, and a memory-logits batch 1 for soft-logit matching. The paper sets the same batch size for 2, 3, and 4, and uses equal weights for the two terms in 5 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
6
with practical perturbation
7
and update
8
For cross-entropy with logits 9, the paper states
$\min_{\bm{\theta} \max_{\|\bdelta \|_2 \leq \rho} L_{total}(\bm{\theta}+\bdelta),$0
where the logit Hessian is positive semidefinite and shrinks as the distribution becomes more one-hot (Singh et al., 4 Feb 2025).
The key decomposition uses $\min_{\bm{\theta} \max_{\|\bdelta \|_2 \leq \rho} L_{total}(\bm{\theta}+\bdelta),$1, where $\min_{\bm{\theta} \max_{\|\bdelta \|_2 \leq \rho} L_{total}(\bm{\theta}+\bdelta),$2 denotes the stacked network outputs over the dataset and $\min_{\bm{\theta} \max_{\|\bdelta \|_2 \leq \rho} L_{total}(\bm{\theta}+\bdelta),$3. The Hessian is split as
$\min_{\bm{\theta} \max_{\|\bdelta \|_2 \leq \rho} L_{total}(\bm{\theta}+\bdelta),$4
with
$\min_{\bm{\theta} \max_{\|\bdelta \|_2 \leq \rho} L_{total}(\bm{\theta}+\bdelta),$5
and $\min_{\bm{\theta} \max_{\|\bdelta \|_2 \leq \rho} L_{total}(\bm{\theta}+\bdelta),$6 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
$\min_{\bm{\theta} \max_{\|\bdelta \|_2 \leq \rho} L_{total}(\bm{\theta}+\bdelta),$7
so that $\min_{\bm{\theta} \max_{\|\bdelta \|_2 \leq \rho} L_{total}(\bm{\theta}+\bdelta),$8 (Singh et al., 4 Feb 2025). In vision, $\min_{\bm{\theta} \max_{\|\bdelta \|_2 \leq \rho} L_{total}(\bm{\theta}+\bdelta),$9 starts near $\bdelta^{\star} \approx \rho \frac{\nabla_{\bm{\theta} \mathcal{L}_{total} (\bm{\theta})}{\|\nabla_{\bm{\theta} \mathcal{L}_{total} (\bm{\theta})\|_2}.$0 and then becomes comparable to $\bdelta^{\star} \approx \rho \frac{\nabla_{\bm{\theta} \mathcal{L}_{total} (\bm{\theta})}{\|\nabla_{\bm{\theta} \mathcal{L}_{total} (\bm{\theta})\|_2}.$1; in NLP on C4, $\bdelta^{\star} \approx \rho \frac{\nabla_{\bm{\theta} \mathcal{L}_{total} (\bm{\theta})}{\|\nabla_{\bm{\theta} \mathcal{L}_{total} (\bm{\theta})\|_2}.$2 for most of training, and $\bdelta^{\star} \approx \rho \frac{\nabla_{\bm{\theta} \mathcal{L}_{total} (\bm{\theta})}{\|\nabla_{\bm{\theta} \mathcal{L}_{total} (\bm{\theta})\|_2}.$3 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 $\bdelta^{\star} \approx \rho \frac{\nabla_{\bm{\theta} \mathcal{L}_{total} (\bm{\theta})}{\|\nabla_{\bm{\theta} \mathcal{L}_{total} (\bm{\theta})\|_2}.$4 and hence $\bdelta^{\star} \approx \rho \frac{\nabla_{\bm{\theta} \mathcal{L}_{total} (\bm{\theta})}{\|\nabla_{\bm{\theta} \mathcal{L}_{total} (\bm{\theta})\|_2}.$5 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
$\bdelta^{\star} \approx \rho \frac{\nabla_{\bm{\theta} \mathcal{L}_{total} (\bm{\theta})}{\|\nabla_{\bm{\theta} \mathcal{L}_{total} (\bm{\theta})\|_2}.$6
Functional-SAM then uses
$\bdelta^{\star} \approx \rho \frac{\nabla_{\bm{\theta} \mathcal{L}_{total} (\bm{\theta})}{\|\nabla_{\bm{\theta} \mathcal{L}_{total} (\bm{\theta})\|_2}.$7
and updates $\bdelta^{\star} \approx \rho \frac{\nabla_{\bm{\theta} \mathcal{L}_{total} (\bm{\theta})}{\|\nabla_{\bm{\theta} \mathcal{L}_{total} (\bm{\theta})\|_2}.$8 through an outer optimizer such as AdamW (Singh et al., 4 Feb 2025). 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 $\bdelta^{\star} \approx \rho \frac{\nabla_{\bm{\theta} \mathcal{L}_{total} (\bm{\theta})}{\|\nabla_{\bm{\theta} \mathcal{L}_{total} (\bm{\theta})\|_2}.$9 is the elementwise inverse of Adam’s second-moment accumulator, the perturbation becomes
0
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 1 directions and exacerbates Logits-SAM.
The combined method, preconditioned Functional-SAM, uses the same preconditioned perturbation inside the Jacobian-scoped update: 2 The paper states that the combination increases robustness and consistently ameliorates Logits-SAM, especially at scale (Singh et al., 4 Feb 2025).
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 3, 4, and 5 are all reduced by SAM-family methods (Singh et al., 4 Feb 2025). The paper emphasizes that vanilla SAM can lower 6 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 7, with target policy 8 and fixed reference policy 9. 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 (Luo et al., 18 Mar 2026).
In the fixed-feature regime, the model uses last-layer parameters 0, fixed features 1, logits 2, probabilities 3, logit gradient 4, and logit Hessian
5
A key proposition states
6
and 7 if 8, so parameter-space second-order effects can be studied entirely in logit space (Luo et al., 18 Mar 2026).
The SAM-corrected logit dynamics are given by
9
and
0
with 1 and 2 (Luo et al., 18 Mar 2026). In the eigenbasis of 3, residual coordinates satisfy
4
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 5, typically the lm_head: 6 Equivalently, with logits 7, one may view the method as applying a normalized perturbation directly in logit space. The implementation computes a pre-perturbation DPO loss, forms 8 on 9, caches the penultimate hidden states 0, recomputes only the final linear projection with 1, 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, 2, “MGSER-SAM obtains 3 testing accuracy in class-IL, which is 4 higher than the ER method,” and that on S-MNIST, class-IL, 5, “MGSER-SAM achieves 6 testing accuracy in class-IL, which is 7 and 8 higher than the ER and DER++ methods respectively” (Li et al., 2024). It also reports 9 task-IL accuracy on S-CIFAR100 with $\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],$0, which is $\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],$1 higher than DER++, and on S-TinyImageNet with $\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],$2 it reports best $\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],$3 class-IL and $\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],$4 task-IL performance. The paper states that MGSER-SAM achieves the lowest forgetting on each benchmark. The cost is approximately $\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],$5 backpropagation relative to ER because SAM requires one pass to compute $\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],$6 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 $\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],$7 worsened performance; optimal $\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],$8 was effectively 0 (no SAM)” (Singh et al., 4 Feb 2025). 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 (Luo et al., 18 Mar 2026). 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 (Singh et al., 4 Feb 2025). 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 (Luo et al., 18 Mar 2026). The continual-learning use is again different: there, soft logits are used to steer the perturbation toward rehearsal-consistent directions (Li et al., 2024).
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 (Li et al., 2024). In the NLP diagnosis, the Gauss–Newton 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],$9
captures a logit-side curvature component that can dominate optimization and anti-align with the functional route (Singh et al., 4 Feb 2025). In DPO, the identity
00
justifies restricting SAM to the output layer, because the curvature structure relevant to squeezing factors through the logit Hessian 01 (Luo et al., 18 Mar 2026).
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 (Li et al., 2024). Functional-SAM and preconditioned SAM preserve roughly the 02 per-step structure of SAM, even if they redirect the regularization path (Singh et al., 4 Feb 2025). 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 03 is too large (Luo et al., 18 Mar 2026).
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 04 schedules; and combinations with objective-level or data-centric interventions such as CHES, IPO, SimPO, or Cal-DPO (Singh et al., 4 Feb 2025, Luo et al., 18 Mar 2026). 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.