SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Table 10: SAM-ON improves over SAM-all for BatchNorm and ResNets on CIFAR-10: Test accuracy for ResNet-like models on CIFAR-10. Bold values mark the better performance between SAM-ON and SAM-all within a SAM-variant, and underline highlights the overall best method per model and augmentation
SAM variant	RN-56	RNxT	WRN-28
all	onlyNorm	all	onlyNorm	all	onlyNorm
basic aug.	SGD	94.28±0.2	95.37±0.1	96.20±0.1
SAM	94.94±0.1	95.18±0.1	96.35±0.2	96.48±0.1	97.08±0.1	97.10±0.0
elem. l2	94.96±0.1	94.94±0.2	96.41±0.1	96.53±0.1	96.98±0.2	97.06±0.0
elem. l2, orig.	95.14±0.1	95.21±0.1	96.40±0.1	96.41±0.1	97.10±0.1	97.07±0.1
elem. l∞	94.93±0.1	94.96±0.0	96.06±0.2	96.22±0.1	96.95±0.2	97.00±0.1
Fisher	95.01±0.1	95.03±0.1	96.31±0.0	96.55±0.0	96.95±0.0	97.13±0.1
layer. l2	94.95±0.2	95.07±0.1	96.07±0.3	96.46±0.1	97.02±0.0	96.96±0.1
+ AA	SGD	94.70±0.1	96.19±0.2	97.01±0.0
SAM	95.25±0.1	95.40±0.1	96.98±0.1	97.22±0.3	97.57±0.1	97.58±0.0
elem. l2	95.12±0.0	94.82±0.2	97.01±0.0	97.21±0.1	97.61±0.0	97.69±0.0
elem. l2, orig.	95.39±0.1	95.60±0.1	97.24±0.0	97.33±0.1	97.60±0.0	97.56±0.0
elem. l∞	95.12±0.1	95.48±0.3	96.70±0.2	96.91±0.2	97.52±0.1	97.62±0.1
Fisher	95.19±0.0	95.38±0.1	96.77±0.0	97.24±0.1	97.53±0.0	97.65±0.1
layer. l2	95.43±0.3	95.28±0.1	96.80±0.1	96.88±0.1	97.60±0.0	97.48±0.1
Table 11: SAM-ON improves over SAM-all for BatchNorm and more ResNet models: Bold values mark the better performance between SAM-ON and SAM-all within a SAM-variant, and underline highlights the overall best method per model and augmentation
SAM variant	VGG-13	VGG-16	VGG-19	DenseNet-100
all	onlyNorm	all	onlyNorm	all	onlyNorm	all	onlyNorm
CIFAR-100	SGD	75.44±0.2	74.43±0.4	73.40±0.2	77.00±0.2
SAM	76.74±0.2	77.57±0.1	75.81±0.2	76.86±0.1	74.08±0.6	75.60±0.1	79.42±0.6	79.90±0.3
elem. l2	76.65±0.1	77.49±0.1	75.95±0.2	76.45±0.2	74.72±0.2	75.12±0.1	78.90±0.2	79.83±0.3
elem. l2, η = 0.01	77.27±0.2	77.37±0.2	76.65±0.1	76.66±0.3	75.00±0.5	75.44±0.2	79.94±0.4	80.14±0.1
elem. l∞	76.82±0.3	77.62±0.2	75.43±0.4	76.68±0.1	72.74±0.2	74.50±0.4	79.47±0.3	79.64±0.2
Fisher, η = 1.	76.76±0.2	77.68±0.4	75.85±0.2	76.99±0.1	74.03±0.2	74.96±0.3	79.68±0.2	80.38±0.3
layer. l2	76.76±0.2	77.91±0.2	75.99±0.2	77.12±0.2	74.65±0.5	75.28±0.2	78.25±0.2	79.86±0.3
CIFAR-10	SGD	94.29±0.0	93.85±0.3	93.82±0.0	94.51±0.1
SAM	94.88±0.1	95.19±0.2	94.96±0.0	95.02±0.1	94.58±0.1	94.81±0.2	95.84±0.2	95.89±0.0
elem. l2	94.97±0.1	95.08±0.0	95.01±0.1	95.02±0.1	94.68±0.0	94.99±0.1	95.76±0.2	95.86±0.2
elem. l2, η = 0.01	94.95±0.0	95.13±0.1	94.87±0.1	95.12±0.1	94.66±0.1	94.87±0.2	95.92±0.3	95.85±0.1
elem. l∞	94.96±0.1	95.06±0.0	94.74±0.2	94.91±0.0	94.68±0.1	94.73±0.1	95.56±0.2	95.91±0.1
Fisher, η = 1.	95.07±0.0	95.17±0.0	94.77±0.0	95.10±0.2	94.55±0.0	94.91±0.1	95.65±0.1	96.00±0.1
layer. l2	94.78±0.1	95.09±0.1	94.54±0.1	95.08±0.1	66.21±48.7	94.96±0.1	95.48±0.2	95.82±0.1
and comparisons with the dynamic sparse perturbation approach (SSAM-D) [42]. We also provide additional results for SSAM-D for a WideResNet-28 on CIFAR-100 data in Table 12. We found optimal performance for SSAM for  $50\%$  sparsity and  $\rho = 0.1$  on CIFAR-10 and  $\rho = 0.2$  on CIFAR-100 (as also observed in [42] for slightly different training settings). We find that although both SSAM approaches can perform on par or even outperform regular SAM, they are less effective than our SAM-ON approach. The generalization gap increases even further when considering the same high sparsity levels as for SAM-ON.
Table 12: Although SSAM-  $F$  and SSAM-  $D$  [42] with different sparsity levels can outperform SAM-all on CIFAR-100 with WRN-28, they are less effective than SAM-ON.
Sparsity	SAM
0%	SAM-ON
99.95%	SAM-rand
99.95%	SSAM-F	SSAM-D
50%	99.95%	50%	99.95%
Accuracy	83.11±0.3	84.19±0.2	80.97±0.2	83.94±0.1	83.14±0.1	83.53±0.1	81.01±0.1
Figure 6: All considered SAM-variants and their SAM-ON counterpart for a WRN-28 on CIFAR-100.
Figure 7: All considered SAM-variants and their SAM-ON counterpart for a ViT-S on CIFAR-100.
Figure 8: SAM-ON outperforms SSAM-F and SSAM-D [42] (with different sparsity levels) and random mask SAM-rand (same sparsity level  $99.93\%$  as SAM-ON) sparse perturbation approaches on CIFAR-10 for ResNet-18.
B.3 Finetuning from ImageNet-21k
Since ViTs are commonly trained on large-scale datasets and then fine-tuned, we investigate this scenario for SAM-ON. In particular, we consider a ViT-S pretrained on ImageNet-21k from [48]. We fine-tune for 9 epochs with SGD for a range of SAM-variants with their respective SAM-ON counterpart. For each setup, we probe three values of  $\rho$  and report the best result in Table 13. We find in this setting that SAM-ON performs on par with SAM-all although there are small differences across SAM variants: for layerwise- $\ell_2$  SAM-ON performs slightly worse, whereas for all other variants SAM-ON performs equally well or slightly better than SAM-all. In all cases SAM-ON outperforms plain SGD.
Table 13: Results for ImageNet-1k fine-tuning of a ViT-S-224 from a ImageNet-21k model.
SGD	SAM
all ON	ASAM elem. l2
all ON	ASAM layer. l2
all ON	ASAM elem. ℓ∞
all ON
81.62	81.75	81.75	81.73	81.75	81.79	81.75	81.84	81.84
B.4 Adversarial robustness
Here, we provide additional results and extend the discussion on adversarial robustness from Section 4.2. In a study by Wei et al. [52] SAM-trained models showed non-trivial robustness to small adversarial perturbations [50]. Since there are several works highlighting the role of normalization layers for adversarial robustness [9, 54], it is interesting to investigate whether the robustness properties of SAM can be preserved when training with SAM-ON instead of SAM-all. In Table 15 we report the adversarial robustness of the ViT-S trained from scratch on ImageNet (as reported in Section 4.2 evaluated with the two white-box attacks from APGD, but for more radii). The SAM-ON models are not only better than the base optimizer, but consistently outperform the SAM-all models by a small margin. For a WRN-28-10 on CIFAR-100 the differences are less pronounced and often within the standard deviation (reported over 3 seeds in Table 14). SAM-ON also improves over SAM-all, but for the ASAM-elementwise- $\ell_{\infty}$  the all-variant is slightly better than the ON-variant. Overall, we find that in order to get SAM-like improvements for adversarial robustness (as shown in [52]) it is enough to only perturb the normalization layers in SAM, illustrating again their special role.
Table 14: Adversarial robustness CIFAR-100: Reported is robust accuracy (in %) for a WRN-28 trained from scratch on CIFAR-100. Adversarial robustness is evaluated with the two whitebox APGD attacks from autoattack [16].
threat model	ε	SGD	SAM	ASAM-el.-l∞
all	ON	all	ON
l2	0.10	18.14±0.11	28.14±1.09	31.28 ±0.50	30.33±0.80	30.16±0.26
l2	0.20	2.33±0.11	5.39±0.34	6.62 ±0.07	6.63 ±0.12	6.10±0.18
l∞	1/255	10.29±0.04	17.96±1.08	19.56 ±0.33	20.69 ±0.81	18.63±0.30
l∞	2/255	0.67±0.01	1.96±0.17	2.16 ±0.07	2.62 ±0.01	2.05±0.17
Clean acc.	80.7±0.2	83.1±0.3	84.2±0.2	83.3±0.2	84.1±0.2
Table 15: Adversarial robustness ImageNet: Reported is robust accuracy (in %) for a ViT-S trained from scratch on ImageNet, as reported in Table 4. Adversarial robustness is evaluated with the two whitebox APGD attacks from autoattack [16].
ε	vanilla	AdamW	vanilla	Lion
SAM-all	SAM-ON	SAM-all	SAM-ON
l2	0.25	19.67±0.47	37.53±0.69	41.16 ±0.24	22.01±0.78	38.52±0.66	43.12 ±0.97
l2	0.50	5.47±0.18	17.71±0.61	22.72 ±0.25	6.63±0.46	19.03±0.92	24.27 ±1.34
l2	1.00	0.43±0.09	3.34±0.36	5.58 ±0.19	0.57±0.07	3.98±0.28	6.64 ±0.69
l∞	0.25/255	33.45±0.80	48.08±0.14	49.34 ±0.08	35.31±0.08	49.57±0.60	51.37 ±0.99
l∞	0.5/255	14.98±0.18	29.68±0.09	32.46 ±0.15	15.86±0.13	31.68±0.62	34.23 ±1.73
l∞	1/255	2.61±0.16	8.64±0.01	10.82 ±0.56	2.93±0.29	10.02±0.56	12.03 ±1.30
Clean acc.	66.89±0.04	71.47±0.12	71.37±0.026	68.20±0.02	71.90±0.19	72.64±0.14
B.5 Machine translation task
To probe the effectiveness of SAM-ON outside the vision domain, we apply it to the IWSLT'14 DE-EN machine translation task, following the setup of Kwon et al. [37]. We report the resulting Bleu scores in Table 16: SAM-all and SAM-ON perform similar (within standard deviations reported over 3 random seeds), both improving over the vanilla optimizer. While being very limited in its scope, this experiment is a first hint that SAM-ON might also be effective outside the vision domain. Proper evaluations, as for instance done in [7], are required to confirm this for large-scale settings.
Table 16: IWSLT-DE-EN Bleu scores. Reported over 3 random seeds.
vanilla	SAM-all	SAM-ON
34.56±0.11	34.83±0.10	34.95±0.16
B.6 Weight distribution after training
In order to get a better understanding of the impact of SAM-ON on  $\gamma$  and  $\beta$  (as defined in Eq. 1), we train a WideResNet-28-10 with different SAM-variants and both SAM-ON and all. We show the distribution of  $|w_i|$ , i.e. the parameter magnitudes, at the end of training for different layer types in Figure 9. Different to the discussion in Section 5.3, we show the  $y$ -axis on log-scale, in order to inspect more nuanced differences. For elementwise  $\ell_2$  there is no strong change in the distribution of the BatchNorm parameters between all and SAM-ON. For elementwise  $\ell_{\infty}$ , layerwise  $\ell_2$  and SAM, however, the magnitude of the BatchNorm parameters shifts clearly towards larger values, especially for the weight parameters. We note that this resembles a pattern we observed when comparing the optimal  $\rho$ -value for all and SAM-ON in Table 7: The optimal  $\rho$  of elementwise  $\ell_2$  did not change much for ResNet architectures, whereas for the other considered methods, it shifted towards larger values for SAM-ON. Additionally and in contrast to the other methods, the elementwise  $\ell_2$  variant showed a strong performance decrease in no-norm (Figure 1), indicating that it implicitly focuses on perturbing the BatchNorm layers already. We note that larger BatchNorm parameters do not necessarily indicate a functionally different network, since there are many reparameterization invariances in ReLU networks, some of which ASAM tries to leverage in its perturbation definition Eq. (4). Nevertheless, the scale of the network still has an impact on the training dynamics, since other methods like e.g. weight decay depend on it. We discuss the impact of weight decay further in Appendix B.9.
B.7 Removing the affine parameters
Frankle et al. [24] found for SGD that fixing the normalization parameters typically decreases the generalization performance of networks. As an ablation, we therefore study the effect of SAM when the normalization weights are non-trainable. This is, we set  $\gamma = 1$  and  $\beta = 0$  and train the remaining parameters with SAM. The results are shown for a WRN-28 in Figure 10, where it can be seen that fixing the normalization parameters (fix-norm) does not lead to a decrease in the performance of SAM. We thus hypothesize that in certain settings, SAM might not leverage the expressive power of the normalization layers, which might contribute to the improved performance of SAM-ON.
B.8 Training BatchNorm and only BatchNorm
The affine parameters of the normalization layers are relatively understudied in the literature. Recently, Frankle et al. [24] were able to obtain surprisingly high performance for ResNet architectures by only training the BatchNorm layers (freezing all other parameters), illustrating their expressive power. We study the effect of SAM in this setting (i.e. when all parameters except for the BatchNorm layers are frozen) for a ResNet-101 and a WRN-28 on CIFAR-10 and find that SAM still aids generalization in this setting (Table 17).
Table 17: Effect of SAM when training only BatchNorm layers, for networks trained on CIFAR-10.
Model	SGD	SAM ρ = 0.01	SAM ρ = 0.05
ResNet-101	78.75	78.63	79.27
WRN-28	63.49	64.48	62.70
B.9 Weight decay and dropout
Here, we explore potential connections of SAM-ON with weight decay and dropout. Since weight decay is sometimes applied to all network parameters, and sometimes normalization layers are omitted, it is worth investigating if the benefits of SAM-ON can be attributed to its interaction with weight decay. To this end, we train a WRN-28 with SGD, SAM-all and SAM-ON, and apply weight decay to either all parameters, all except the normalization layers, or not at all (Figure 11, right). For each setting SAM-ON outperforms SAM, outlining that its success should not be attributed to the interaction with weight decay.
We further test if SAM-ON-like performance can be achieved by simply applying stronger regularization and stochasticity to the normalization parameters. To this end, we apply dropout solely on the normalization layers (Figure 11, left) and find that this is not the case.
Figure 9: SAM-ON leads to a shift in the distribution of  $\gamma$ .
Figure 10: When training with SAM, fixing  $\gamma = 1$ ,  $\beta = 0$  (fix-norm) barely changes the performance of the network. WRN-28, CIFAR-100.
Figure 11: Left: Applying dropout only to the normalization layers (blue/red) performs worse than SAM-ON. Right: SAM-ON improves over SAM-all irrespective of whether weight decay is applied to all parameters (green), all except the normalization layers (blue) or not at all (yellow).
B.10 Details on sharpness evaluation
For the following discussion, we note that the term generalization is sometimes used as the difference between train and test error, while in other cases people use it as a synonym for test error. Since in CIFAR settings the models achieve train error close to zero, the two definitions become equivalent.
Many studies have attempted to better understand the possible connection between the generalization of deep neural networks and the flatness of the loss-surface [26, 32, 22, 34, 18]. Recently, Andriushchenko et al. [4] conducted a large-scale study for a range of models, datasets, and sharpness-definitions, finding that "while there definitely exist restricted settings where correlation between sharpness and generalization is significantly positive (e.g., for ResNets on CIFAR-10 with a specific combination of augmentations and mixup) it is not true anymore when we compare all models jointly" and concluding "that one should avoid blanket statements like flat minima generalize better". In order to evaluate sharpness, we therefore adopt their setup and choose the best-performing sharpness measure for CIFAR from their study, which is logit-normalized elementwise-adaptive worst-case- $\ell_{\infty}-m$ -sharpness. This is,  $m$ -sharpness  $s_w^m$  is defined as the largest possible change in loss within the adaptive perturbation model defined in 4,
$$
s _ {w} ^ {m} = \mathbb {E} _ {x, y \sim D _ {m}} \max  _ {\left\| T _ {w} ^ {- 1} \epsilon \right\| _ {p} \leq \rho} L (\mathbf {w} + \epsilon) - L (\mathbf {w}) \tag {6}
$$
where  $T_{w}^{i} = |w_{i}|$ ,  $p = \infty$  and  $D_{m}$  returns data batches of size  $m$ .  $\rho$  here denotes the size of the ball over which sharpness is evaluated and is not to be confused with the  $\rho$  from the SAM-algorithm. Like for ASAM [37], the motivation behind adaptive sharpness measures is to make them invariant to reparameterizations of the network. Further, the logit-outputs of the network are normalized with respect to their  $\ell_{2}$ -norm in order to mitigate the scale-sensitivity of classification losses. In practice, Andriushchenko et al. [4] compute  $s_w^m$  over a subset of the train set of size 1024 and use  $m = 128$ , i.e. average 8 batches. We use a subset of size 2048 in order to obtain more reliable sharpness estimates, and adopt  $m = 128$ . The maximization in (6) is performed with AutoPGD [16], a hyperparameter-free method designed for accurate estimation of adversarial robustness. It is to note that except for the logit-normalization, the sharpness definition reported in Table 6 corresponds exactly to the perturbation model that ASAM elementwise  $\ell_{\infty}$  uses, and hence the 1-step sharpness reported should be fairly close the the objective that ASAM elementwise  $\ell_{\infty}$  actually minimizes during training. While ASAM elementwise  $\ell_{\infty}$  yields slightly smaller sharpness values than the conventional SAM algorithm, the differences are rather small when compared to the significantly sharper SAM-ON models. For the results in Table 6 in the main paper we tuned the sharpness radius  $\rho$  such that we obtain sharpness values similar to those reported to yield the highest correlation in Andriushchenko et al. [4]. In Table 18 we report sharpness values for a ResNeXt-model, in addition to the WRN-28 from the main paper. In all cases the SAM-ON models are sharper than the SAM-all models yet generalize better. In Table 19 we further report other sharpness measures without logit-normalization for a WRN-28. SAM-ON is sharper than SAM-all with respect to most metrics, although there exist some exceptions. It should however be stressed that many of those metrics did not show good correlation with generalization in the study by Andriushchenko et al. [4].
Table 18: Sharpness evaluation of both a WRN-28 and a ResNeXt. SAM-ON is sharper than SAM-all in all cases. Shown is 20-step logit-normalized  $\ell_{\infty}$  sharpness from [2], averaged over three models per method. Dataset considered is CIFAR-100.
SGD	SAM	ASAM-el.-\( \ell_{\infty} \)
all	ON	all	ON
WRN-28	Test Accuracy (%)	\( 80.71^{\pm 0.2} \)	\( 83.11^{\pm 0.3} \)	\( 84.19^{\pm 0.2} \)	\( 83.25^{\pm 0.2} \)	\( 84.14^{\pm 0.2} \)
\( \ell_{\infty} \)-sharpness, \( \rho = 0.003 \)	\( 0.071^{\pm 0.000} \)	\( 0.048^{\pm 0.001} \)	\( 0.090^{\pm 0.005} \)	\( 0.048^{\pm 0.001} \)	\( 0.078^{\pm 0.004} \)
\( \ell_{\infty} \)-sharpness, \( \rho = 0.005 \)	\( 0.201^{\pm 0.001} \)	\( 0.139^{\pm 0.004} \)	\( 0.296^{\pm 0.018} \)	\( 0.124^{\pm 0.002} \)	\( 0.283^{\pm 0.011} \)
\( \ell_{\infty} \)-sharpness, \( \rho = 0.007 \)	\( 0.433^{\pm 0.002} \)	\( 0.309^{\pm 0.011} \)	\( 0.585^{\pm 0.018} \)	\( 0.255^{\pm 0.005} \)	\( 0.580^{\pm 0.020} \)
ResNeXt	Test Accuracy (%)	\( 80.16^{\pm 0.3} \)	\( 81.79^{\pm 0.4} \)	\( 82.22^{\pm 0.2} \)	\( 81.02^{\pm 0.6} \)	\( 82.38^{\pm 0.3} \)
\( \ell_{\infty} \)-sharpness, \( \rho = 0.001 \)	\( 0.036^{\pm 0.001} \)	\( 0.029^{\pm 0.000} \)	\( 0.034^{\pm 0.000} \)	\( 0.026^{\pm 0.002} \)	\( 0.034^{\pm 0.001} \)
\( \ell_{\infty} \)-sharpness, \( \rho = 0.003 \)	\( 0.164^{\pm 0.005} \)	\( 0.117^{\pm 0.004} \)	\( 0.140^{\pm 0.002} \)	\( 0.099^{\pm 0.010} \)	\( 0.147^{\pm 0.001} \)
\( \ell_{\infty} \)-sharpness, \( \rho = 0.005 \)	\( 0.383^{\pm 0.011} \)	\( 0.252^{\pm 0.008} \)	\( 0.291^{\pm 0.005} \)	\( 0.203^{\pm 0.021} \)	\( 0.312^{\pm 0.001} \)
Table 19: Additional sharpness measures. WRN-28 (no logitnorm).
adaptive	SGD	SAM	ASAM-el.-\( \ell_{\infty} \)
all	ON	all	ON
Test Accuracy (%)	80.71±0.2	83.11±0.3	84.19±0.2	83.25±0.2	84.14±0.2
\( \ell_2avg, \rho=0.005 \)	False	1.358±0.049	0.515±0.020	2.372±0.071	0.569±0.012	2.141±0.045
\( \ell_2avg, \rho=0.1 \)	True	0.042±0.001	0.019±0.001	0.022±0.001	0.040±0.001	0.019±0.001
\( \ell_{\infty}avg, \rho=0.01 \)	False	2.643±0.097	1.264±0.028	3.455±0.050	1.304±0.007	3.259±0.031
\( \ell_{\infty}avg, \rho=0.2 \)	True	0.078±0.001	0.035±0.001	0.034±0.004	0.068±0.003	0.031±0.001
\( \ell_{2}-worst, \rho=0.05 \)	False	0.501±0.048	0.655±0.277	0.701±0.057	0.768±0.141	0.313±0.044
\( \ell_{2}-worst, \rho=0.25 \)	True	0.065±0.008	0.033±0.004	0.037±0.017	0.056±0.006	0.062±0.001
\( \ell_{\infty}-worst, \rho=1e-05 \)	False	0.149±0.003	0.055±0.002	0.144±0.005	0.050±0.002	0.123±0.007
\( \ell_{\infty}-worst, \rho=0.004 \)	True	0.537±0.023	0.262±0.009	0.600±0.053	0.255±0.011	0.505±0.027
C Convergence Analysis
We provide in this section a convergence analysis for SAM-ON in the non-convex setting. Using standard assumptions we obtain a theorem which resembles findings for closely related methods such as found in [2, 42].
Our assumptions:
Assumption C.1. We assume function  $f: \mathbb{R}^n \to \mathbb{R}$  to be  $L$ -smooth: there exists  $L > 0$  such that
$$
\| \nabla f (v) - \nabla f (w) \| _ {2} \leq L \| v - w \| _ {2}, \forall v, w \in \mathbb {R} ^ {n}. \tag {7}
$$
Assumption C.2. There exists  $M > 0$  for any sample  $x_{i}$  such that
$$
\left\| \nabla f _ {x _ {i}} (w) \right\| _ {2} ^ {2} \leq M, \forall w \in \mathbb {R} ^ {n}. \tag {8}
$$
Remark C.3. If Assumption C.1 holds ( $L$ -smoothness), then  $\forall v, w \in \mathbb{R}^n$ :
$$
| f (v) - \left(f (w) + \nabla f (w) ^ {T} (v - w)\right) | \leq \frac {L}{2} \| v - w \| _ {2} ^ {2}. \tag {9}
$$
This well-known result can be derived using the fundamental theorem of calculus and Cauchy-Schwartz.
Remark C.4. Assumption C.2 guarantees that the variance of the stochastic gradient is less than  $M$ .
SAM-ON. In the following we shall denote the true gradient as  $\nabla f(w)$  and the noisy observation gradient as  $g(w)$ . The gradient of the loss of the  $i$ th training example is denoted as  $g_{x_i}(w)$ . We partition the neural network parameters layer-wise as  $w = \{w_N, w_A\}$ , with  $w_N \in \mathbb{R}^{n_N}$ ,  $w_A \in \mathbb{R}^{n_A}$ ,  $n = n_N + n_A$ , where  $w_N$  represent the normalization layer parameters and  $w_A$  all other layers. The
iteration for  $w_{N}$  is:
$$
w _ {N} ^ {t + 1 / 2} = w _ {N} ^ {t} + \rho \frac {g _ {N , x _ {i}} (w ^ {t})}{\| g _ {N , x _ {i}} (w ^ {t}) \|}
$$
$$
w _ {N} ^ {t + 1} = w _ {N} ^ {t} - h g _ {N, x _ {i}} \left(w ^ {t + 1 / 2}\right) \tag {10}
$$
and for  $w_{A}$  is:
$$
w _ {A} ^ {t + 1 / 2} = w _ {A} ^ {t}
$$
$$
w _ {A} ^ {t + 1} = w _ {A} ^ {t} - h g _ {A, x _ {i}} \left(w ^ {t + 1 / 2}\right). \tag {11}
$$
Theorem C.5. Assuming  $C.1$  and  $C.2$ ,  $h \leq 1 / L$ , we obtain:
$$
\frac {1}{T} \sum_ {t = 0} ^ {T - 1} \mathbb {E} \left[ \| \nabla f (w ^ {t}) \| ^ {2} \right] \leq \frac {2 \left(f \left(w ^ {0}\right) - f \left(w ^ {*}\right)\right)}{h T} + 2 L h M + L ^ {2} \rho^ {2} (1 + L h), \tag {12}
$$
with  $w^{*}$  the optimal solution to  $f(w)$ .
Proof. From Assumption C.1 and thus Remark C.3 it follows that:
$$
\begin{array}{l} f \left(w ^ {t + 1}\right) \leq f \left(w ^ {t}\right) + \nabla f \left(w ^ {t}\right) \cdot \left(w ^ {t + 1} - w ^ {t}\right) + \frac {L}{2} \| w ^ {t + 1} - w ^ {t} \| ^ {2} (13) \\ \leq f \left(w ^ {t}\right) - h \nabla f \left(w ^ {t}\right) \cdot g _ {x _ {i}} \left(w ^ {t + 1 / 2}\right) + \frac {h ^ {2} L}{2} \left\| g _ {x _ {i}} \left(w ^ {t + 1 / 2}\right) \right\| ^ {2} (14) \\ = f \left(w ^ {t}\right) - h \nabla f \left(w ^ {t}\right) \cdot g _ {x _ {i}} \left(w ^ {t + 1 / 2}\right) \\ + \frac {h ^ {2} L}{2} \left(\| \nabla f (w ^ {t}) - g _ {x _ {i}} \left(w ^ {t + 1 / 2}\right) \| ^ {2} - \| \nabla f (w ^ {t}) \| ^ {2} + 2 \left(\nabla f (w ^ {t}) \cdot g _ {x _ {i}} \left(w ^ {t + 1 / 2}\right)\right)\right) \\ = f (w ^ {t}) - \frac {L h ^ {2}}{2} \| \nabla f (w ^ {t}) \| ^ {2} + \frac {L h ^ {2}}{2} \| \nabla f (w ^ {t}) - g _ {x _ {i}} (w ^ {t + 1 / 2}) \| ^ {2} \\ - (1 - L h) h \left(\nabla f \left(w ^ {t}\right) \cdot g _ {x _ {i}} \left(w ^ {t + 1 / 2}\right)\right) (15) \\ \leq f (w ^ {t}) - \frac {L h ^ {2}}{2} \| \nabla f (w ^ {t}) \| ^ {2} + L h ^ {2} \| \nabla f (w ^ {t}) - g _ {x _ {i}} (w ^ {t}) \| ^ {2} \\ + L h ^ {2} \left\| g _ {x _ {i}} \left(w ^ {t}\right) - g _ {x _ {i}} \left(w ^ {t + 1 / 2}\right) \right\| ^ {2} - (1 - L h) h \left(\nabla f \left(w ^ {t}\right) \cdot g _ {x _ {i}} \left(w ^ {t + 1 / 2}\right)\right). (16) \\ \end{array}
$$
Taking the double expectation gives (because unbiased gradient and Assumption C.2 and Remark C.4):
$$
\begin{array}{l} \mathbb {E} [ f (w ^ {t + 1}) ] \leq \mathbb {E} [ f (w ^ {t}) ] - \frac {L h ^ {2}}{2} \mathbb {E} \| \nabla f (w ^ {t}) \| ^ {2} + L h ^ {2} M \\ + \underbrace {L h ^ {2} \| g \left(w ^ {t}\right) - g \left(w ^ {t + 1 / 2}\right) \| ^ {2}} _ {\mathcal {A}} - (1 - L h) h \underbrace {\mathbb {E} \left[ \nabla f \left(w ^ {t}\right) \cdot g \left(w ^ {t + 1 / 2}\right) \right]} _ {\mathcal {B}}. \tag {17} \\ \end{array}
$$
For term  $\mathcal{A}$  we obtain using Assumption C.1:
$$
\mathcal {A} \leq L ^ {3} h ^ {2} \| w ^ {t} - w ^ {t + 1 / 2} \| ^ {2} = L ^ {3} h ^ {2} \rho^ {2}. \tag {18}
$$
For term  $\mathcal{B}$  we obtain: