Align-SAM: Seeking Flatter Minima for Better Cross-Subset Alignment
Authors: Van-Anh Nguyen
Mehrtash Harandi
Thanh-Toan Do
Linh Ngo Van
Dinh Phung
Trung Le
Venue: ICLR 2026 Poster (2026)
Keywords: Sharpness-aware
gradient alignment
Primary area: optimization

Abstract: Sharpness-Aware Minimization (SAM) has proven effective in enhancing deep neural network training by simultaneously minimizing the training loss and the sharpness of the loss landscape, thereby guiding models toward flatter minima that are empirically linked to improved generalization. From another perspective, generalization can be seen as a model’s ability to remain stable under distributional variability. In particular, effective learning requires that updates derived from different subsets or resamplings of the same data distribution remain consistent. In this work, we investigate the connection between the flatness induced by SAM and the alignment of gradients across random subsets of the data distribution, and propose \textit{Align-SAM} as a novel strategy to further enhance model generalization. Align-SAM extends the core principles of SAM by promoting optimization toward flatter minima on a primary subset (the training set), while simultaneously enforcing low loss on an auxiliary subset drawn from the same distribution. This dual-objective approach leads to solutions that are not only resilient to local perturbations but also robust against distributional shifts in each training iteration. Empirical evaluations demonstrate that Align-SAM consistently improves generalization across diverse datasets and challenging settings, including scenarios with noisy labels and limited data availability.

Published as a conference paper at ICLR 2026
ALIGN-SAM: SEEKING FLATTER MINIMA FOR BETTER -ROSS-SUBSET ALIGNMENT
Van-Anh Nguyen
Trung Le
Department of Data Science and AI Monash University, Australia
Department of Data Science and AI Monash University, Australia grunglm@monash.co edu
van-anh, van-anh.nguuenn@eene
Mehrtash Harandi
Thanh-Toan Do
Department of Electrical and Computer Systems Engineering Monash University, Australia mehrtash, rand leeennd@@moonnn @monash. edu
Department of Data Science and Al Monash University, Australia toan www.oooomoaeeeooo edu
Linh Ngo Van
Dinh Phung
Computer Science
Department of Data Science and AI Monash University, Australia dinh. pwwnggmonash. edu
Hanoi University of Science and Technology, Vietnam linhnvesoict.o saanveessic...utttt..... edu.
ABSTRACT
Sharpness-Aware Minimization (SAM) has proven effective in enhancing deep neural network performance by simultaneously minimizing the training loss and the sharpness the loss landscape, thereby guiding models toward flatter minima that are empirically linked to improved generalization. From another perspective, generalization can be seen as a model's ability to remain stable under distributional variability.. In particular, effective learning requires that updates derived from different subsets or resamplings of the same data distribution remain consistent. In this work, we investigate the connection between the flatness induced by SAM and the alignment of gradients across random subsets of the data distribution, and propose Align-SAM
novel strategy further enhance model generalization. Align-SAM extends the core principles of SAM by promoting optimization toward flatter minima primary subset (the training set), while simultaneously enforcing low loss on an auxiliary subset that is drawn from the same distribution. This dual- objective approach leads solutions that are only resilient local perturbations but also robust against distributional shifts in each training iteration. Empirical evaluations demonstrate that Align-SAM consistently improves generalization across diverse datasets and challenging settings, including scenarios with noisy labels and limited data availability.
INTRODUCTION
Deep neural networks have emerged as the dominant approach for solving complex tasks such as classification, often outperforming traditional machine learning models. These models learn by adjusting a vast number of parameters minimize prediction errors or maximize task-speeific performance. In practice, training is conducted on finite dataset S, sampled from an unknown underlying distribution significantly impact model efficiency and performance Hestness et al. (2017); Kaplan et al. (2020). Despite their ability to learn complex patterns, deep learning models can also capture noise or random fluctuations in training data, leading overfitting Arpit al. (2017); Zhang et al. (016); Liu et al. (200)). This results in excellent performance on training data but poor predictions on new, unseen data, especially with domain shifts. Generalization McAllester (1999); Dziugaite & Roy (2017b). measured by comparing prediction errors on and D. becomes crucial. Balancing a model 's ability to fit training data with its risk of overfitting is key challenge in machine learning.
The quality and alignment of this dataset with the target distribution

---

Published as a conference paper at ICLR 2026
Several studies have been done this problem, both from theoretical and practical perspectives. Statistical learning theory has proposed different complexity measures that are capable of controlling
generalization errors [Vapnik] [9998] [Bartlett & Mendelson 2003; Mukherjee al. 2002 Bousquet & Hlisseeff 2002 Poggio et al. (004). In general, they develop bound for the general error on D. Theory suggests that minimizing the intractable general error on 'D is equivalent minimizing the empirical loss on S with some constraints to the complexity of models and training size (Alquier) ct al. 2016b). An alternative strategy for mitigating generalization errors involves the utilization of an optimizer to learn optimal parameters for models with specific local geometry. This approach enables models to find wider local minima (i.e., flat minima), which makes them more robust against data shift between training and testing sets (Jiang et al. 2020: Petzka et al. 2021 Huang et al. 2025).
The connection between a model's generalization and the width of minima has been investigated theoretically and empirically in many studies, notably (Hochreiter & Schmidhuber 1994, Neyshabur al., 2017: Dinh et al., 2017, Fort & Ganguli| 2019). specific method within this paradigm Sharpness-Awerr Aware Minimisation (SAM) (Foret et al. 2021). which has emerged as effective technique for enhancing the generalization ability of deep learning models. SAM seeks a perturbed model within the vicinity of a current model that maximizes the loss over training set. Eventually, SAM leads the model to the region where both the current model and its perturbation model have low loss values, ensuring flatness. The success of SAM and its variants (Kwon et al. [2021] Kim et al., 2022) has inspired further investigation into its formulation and hehavior, as evidenced by recent
works such as (Kaddour et al. 2022 Möllenhoff & Khan 2022 Andriushchenko & Flammarion
2022a Ji et al. 2024).
Additionally, inspired by the PAC-Bayes theorem Alquier 20233) SAM provides an upper bound on the generalization loss over the data distribution by considering the loss of perturbed models trained on a random dataset within a neighborhood of parameters. Such framework has been shown to encourage convergence to flatter minima on the random dataset S, which in turn promotes improving generalization. In practice, however, directly minimizing sharpness over the entire training set S would require forward and backward passes on all of S, which is computationally infeasible with large datasets. Instead, standard stochastic optimizers, such as SGD or Adam, are applied that rely on randomly sampled mini- batches, thereby reducing sharpness on each subset and stochastically approximating minimization of sharpness over the full dataset S.
D. This guides optimization toward minimizing the worst-case loss
From a complementary perspective, generalization can be viewed as the model's 's ability to remain reliable on new subsets drawn from the same data distribution D. Specifically, we adopt the viewpoint that a model exhibits strong generalization if, although optimized primarily on one random subset S, it can also perform well another independently drawn auxiliary subset Sa where S ~ D. Motivated by this, we formulate our objective as finding models that minimize sharpness and loss on the primary subset while simultaneously maintaining low loss on the auxiliary subset, thereby ensuring stability across resamplings of the data distribution. To achieve this, we propose Align-SAM, a novel method that updates model parameters toward solutions that are both flat and low-loss on the primary subset, while maintaining robust performance across auxiliary subsets, thus implicitly promoting stronger generalization to the full distribution.
In summary, our contributions in this work are as follows:
We approach generalization from novel perspective by framing it an alignment across random subsets drawn from the same data distribution. Building on this viewpoint, we propose Align-SAM,, method designed to enhance both model flatness and stability under distributional variability. Align-SAM primarily updates model parameters by guiding them toward regions in parameter space that minimize sharpness and loss on a primary subset, while simultaneously encouraging strong performance on an auxiliary subset sampled independently from the same distribution. This is achieved by leveraging a combination of gradients computed on both subsets during optimization.
We demonstrate the effectiveness of Align-SAM in enhancing generalization performance across variety of settings. Our evaluation begins with image classification tasks, covering both training from scratch and transfer learning on datasets ranging from small to large scale. We further assess its robustness under noisy label conditions with varying noise levels. Additionally, we extend our experiments to meta-learning scenarios to evaluate Align-SAM's ability generalize beyond meta-training tasks and adapt across diverse
2

---

