SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Theorem 3 demonstrates that VaSSO alleviates the friendly adversary problem by promoting stability. Qualitatively, VaSSO is roughly  $\sqrt{\theta} \in (0,1)$  times more stable relative to SAM, since the term in big  $\mathcal{O}$  notation is negligible given a sufficiently large  $T$ . Theorem 3 also guides the choice of  $\theta$  - preferably small but not too small, otherwise the term in big  $\mathcal{O}$  is inversely amplified.
3.2 Additional perspectives of VaSSO
Having discussed about the stability, this subsection proceeds with other aspects of VaSSO for a thorough characterization.
Convergence. Summarized in the following corollary, the convergence of VaSSO can be pursued as a direct consequence of Theorem 1. The reason is that  $\pmb{\epsilon}_t\in \mathbb{S}_{\rho}(\mathbf{0})$  is satisfied by (4).
Corollary 1 (VaSSO convergence). Suppose that Assumptions 1 - 3 hold. Choosing  $\eta_t$  and  $\rho$  the same as Theorem 1, then for any  $\theta \in (0,1)$ , VaSSO ensures that
VaSSO better reflects sharpness around optimum. Consider a near optimal region where  $\| \nabla f(\mathbf{x}_t)\| \to 0$ . Suppose that we are in a big data regime where  $\mathbf{g}_t(\mathbf{x}_t) = \nabla f(\mathbf{x}_t) + \boldsymbol{\zeta}$  for some Gaussian random variable  $\zeta$ . The covariance matrix of  $\zeta$  is assumed to be  $\sigma^2\mathbf{I}$  for simplicity, but our discussion can be extended to more general scenarios using arguments from von Mises-Fisher statistics (Mardia and Jupp, 2000). SAM has difficulty to estimate the flatness in this case, since  $\epsilon_t \approx \rho \zeta / \| \zeta\|$  uniformly distributes over  $\mathbb{S}_{\rho}(\mathbf{0})$  regardless of whether the neighboring region is sharp. On the other hand, VaSSO has  $\epsilon_t = \rho \mathbf{d}_t / \| \mathbf{d}_t\|$ . Because  $\{\mathbf{g}_{\tau}(\mathbf{x}_{\tau})\}_{\tau}$  on sharper valley tend to have larger magnitude, their EMA  $\mathbf{d}_t$  is helpful for distinguishing sharp with flat valleys.
Memory efficient implementation. Although at first glance VaSSO has to keep both  $\mathbf{d}_t$  and  $\epsilon_t$  in memory, it can be implemented in a much more memory efficient manner. It is sufficient to store  $\mathbf{d}_t$  together with a scalar  $\| \mathbf{d}_t\|$  so that  $\epsilon_{t}$  can be recovered on demand through normalization; see (4b). Hence, VaSSO has the same memory consumption as SAM.
CIFAR10	SGD	SAM	ASAM	FisherSAM	VaSSO
VGG-11-BN	93.20±0.05	93.82±0.05	93.47±0.04	93.60±0.09	94.10±0.07
ResNet-18	96.25±0.06	96.58±0.10	96.33±0.09	96.72±0.03	96.77±0.09
WRN-28-10	97.08±0.16	97.32±0.11	97.15±0.05	97.46±0.18	97.54±0.12
PyramidNet-110	97.39±0.09	97.85±0.14	97.56±0.11	97.84±0.18	97.93±0.08
Table 1: Test accuracy (%) of VaSSO on various neural networks trained on CIFAR10.
Extensions. VaSSO has the potential to boost the performance of other SAM family approaches by stabilizing their stochastic linearization through variance suppression. For example, adaptive SAM methods (Kwon et al., 2021; Kim et al., 2022) ensure scale invariance for SAM, and GSAM (Zhuang et al., 2022) jointly minimizes a surrogated gap with (1). Nevertheless, these SAM variants leverage stochastic linearization in (2). It is thus envisioned that VaSSO can also alleviate the possible friendly adversary issues therein. Confined by computational resources, we only integrate VaSSO with GSAM in our experiments, and additional evaluation has been added into our research agenda.
4 Numerical tests
To support our theoretical findings and validate the powerfulness of variance suppression, this section assesses generalization performance of VaSSO via various learning tasks across vision and language domains. All experiments are run on NVIDIA V100 GPUs.
4.1 Image classification
Benchmarks. Building on top of the selected base optimizer such as SGD and AdamW (Kingma and Ba, 2014; Loshchilov and Hutter, 2017), the test accuracy of VaSSO is compared with SAM and two adaptive approaches, ASAM and FisherSAM (Foret et al., 2021; Kwon et al., 2021; Kim et al., 2022).
CIFAR10. Neural networks including VGG-11, ResNet-18, WRN-28-10 and PyramidNet-110 are trained on CIFAR10. Standard implementation including random crop, random horizontal flip, normalization and cutout (Devries and Taylor, 2017) are leveraged for data augmentation. The first three models are trained for 200 epochs with a batchsize of 128, and PyramidNet-110 is trained for 300 epochs using batchsize 256. Cosine learning rate schedule is applied in all settings. The first three models use initial learning rate 0.05, and PyramidNet adopts 0.1. Weight decay is chosen as 0.001 for SAM, ASAM, FisherSAM and VaSSO following (Du et al., 2022a; Mi et al., 2022), but 0.0005 for SGD. We tune  $\rho$  from \{0.01, 0.05, 0.1, 0.2, 0.5\} for SAM and find that  $\rho = 0.1$  gives the best results for ResNet and WRN,  $\rho = 0.05$  and  $\rho = 0.2$  suit best for and VGG and PyramidNet, respectively. ASAM and VaSSO adopt the same  $\rho$  as SAM. FisherSAM uses the recommended  $\rho = 0.1$  (Kim et al., 2022). For VaSSO, we tune  $\theta = \{0.4, 0.9\}$  and report the best accuracy although VaSSO with both parameters outperforms SAM. We find that  $\theta = 0.4$  works the best for ResNet-18 and WRN-28-10 while  $\theta = 0.9$  achieves the best accuracy in other cases.
It is shown in Table 1 that VaSSO offers 0.2 to 0.3 accuracy improvement over SAM in all tested scenarios except for PyramidNet-110, where the improvement is about 0.1. These results illustrate that suppressed variance and the induced stabilized adversary are indeed beneficial for generalizability.
CIFAR100. The training setups on this dataset are the same as those on CIFAR10, except for the best choice for  $\rho$  of SAM is 0.2. The numerical results are listed in Table 2. It can be seen that SAM has significant generalization gain over SGD, and this gain is further amplified by VaSSO. On all tested models, VaSSO improves the test accuracy of SAM by 0.2 to 0.3. These experiments once again corroborate the generalization merits of VaSSO as a blessing of the stabilized adversary.
ImageNet. Next, we investigate the performance of VaSSO on larger scale experiments by training ResNet-50 and ViT-S/32 on ImageNet (Deng et al., 2009). Implementation details are deferred to Appendix C. Note that the baseline optimizer is SGD for ResNet and AdamW for ViT. VaSSO is also integrated with GSAM  $(\mathrm{V} + \mathrm{G})$  to demonstrate that the variance suppression also benefits other SAM type approaches (Zhuang et al., 2022). For ResNet-50, it can be observed that vanilla VaSSO
CIFAR100	SGD	SAM	ASAM	FisherSAM	VaSSO
ResNet-18	77.90±0.07	80.96±0.12	79.91±0.04	80.99±0.13	81.30±0.13
WRN-28-10	81.71±0.13	84.88±0.10	83.54±0.14	84.91±0.07	85.06±0.05
PyramidNet-110	83.50±0.12	85.60±0.11	83.72±0.09	85.55±0.14	85.85±0.09
Table 2: Test accuracy (%) of VaSSO on various neural networks trained on CIFAR100.
ImageNet	vanilla	SAM	ASAM	GSAM	VaSSO	V+G
ResNet-50	76.62±0.12	77.16±0.14	77.10±0.16	77.20±0.13	77.42±0.13	77.48±0.04
ViT-S/32	68.12±0.05	68.98±0.08	68.74±0.11	69.42±0.18	69.54±0.15	69.61±0.11
outperforms other SAM variants, and offers a gain of 0.26 over SAM.  $\mathrm{V + G}$  showcases the best performance with a gain of 0.28 on top of GSAM. VaSSO and  $\mathrm{V + G}$  also exhibit the best test accuracy on ViT-S/32, where VaSSO improves SAM by 0.56 and  $\mathrm{V + G}$  outperforms GSAM by 0.19. These numerical improvement demonstrates that stability of adversaries is indeed desirable.
4.2 Neural machine translation
Having demonstrated the benefits of a suppressed variance on vision tasks, we then test VaSSO on German to English translation using a Transformer (Vaswani et al., 2017) trained on IWSLT-14 dataset (Cettolo et al., 2014). The faireq implementation is adopted. AdamW is chosen as base optimizer in SAM and VaSSO because of its improved performance over SGD. The learning rate of AdamW is initialized to  $5 \times 10^{-4}$  and then follows an inverse square root schedule. For momentum, we choose  $\beta_{1} = 0.9$  and  $\beta_{2} = 0.98$ . Label smoothing is also applied with a rate of 0.1. Hyperparameter  $\rho$  is tuned for SAM from  $\{0.01, 0.05, 0.1, 0.2\}$ , and  $\rho = 0.1$  performs the best. The same  $\rho$  is picked for ASAM and VaSSO as well.
The validation perplexity and test BLEU scores are shown in Table 4. It can be seen that both SAM and ASAM have better performance on validation perplexity and BLEU relative to AdamW. Although VaSSO with  $\theta = 0.9$  has slightly higher validation perplexity, its BLEU score outperforms SAM and ASAM. VaSSO with  $\theta = 0.4$  showcases the best generalization performance on this task, providing a 0.22 improvement on BLEU score relative to AdamW. This aligns with Theorems 2 and 3, which suggest that a small  $\theta$  is more beneficial to the stability of adversary.
4.3 Additional tests
Additional experiments are conducted to corroborate the merits of suppressed variance and stabilized adversary in VaSSO. In particular, this subsection evaluates several flatness related metrics after training a ResNet-18 on CIFAR10 for 200 epochs, utilizing the same hyperparameters as those in Section 4.1.
Table 3: Test accuracy (%) of VaSSO on ImageNet, where  $\mathrm{V} + \mathrm{G}$  is short for VaSSO + GSAM.
SGD	SAM	VaSSO
λ1	82.52	26.40	23.32
λ1/λ5	16.63	2.12	1.86
Table 5: Hessian spectrum of a ResNet-18 trained on CIFAR10.
Hessian spectrum. We first assess Hessian eigenvalues of a ResNet-18 trained with SAM and VaSSO. We focus on the largest eigenvalue  $\lambda_{1}$  and the ratio of largest to the fifth largest eigenvalue  $\lambda_{1} / \lambda_{5}$ . These measurements are also adopted in (Foret et al., 2021; Jastrzebski et al., 2020) to reflect the flatness of the solution, where smaller numbers are more preferable. Because exact calculation for Hessian spectrum is too expensive provided the size of ResNet-18, we instead leverage Lanczos algorithm for approximation (Ghorbani et al., 2019). The results can be found in Table 5. It can be seen that SAM indeed converges to a much flatter solution compared with SGD, and VaSSO further improves upon SAM. This confirms that the friendly adversary issue is indeed alleviated by
AdamW	SAM	ASAM	VaSSO (θ = 0.9)	VaSSO (θ = 0.4)
val. ppl.	5.02±0.03	5.00±0.04	4.99±0.03	5.00±0.03	4.99±0.03
BLEU	34.66±0.06	34.75±0.04	34.76±0.04	34.81±0.04	34.88±0.03
Table 4: Performance of VaSSO for training a Transformer on IWSLT-14 dataset.
SAM	VaSSO (θ = 0.9)	VaSSO (θ = 0.4)	VaSSO (θ = 0.2)
25% label noise	96.39±0.12	96.36±0.11	96.42±0.12	96.48±0.09
50% label noise	93.93±0.21	94.00±0.24	94.63±0.21	94.93±0.16
75% label noise	75.36±0.42	77.40±0.37	80.94±0.40	85.02±0.39
Table 6: Test accuracy (%) of VaSSO on CIFAR10 under different levels of label noise.
the suppressed variance in VaSSO, which in turn boosts the generalizability of ResNet-18 as shown earlier in Section 4.1.
Label noise. It is known that SAM holds great potential to harness robustness to neural networks under the appearance of label noise in training data (Foret et al., 2021). As the training loss landscape is largely perturbed by the label noise, this is a setting where the suppressed variance and stabilized adversaries are expected to be advantageous. In our experiments, we measure the performance VaSSO in the scenarios where certain fraction of the training labels are randomly flipped. Considering  $\theta = \{0.9, 0.4, 0.2\}$ , the corresponding test accuracies are summarized in Table 6.
Our first observation is that VaSSO outperforms SAM at different levels of label noise. VaSSO elevates higher generalization improvement as the ratio of label noise grows. In the case of  $75\%$  label noise, VaSSO with  $\theta = 0.4$  nontrivially outperforms SAM with an absolute improvement more than 5, while VaSSO with  $\theta = 0.2$  markedly improves SAM by roughly 10. In all scenarios,  $\theta = 0.2$  showcases the best performance and  $\theta = 0.9$  exhibits the worst generalization when comparing among VaSSO. In addition, when fixing the choice to  $\theta$ , e.g.,  $\theta = 0.2$ , it is found that VaSSO has larger absolute accuracy improvement over SAM under higher level of label noise. These observations coincide with Theorem 3, which predicts that VaSSO is suitable for settings with larger label noise due to enhanced stability especially when  $\theta$  is chosen small (but not too small).
5 Other related works
This section discusses additional related work on generalizability of DNNs. The possibility of blending VaSSO with other approaches is also entailed to broaden the scope of this work.
Sharpness and generalization. Since the study of Keskar et al. (2016), the relation between sharpness and generalization has been intensively investigated. It is observed that sharpness is closely correlated with the ratio between learning rate and batchsize in SGD (Jastrzejski et al., 2017). Theoretical understandings on the generalization error using sharpness-related measures can be found in e.g., (Dziugaite and Roy, 2017; Neyshabur et al., 2017; Wang and Mao, 2022). These works justify the goal of seeking for a flatter valley to enhance generalizability. Targeting at a flatter minimum, approaches other than SAM are also developed. For example, Izmailov et al. (2018) proposes stochastic weight averaging for DNNs. Wu et al. (2020) studies a similar algorithm as SAM while putting more emphases on the robustness of adversarial training.
Other SAM type approaches. Besides the discussed ones such as GSAM and ASAM, (Zhao et al., 2022a) proposes a variant of SAM by penalizing the gradient norm based on the observation where sharper valley tends to have gradient with larger norm. Barrett and Dherin (2021) arrive at a similar conclusion by analyzing the gradient flow. Exploiting multiple (ascent) steps to find an adversary is systematically studied in (Kim et al., 2023). SAM has also been extended to tackle the challenges in
domain adaptation (Wang et al., 2023). However, these works overlook the friendly adversary issue, and the proposed VaSSO provides algorithmic possibilities for generalization benefits by stabilizing their adversaries. Since the desirable confluence with VaSSO can be intricate, we leave an in-depth investigation for future work.
Limitation of VaSSO and possible solutions. The drastically improved generalization of VaSSO comes at the cost of additional computation. Similar to SAM, VaSSO requires to backpropagate twice per iteration. Various works have tackled this issue and developed lightweight SAM. LookSAM computes the extra stochastic gradient once every a few iterations and reuses it in a fine-grained manner to approximate the additional gradient (Liu et al., 2022). ESAM obtains its adversary based on stochastic weight perturbation, and further saves computation by selecting a subset of the minibatch data for gradient computation (Du et al., 2022a). The computational burden of SAM can be compressed by switching between SAM and SGD following a predesigned schedule (Zhao et al., 2022b), or in an adaptive fashion (Jiang et al., 2023). SAF connects SAM with distillation for computational merits (Du et al., 2022b). It should be pointed out that most of these works follow the stochastic linearization of SAM, hence can also encounter the issue of friendly adversary. This opens the door of merging VaSSO with these approaches for generalization merits while respecting computational overhead simultaneously. This has been included in our research agenda.
This contribution demonstrates that stabilizing adversary through variance suppression consolidates the generalization merits of sharpness aware optimization. The proposed approach, VaSSO, provably facilitates stability over SAM. The theoretical merit of VaSSO reveals itself in numerical experiments, and catalyzes model-agnostic improvement over SAM among various vision and language tasks. Moreover, VaSSO nontrivially enhances model robustness against high levels of label noise. Our results corroborate VaSSO as a competitive alternative of SAM.
This research is supported by NSF grants 2128593, 2126052, 2212318, 2220292, and 2312547. The authors would also like to thank anonymous reviewers for their feedback.
Maksym Andriushchenko and Nicolas Flamarion. Towards understanding sharpness-aware minimization. In Proc. Int. Conf. Machine Learning, pages 639-668, 2022.
David GT Barrett and Benoit Dherin. Implicit gradient regularization. In Proc. Int. Conf. Learning Representation, 2021.
Peter L Bartlett, Philip M Long, and Olivier Bousquet. The dynamics of sharpness-aware minimization: Bouncing across ravines and drifting towards wide minima. arXiv:2210.01513, 2022.
Léon Bottou, Frank E Curtis, and Jorge Nocedal. Optimization methods for large-scale machine learning. arXiv preprint arXiv:1606.04838, 2016.
M Cettolo, J Niehues, S Stker, L Bentivogli, and M Federico. Report on the 11th iwslt evaluation campaign, iwslt 2014. 2014.
Pratik Chaudhari, Anna Choromanska, Stefano Soatto, Yann LeCun, Carlo Baldassi, Christian Borgs, Jennifer Chayes, Levent Sagun, and Riccardo Zecchina. Entropy-SGD: Biasing gradient descent into wide valleys. In Proc. Int. Conf. Learning Representation, 2017.
Xiangning Chen, Cho-Jui Hsieh, and Boqing Gong. When vision transformers outperform resnets without pre-training or strong data augmentations. In Proc. Int. Conf. Learning Representation, 2022.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. ImageNet: A large-scale hierarchical image database. In Proc. Conf. Computer Vision and Pattern Recognition, pages 248-255, 2009.
Dongxian Wu, Shu-Tao Xia, and Yisen Wang. Adversarial weight perturbation helps robust generalization. In Proc. Adv. Neural Info. Processing Systems, volume 33, pages 2958-2969, 2020.
Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learning (still) requires rethinking generalization. Communications of the ACM, 64(3):107-115, 2021.
Zhiyuan Zhang, Ruixuan Luo, Qi Su, and Xu Sun. GA-SAM: Gradient-strength based adaptive sharpness-aware minimization for improved generalization. In Proc. Conf. Empirical Methods in Natural Language Processing, 2022.
Yang Zhao, Hao Zhang, and Xiuyuan Hu. Penalizing gradient norm for efficiently improving generalization in deep learning. In Proc. Int. Conf. Machine Learning, pages 26982-26992, 2022a.
Yang Zhao, Hao Zhang, and Xiuyuan Hu. SS-SAM: Stochastic scheduled sharpness-aware minimization for efficiently training deep neural networks. arXiv:2203.09962, 2022b.
Juntang Zhuang, Boqing Gong, Liangzhe Yuan, Yin Cui, Hartwig Adam, Nicha Dvornek, Sekhar Tatikonda, James Duncan, and Ting Liu. Surrogate gap minimization improves sharpness-aware training. In Proc. Int. Conf. Learning Representation, 2022.
Supplementary Document for "Enhancing Sharpness-Aware Optimization Through Variance Suppression"
A Linking SAM adversary with stochastic Frank Wolfe
A.1 Stochastic Frank Wolfe (SFW)
We first briefly review SFW. Consider the following nonconvex stochastic optimization
$$
\max  _ {\mathbf {x} \in \mathcal {X}} h (\mathbf {x}) := \mathbb {E} _ {\xi} [ h (\mathbf {x}, \xi) ] \tag {6}
$$
where  $\mathcal{X}$  is a convex and compact constraint set. SFW for solving (6) is summarized below.
Algorithm 2 SFW (Reddi et al., 2016)
1: Initialize:  $\mathbf{x}_0\in \mathcal{X}$
2: for  $t = 0,1,\ldots ,T - 1$  do
3: draw iid samples  $\{\xi_t^b\}_{b = 1}^{B_t}$
4: let  $\hat{\mathbf{g}}_t = \frac{1}{B_t}\sum_{b = 1}^{B_t}\nabla h(\mathbf{x}_t,\xi_t^b)$
5:  $\mathbf{v}_{t + 1} = \arg \max_{\mathbf{v}\in \mathcal{X}}\langle \hat{\mathbf{g}}_t,\mathbf{v}\rangle$
6:  $\mathbf{x}_{t + 1} = (1 - \gamma_t)\mathbf{x}_t + \gamma_t\mathbf{v}_{t + 1}$
7: end for
It has been shown in (Reddi et al., 2016, Theorem 2) that one has to use a sufficient large batch size  $B_{t} = \mathcal{O}(T), \forall t$  to ensure convergence of SFW. This is because line 5 in Alg. 2 is extremely sensitive to gradient noise.
A.2 The adversary of SAM
By choosing  $h(\epsilon) = f(\mathbf{x}_t + \epsilon)$  and  $\mathcal{X} = \mathbb{S}_{\rho}(\mathbf{0})$ , it is not hard to observe that 1-iteration SFW with  $\gamma_0 = 1$  gives equivalent solution to the stochastic linearization in SAM; cf. (2) and (3). This link suggests that the SAM adversary also suffers from stability issues in the same way as SFW. Moreover, what amplifies this issue in SAM is the adoption of a constant batch size, which is typically small and far less than the  $\mathcal{O}(T)$  requirement for SFW.
Our solution VaSSO takes inspiration from modified SFW approaches which leverage a constant batch size to ensure convergence; see e.g., (Mokhtari et al., 2020; Li et al., 2021). Even though, coping with SAM's instability is still challenging with two major obstacles. First, SAM uses one-step SFW, which internally breaks nice analytical structures. Moreover, the inner maximization (i.e., the objective function to the SFW) varies every iteration along with the updated  $\mathbf{x}_t$ .
A.3 The three dimensional example in Fig. 2
Detailed implementation for Fig. 2 is listed below. We use  $\nabla f(\mathbf{x}) = [0.2, -0.1, 0.6]$ . The stochastic noise is  $\pmb{\xi} = [\xi_1, \xi_2, \xi_3]$ , where  $\xi_1, \xi_2, \xi_3$  are iid Gaussian random variables with variance scaling with  $0.2, 1, 2$ , respectively. We scale the variance to change the SNR. We generate 100 adversaries by solving  $\arg \max_{\| \epsilon \| \leq \rho} \langle \nabla f(\mathbf{x}) + \pmb{\xi}, \pmb{\epsilon} \rangle$  for each choice of SNR. As shown in Fig. 2, the adversaries are unlikely to capture the sharpness information when the SNR is small, because they spread indistinguishably over the sphere.
B More on  $m$ -sharpness
$m$ -sharpness can be ill-posed. Our reason for not studying  $m$ -sharpness directly is that its formulation (Andriushchenko and Flammarion, 2022, eq. (3)) may be ill-posed mathematically due to the lack of a clear definition on how the dataset  $\mathcal{S}$  is partitioned. Consider the following example, where the same notation as (Andriushchenko and Flammarion, 2022) is adopted for convenience. Suppose that the loss function is  $l_{i}(w) = a_{i}w^{2} + b_{i}w$ , where  $(a_{i},b_{i})$  are data points and  $w$  is the
parameter to be optimized. Let the dataset have 4 samples,  $(a_{1} = 0,b_{1} = 1)$  .  $(a_{2} = 0,b_{2} = -1)$ $(a_{3} = -1,b_{3} = 0)$  ; and,  $(a_4 = 1,b_4 = 0)$  . Consider 2-sharpness.
- If the data partition is  $\{1,2\}$  and  $\{3,4\}$ , the objective of 2-sharness i.e., equation (3) in (Andriushchenko and Flamarion, 2022), becomes  $\min_{w} \sum_{i=1}^{2} \max_{||\delta|| < \rho} 0$ .
- If the data partition is  $\{1,3\}$  and  $\{2,4\}$ , the objective is  $\min_{w} \sum_{i=1}^{2} \max_{||\delta|| < \rho} f_i(w, \delta)$ , where  $f_1$  is the loss on partition  $\{1,3\}$ , i.e.,  $f_1(w, \delta) = -(w + \delta)^2 + (w + \delta)$ ; and  $f_2(w, \delta) = (w + \delta)^2 - (w + \delta)$  is the loss on partition  $\{3,4\}$ .
Clearly, the objective functions are different when the data partition varies. This makes the problem ill-posed - different manners of data partition lead to entirely different loss curvature. In practice, the data partition even vary with a frequency of an epoch due to the random shuffle.
C Details on numerical results
CIFAR10 and CIFAR100. For these small resolution datasets, we slightly change the first convolution layer of ResNet18 and WRN-28-10 to one with  $3 \times 3$  kernel size, 1 stride and 1 padding following Mi et al. (2022). The results on SGD and SAM demonstrate that the accuracy is almost identical to the vanilla model.