Published as a conference paper at ICLR 2026
domains. The consistent performance gains across these experiments show that Align- SAM not only improves robustness to label noise and generalization across tasks but also promotes more stable and reliable predictions in varied settings.
2 RELATED WORKS
Sharpness-Awan Aware Minimization. The correlation between the wider minima and the generalization capacity has been extensively explored both theoretically and empirically various studies Tan et al. (2025); Diang et al. (2020); Petzka et al. (202T): Dziugaite & Roy (2017a) Zhuang et (0022) K won ctt al. 2021]. Many works suggested that finding flat minimizers might help to reduce generalization error and increase robustness to data distributional shift problems in various settings Jiang al. (2020); ; PPtzkk et al. (202T); :: Huang et al. 2025). There are multiple works have explored the impact of different training parameters, including batch size, learning rate, gradient covariance, and dropout, on the flatness of discovered minima such as Keskar et al. (2017): Jastrzebski et al. (2017) Wei et al.
(2020): Deng al. 2025).
Sharpness-aware minimization (SAM) (Foret et al. 2021 is a recent optimization technique designed to improve the generalization error of neural networks by considering the sharpness of the loss land- scape during training. SAM minimizes the worst-case loss around the current model and effectively updates models towards flatter minima to achieve low training loss and maximize generalization performance on new and unseen data. SAM has been successfully applied to various tasks and domains, such as vision models (Chen et al. 2021), language models (Bahri al. [2022). federated learning (Qu et al.) 2022; Xing et al.] (025)) Bayesian Neural Networks (Nguyen et al.) (023). domain generalization (Cha et al.
(021) 1,, multi-task learning (Phan et al. 2022) and meta-learning bi-level optimization (Abbas et al. 2022). Multiple varieties of SAM have been developed to address limitations of the original method, including ASAM Kwon et al., 2021). Friendly-SAM (Li et al.)
2024), GSAM (Du et al. 2022). VASSO (Li & Giannakis 2024). )) and other curvature- or alignment- aware extensions. Efficiency-oriente approaches such as SAF Du et al. (2022) approximate SAM's perturbation to reduce computational overhead, GNAM Zhang et al. (2023) promotes first-order flatness through gradient-norm regularization. LookS A. Liu et al. (2022) further improves efficiency by reusing the same perturbation direction across consecutive iterations, significantly lowering SAM's computation while maintaining its flatness- inspired by the Lookahead optimizer to deploy multi-step strategies to explore flatter regions of the loss landscape: Lookahead-SAAl (Yu et al.) 2024) integrates Lookahead's ' extrapolation-interpllaiio mechanism with SAM to reach wider minima, whereas Lookhehind-SAr AAM (Mordido et al.) 2024) incorporates backward steps before the SAM update to better navigate sharp regions.
behaviour. Recent work has also been
Implicit Biases and Behaviors of SAM. SAM was inspired by the PAC-Bayes theorem /Alquier et al. (2016a); Alquier (2023) Alquier et al. (2016b), which provides an upper bound on generalization loss and motivates the pursuit of flat minima. Several works have since sought to better understand and improve Sharpness- Aware Minimization (SAM). Andriushchenko and Flammarion Andriushchenko & Flammarion (0226)) offer theoretical insights into SAM's 's optimization dynamics, emphasizing its implicit regularization effects. Compagnoni et al. Monzio Compagnoni et al.) (2023) model SAM using a stochastic differential equation (SDE), providing a continuous time interpretation of its behavior. Wen et al. (Wen et al. 2003b) study how SAM reduces sharpness during optimization, while Chen et al. Chen et al. (2023) show that SAM generalizes better than SGD by avoiding sharp minima. To further improve SAM, Luo al. (Luo et al. 2024) introduce explicit eigenvalue regularization to control curvature during training. Finally, Wen et al. (Wen et al. (223a) argue that the generalization benefits of sharpness-awe aware algorithms stem not only from sharpness minimization but also from other favorable inductive biases.
complementary strand of work examines data distribution, class imbalance, and long-tailed regimes, where sharpness varies significantly across classes. Nguyen et al. (2024) analyze the features learned under SAM and show that SAM encourages models to rely on both simple and complex features. They further demonstrate that modifying the training data distribution to reduce simplicity bias improves in- distribution generalization. In inbalaiieal settings, several methods adapt SAM to better handle head-tail disparities. ImbSAM Zhou et al. (2023a) incorporates imbalance-awart perturbations to prevent minority classes from being overshadowed by head-class gradients. Class-Condition Conditional SAM Zhou et al. (2023b) aligns SAM's perturbation with class-specific curvature to reduce sharpness for
3

---

Published as a conference paper at ICLR 2026
tail classes, improving long-tailed accuracy. Focal-SAM et al. (2025) integrates focal reweighting with SAM to emphasize hard or misclassified instances while mantainiigg stable optimization.
PROPOSED METHOD
Notions. We start by introducing the notions used throughout our paper. We denote D as the data/label distribution to generate pairs of data/label (x, y). Given model with the model parameter 8, we denote the per-sample loss induced by (z, y) as É(2, y; 0). Let be a random subset drawn from the distribution D. We denote the empirical and generalization losses as Ls (0) = Es [e (I, y: 0)] and LD (0) = ED [e (xr y; 8)] respectively. We define Lp [0 S) as an upper bound defined over S of the general loss LD (0). Note that inspired by SAM (Foret et al., 2021).
use the sharpness over S
to define LD (0 5) (see Theorem D. Finally, we use |A] to denote the cardinality of set A.
3.1 PROBLEM FORMULATION
Given random subset 5 whose examples are sampled from D (i.e., St ~ DN., with N, = |5")) we use LD (0 (S)) to train models. St is known as the training set. Among the models that minimize this loss, we select the one that minimizes the general loss as follows:
min LD (0") s.t. 0t E AD (5*) argmin, (0)
We note that Ap (S*) returns the optimal models 8" that minimizes the upper bound CD (# S*). Among the set of such minimizers 0*, we select the one that further minimizes the true generalization loss LD. The reason for the formulation in ( is that although Lp (0 5t) is an upper bound of the general loss LD (0), there always exists gap between them. Therefore, the additional outer minimization helps to refine the solutions.
We now denote St (i.e., S NN with Na 18 as other random subset sampled from D, Se called the auxiliary set. With respect to this auxiliary set, we have the following theorem.
Theorem 1. Under conditions Cp (0) I Ee., CO + E) similar SAM Foret et al. 2021).
with probability greater than 1 - (i.e., E [0, 11) over the choice of S DNo,, we then have for any optimal models 0* AD))
8L
N - k
4L VNa
(D((0)= < LD((') (()) +
log
0(1)
VNa
8
4L
| 112
Na
k log
1+
1
log
(2)
NNa
P
k
where denote LD (0+ S) == maxe: 10 Ls (0') for any random subset S DN (i.e
St S*,S), and L is the upper bound of the loss function (i.e., E (I, y: 0) S L, Vx, y, 8.0) 0), le is the model size as the length of vector 9, and > 0 is the perturbation radius.
S) max 11 Ls (8') can be viewed as an upper bound of the generalization loss LD (0*). up to constant difference. Moreover, our theorem l (see Appendix for proof) can be viewed as an extension of Theorem in Foret et al. (002T), where we apply the PAC-Bayes theorem from Alquier et al. (2016a) to prove an upper bound for the generalization loss of any bounded loss, instead of the 0-l loss in Foret et al. 2021). We can generalize this proof for S
According Theorem LD (0"
explain why we use LD (0 5t) = max -1las Ls (6") as an objective to minimize, as in
Based on Theorem we can rewrite the objectives in I as:
min Lo (8" (S)) s.t. E' Ap (st) == argmin, Co st),
(3)
where LD (0) (S) max | م|| Ls (0"). Among all models that minimize the upper general- ization bound on random training subset St. select the that further minimizes the upper
4

---

Published as a conference paper at ICLR 2026
generalization bound independently drawn auxiliary subset S". In other words, the optimal solution is the one that achieves low sharpness and loss on the primary subset while simultaneously maintaining low loss the auxiliary subset.
Our theory works for S S ~ D, where D is the distribution to generate (z, y). In the practical version of Algorithm 1, replace DD by the empirical distribution S and at each iteration, sample two mini-batches B'; B ~ S. Because when the training size (i.e. |S|) approaches oc, the distribution S asymptotically converges the distribution D. Using stochastic optimization, we reformulate this objective into an iterative update scheme, where the model is trained with two independently drawn batihe, B" and B", such that each update encourages convergence toward flat minima while aligning performance across subsets in every iteration.
3.2
OUR SOLUTION
Our motivation here to primarily optimize the loss over the training set St, while using S" to further enhance the generalization ability, where both S and 5 are random subsets drawn from the same data distribution. Our formulation in (3) has the same form as a bi-level optimization problem similar to MAML (Finn et al. 2017), developed for meta-learning. Inspired by MAML, naive approach would be to consider f()) = argmin, (9 5²) (i.e., f(0) = 8 - V/Vp (0 S*)) and finding 0 == argmin, Cp (f()) (a)) with respect to 8. However, this naive approach does not align with our objective, as it mainly focuses on optimizing the loss L. (f(0) .S5) over the auxiliary set 88, in which, the auxiliary set acts like the validation set in MAMI.. Here we note that in Franceschil let al. (2018), bi-level optimization was employed to learn optimal hyperparameters, (e.g., the weight of a regularizer) by finding hyperparameters such that model trained on a training set performs well on a validation set. This is fundamentally different from our aim, which is to study how to achieve flat minima that align two independent random subsets in every update step. Moreover, directly adapting the MAML li-ee formulation would require distinct training and validation sets, which are often unavailable in most scenarios. For these reasons, both our theoretical framework and technical approach differ substantially from those in Franceschi et al. (2018).
Using stochastic optimization, we reformulate the objective (3) into an iterative update scheme, where the model is trained using two independently drawn batches sampled independently from the training set S. The Align-SAM is presented as follows: at each iteration, our primary objective is to optimize L (0 B²). primarily based its gradi- ents, in such a way that future models are able to implicitly perform well on B". To achieve this, similar to SAM (Foret et al. 2021), ) we approximate L (0 B') max -al Sp Sp Lz (6") z Cg: (0 + 1/1 V g (0) for a sufficient small learning rate 171 > 0 (i.c., m VV (0) || p) and C (0) Bd) = max 01 Sp CB= (0') = LB (0 + 72 gB (0)) for a sufficient small learning rate 1/2 > (i.e., m/ll I VB (0) S p).
Be and 8" with both
At each iteration, while primarily using the gradients of (0) | B¹) for optimization, we also utilize the gradient of L (0) Ba) to ensure congruent behavior between these two gradients. Specifically, at the 1-th iteration, we update as follows:
0J = 0: + m2 V VeLB- (01)
(4)
8 = 8: + m V Fz (8t) 1/2 VoL (0))
(5)
0111 = 0: v CL
where 7h > 0, 72 > 0, and 17 > 0 are the learning rates, while LB (0t) and Lg (6)) represent the empirical losses over the mini-batches Bt B ~ 5t respectively.
According to (6) (i.e., 0,+1 = 0:- 11--VVLL V (0)) (O+)). (+)).04+) 0.41 is updated to minimize L 0
We now
do first-order Taylor expansion for L
01 as
Lg. (07) ~ LB: (0t) + m LL (0t) 11 72 VoLg (0r) VoL (0r)
where specifies the dot product.
5

---

Published as a conference paper at ICLR 2026
From 7 reach the conclusion that the update in (6 aims to minimize simultaneously (i) Lg (0:).
(ii) Vo (0t) l1 and maximize (iii) VoL (61) Vo
6f) While the effects in (i) and (ii)
V&CB= (0r) encourages two
similar to SAM (Foret et al. 2021 maximizing Be (0t)
gradients of the losses over B3 and Bd to become more congruent. The following theorem shows that, during training, the two gradients of interest become more congruent.
(olge(01) (of) and m2 S
Theorem 2. For sufficiently small learning rates mn I
12 VoCo (6:)
H (n((O//Vo
Vo (9t) VoCge (0F)
Voh (0)) VoC (0))
min G VoLan (f?) (gr( (6r)
we have
VoLoa (6r) H (0?) VO (01)
VoL (0t) VoLB" 09
ifVoL (01)
VeLB= (0r) z
VOLB (01).v
VoLB" (8f).VoLD'
09)
VoLg (0i) VoLga
0p)
otherwise
(8)
Theorem (see Appendix A.1 for proof) indicates that two gradients VoLBi (8) and VoLB" 08 encouraged to be more congruent since update aims maximize its lower bound x VoCze (0₂) VoLBa (0g) (i.e., c = 0.5 or c 1.5).
Practical Algorithm. Inspired by SAM (Foret et al. 2021 we set 7)2 VaLne (6:) and ?h X VoC al where are perturbation radius and À trade-off coefficient for combining gradient from Be and Ba In practice, we observe that setting X > 1, which prioritizes the gradient from the training mini-batch B", results in improved performance. This trade-off is discussed in Section A.3
The pseudo-code of Align-SAM is presented in Algorithm II Compared standard SAM, our method requires additional forward and backward passes due to the use of an auxiliary batch. To reduce computational overhead, we substitute the gradient of purturbation model on the auxiliary set
Bft with the gradient of the current model VoLB- (0)) in Equation This reuse of the VoLo"
gradient maintains the primary objective of AlignSAM: maximizing loss on the training set while simultaneously minimizing on the auxiliary set. Additionally, we the auxiliary batch size |B"| significantly smaller than the primary training batch size |B*,, ensuring that most computation is devoted to the main update step. As result, Align-SAM only marginally slower than standard Further details are provided in the Appendix.
SAM, as reported in Table 9
3.3 CONVERGENCE ANALYSIS
is well known that the normalized (practical) version of SAM does not converge to the minimizer of the training loss, rigorously demonstrated in (Si & Yun 2023) (Theorem 4.6), of the most comprehensive analyses of SAM's convergence behavior. Our proposed approach shares the same convergence rate as standard SAM, established in Si & Yun (2023) (Theorem 4.6). Details in Appendix A.2
4 EXPERIMENTS
In this section, we present the results of various experiments to evaluate the effectiveness of our Align-SAM, including training from seratch, transfer learning on different dataset sizes, learning with noisy labels, and a meta-learning setting.
4.1 IMAGE CLASSIFICATION FROM SCRATCH
We first conduct experiments on ImageNet, Food101, and CIFAR datasets with standard image classification settings trained from scratch. The performance is compared with baseline models trained

---

Published as a conference paper at ICLR 2026
Algorithm 1 Pseudo-code of Align-SAM
1: Input: p. X, 1J. the number of iterations T, and the training set S
2: Output: the optimal model Or.
3: for 1 to T do
4: Sample
B", B', _
S.
5: Ya VaLB= (0t)
6: g: VgL (0t)
7: Compute Of c 0: + p 11oel12
1laalla
8: Compute 0141 t 01 ---11Vool VoL
0i)
9:
end for
with the SGD, SAM, ASAM, and the integration ASAM and Align-SAM. For all experiments of Align-SAM, consistently set and discuss the effect of this trade-off in Section A.3
ImageNet dataset. We use R and ResNet34 models for experiments on the ImageNet dataset, with input size of 224 x 224. For all experiments with Align-SAM, we consistently set À while the perturbation radius is contigured according to the SAM method. Specifically, in this experiment, we set p 0.1 for both SAM and Align-SAM. The models are trained for 200 epochs with basic data augmentations (random cropping, horizontal flipping, and normalization). We use an initial learning rate of 0.1, batch size of 2048 for the training bati-batches, following a cosine learning schedule across all experiments in this paper. We extend this experiment to the mid-sized Food101 dataset using the same settings, except for batch size of 128 for the training and 32 for the auxiliary mini-batches. Performance results are detailed in Table [
and 512 for the auxiliary
Table 1: Classification accuracy on the ImageNet and Food101 datasets. All models are trained from scratch with 200 epochs.
Resnet18
Resnet34
Dataset
Method
Top-1
Top-5
Top-1
Top-5
ImageNet
SAM
62.46
84.19
63.73
84.95
Align-SAM
63.64
85.22
65.89
86.84
Food101
SAM
73.15
89.85
73.87
90.84
Align-SAM
73.45
90.35
74.47
91.27
CIFAR dataset. We used three architectures: WideResner28x110 10, Pyramid101, and Densenet122 with an input size of 32 32 for CIFAR datasets. To replicate the baseline experiments, we followed the hyperparameters provided in the original papers. Specifically, for CIFAR- 100, we set p 0.1, and for CIFAR-10, we used p 0.05 for SAM, VASSO and Align-SAM. The same procedure and settings were applied to ASAM and Alin-ASAM, with the perturbation radius 1.0 for CIFAR- 100 and p 0.5 for CIFAR-10. 10. Other training configurations are consistent with those used in the ImageNet experiments, except for data augmentations (horizontal flipping, four-pixel padding, and random cropping). We use / 0.9 as the default parameter for VASSO. The results are reproduced and reported in Tables 2 while the SGD results referenced from Foret et al. (2021).
Our proposed method outperforms the baselines across various settings. On both ImageNet and Fo1,101, it significantly surpasses the baselines, with notable improvement in both Top-1 and Top-5 accuracy. For CIFAR-10, performance close to the saturation point, making further improvements challenging. Nevertheless, Align- SAM achieves slight enhancements across all cases. On CIFAR- 100, where models are more prone to overfitting compared to CIFAR- 10, Align-SAM still delivers competitive results.
7

---

Published as a conference paper at ICLR 2026
Table Classification accuracy on the CIFAR datasets, All models are trained from scratch three times with different random seeds and we report the mean and standard deviation of accuracies.
Method
W
Pyramid101
Densenet121
Dataset CIFAR-100
Foret
2021
81.20 # 0.200
80.30 0.300
SAM Foret et al.
2021
83.00 = 0.035
81.99 0.636
68.72 0.409
VASSO Giannakis 2024
83.11 0.063
82.04 82.04
0.261
GSAM Ghua al. 2222
83.13 0.099
81.87 0.143
68.88 0.201
LookSAM (K=1) 1) al.
2022
82.89 #
82.25 0.273
69.05 0.182
Align-SAM
83.72 = 0.049
82.53 0.282
69.10 0.311
ASAM
83.16 - 0.296
82.02 0.134
69 62 0.120
Align-ASAM
83.88 + 0.042
82.31 0.183
69.71 0.339
Dataset CIFAR-10
SGD Forct 2021
96.50 + 0.100
96.00 0.100
Force FoM
96.87 0.027
96.17 0.174
91.28 0.241
VASSO
2024
96.84 + 0:014
+
91.18 0.063
GSAM Zhuang 2022
96.91 0.020
91.50 0.109
LookSAMIIK- Lm
2022
97.00 0.181
0.200
91.10 0.196
Align-SAM
96.91 0.007
96.47
0.219
91.54 0.307
ASAM Kwon al. 2021
96.91 + 0.063
96.45 0.042
92.04 + 0.240
Align-ASAM
97.15 + 0.063
96.56 + 0.261
92.02 + 0.000
4.2 TRANSFER LEARNING
In this subsection, we further evaluate Align-SAM in the transfer learning setting using the ImageNet pre- rreitrained models fine-tune both small-size, mid-size, and large-size datasets. All initialized weights are available on the Pytorch library.
Table Transfer learning on ImageNet with Resnet models.
Top-1 Acc
Top-5 Acc
Model
SAM Foret 2021
Align-SAM
SAM et 2021
Align SAM
Resnet18
70.44 0.12
70.92 + 0.05
89.63 +: 0.04
89,90 : 0.04
Resnet34
73.06 0.48
73.94 0.14
91.29 + 92.58 + 0.01
Resnet50
75.17 0.23
75.91 0.19
92.83
0.01
VIT-Adapter-S
79.96 0.30
80.04 + 0.32
95.35 + 0.01
95.39 = 0.03
First, conduct experiments ImageNet using three models from the ResNet family and ViT-Adapter-S (which incorporates lightweight Adapter modules with a plain ViT-Small backbone). The ResNet models are pre-trained on ImageNet, while the backbone ViT-Small of ViT-Adapter-S is pre-trained on ImageNet-2Ik. Each model is then fine-tuned for 50 epochs using either SAM or Align-SAA SAM with learning rate of 0.01. We set & 0.05 for SAM and Align-SAM, basic augmentation techniques are the same training from the scratch setting. Results reported in Table show that our methods outperform baselines with significant gap in both top-1 and top-5 accuracies.
Next, we examine this setting small and mid-sized datasets on three models of the EfficientNet family. We fine-tune with learning rate of 0.05 in epochs and use p 0.1 for all experiments of SAM, VASSO (with 0 = 0.9 as the default) Li & Giannakis (2024), and Align-SAM. In Table Align-SAM achieves a noticeable improvement compared to most the baselines on small-size, mid-size, and large-size datasets, demonstrating its robustness and stability across various experiment settings.
4.3 TRAIN WITH NOISY LABEL
In addition to mitigating data shifts between training and testing datasets, evaluate the robustness of Align-SAM against noisy labels on the standard training procedure. Specifically, we adopt a classical noisy-label setting for CIFAR-10 and CIFAR-100, in which portion of the training set's
8

---

Published as a conference paper at ICLR 2026
Table Transfer learning accuracy of small and medium datasets. All models are fine-tuned from pre-trained weights on ImageNet.
Dotaset
Top-1 Ace
Top-5 Acc
Align-SAM
Align-SAM
Stanford Cars
89.14:0..11
89.68 1 0.17
89.91 2: 0.24
90.39
1: 0.07
97.60
0.20
98.04 1: 0.07
98.03 1: 0.09
98.30 1: 0.09
RGVC-Aircrft
65,83 0.23
86.25
86.23 +.20+0..2
87.22
0.27
98.72
95.87
0.06
95.94 0.05
95.94 0.03
Oxford IITI Pets
0.11
92.34
0.11
92.40
0.16
0.17
99.23
99.35
002
0.01
Flower102
95,06 0.01
95.22
1 0.14
95.37
0.11
95.43
1: 0.10
99.08
0.18
0:019
99.23
1: 0.09
99.18 1: 0.02
Food101
83,50 0.01
85,12 +: 0:07
84.65
= 0,03
85.51
:: 0.02
96.10
0.32
0.08 96.83
96.60 :: 0.09
97.14 0.00
Country211
0.03
KMlilentNNe-833
Suanford Cars
89.01 0.19
0.09 85.19
89.40
89.55
0.12
89.86
1: 0.14
97.73 =: 0.21
98.03
0.07
98.01 1: 0.05
98.10
0.01
RGVC-Airerff
84.88
85.15
85.78
0.25
95.53
0.12
95.67
0.00
95.67
0.02
96.26
Oxford LITT Pets
92.68
92.58
0.02
92.64
0.06
99.00
0.01
99.19
99.23
0.10
99.26
Flower102
94.59
0.10
94.73
1: 0.14
94.94
2 0.17
95.32 -: 0.26
98.95
0.08
99.12 11221:.11
99.13 1:33::00 0.07
99.26
0.07
Food101
83.75 =: 0.12
85.79
4: 0.13
85.69
0.14
85.95
0.13
96.22
0.02
97.12 0.00
97.07
4: 0.09
97.33 0.00
Country211
12.96 0.01
13.38
0.09
13.36
25.78
008
25.91
0.05
KflclentNet-84
Stanford Cars
84.72
0.04
85.08 + 0.16
85.00 + 0.07
85.46
0.32
96.41
0.07
96.45
+0001
96:53 + 0..04
96.81 01.00
FGVC-Aircraf)
Oxford LITT Pots
91.89
0:13
92.02
0.23
92.04
92.07
0.00
99.43
0.07
0.08
99.44
Flower102
92.73 = 0.04
93.02
1: 0.14
93.02
2 0.16
93.07 1: 0.16
98.49
0.07
98.68 1: 0.02
98.73 1: 0.07
98.63 1: 0.05
Food101
84,55
0.14
86.13
0.06
86.18 0.10
86.40
0.44
96.31
0.03
97.07 0.01
97.07 4: 0.03
97.31 0.02
Country211
14.63
0.13
14.77
28.09
1.77
0.18
labels symmetrically flipped with noise fractions {0.2, 0.4, 0.6, 0.8), while the testing set's labels remain unchanged.
Table 5: Results under label noise CIFAR datasets with ResNet32. Each experiment is conducted three times using different random seeds, and we report their averages and standard deviations.
Noise (%)
Method
0.2
0.4
0.6
0.8
Dataset CIFAR-100
SGD
66.22 + 0.355
59.26 + 0.045
46.77 + 0.020
26.49 + 0.640
SAM
66.16 0.721
59.95
0.622
50.81 0.353
24.26 1.209
ESAM
65.73 0.219
58.96 + 0.381
49.36 1.103
25.92 1.173
VASSO 0.9)
6.252
59.67 0.318
50.09 + 0.353
20.85 0.077
VASSO 0.2)
0.042
$9.07
0.820
48.35 - 1.046
28.49 0.551
Align-SAM SAg
66.78
0.657
60.78 + 0.636
51.03 + 0.502
27.66 - 1.265
ASAM
66.88
0.593
61.53
0.487
$2.77 1 0.561
30.33 1.788
Align-ASAM
67.38 + 0.106
62.72 0.304
54,58 + 0.572
32.77 0.388
Dataset CIFAR-10
SGD
89.98 + 0.070
84.83 +: 0.085
75.06 + 0.385
54.47 + 1.265
SAM
91.26 - 0.007
88.19 1.060
83.43 0.622
61.69 0.289
FSAM
91.35 0.318
87.58 + 0.353
82.78 + 2.057
58.09 2.276
VASSO (0 0.9)
91.47 0.487
88.17 0.890
83.75 0.480
67.71 4.129
VASSO 0.2)
90.45 0.855
86.28 0.997
77.33 - 0.806
70.95 0.770
Align-SAM
92.38 0.007
90.20 0.318
85,33 + 0.268
70.02 0.403
ASAM
91.98 0.007
89.24
0.572
84.39
0.445
64.82 6.880
Align-ASAM
92.06 0.367
90.01 0.282
86.09 0.657
73.25 0.353
All experiments are conducted using the ResNet32 architecture, with models trained from scratch for 200 epochs. The batch size set to 5 512 for the training mini-batches and 128 for the auxiliary mini-batches. Following Foret et al. (2021).
set 0.1 SAM, FSAM al. (2024), VASSO Li Giannakis (2024). and Align-SA.M p = 1.0 for ASAM and Align-ASAM when training with all noise levels. Exceptionally. for the setting with 80% noisy labels, the perturbation radius for SAM, FSAM, Align-SAM ASAM, and Align-ASAM.
reduced by half to ensure more stable convergence. This observation is also noted in Li al. (2044 and Foret et al. (021)

---

Published as a conference paper at ICLR 2026
In line with Li et al. (2024), we apply additional cutout techniques along with the basic augmentations
outlined in Section 4.1 We report the results of VASSO with 0 = {0.2, 0.9), as presented in their paper & Giannakis (024), where 0 = 0.2 is expected to yield better performance. However, observe that = 0.2 performs better only in the setting with 80% noisy labels, while in other noisy label settings, it gives a lower result compared to 0.9. Each experiment is repeated three times with different random seeds, and we report the average and standard deviation of the results in Table Note that training with SGD is prone overfitting as the number of epochs increases. Therefore, we present the best results for SGD training at both 200 and 400 epochs.
5
ABLATION STUDY
5.1 COSINE SIMILARITY OF GRADIENTS
We present the cosine similarity of gradients before and update the model using SAM and Align-SAM in Figure
Detail analysis is presented in Appendix A.3
SAM
Align-SMM
305
simberry
Convine simberrry
Cosine simbarity
Epochs
Epochs
(a) Before updating model
(b) After updating model
(c) Improvement after updating
Figure 1: Cosine similarity of two gradients VoLge (0:) and VoL..
0p
(a) before updating model,
(b) after updating model and (c) the improvement of this similarity.
The size of the auxiliary subset | and complexity, the sensitivity of trade-off coefficient X, the analysis of loss landscape, and details of these experiments are presented in Appendix A.3
CONCLUSION AND LIMITATION
In conclusion, this work revisits Sharpness-Aware Minimization (SAM) through the lens of cross- subset alignment, offering a fresh perspective on generalization. While SAM encourages flat minima to improve generalization, we argue that effective generalization also hinges on the alignment between two independently sampled subsets from the same data distribution. Building on this insight, introduce Align-SAM find models that minimize sharpness and loss on the primary subset while simultancously maintaining low loss on the auxiliary subset, thereby ensuring stability across resamplings of the data distribution. By explicitly aligning the optimization process across both subsets, Align-SAM produces models that are not only robust to perturbations but also more resilient to distributional shifts. Extensive experiments confirm that Align-SAM delivers consistent gains in generalization, particularly under challenging conditions such as label noise and data scarcity. One limitation to note that using an additional auxiliary subset in each training iteration may increase training time (depending on the size of the auxiliary sets). We view this as trade-off between performance and training complexity. However, this issue could potentially be mitigated by reusing the gradients from the previous steps. We leave this as direction for future work to reduce training complexity while maintaining performance.
ACKNOWLEDGMENT
Trung Le, Mehrtash Harandi, and Dinh Phung were supported by the ARC Discovery Project grant DP250100262. Additionally, Dinh Phung further acknowledged the support from the Australian Research Council (ARC) Discovery Project OP230101176. Trung Le and Mehrtash Harandi were also supported by the Air Force Office of Scientific Research under award number (09550-3--S-0000 This research/worl supported by Monash eResearch capabilities, including M3.
10

---

Published as a conference paper at ICLR 2026
REFERENCES
Momin Abbas, Quan Xiao, Lisha Chen, Pin- Yu Chen, and Tianyi Chen. Sharp-maml: Sharpness- aware model-agnostic meta learning. arXiv preprint arXiv:2206.03990 03996, 2022.
Pierre Alquier. User-friendly introduction to pac-bayes bounds, 2023. URL https://a/xxx 11216
Pierre Alquier, James Ridgway, and Nicolas Chopin. On the properties of variational approximations of gibbs posteriors. Journal of Machine Learning Research, 17(236), 2016a. URL http: //jmlr.org// wttp:ooppppppps///..ot
1900html
Pierre Alquier, James Ridgway, and Nicolas Chopin. On the properties variational approximations of gibbs posteriors. The Journal of Machine Learning Research, 17(1):8374 8414, 2016b.
Maksym Andriushchenko and Nicolas Flammarion. Towards understanding sharpness-aware mini- mization. In International Conference Machine Learning, pp. 639-666-6 PMLR, 2022a.
Maksym Andriushchenko and Nicolas Flammarion. Towards understanding sharpness-aware mini- mization. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato (eds.), Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pp. 639-668. PMLR, 17-23 Jul 2022b. http://ddiissshhhee2222mm html
URL https:://pooceeeingggg mmr..ngssm.r.......mmm.mmmmm.mmm mlr www.pressssee wtw.:/anww..............ee.e.e.e
Devansh Arpit, Stanistaw Jastrzebski, Nicolas Ballas, David Krueger, Emmanuel Bengio, Maxinder Kanwal, Tegan Maharaj. Asja Fischer, Aaron Courville, Yoshua Bengio, et al. A closer look at memorization in deep networks. In International conference on machine learning, pp. 233-242.
PMLR, 2017.
Dara Bahri, Hossein Mobahi, and Yi Tay. Sharpness-aware minimization improves language model generalization. In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 7360-7371, Dublin, Ireland, May 2022. Association for Computational Linguisties. doi: 10. 10.186653/11/2022aalloooo htttps///eelaattooooo...ogg22222
URL
Peter L. Bartlett and Shahar Mendelson. Rademacher and gaussian complexities: Risk bounds and structural results. J. Mach. Learn. Res., 3(mull):4633445 mar 2003. ISSN 1532-4435.
Olivier Bousquet and André Elisseeff. Stability and generalization. The Journal of Machine Learning Research, 2:499-526, 2002.
Junbum Cha, Sanghyuk Chun, Kyungjae Lee, Han-Cheol Cho, Seunghyun Park, Yunsung Lee, and Sungrae Park. Swad: Domain generalization by seeking flat minima. Advances in Neural Information Processing Systems, 34:22405 -222405-22488,, 2021.
Xiangning Chen, Cho-Jui Hsieh, and Boqing Gong. When vision transformers outperform resnets without pre- pre-training or strong data augmentations. arXiv preprint arXiv:2106.01/ 01 548, 2021.
Zixiang Chen, Junkai Zhang, Yiwen Kou, Xiangning Chen, Cho-Jui Hsieh, and Quanquan Gu. Why does sharpness-awar minimization generalize better than sgd? In Proceedings of the 37th International Conference on Neural Information Processing Systems, NIPS '23, Red Hook, NY, USA, 2023. Curran Associates Inc.
Jiaxin Deng, Junbiao Pang, Baochang Zhang, and Guodong Guo. Asymptotic unbiased sample sampling to speed up sharpness-awan awar minimization. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 39, pp. 16208-16216, 2025.
Laurent Dinh, Razvan Pascanu, Samy Bengio, and Yoshua Bengio. Sharp minima can generalize for deep In International Conference on Machine Learning, pp. 1019-1028. PMLR, 2017.
Jiawei Du, Daquan Zhou, Jiashi Feng, Vincent Tan, and Joey Tianyi Zhou. Sharpness- training for free. Advances in Neural Information Processing Systems, 35:23439-23451, 2022.
2022.
11

---

Published as a conference paper at ICLR 2026
Gintare Karolina Dziugaite and Daniel M. Roy. Computing nonvacuous generalization bounds for deep (stochastic) neural networks with many more parameters than training data. UAL. AUAI Press,2017a.
Gintare Karolina Dziugaite and Daniel M Roy. Computing nonvacuous generalization bounds for deep (stochastic) neural networks with many more parameters than training data. arXiv preprint arXiv: 1008, 2017b.
Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adaptation of deep networks. Doina Precup and Yee Whye Teh (eds.). Proceedings of the 34th International Conference Machine Learning, volume 70 of Proceedings of Machine Learning Research, pp. PMLR, 06-11 Aug 2017.
Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware minimization for efficiently improving generalization. In International Conference on Learning Representations, 2021. URL URL/attss//
fortmm/ormmmmmm d 6T11mpssrMM
Stanislav Fort and Surya Ganguli. Emergent properties of the local geometry of neural loss landscapes. arXiv preprint arXiv:1910 2019.
Luca Franceschi, Paolo Frasconi, Saverio Salzo, Riccardo Grazzi, and Massimiliano Pontil. Bilevel programming for hyperparameter optimization and meta-learning In Jennifer Dy and Andreas Krause (eds.), Proceedings of the 35th International Conference Machine Learning, volume 80 of Proceedings of Machine Leurning Research, pp. 1577. PMLR, 10-15 Jul 2018. URL https: https::/prroeeeingg mlr press www...ppeaa/////raaaaaeeeca wwa0/fraaneecchllee. html
Joel Hestness, Sharan Narang, Newsha Ardalani, Gregory F. Diamos, Heewoo Jun, Hassan Kianinejad, Md. Mostofa Ali Patwary, Yang Yang. and Yanqi Zhou. Deep learning scaling is predictable, empirically. CoRR, abs/1712.004 2017.
Sepp Hochreiter and Jürgen Schmidhuber. Simplifying neural nets by discovering flat minima. In NIPS, pp. 529-536. MIT Press, 1994.
Bin Huang, Ying Xie, and Chaoyang Xu. Learning with noisy labels via clean aware sharpness aware minimization. Scientific Reports, 15(1):1350) 2025.
Stanislaw Jastrzebski, Zachary Kenton, Devansh Arpit, Nicolas Ballas, Asja Fischer, Yoshua Bengio, and Amos Storkey, Three factors influencing minima in sgd. ArXiv, abs/1711.04623 2017.
Jie Ji, Gen Li, Jingjing Fu, Fatemeh Afghah, Linke Guo, Xiaoyong Yuan, and Xiaolong Ma. A single- step, sharpness-aware minimization all you need to achieve efficient and accurate sparse training. In A. Globerson, L. Mackey, D. Belgrave, A. Fan, U. Paquet, J. Tomczak, and C. Zhang (eds.), Advances in Neural Information Processing Systems, volume 37, pp. 44269-44290. Curran Asso- ciates, Inc., 2024. URL oBtp:///hh..aaacccceeeeeeeeeeee2
wtwp:riiss.ccppppppeeeeeeeepepee files/paper
Yiding Jiang, Behnam Neyshabur, Hossein Mobahi, Dilip Krishnan, and Samy Bengio. Fantastic generalization measures and where to find them. ICLR. OpenReview.net, 2020.
Jean Kaddour, Linqing Liu, Ricardo Silva, and Matt J Kusner. A fair comparison of two popular flat minima optimizers: Stochastic weight averaging vs. sharpness-aware minimization. arXiv preprint 05t::22022006661 1, 2022.
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B. Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. CoRR, abs/2001.08366 2020.
Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. On batch-batch training for deep learning: Generalization gap and sharp minima. In ICLR. OpenReview.net, 2017.
12

---

Published as a conference paper at ICLR 2026
Minyoung Kim, Da Li, Shell X Hu, and Timothy Hospedales. Fisher SAM: Information geometry and sharpness aware minimisation. In Kamalika Chaudhuri, Stefanie Jegelka. Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato (eds.), Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pp. 11148- 11161. PMLR, 17-23 Jul 2022.
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 Conference on Machine Learning, pp. 5905-5914, PMLR, 2021.
Bingcong Li and Georgios Giannakis. Enhancing sharpness-aware optimization through variance suppression. Advances in Neural Information Processing Systems, 36, 2024.
Sicong Li, Qianqian Xu, Zhiyong Yang, Zitai Wang, Linchao Zhang, Xiaochun Cao, and Qingming Huang. Focal-SAN, Focal sharpness-awar minimization for long-tailed classification. In Forty- second International Conference on Machine Learning, 2025, URL httt/tttt/ppeeetttt //tp:::e///peeeevieww net htte/forum?ii 1it.ckk2ccckKkkt
Tao Li, Pan Zhou, Zhengbao He, Xinwen Cheng, and Xiaolin Huang. Friendly sharpness-aware minimization. In Proceedings of the IEEF/CVF Conference on Computer Vision and Pattern Recognition, pp. 5631-5640 2024.
Sheng Liu, Jonathan Niles- Weed, Narges Razavian, and Carlos Fernandez- Granda. Early-learning regularization prevents memorization of noisy labels. Advances neural information processing systems, 33:20331 2020.
Yong Liu, Siqi Mai, Xiangning Chen, Cho-Jui Hsieh, and Yang You. Towards efficient and scalable sharpness- aware minimization. In Proceedings of the IEEE/CVF Conference on Computer Vision and Puttern Recognition, pp. 12360-12370, 2022.
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 on Neural Information Processing Systems, 2024. httsttts: /openrreie/eeeeee.eee nnt// JF
David McAllester. Pac-bayesian model averaging. In Proceedings of the Twelfth Annual Confer- ence on Compitational Learning Theory, COLT '99, pp. 164-170, New York, NY, USA, 1999, Association for Computing Machinery. ISBN 131674. doi: 10. 00.1145/377400.3074533 URL
https: /doi.. www.oor....... 10......111/75 1145/30 740444007777733
Thomas Möllenhoff and Mohammad Emtiyaz Khan. Sam as optimal relaxation of bayes. arXiv
preprint arXiv:2210.01620, 2022.
Enea Monzio Compagnoni, Luca Biggio, Antonio Orvieto, Frank Norbert Proske, Hans Kersting, and Aurelien Lucchi. An SDE for modeling SAM: Theory and insights. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato. and Jonathan Scarlett (eds.). Proceedings of the 40th International Conference on Machine Learning. volume 202 of Proceedings of Machine Learning Research, pp. 115209-252533 PMLR, 23-29 Jul 2023. URL hittps://procest ngs ml www.prrss/s.... /202s http///2222/mmmcio/ monzio compagnon 23a html
Gonçalo Mordido, Pranshu Malviya, Aristide Baratin, and Sarath Chandar. Lookbehind-SAM, k steps back, step forward. In Forty-first International Conference on Machine Learning, 2024. URL
net forum? id= 1wCWWE
Sayan Mukherjee, Partha Niyogi, Tomaso A. Poggio, and Ryan M. Rifkin. Statistical learning: Stabil- ity is sufffciiett for generalization and necessary and sufficient for consistency of empirical risk min- imization. 2002. URL htttps://aa
www.haannticccccooo wwwp://raa...ooooouuuo.uuuu. 17:7803 7478036
Behnam Neyshabur, Srinadh Bhojanapalli, David McAllester, and Nati Srebro. Exploring generaliza- tion in deep learning. Advances neural information processing systems, 30, 2017.
Tuan H Nguyen, Paymon Haddad, Eric Gan, and Baharan Mirzasoleiman. Changing the training data distribution to reduce simplicity bias improves in distribution generalization. Advances Neural Information Processing Systemx, 37:68854 00668854-68866 2024.
13

---

Published as a conference paper at ICLR 2026
Van-Anh Nguyen, Tung-Long Vuong, Hoang Phan, Thanh-Toan Do, Dinh Phung, and Trung Le. Flat seeking bayesian neural network. In Advances in Neural Information Processing Systems, 2023.
Henning Petzka, Michael Kamp, Linara Adilova, Cristian Sminchisescu, and Mario Boley. Relative flatness and generalization. In NeurIPS, pp. 18420-18432, 2021.
Hoang Phan, Ngoc Tran, Trung Le, Toan Tran, Nhat Ho, and Dinh Phung. Stochastic multiple target sampling gradient descent. Advances in neural information processing systems, 2022.
Tomaso Poggio, Ryan Rifkin, Sayan Mukherjee, and Partha Niyogi. General conditions for predictiv- ity in learning theory. Nature, 428(6981 (1:881)):111222 2004.
Zhe Qu Xingyu Li, Rui Duan, Yao Liu, Bo Tang, and Zhuo Lu. Generalized federated learning via sharpness aware minimization. arXiv preprint arXiv:2206.0261 2022.
Dongkuk Si and Chulhee Yun. Practical sharpness-aware minimization cannot converge all the way to optima. In Alice Oh, Tristan Naumann, Amir Globerson, Kate Saenko. Moritz Hardt, and Sergey Levine (eds.), Advances in Neural Information Processing Systems 36: Annual Conference on Neural Information Processing Systems 2023, NewrIPS 2023, New Orleans, IA, USA, December 10 16, 2023, 2023.
Chengli Tan, Jiangshe Zhang, Junmin Liu, Yicheng Wang, and Yunda Hao. Stabilizing sharpness- aware minimization through a simple renormalization strategy. Journal of Machine Learning Research, 26(68):1-35 2025.
Vladimir Naumovich Vapnik. Statistical learning theory. In Adaptive and Learning Systems for Signal Processing, Communicutions, and Control. Wiley, 1998.
Colin Wei, Sham Kakade, and Tengyu Ma. The implicit and explicit regularization effects of dropout. In International conference on machine learning, pp. 10181-10192. PMLR, 2020.
Kaiyue Wen, Zhiyuan Li, and Tengyu Ma. Sharpness minimization algorithms do not only minimize sharpness achieve better generalization. In Thirty-seventt Conference on Neural Information Processing Systems, 2023a. / wttp:/eeetttoootndden.eeeoe
Kaiyue Wen, Tengyu Ma, and Zhiyuan Li. How sharpness-awand minimization minimizes sharpness? In The Eleventh International Conference Learning Representations, 2023b. URL https
/ppeneeview net fow.toidddosppgeeeeee
Xinda Xing, Qiugang Zhan. Xiurui Xie, Yuning Yang, Qiang Wang, and Guisong Liu. Flexible sharpness-aware personalized federated learning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 39, pp. 21707-21715, 2025.
Runsheng Yu, Youzhi Zhang, and James Kwok. Improving sharpness-aware minimization by lookahead. In Forty-first International Conference on Machine Learning, 2024.
Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learning requires rethinking generalization. arXiv preprint arXiv: 16/111111.00 1.03530, 2016.
Xingxuan Zhang, Renzhe Xu, Han Yu, Hao Zou, and Peng Cui. Gradient norm aware minimization seeks first-order flatness and improves generalization. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 20247-20251 2023.
Yixuan Zhou, Yi Qu, Xing Xu, and Hengtao Shen. Imbsam: A closer look sharpness-aware minimization in class-imbalanced recognition. In Proceedings of the IEEE/VV International Conference on Computer Vision, pp. 11345-11355, 2023a.
Zhipeng Zhou, Lanqing Li, Peilin Zhao, Pheng-Ann Heng, and Wei Gong. Class-conditional sharpness-awen aware mininization for deep long-tailed recognition. In Proceedings of the IEEF/CVF conference on computer vision and paltern recognition, pp. 3499-3509, 2023b.
Juntang Zhuang, Boqing Gong, Liangzhe Yuan, Yin Cui, Hartwig Adam, Nicha C. Dvomek, Sekhar Tatikonda, James S. Duncan, and Ting Liu. Surrogate gap minimization improves sharpness aware training. In International Conference Learning Representations (ICLR), 2022. URL https://arxiv: wwwpriii....aaaaaa
14

---

Published as a conference paper at ICLR 2026
A APPENDIX / SUPPLEMENTAL MATERIAL
In this appendix, we present the proof' in our paper and additional experiments. We open-source our code and provide instruction, scripts, and log files to reproduce experiments at https:// anonymous 4open.
http:////AA111aA---000000 README. md
A.1
ALL PROOPS
Proof of Theorem
Proof. We use the PAC-Bayes theory in this proof. In PAC-Bayes theory, could follow a distribution, says P. thus we define the expected loss over 0 distributed by P as follows:
CD P) Eg [C
Ls P) Eg~p [s){
For any distribution P = (0, o and Q = o(0.02Ik) over 0 E Rt. where P is the prior distribution and is the posterior distribution, use the PAC-Bayes theorem in Alquier al. (2016a). for all B > 0, with probability at least 1 - 8, we have
CD( Q) Ls( + 1
1 + v(S,N)). KL(QIIP) log
where V is defined as
L (3, NN log Epp
exp -
When the loss function is bounded by L, then
3²12
(3,N) S
8N
v8N KL( P)
The task is to minimize the second term of RHS of 19 thus choose Then the second term of RHS of ( is equal to
KL P) + log
x L.
2N
The KL divergence between Q and P, when they are Gaussian, is given by formula
ko 2 + 1012
of 0²
(L( 2
k + l log
ofp
For given posterior distribution Q with fixed o², to minimize the KL term, the of should be equal to x² + ||0|| //k. In this case, the KI. term is no less than
klog (1 (1+
|10011 ko
Thus, the second term of RHS is
KL( + log
k: log (1+ +
|10112
L2
kaz
L L>L
2N
4N
when | > o²{ exp((NNkk) 1}. Hence, for any |10|| > o²{ exp(4N/k) - 1]. we have the RHS greater than the LHS, and the inequality is trivial. this work, we only consider the case:
|1111 < o² (exp((AAN/) 1)-
(10)
Distribution P is Gaussian centered around 0 with variance of =2 + /k/ which unknown the time up the inequality, since is unknown. Meanwhile, we have specify P in advance,
15

---

Published as a conference paper at ICLR 2026
since P is the prior distribution. To deal with this problem, could choose family of P such that its means cover the space of I satisfying inequality (TO). We set
0² + exp(4N/R))
Pj = (0, cexp 1- 1 Tla) p : (P, j 1.2....)
k
Then the following inequality holds for a particular distribution Pj with probability - 6j with 6j
Eg ) I ( 1 KL(QIIP) + log 1 ōj 3
+ ((B
Use the well-known equation: holds with every j. We pick
1==1 fs
x²
then with probability 1 ô, the above inequality
0² + 11012//
0² + 100 10/
5 = 1 A: log
ks log ²² exp (4N/E))
Therefore,
0² + 11011//k
1 8
k klog
0² + |1011 /1:
1 j-j
0² + 1100112/k 1
log
S log
C
k
c
k
²² + 11012/kk S cexp {
[1-j"
S exp(11kk)
-10111/k]
k
0² + 11012/k - or I <xpp(1/k) + /101///6].
Thus the KL term could be bounded as follow
ko + 11012
pg.
KL(QIIP3-) =
ts + ls log
2
of
k(o² + 11012/k) - k + k log
exp(L/k)(o 1011//k)
< 2
0² + 11012/kk
g²
1
exp(1/k)(o22 + (0110)(22+1101)/4))
llo
2
0²
1
k log (1+ 1100112 }
2
ko
For the term log o5= with recall that c o2= (1 + exp(4N/k)) and
x² 1100
jr - k log oz [1+ expt (NN771
we have
7
1 + log
π
log 85-
= log
log g
+ 2 log(i")
65
6
log
1
+
log
2²
2lo
[ + k klo
0² (1 + exp(4N//))
8
0² + 10112/
2²
log
log
2 2og
k (1 + exp(4N/k)))
8
n
4N (1+ + k
< log 8
+ log
+ 2 log
6
1
π²
< log
+ log
+ log(1 + k + 4N).
16

---

Published as a conference paper at ICLR 2026
Hence, the inequality
KL(QI| + log o,
Co ,N0, (((,00111)) s Ls Lo (I.N((,, +
x L
2N
Ls Ls (5,00, 021)
L
11012) ko2
1 11
k
1
+ 2 log
+ +4log(N k)
2VN
66
$ Ls (6',N(0,,
L
|1013 + 0(1) + 2og z
1
+ 2/N
k log (1
4 log(N + k).
ko²
Since ||0' - 0|2 is & chi- chi-square distribution, for any positive we have
P( 01 - ko² > 20² 2to²)) I exp(-t).
By choosing t = 1/2 (N), with probability 1 - 1-1 we have
||0' 011 v 0² log(NN ko o² 2k log(N)) S ko (1 l)g(NN) 2
log( N) k
By setting o = px x (Vl: + vlog(N)) we have 100 - -112 S p². Hence, we get
(s(0',N(0.00t1)) = ) [fo]] = 110000150
Es [dN
/
Es (0,00.011
<
max Ls(f") L
VN
100 10'S
N
Ls(0") + 2L
< max
00-00
VN
It follows that
4L
|1012
CD( <
max
Ls(l') +
k
[
(1+ (log(N)/k)3)
VN
x²
N + k
+ 2
log
+ +O(1))
8
4L VN
|1012 (1 (1+
Lo(
()+
k klo
p² (1 /log(N)/k)))
N + k + 2/ log
+ 001)
By choosing 0 = 0*, which is the solution on random subset S and S S", which another subset, St 5t ~ D, hence N = Na, we reach the conclusion.
Proof of Theorem 2
Proof. We have
L 01 = Ln, (01) + Vn/lVoL. (0)) 1/2 72 VoLn (0:) VoL 08
17

---

Published as a conference paper at ICLR 2026
This follows that
VOLB' 01
VoLB, (01) + H (81) VBL (01)
122 HB (0t) VoL.= (0t) + HB
( (8i)
where HB: (0)) = V2LB, (01) and Hga (8) = V2LBaa (0r) are the Hessian matrices.
VOL (0t)
VoLB' (0:) VOLB, (0i) VoLnn
0f
+ 2m VoLB=
(f))
H (0:) VeLge (0i)
1/2 VoLB
(8))
T
Hg. (0:) VoLB" (0p)
1/2 VoLoa
0
T
H
(if) VoL. (01)
We now choose 7/1 S VoCo, VoC (f
we then have
12 volnn (0r) (0:) Fo
?i VoLno
(0+) HB: (01) VOL (0:)
VoLo, (01) VoLB= (0F)
12
This further implies
T (0t)² Hg: (0:) VoLg (0i) 2
VoLn, (0i) VaLna (or)
VaLno
12
(0, (o))
VoLn, (0.) (or)
Next we choose 7/2 min
Vo (0²) " ( (o+)
Volle (6f) 18 (0 2 (0))
then have
VoLB, (0)) VoLBa (8
(0:)
T H (0:) VoLB= (0r)
1/2 VOLnu
VoLn, (01) VoL 0p)
72 VOLB
(8
T
Hg (fi) VoLB" (0P) 2
6
VoLB, (0) VoLBo (0+)
1/2 VoL
(0+)
T HB (0p)
VoL.g (0r) S
VoLn, (0))
VoL na 0f)
12 VoLB (0:) T
HB (8)) VoLn: (0:)
6
Finally, we yield
VoLB (09)
VOLB
07) 2 VOLB, (0:) VoLBr (0t)
1 VOLB, (0)) |VoLB
1
VoL (or)
2
Proof of Theorem
Proof. We first denote 0: = 0 + (p1 p22 vc((0))
[VL:( Using the 8-smoothness, we have
Ls (0 < L. (0:) + vLs (0))- (0t41 - 0,) + 3
2 |0 - 0,1/
< Ls (0)) V (0)) vL (o:) + B1²
2 V (0: 12
18

---

Published as a conference paper at ICLR 2026
Taking the expectation, we gain
nB V Ls (8)) BL (Bt) 808² E |VL 88) [1a
E [Css (Bt+1)) SB [Cs (0:)
SB
[Cs (0))
E. [[VV (0)) [1a] - B VLs (81) . [ (8)) vLs (8:)
Gn
+
E
L
0 ||2]
2
[I| V L. (0t) [1 - 7 VLs (0)) [VLg: (0)) - VL (8))
=E
[Cs (0:)]
+
Bn2
E
|VL
88 11 | - n VLs (0) . [vCs (0:) - vLs (0))
2
<B [Cs (0))
nE
[ (0)) [P]
+
7 E
[| VLs (0:) |12]
E
V (8) VLz. (0)) 112]
2
n
|VLs (0))
0: 01
[Vcs (8)) vCs (8))
+
_nr² E
|VL
(8) 12
p1 - P2
8
fr² 2
<B (00)
E [|| Ls Ls 11]
+
[10] - 0[11
+
EE
IVLB LB
0
[|2]
2
2
| (0t) ||
100 0.112
+ 3y
P1
P2
<E LCs (01)
"E
V Ls (8)) 11a]
2
V BB (0))
8 E
BL Bd (0))
vLs (6))
+
2
|
LB Bd (0:)
P2
09) ||
(p1 p2)
|LL (0)) 1|3
3n7
[ (0t) || |00 8/112)
+
E
VL
(8) |1
+ BTB
2
P1 P2
<E
[Cs (01)
V Ls (0) [1 + +383 (p) + 2 P1P2)
+ (p1 (2) Bnl L (0)) +| 377
E LB 01 [2]
2
<B(Cs [6]) -" B
[IVLs (0)
+
33 n
(pi + 2/ P1P2)
+ -22) B
[ [0:]]
+ 3,1² [ L
(0t) vLs (0)) [|a] + B E IIVLs (ot) [2]]
<E|Cs (0:) - E
|| Ls (0) [1 + 3.8 + 22- P1P2)
+ (p1 (2) Bn V (0)) |||+ 11+ (0" + GO)
< ECss (0)) "
2 (E 11 VLs (0) ||| - 2p8)2 + $1 p2
+ +333 n (pi + p - p1P2) Br² (02 +
Rearrange the terms, we obtain
E VLs (f)) ||| - A0s)² 2 (E [Ls (8)) - E [Cs (8:-1))
17
1
~ + 3 (oî -p1P2) + B1B (0² + G²2
19

---

Published as a conference paper at ICLR 2026
Take sum / from 1 to T, we reach
2 n/
T B ||| VLs (0)) III - 2p3)² < (B Cs (0o) - E[s (0r+1))) + Ap?
+ 63 (oi + 03 - + 23r (02 G²)
+
2
(E LGs (0o)) 1.) Ap2 6,3 (7pl + 7 8p1
1/
2 (o² + +G²)
2A
6B² (7pî + 7 8p1p2) + + (o² + G²)
nT
Substitute 7j VBT(o=+07) we arrive at
4v 3A (02 + G2)
LB (0)) ||I Ap3)
+ +6,32 (7pî 7pp 8p1p2)
VT
Proof of Theorem 3.1
Proof. We have
T
min O....,
L (0)) |II Ap31 -
E El (0,) III Ap8/
T
[IIVV (0:) |1 AD3)²
T 1=1
2 + G2) 1/4
2 |a
+ 6 + - 8p1p2)
T1/4
This implies that
2 203 (02 + G2) 1/4
min L (0:)
+ B 6 (7 77p -81122)
71/1
+ Ap3.
A.2 CONVERGENCE ANALYSTS
To do convergence analysis for Align-SAM, we make the following assumptions (Si & Yun 2023
A1 (G-Lipchitz).
The loss function Cs is G-Lipchitz, i.c., (Cs(0) Ls((') - Gll@ 0'|l- A2 (8-smoothness). The loss function Ls is 8-smooth, if [[V (0) - vLs (0') || < 3|10 0'll for
all 0, 0'.
A3 (Bounded variance). For any batch B ~ S, EB ||| CB (0) - Ls (0) [| < 0² for all
Theorem 3. Assume that the loss function Ls satisfies the assumptions AI,A2, and A3, and L*
fo -c. Under Alig starting from Og and the learning rate d /
V
have
E [ (8:) || 111-Ap1 App)) < 4
4/3A 3A (0 G²)
+ 6 (7p² + 7 8p1P2)
T
VT
where Ap P p2 and A L.(( TT"
20

---

Published as a conference paper at ICLR 2026
Table Domain Generalization setting, All models are trained on ImageNet- 1k and then evaluated on ImageNet-Il Ik (clean validation set), Imagenet-R, R, and Imagenet-C datasets.
Top-1 Ace
Top-5 Ace
Method
Imagenet
ImageNet-R
ImageNet-C
Imagenet
ImageNet-R
ImageNet-C'
ResNet18
Transfer Learning
SAM
70.52
34.18 34.38
48.27
89.60 89,94
52.82 53.23
72.17 72.61
Ali
70.88
48.69
ResNet18 From Seratch
SAM
62.46
25.86 26.20
32.96 34,06
73.15
55.72
Align-SAM
63.64
73,45
43,99
57.42
Corollary 3.1. Under the assumptions as in Theorem
have
2 2 (o + G2)] 1/4
+ (Vo +7p - + Ap).
E VI (I) 0.....
min
71/4
is well known that the normalized (practical) version of SAM does not converge the minimizer of the training loss, as rigorously demonstrated in [ii & Yun 2023) (Theorem 4.6), one of the most comprehensive analyses of SAM's convergence behavior. Our proposed approach shares the same convergence rate as standard SAM, as established in Si & Yun 2023) (Theorem 4.6).
A.3 ADDITIONAL EXPERIMENTS
Experiments of Domain Generalization To strengthen the claims about Align-SAM, we evaluated Align-SAM's robustness under domain shifts by training on ImageNet- 1K and testing on ImageNet- 1K (clean test set), ImageNet-R R (artistic renditions), and ImageNet-C (with various corruptions). These shifts demonstrate Align-SAM's robustness, it consistently outperforms SAM across all setups. The results are shown in Table
Experiments of meta-learning setting The concept of Align-SAM is inspired by the agnostic approach in the MAML setting, where the meta-model optimized the meta-training set but aims to minimize loss on the validation set, assuming both the training and validation sets are from the same data distribution. Different from this original idea, Align-SAM uses the gradient from another auxiliary set as an indicator to close the generalization gap between the training and testing sets. Despite this difference, both approaches share the same underlying objective, making it reasonable to expect that applying Align-SAM in the MAML setting will result in improved generalization performance.
Table 7: Meta-learning results on Mini-Imagenet dataset. All baseline results are taken from Abbas al. (2022)
Accuracy
Method
5 ways shot 5 ways 5 shots
MAML SHARP-MAMLA
47.13
62.20
49.72
63.18
Align-SAMtow
50.08
64.29
We compare our approach with standard MAMI. and Sharp-MAME (Abbas et al., 2022 which also address the loss landscape flatness in bilevel models. MAML is typically framed as a bilevel optimization problem, consisting of meta-update step to learn a shared model initialization and a fine-tuning step to adapt task-specific models. Sharp-MAML analyzes the geometry of MAML's loss landscape and introduces the use of SAM to avoid sharp local minima in MAML loss functions. Sharp-MAML proposes three variants: Sharp-MAMLoo (applying SAM only to the fine-tuning step), (applying SAM to both steps). Since our Align-SAM shares the same objective SAM to improve model
Sharp- Searp-MAML,p (applying SAM only to the meta-update step), and Charp-MAMAL
21

---

Published as a conference paper at ICLR 2026
Table Meta-learning results on Omniglot dataset. All baseline results are taken from Abbas et al. (0022)
Accuracy
Method
20 ways I shot
20 ways 5 shots
MAML
91.77
96. 16
SHARP-MAMI
92.89
96.59
92.66
97.28
generalization it can replace SAM optimization in both the meta-update and fine-tuning steps of the MAML model.
The experiments follow the setup from Abbas et al. (2022), specifically using the Sharp-MAMLIs,
variation, which focuses minimizing the sharpness of meta-models fine-tuned the meta-training set. For Align-SAM, we set À 2 and follow the setup exactly for p, inner gradient steps, and step size in Abbas al. (2022) for both Sharp-MAMLion and Align-SAM The results are reported in Tables 7 and Our method outperforms most baselines with significant improvements, demonstrating the effectiveness of Align-SAM and flexibility across various settings.
Cosine similarity of gradients Theorem we prove that minimizing the loss function Lg
could encourage two gradients VoL 05 and VoL (9 to be more congruent since our
update aims maximize its lower bound, which is VOLB (0t) VoLB- (0f) In this sec- tion, we measure the cosine similarity between two gradients VoLn: (0t) and VoL (0p) be- fore (denoted cosines in Figure Tap and after (denoted as cosine, in Figure Ib) updating the model and measure the change these two score (denoted as change). Mathematically, we (6
VeCnt (o CoCno (6g)
Volne (8.41) Vol cosine, = [V Vo
define cosines =
and
change cosine cossines
cosineg
As shown in Figure Tc both SAM and Align-SAM improve the similarity after updating the model, this improvement also increases across training epochs. However, the similarity score of our Align- SAM is always higher than SAM across the training process, both before and after updating the model. It is evident that our Align-SAM encourages gradients of the training subset and the auxiliary subset to be more similar during the training process.
Batch size _ and complexity Our method, as presented in Algorithm | use a gradient on the auxiliary subset helper indicator to lead the model to wider local minima while maintaining low loss on a random subset from the same distribution, and the model should be updated mainly using training samples. Increasing the batch size | could potentially increase performance and training time. In Table we present the results of Align-SAM with various sizes |Ba of CIFAR- -00 with Resnet32 while maintaining a fixed training batch size [B*| 512. We consider performance and training complexity to be the rafe-offf of Align-SAM and find that setting [B| = 4 B" works well for all experiments.
Additionally, we examine different numbers of epochs with SAM, Align-SAM, and Lehine SAM Mordido al. (2024). As presented in Table 9 Align-SAM outperforms both SAM and Lookbehind under the same training time.
Sensitivity of trade-off coefficient X Throughout this paper, we used consistent setting of = 2, which the trade-off coefficient for combining gradients from Be and 8". While this hyperparameter could be optimized for each experiment individually, we find that this configuration delivers good performance across most experiments. By setting À > 1, we ensure that the perturbed model prioritizes maximizing the loss on the training mini-batch B* rather than minimizing it the auxiliary mini-batch B". This approach encourages the model to focus primarily on minimizing sharpness during the actual update step in Formula 6
22

---

Published as a conference paper at ICLR 2026
Table 9: Experiments on different sizes of auxiliary mini-batch with fixed size of training mini-batch is 51 512 samples
Auxiliary batch-size
Accuracy
Training time per epoch
Method
Epoch
Total
200
70.31 = 0.233
11s
36.6m
220
70.45 + 0.303
11s
40.3m
SAM
250
70.06 + 0.078
11s
45.8m
270
71.15 + 0.293
11s
49.5m
16
70.58 + 0.219
11s
37.5m
71.07 0.172
12s
40.0m
Align-SAM
200
64
70.67 + 0.049
13s
43.3m
128
71.21 0.056
14s
46.6m
256
71.04 = 0.207
15s
50.0m
SAM
220
70.45 + 0.303
11s
40.3m
Align-SAM
170
128
70.83 + 0.209
14s
39.6m
Performance fixed
Performance of various with fixed = 0.1
955
Accuraoy
70
455
08
Align-Si SAM
AIgn-SAM
SAM
54
0.25
0.5
1.5
025
05
2.5
3.5
Trade-off 2
Trade-off
(a) Experiments with CIFAR- 100 on Resnet32
(b) Experiments with Flower102 on EfficientNet- -B
Figure Experiments of various trade-off off with fixed perturbation radius p
To verify the impact of this hyperparameter on model performance, we conduct experiments with varying values of trade-off > and present the results in Figure Notably, the configuration where
A > 1 consistently yields higher accuracy compared to the setting where À < 1. When decreasing X, the model places more emphasis on minimizing the auxiliary loss, rather than sharpness on the training set during the actual update step in Formula 6 ultimately reducing performance.
Analysis of loss landscape and eigenvalues of the Hessian matrix We demonstrate the effective- ness of Align-SAM in guiding models toward flatter regions of the loss landscape, as compared to both SAM and SGD, in Figures 3 andd 4 The loss landscapes are visualized with the same setting, the blue areas represent lower loss values, while the red areas indicate higher loss values. Although SAM is shown to lead the model to flatter region than SGD, Align-SAN achieves an even smoother and significantly flatter loss landscape, especially experiments with EfficientNet-B22 in Figure 3
To further validate that Align-SAM successfully locates minima with low curvature, we compute the Hessian the loss landscape and report the five largest eigenvalues, sorted from \ to As. in T These eigenvalues provide insight into the curvature of the model at the optimized parameters. Larger eigenvalues indicate steeper curvature, meaning the model is more sensitive to small changes in its parameters. Conversely, smaller eigenvalues suggest flatter minima, which are typically associated with improved robustness, better generalization, and reduced sensitivity to overfitting. Negative eigenvalues indicate non-convex curvature in certain directions.
23

---

Published as a conference paper at ICLR 2026
Figure 3: Loss landscape of EffecientNett-22 trained on Flower102 dataset with (left) SGD, (middle) SAM, and (right) Align-SAM.
Figure 4: Loss landscape of ResNet32 trained (left) SGD, (middle) SAM, and (right) Align-SAM on Cifar100 dataset.
Ratio of top-5 largest eigenvalues of Hessian matrix
Methods
A2
Ag
X4
As
EfficientNet-B2 on Flower 102
SGD
2.05 10a
4.55
7.88
-4.36
-4.88
SAM
1.61 x 103
1.27
1.31
1.54
-1.65
Align-SAM 0.61 x 103
1.48
1.64
1.90
1.96
Resnet32 on Cifar100
SGD
3.07 x 10
1.27
1.46
1.87
2.13
SAM
1.42 x 10
1.47
1.65
1.79
1.89
Align-SAM
1.11 X 10.
1.32
1.68
1.85
2.09
Table 10: Eigenvalues of Hessian matrix
As shown in Table 10 Align-SAM consistently achieves positive and lower cigenvalues compared to the baseline methods, suggesting that it effectively leads the model toward flatter regions the loss landscape. These results further support the efficacy of Align-SAM in optimizing for smoother and more stable solutions across a variety of architectures and tasks.
A.4 ETHICAL STATEMENT AND THE USE OF LARGE LANGUAGE MODELS
We used large language model (ChatGPT) to help with editing this paper. It was only used for simple tasks such as fixing typos, rephrasing sentences for clarity, and improving word choice. All ideas, experiments, and analyses were done by the authors, and use of LLMs does not affect the reproducibility of our work.
24


Reviews:
[{"review_id": "TCRuEcixeU", "reviewer": "Reviewer_GCo1", "summary": "This paper proposes Align-SAM, a new optimization method based on Sharpness-Aware Minimization (SAM). The authors introduce the idea of \"cross-subset alignment,\" arguing that good generalization requires a model trained on one subset of data to also perform well on an auxiliary, independently drawn subset. The proposed Align-SAM method extends SAM by optimizing for flat minima on a primary training batch while simultaneously enforcing low loss on this auxiliary batch. This is implemented as a modification to the SAM update rule that incorporates gradients from both batches, with the goal of encouraging gradient alignment. The authors provide empirical results across various settings, including image classification, transfer learning, and training with noisy labels, claiming consistent generalization improvements.", "questions": "1. Should the authors clarify that this paper focus on *in-distribution* generalization, given that *out-of-distribution* generalization is a distinct concept?\n2. What are the meaning of $\\nabla\\_\\theta \\mathcal{L}\\_{B\\_t}(\\tilde{\\theta}\\_l^t)$ and $\\nabla\\_\\theta \\mathcal{L}\\_{B\\_a}(\\tilde{\\theta}\\_l^a)$? Do we want to align the gradient of $B\\_t$ and $B\\_a$ at $\\theta\\_l$, i.e., $\\nabla\\_\\theta \\mathcal{L}\\_{B\\_t}({\\theta}\\_l)$ and $\\nabla\\_\\theta \\mathcal{L}\\_{B\\_a}({\\theta}\\_l)$, instead?\n3. The comparison appears to be unfair because Align-SAM effectively operates with a larger batch size than the baseline methods.\n4. Could you please add standard deviations to Table 1 and Table 3 for a more rigorous comparison? Also, given its importance, should the analysis of Figure 1 (currently in Appendix A.3) be moved to the main text?\n5. How does Align-SAM compare to a simpler baseline that adds a direct (weighted) gradient alignment regularizer to the standard SAM loss?\n6. Could the authors clarify the intuition for why Align-SAM improves performance under noisy labels, especially when the test dataset follows a different distribution?\n7. Missing related work about SAM\n- Yu, Runsheng, Youzhi Zhang, and James Kwok. \"Improving sharpness-aware minimization by lookahead.\" Forty-first International Conference on Machine Learning. 2024.\n- Nguyen, Tuan H., et al. \"Changing the training data distribution to reduce simplicity bias improves in-distribution generalization.\" Advances in Neural Information Processing Systems 37 (2024): 68854-68896.", "limitations": "", "rating": 2, "confidence": 4, "soundness": 2, "presentation": 1, "contribution": 1, "strengths": "**Thorough Evaluation:** The paper presents a wide range of experiments, evaluating the method on standard classification, transfer learning, domain generalization, and noisy labels.", "weaknesses": "- **Marginal Gains:** The central claim of \"consistent improvement\" is overstated. In many key experiments (e.g., CIFAR classification and transfer learning), the performance gains over strong baselines like SAM and ASAM are marginal and often fall within the reported standard deviation. These minor benefits do not appear to justify the significant computational overhead of an extra forward/backward pass per step.\n- **Lack of Novelty and Missing Discussion:** The novelty of the update rule is questionable, as it resembles previous works like Lookahead SAM [1] and Lookbehind SAM [2]. The authors fail to adequately discuss or compare against these highly relevant methods in the main text, making it difficult to assess the originality and relative performance of the proposed contribution.\n\n[1] Yu, Runsheng, Youzhi Zhang, and James Kwok. \"Improving sharpness-aware minimization by lookahead.\" *Forty-first International Conference on Machine Learning*. 2024.\n\n[2] Mordido, Gonçalo, et al. \"Lookbehind-SAM: k steps back, 1 step forward.\" *arXiv preprint arXiv:2307.16704* (2023).", "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": "3QNEDaJaWq", "reviewer": "Reviewer_Engx", "summary": "This paper revisits Sharpness-Aware Minimization (SAM) from the perspective of cross-subset alignment. The authors argue that good generalization requires not only finding flat minima on the training subset, but also maintaining stable performance across independently sampled subsets from the same data distribution. To achieve this, they propose Align-SAM, which incorporates an auxiliary mini-batch and encourages gradient alignment between the primary and auxiliary subsets during training. The method is theoretically supported by a PAC Bayes generalization bound and a gradient congruence theorem. Extensive experiments validate the effectiveness of the proposed method.", "questions": "Please see above.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper presents a clear motivation. It aims to reduce sharpness and loss on the primary subset while maintaining low loss on the auxiliary subset, aligning well with the goal of improving generalization.\n- The paper is theoretically grounded. It builds on PAC-Bayesian theory to formalize how cross-subset alignment reduces generalization error, and provides mathematical guarantees showing that the proposed update rule enhances gradient congruence across subsets, thereby offering a solid theoretical foundation for the method.\n- The experiments span a wide range of tasks, including standard classification, transfer learning, and learning with label noise, demonstrating the robustness and effectiveness of the proposed method across diverse scenarios.", "weaknesses": "- According to Algorithm 1, the proposed method requires four backward passes for each parameter update. Although the authors state that using a much smaller auxiliary batch size $|B^a|$ than the target batch size $|B^t|$ can mitigate the computational overhead, this strategy may necessitate a large overall batch size (e.g., Table 7 uses a batch size of 512), and large-batch training may potentially harm performance.\n\n- The experiments primarily evaluate CNN-based architectures (e.g., WideResNet, PyramidNet, DenseNet). Including results on transformer-based models such as ViT would further strengthen the paper, given their broad adoption and distinct optimization characteristics.\n\n- More representative SAM-based methods, such as SAF [1], GAM [2], ImbSAM [3], CC-SAM [4], and Focal-SAM [5] should be included in the related work for a more comprehensive review.\n- There are some typos, such as in line 307, where \"Asgnostic-SAM\" should be corrected to \"Agnostic-SAM\"\n\n-----\n\n[1] Sharpness-Aware Training for Free, NeurIPS 2022\n\n[2] Gradient Norm Aware Minimization Seeks First-Order Flatness and Improves Generalization, CVPR 2023\n\n[3] ImbSAM: A Closer Look at Sharpness-Aware Minimization in Class-Imbalanced Recognition, ICCV 2023\n\n[4] Class-Conditional Sharpness-Aware Minimization for Deep Long-Tailed Recognition, CVPR 2023\n\n[5] Focal-SAM: Focal Sharpness-Aware Minimization for Long-Tailed Classification, ICML 2025", "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": "QfTaVWzNhd", "reviewer": "Reviewer_2St2", "summary": "The paper proposes Align-SAM, a novel optimization method that extends SAM to further enhance generalization under different subsets. Align-SAM guides the model toward solutions that minimize sharpness and loss on a primary training subset while simultaneously maintaining low loss on an independently drawn auxiliary subset sampled from the same distribution. This approach views generalization through the lens of cross-subset alignment and is designed to find minima that are not only resilient to local perturbations but also robust against distributional variability in each training iteration. Empirical results demonstrate consistent performance gains over SAM and its variants across diverse and challenging settings.", "questions": "1. Could you elaborate on the reasoning why, in the statement of Theorem 1 (L172), $$\\mathcal{L}_{\\mathcal{D}} (\\theta^* | \\mathcal{S}^{a})$$ is used as the objective to minimize (Equation 3), rather than $\\mathcal{S})$.\n\n2. In line with W3(a), could you briefly explain the core mechanisms of baseline methods and explicitly state whether they require an auxiliary set?", "limitations": "", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. **Theoretical Foundation and Methodological Proposal**: The authors approach generalization from a novel perspective, framing it as an alignment across random subsets drawn from the same data distribution. Building on this viewpoint, the methodology is supported by Theorem 1, an extension of the PAC-Bayes theorem, which establishes an upper bound on generalization loss using an independently drawn auxiliary subset. Furthermore, Theorem 2 provides theoretical justification that the proposed iterative update scheme (Equations 4–6) successfully maximizes the dot product of the gradients computed on the primary batch and the auxiliary batch, thereby encouraging these gradients to become more congruent.\n\n2. **Extensive Experiments**: Align-SAM is evaluated extensively across diverse and challenging settings, demonstrating consistent effectiveness. The experiments include: training from scratch on large-scale datasets (ImageNet, Food101), mid-sized datasets (CIFAR-100/10), transfer learning using various architectures (ResNet, EfficientNet, ViT-Adapter-S), and evaluation of robustness against noisy labels (CIFAR-10/100). Additionally, domain generalization tests show robustness under domain shifts (ImageNet-R and ImageNet-C). Ablation studies reinforce the core mechanism, showing that Align-SAM achieves significantly flatter loss landscapes and lower Hessian eigenvalues compared to SAM and SGD.", "weaknesses": "# Major Concerns\n\n> 1. **Insufficient Explanation of Preliminaries**:\n> The theoretical foundation and the proposed Align-SAM heavily rely on familiarity with SAM, making the manuscript very difficult to follow for readers unfamiliar with SAM's specifics. Specifically, in Theorem 1, rather than stating \"Under some mild conditions similar to SAM (Foret et al., 2021),\" the specific conditions necessary for the theorem's validity should be explicitly mentioned, as is standard for a formal theorem.\n\n> 2. **Necessity of the Auxiliary Batch Set**: \n> Align-SAM introduces an auxiliary batch requiring additional computational overhead due to extra forward and backward passes compared to standard SAM. Even though authors report Table 7 to alleviate the issues with this, there is still a question for the efficiency  of Align-SAM (please refer to W3(b) for details).\n\n\n> 3. **Questions Regarding Fair Comparison**:\n>\n>> a. **Baseline Explanation**: The experimental baselines require further description. The paper compares Align-SAM against VASSO, GSAM, and LookSAM. Information should be provided (briefly) on the core methods utilized by each of these techniques, and crucially, whether they require an auxiliary set. This information is necessary for the reader to assess the fairness of the comparison and the complexity trade-offs.\n>\n>> b. **Practicality of Training Time**: Table 7 reports training time overhead based on CIFAR100 experiments. However, the large-scale ImageNet experiments utilize batch sizes of $\\|B_t \\|=2048$ and $\\|B_a \\|=512$. Given the scale of ImageNet and the fourfold increase in batch sizes, it is likely that the actual training time overhead is significantly larger. This raises doubts about the practical relevance of Table 7's results for large-scale implementations like ImageNet.\n\n\n# Minor Concerns\n> 4. **The manuscript requires general refinement:**\n>\n>> a. **Reference citation formats appear incorrect** or duplicated in the \"Related Works\" section.\n>\n>> b. **Notation is sometimes omitted**, e.g. $L$ in Theorem 1.\n>\n>> c. In Line L308, the text refers to \"Asgnostic-SAM.\" Please confirm if this is a **typo** and should be \"Align-SAM\" or \"Align-ASAM\".\n>\n>> d.  In the formulation of the Taylor expansion leading to Equation (7) (L235), the step often involves an equality (=) based on a first-order  (L221). It is suggested to clarify whether the equality should be changed to an approximation (≈), and if so, discussion regarding the magnitude of the error term should be included for completeness.", "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": "DMf0PL3rkR", "reviewer": "Reviewer_B43K", "summary": "The paper propses Align-SAM, a variation of SAM which modifies the SAM update to better align gradients between different subsets of the training set. The paper demonstates theoretically that indeed their algorithm leads to better alignment as well as the same convergence rate as SAM. Empirically, the method outperforms SAM variants.", "questions": "- What is the behavior of the method when the tradeoff coefficient is 0?\n- Is there a study on how changing \\rho affects performance?\n- See minor comments above", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The main strength of this paper in my view are the strong empirical results. Align-SAM's strong empirical performance is demonstrated on numerous datasets, architectures and against several baseliens. The method is also intuitive and is theoretically motivated.", "weaknesses": "The paper doesn't too have many weaknesses in my view. First, in Figure 2, the tradeoff coefficient range cuts off on the lower side at 0.25. Is it possible to consider a tradeoff coefficient of zero? Relatedly, is there a study varying the size of \\rho? Second, several of the figures and tables are too small- I hope the authors use the extra page to enlarge these (e.g. Table 2, 3, 5, Figure 1).\n\nMinor comments\n- Is \"Agnostic-SAM\" in Figure 2 another name for Align-SAM?", "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)