Original: Improved Deep Neural Network Generalization Using m-Sharpness-Aware Minimization
Revised: Improved Deep Neural Network Generalization Using m-Sharpness-Aware Minimization

Original: author names withheld

Original: Under Review for OPT

Original: Abstract
Revised: Abstract

Original: Modern deep learning models are over-parameterized, where the optimization setup strongly affects the generalization performance.
Revised: Modern deep learning models are over-parameterized, where the optimization setup strongly affects the generalization performance.

Original: A key element of reliable optimization for these systems is the modification of the loss function.
Revised: A key element of reliable optimization for these systems is the modification of the loss function.

Original: Sharpness-Aware Minimization (SAM) modifies the underlying loss function to guide descent methods towards flatter minima, which arguably have better generalization abilities.
Revised: Sharpness-Aware Minimization (SAM) modifies the underlying loss function to guide descent methods towards flatter minima, which arguably have better generalization abilities.

Original: In this paper, we focus on a variant of SAM known as mSAM, which, during training, averages the updates generated by adversarial perturbations across several disjoint shards of a mini-batch.
Revised: In this paper, we focus on a variant of SAM known as mSAM, which, during training, averages the updates generated by adversarial perturbations across several disjoint shards of a mini-batch.

Original: Recent work suggests that mSAM can outperform SAM in terms of test accuracy.
Revised: Recent work suggests that mSAM can outperform SAM in terms of test accuracy.

Original: However, a comprehensive empirical study of mSAM is missing from the literature— previous results have mostly been limited to specific architectures and datasets.
Revised: However, a comprehensive empirical study of mSAM is missing from the literature— previous results have mostly been limited to specific architectures and datasets.

Original: To that end, this paper presents a thorough empirical evaluation of mSAM on various tasks and datasets.
Revised: To that end, this paper presents a thorough empirical evaluation of mSAM on various tasks and datasets.

Original: We provide a flexible implementation of mSAM and compare the generalization performance of mSAM to the performance of SAM and vanilla training on different image classification and natural language processing tasks.
Revised: We provide a flexible implementation of mSAM and compare the generalization performance of mSAM to the performance of SAM and vanilla training on different image classification and natural language processing tasks.

Original: We also conduct careful experiments to understand the computational cost of training with mSAM, its sensitivity to hyperparameters and its correlation with the flatness of the loss landscape.
Revised: We also conduct careful experiments to understand the computational cost of training with mSAM, its sensitivity to hyperparameters and its correlation with the flatness of the loss landscape.

Original: Our analysis reveals that mSAM yields superior generalization performance and flatter minima, compared to SAM, across a wide range of tasks without significantly increasing computational costs.
Revised: Our analysis reveals that mSAM yields superior generalization performance and flatter minima, compared to SAM, across a wide range of tasks without significantly increasing computational costs.

Original: 1. Introduction
Revised: 1. Introduction

Original: In recent years, overparameterized deep neural networks (DNNs) have become a staple of modern machine learning advancements and crucial to many domains.
Revised: In recent years, overparameterized deep neural networks (DNNs) have become a staple of modern machine learning advancements and crucial to many domains.

Original: In particular, overparameterization has been successful in furthering the state-of-the-art performance for domains like image understanding [11, 16, 23], natural language processing (NLP) [4, 17, 24] and recommender systems [9, 20].
Revised: In particular, overparameterization has been successful in furthering the state-of-the-art performance for domains like image understanding [11, 16, 23], natural language processing (NLP) [4, 17, 24] and recommender systems [9, 20].

Original: Training DNNs requires minimizing complex and non-convex loss functions with an abundance of local minima.
Revised: Training DNNs requires minimizing complex and non-convex loss functions with an abundance of local minima.

Original: Interestingly, different local minima can have varying loss values and generalizability on unseen data.
Revised: Interestingly, different local minima can have varying loss values and generalizability on unseen data.

Original: Thus, it is essential to carefully choose an optimization scheme that can seek out minima that yield strong generalization performance.
Revised: Thus, it is essential to carefully choose an optimization scheme that can seek out minima that yield strong generalization performance.

Original: In recent years, a wide range of optimization algorithms has been developed for various domains, e.g. stochastic gradient descent (SGD), heavy-ball momentum [22], Adam [14], LAMB [30] among others.
Revised: In recent years, a wide range of optimization algorithms has been developed for various domains, e.g. stochastic gradient descent (SGD), heavy-ball momentum [22], Adam [14], LAMB [30] among others.

Original: These methods help yield strong generalization performance when coupled with suitable regularisation techniques.
Revised: These methods help yield strong generalization performance when coupled with suitable regularization techniques.

Original: Recently, extensive work has been done to study the correlation between the geometry of the loss landscape and generalization performance [6, 10, 13, 21, 28, 29].
Revised: Recently, extensive work has been done to study the correlation between the geometry of the loss landscape and generalization performance [6, 10, 13, 21, 28, 29].

Original: The newly proposed Sharpness- © . Aware Minimization (SAM) algorithm [7] leverages this correlation between the flatness of minima and generalization performance by seeking to obtain flatter solutions during training, leading to superior generalization for a panoply of tasks and domains.
Revised: The newly proposed SharpnessAware Minimization (SAM) algorithm [7] leverages this correlation between the flatness of minima and generalization performance by seeking to obtain flatter solutions during training, leading to superior generalization for a panoply of tasks and domains.

Original: In this paper, we focus on a variant of
Revised: In this paper, we focus on a variant of

Original: SAM called mSAM [7].
Revised: SAM called mSAM [7].

Original: Instead of using a single adversarial perturbation for the entire mini-batch mSAM executes SAM-like training by averaging the updates generated by adversarial perturbations across several disjoint shards of a mini-batch.
Revised: Instead of using a single adversarial perturbation for the entire mini-batch mSAM executes SAM-like training by averaging the updates generated by adversarial perturbations across several disjoint shards of a mini-batch.

Original: The name mSAM stems from the usage of m such disjoint shards.
Revised: The name mSAM stems from the usage of m such disjoint shards.

Revised: SAM represents flatness by looking at what happens in one particular adversarial direction, whereas mSAM leverages several such directions; which may better represent flatness.

Original: Related Work:
Revised: Related Work:

Original: Although the sharpness of the loss landscape can be calculated using several different measures (for example, the largest eigenvalue [27] or trace [12] of the Hessian of the loss), most of these measures are computationally too expensive for practical purposes.
Revised: Although the sharpness of the loss landscape can be calculated using several different measures (for example, the largest eigenvalue [27] or trace [12] of the Hessian of the loss), most of these measures are computationally too expensive for practical purposes.

Original: The main idea behind the SAM [7] algorithm is to encourage the network to seek regions where the worst loss value in a local neighbourhood is not too large (see Section 2 for more details).
Revised: The main idea behind the SAM [7] algorithm is to encourage the network to seek regions where the worst loss value in a local neighbourhood is not too large (see Section 2 for more details).

Original: This proxy of sharpness lends itself to easy computation, unlike the measures of sharpness described earlier.
Revised: This proxy of sharpness lends itself to easy computation, unlike the measures of sharpness described earlier.

Original: SAM has sparked interest in sharpness-aware training, resulting in several variants [5, 19, 32].
Revised: SAM has sparked interest in sharpness-aware training, resulting in several variants [5, 19, 32].

Original: In this paper, we focus on a less-understood variant of SAM, known as mSAM.
Revised: In this paper, we focus on a less-understood variant of SAM, known as mSAM.

Original: While SAM uses a single adversarial perturbation for a mini-batch, mSAM divides the mini-batch into disjoint shards and computes updates using adversarial perturbations applied to each shard.
Revised: While SAM uses a single adversarial perturbation for a mini-batch, mSAM divides the mini-batch into disjoint shards and computes updates using adversarial perturbations applied to each shard.

Original: The updates are then averaged, resulting in a single mSAM update.
Revised: The updates are then averaged, resulting in a single mSAM update.

Original: In [7], mSAM is used implicitly to reduce the computational cost of SAM by avoiding synchronization across multiple GPUs (referred to as “accelerators” hereinafter).
Revised: In [7], mSAM is used implicitly to reduce the computational cost of SAM by avoiding synchronization across multiple GPUs (referred to as “accelerators” hereinafter).

Original: Thus, each accelerator’s shard of data ends up with its own adversarial perturbation.
Revised: Thus, each accelerator’s shard of data ends up with its own adversarial perturbation.

Original: Recently, it has been observed via limited experimentation that mSAM results in better generalization performance [2, 3, 7].
Revised: Recently, it has been observed via limited experimentation that mSAM results in better generalization performance [2, 3, 7].

Original: Although the authors of [1] present mathematical expressions for mSAM, their analysis is primarily focused on a particular version of mSAM (see Section 2 for more details).
Revised: Although the authors of [1] present mathematical expressions for mSAM, their analysis is primarily focused on a particular version of mSAM (see Section 2 for more details).

Original: The experiments are also limited to image classification tasks on small architectures.
Revised: The experiments are also limited to image classification tasks on small architectures.

Original: As a result, the current understanding of mSAM is limited, and a comprehensive empirical study of its performance on state-of-the-art architectures and datasets is missing from the literature.
Revised: As a result, the current understanding of mSAM is limited, and a comprehensive empirical study of its performance on state-of-the-art architectures and datasets is missing from the literature.

Original: Our Contributions:
Revised: Our Contributions:

Original: We conduct a thorough empirical study of the mSAM algorithm and perform the following explorations that extol the value of mSAM.
Revised: We conduct a thorough empirical study of the mSAM algorithm and perform the following explorations that extol the value of mSAM.

Original: (i)
Revised: (i)

Original: Starting from the mathematical description of mSAM, we present an explicit flexible implementation of mSAM that does not rely on accelerator synchronization and is compatible with any single/multi-accelerator setup.
Revised: Starting from the mathematical description of mSAM, we present an explicit flexible implementation of mSAM that does not rely on accelerator synchronization and is compatible with any single/multi-accelerator setup.

Original: (ii) We conduct extensive experiments on a wide variety of tasks using images and NLP data, leveraging architectures like Convolutional Neural Networks (CNNs) and transformers.
Revised: (ii) We conduct extensive experiments on a wide variety of tasks using images and NLP data, leveraging architectures like Convolutional Neural Networks (CNNs) and transformers.

Original: In our experiments, mSAM consistently outperforms SAM and vanilla training.
Revised: In our experiments, mSAM consistently outperforms SAM and vanilla training.

Original: (iii)
Revised: (iii)

Original: We also conduct careful experiments to understand the computational cost of training with mSAM, its sensitivity to hyperparameters like m and its correlation with the flatness of the loss landscape.
Revised: We also conduct careful experiments to understand the computational cost of training with mSAM, its sensitivity to hyperparameters like m and its correlation with the flatness of the loss landscape.

Original: Our thorough empirical investigation reveals that mSAM substantially improves the generalization performance of SAM without significantly increasing the computational cost of training.
Revised: Our thorough empirical investigation reveals that mSAM substantially improves the generalization performance of SAM without significantly increasing the computational cost of training.

Original: 2. Algorithm
Revised: 2. Algorithm

Original: mSAM is based on the SAM algorithm that aims to obtain flat solutions to the empirical loss function.
Revised: mSAM is based on the SAM algorithm that aims to obtain flat solutions to the empirical loss function.

Original: In particular, SAM tries to find a solution that minimizes the worst-case loss in a ball around the solution.
Revised: In particular, SAM tries to find a solution that minimizes the worst-case loss in a ball around the solution.

Original: Mathematically, let S = { ( x i , y i ) , i ∈
Revised: Mathematically, let S = { ( x i , y i ) , i ∈

Original: [ n ] : x i ∈ X , y i ∈ Y} be a dataset of n samples, where X is the set of features and Y is the set of outcomes..
Revised: [ n ] : x i ∈ X , y i ∈ Y} be a dataset of n samples, where X is the set of features and Y is the set of outcomes.

Original: Moreover, let ℓ : R d × X × Y (cid:55)→ R be a differentiable loss function, where d is the number of model parameters.
Revised: Moreover, let ℓ : R d × X × Y (cid:55)→ R be a differentiable loss function, where d is the number of model parameters.

Original: The empirical loss over the dataset S is defined as L S ( w ) =
Revised: The empirical loss over the dataset S is defined as L S ( w ) =

Original: (cid:80) ni
Revised: (cid:80) ni

Original: =1 ℓ ( w ; x i , y i ) /n , where w parameterizes the neural network.
Revised: =1 ℓ ( w ; x i , y i ) /n , where w parameterizes the neural network.

Original: With this notation in place, the SAM loss function is defined as [7]:
Revised: With this notation in place, the SAM loss function is defined as [7]:

Original: [Equation]
Revised: [Equation]

Original: for some p ≥ 1 .
Revised: for some p ≥ 1 .

Original: In this work, we use p = 2 .
Revised: In this work, we use p = 2 .

Original: In practice, however, the maximization step in (1) cannot be done in closed-form.
Revised: In practice, however, the maximization step in (1) cannot be done in closed-form.

Original: Hence, authors in [7] use a first-order approximation to L S to simplify (1) as 
Revised: Hence, authors in [7] use a first-order approximation to L S to simplify (1) as [Equation]

Revised: It is easy to see that the maximum in Problem (2) is achieved for

Original: [Equation]
Revised: [Equation]

Revised: As a result, L SAM S ≈ L S ( w + ˆ ϵ ) .

Original: The gradient of the SAM loss,keeping only the first order terms, is given by:
Revised: This leads to the gradient ∇L [Equation] However, calculating ∂ ( w +ˆ ϵ ) /∂w in (4) involves second order terms that require access to Hessian, which can be computationally inefficient in practice. Thus, by ignoring second order terms in (4), gradient of the SAM loss, is approximated as [7]:

Original: [Equation]
Revised: [Equation]

Original: which is used in the SAM algorithm (for example, in conjunction with SGD).
Revised: which is used in the SAM algorithm (for example, in conjunction with SGD).

Original: We refer to [7] for more details and intuitions about SAM.
Revised: We refer to [7] for more details and intuitions about SAM.

Original: We call the inner gradient calculations on the right-hand side of (3) as the SAM ascent step and the outer gradient calculations as the gradient step.
Revised: We call the inner gradient calculations on the right-hand side of (5) as the SAM ascent step and the outer gradient calculations as the gradient step.

Original: mSAM [7] is a variation of the SAM algorithm.
Revised: mSAM [7] is a variation of the SAM algorithm.

Original: In general, for mSAM, a minibatch of data S is further divided into m smaller disjoint shards (aka “micro-batches”), such as S 1 , · · · , S m where ∪ mi =1 S
Revised: In general, for mSAM, a minibatch of data S is further divided into m smaller disjoint shards (aka “micro-batches”), such as S 1 , · · · , S m where ∪ mi =1 S

Original: i
Revised: i

Original: =
Revised: =

Original: S .
Revised: S .

Original: For simplicity, we assume |S 1 | = · · · = |S m | = |S| /m
Revised: For simplicity, we assume |S 1 | = · · · = |S m | = |S| /m

Original: although such an assumption is not necessary in general.
Revised: although such an assumption is not necessary in general.

Original: The mSAM loss is a variation of the SAM loss, defined as:
Revised: The mSAM loss is a variation of the SAM loss, defined as:

Original: [Equation]
Revised: [Equation]

Revised: [Equation]

Revised: [Table]

Revised: [Equation]

Revised: [Table]

Revised: [Equation]

Original: Intuitively, mSAM is a version of SAM where the ascent step (or the weight perturbation) of SAM is done independently on each micro-batch using different ϵ ( i ) , instead of using an average perturbation such as ϵ for all micro-batches.
Revised: [Table] [Equation] [Table] [Equation]

Revised: [Table]

Revised: [Equation]

Revised: [Table]

Revised: Intuitively

Revised: [Table]

Revised: mSAM is a version of SAM where the ascent step (or the weight perturbation) of SAM is done independently on each micro-batch using different ϵ ( i ) , instead of using an average perturbation such as ϵ for all micro-batches.

Original: The mSAM gradient can thereby be derived as:
Revised: The mSAM gradient can thereby be derived as:

Original: [Equation]
Revised: [Equation]

Original: where (5) is a first-order approximation to the gradient of (4).
Revised: where (7) is a first-order approximation to the gradient of (6).

Original: We also note that the loss (4) is related to the mSAM definition of [1].
Revised: We also note that the loss (6) is related to the mSAM definition of [1].

Original: See Table 1 for a side-by-side comparison of SAM and mSAM, and their different implementations.
Revised: See Table 1 for a side-by-side comparison of SAM and mSAM [Table] and their different implementations.

Original: An important distinction between our work and prior work is that we treat m as a model hyperparameter to improve generalization.
Revised: An important distinction between our work and prior work is that we treat m as a model hyperparameter to improve generalization.

Original: In particular, in mSAM implementation of [7], the value of m is fixed to the number of hardware accelerators, micro-batch
Revised: In particular, in mSAM implementation of [7], the value of m is fixed to the number of hardware accelerators, micro-batch

Original: i is the part of the data that is loaded onto accelerator i , and each accelerator uses a separate perturbation, simulating the effect of mSAM.
Revised: i is the part of the data that is loaded onto accelerator i , and each accelerator uses a separate perturbation, simulating the effect of mSAM.

Original: With this implementation, m is an artefact of the hardware setup.
Revised: With this implementation, m is an artefact of the hardware setup.

Original: On the other hand, the analysis of [1] mostly concerns the value m = |S| .
Revised: On the other hand, the analysis of [1] mostly concerns the value m = |S| .

Original: In contrast, we consider a wide range of values for m in our experiments and choose the  value of m that leads to the best generalization performance.
Revised: In contrast, we consider a wide range of values for m in our experiments—this offers the flexibility to choose an appropriate value of m that leads to a better generalization performance.

Original: Moreover, our implementation supports any single/multi-accelerator setup and allows the user to set an appropriate value of m .[Table] [Table]
Revised: Moreover, our implementation supports any single/multi-accelerator setup and allows the user to set an appropriate value of m .

Original: 3. Numerical Experiments
Revised: 3. Numerical Experiments

Original: This section compares mSAM to SAM and vanilla optimization methods  on various model architectures and datasets.
Revised: This section compares mSAM to SAM and vanilla optimization methods (i.e. without sharpnessaware modification) on various model architectures and datasets.

Original: We report the average and standard deviation of accuracy on the test data over five independent runs.
Revised: We report the average and standard deviation of accuracy on the test data over five independent runs.

Revised: We also note that we use the same values of hyper-parameters for all algorithms, where ρ is chosen based on the best validation performance for SAM and other hyper-parameters are chosen based on the best validation error for vanilla methods.

Revised: Moreover, although it is possible to use different values of ρ for each micro-batch in mSAM, doing so requires tuning numerous hyper-parameters (as we usually take m to be large) which is computationally infeasible.

Revised: Therefore, we use the same value of ρ for all micro-batches.

Revised: [Table]

Original: 3.1. Image Classification
Revised: 3.1. Image Classification

Original: In our first set of experiments on image classification datasets, we compare the performance of mSAM, SAM and vanilla methods with multiple CNN architectures such as ResNets [11] and
Revised: In our first set of experiments on image classification datasets [Table] we compare the performance of mSAM, SAM and vanilla methods with multiple CNN architectures such as ResNets [11] and

Original: WideResNet [31].
Revised: WideResNet [31].

Original: We use CIFAR10/100 [15] datasets as our test bed.
Revised: We use CIFAR10/100 [15] datasets as our test bed.

Original: The average accuracies for the experiments with CIFAR data are reported in Table 2.
Revised: The average accuracies for the experiments with CIFAR data are reported in Table 2.

Original: These experiments are done on four NVidia V100 GPUs, with an effective batch size of 512 .
Revised: These experiments are done on four NVidia V100 GPUs [Table] with an effective batch size of 512 and the value of m = 32 for mSAM.

Original: Other hyperparameters used to produce these results can be found in Appendix B. Overall, mSAM leads to better accuracy than SAM and vanilla methods in all cases reported in Table 2.
Revised: Other hyper-parameters used to produce these results can be found in Appendix B. Overall, mSAM leads to better accuracy than SAM and vanilla methods in all cases reported in Table 2.

Original: In particular, mSAM achieves the best improvement in CIFAR100 data (about 2% on ResNet50 and about 1% on WRN28-10) which is a more challenging dataset compared to CIFAR10.
Revised: In particular, mSAM achieves the best improvement in CIFAR100 data (about 2% on ResNet50 and about 1% on WRN-28-10) which is a more challenging dataset compared to CIFAR10.

Original: Following the recent results that suggest that sharpness-aware optimization can improve the training quality of ViTs [3], we conduct additional experiments on a ViT architecture.
Revised: Following the recent results that suggest that sharpness-aware optimization can improve the training quality of ViTs [3] [Table] we conduct additional experiments on a ViT architecture.

Original: In particular, we use the pre-trained ViT-B/16 checkpoint from [26] and fine-tune the model on CIFAR10/100 data independently.
Revised: In particular, we use the pre-trained ViT-B/16 checkpoint from [26] and fine-tune the model on CIFAR10/100 data independently.

Original: The average accuracy results for ViT fine-tuning are reported in Table 2.
Revised: The average accuracy results for ViT fine-tuning are reported in Table 2.

Original: Similar to the CNN case, mSAM leads to better results than SAM that improves upon vanilla optimization further.ResNet
Revised: Similar to the CNN case, mSAM leads to better results than SAM that improves upon vanilla optimization further.

Original: WRN-28- [Figure] 10m      1010m 81.
Revised: We note that in these experiments, m = 32 is fixed for mSAM.

Original: 3.2. NLP Fine-tuning
Revised: 3.2. NLP Fine-tuning

Original: Our next set of experiments is based on NLP analysis.
Revised: Our next set of experiments is based on NLP analysis.

Original: We select four tasks from GLUE benchmark [25].
Revised: We select four tasks from GLUE benchmark [25].

Original: In particular, we choose COLA and MRPC as two small datasets and SST-2 and QQP as two larger datasets for empirical evaluation.
Revised: In particular, we choose COLA and MRPC as two small datasets and SST-2 and QQP as two larger datasets for empirical evaluation.

Original: The fine-tuning experiments with the RoBERTa-base model [18] are performed on four NVidia V100 GPUs with an effective batch size of 32 .
Revised: The fine-tuning experiments with the RoBERTa-base model [18] are performed on four NVidia V100 GPUs with an effective batch size of 32 .

Original: For the ease of reproduction of the results, we tabulate all the hyper-parameters used in Appendix C. For the fine-tuning experiments, we report the average value of Matthews Correlation Coefficient for
Revised: For the ease of reproduction of the results, we tabulate all the hyper-parameters used in Appendix C. For the fine-tuning experiments, we report the average value of Matthews Correlation Coefficient for

Original: COLA, and average accuracy for other datasets in Table 3.
Revised: COLA, and average accuracy for other datasets in Table 3.

Original: Overall, mSAM performs better than the baseline methods on these datasets.
Revised: Overall, mSAM performs better than the baseline methods on these datasets.

Original: However, the variance among different runs is comparably high for smaller datasets such as COLA and MRPC.
Revised: However, the variance among different runs is comparably high for smaller datasets such as COLA and MRPC.

Original: On the other hand, the results on larger data such as
Revised: On the other hand, the results on larger data such as

Original: SST-2 and QQP are expectedly more robust across different runs. 
Revised: SST-2 and QQP are expectedly more robust across different runs. ResNet WRN-28- [Figure]

Revised: [Table]

Revised: [Table]

Original: 4. A Deeper Investigation of mSAM
Revised: 4. A Deeper Investigation of mSAM

Original: To further understand the mSAM algorithm, we design and report the experiments in this section.
Revised: To further understand the mSAM algorithm [Table] we design and report some experiments in this section.

Original: Additional experimental results are moved to the Appendix A.
Revised: Additional experimental results are moved to the Appendix A.

Original: Effect of varying m :
Revised: Effect of varying m :

Original: In our experiments, we have observed that a larger value of m often leads to better test accuracy.
Revised: In our experiments, we have observed that a larger value of m often leads to better test accuracy.

Original: We recover SAM by setting m = 1 , which produces inferior results.
Revised: We recover SAM by setting m = [Table] which produces inferior results.

Original: To test this hypothesis, we set up the experiments with the CIFAR100 dataset on two CNNs, ResNet and WRN-28-10 in the same setup as in Section 3.1.
Revised: To test this hypothesis [Table] we set up the experiments with the CIFAR100 dataset on two CNNs, ResNet and WRN-28-10 in the same setup as in Section 3.1.

Original: We run mSAM for different values of m ∈ { 4 , 8 , 16 , 32 , 64 } .
Revised: We run mSAM for different values of m ∈ { 4 , 8 , 16 , 32 , 64 } .

Original: The accuracy results for these experiments are shown in Figure 1.
Revised: The accuracy results for these experiments are shown in Figure 1.

Original: Increasing m improves the performance up to m ≈ 32 .
Revised: Increasing m improves the performance up to m ≈ 32 .

Original: However, a value of m larger than this threshold either leads to worse performance or marginal improvements, so increasing m does not necessarily result in better generalization.
Revised: However, a value of m larger than this threshold either leads to worse performance or marginal improvements, so increasing m does not necessarily result in better generalization.

Original: Intuitively speaking, when the micro-batch is too small, the perturbation derived according to the micro-batch might not be a good estimate of the actual SAM perturbation, leading to worse performance.
Revised: Intuitively speaking, when the micro-batch is too small, the perturbation derived according to the micro-batch might not be a good estimate of the actual SAM perturbation, leading to worse performance.

Original: We leave the theoretical analysis of such a phenomenon an interesting direction for future research.
Revised: We leave the theoretical analysis of such a phenomenon an interesting direction for future research.

Original: [Table] [Table]
Revised: We also note that understanding how the optimal value of m and batch size interact is an open question left for future work.

Original: Are mSAM solutions flat?
Revised: Are mSAM solutions flat?

Original: The SAM algorithm hypothesizes that flat solutions generalize better.
Revised: The SAM algorithm hypothesizes that flat solutions generalize better.

Original: Since mSAM consistently outperforms SAM [Table] it is worth investigating if mSAM settles for even flatter solutions.
Revised: Since mSAM consistently outperforms SAM, it is worth investigating if mSAM settles for even flatter solutions.

Original: To that end and to quantify sharpness [Table] we calculate the largest eigenvalue of the
Revised: To that end and to quantify sharpness, we calculate the largest eigenvalue of the

Original: Hessian of the loss function L S ( w ) at the final solution, denoted as λ max .
Revised: Hessian of the loss function L S ( w ) at the final solution, denoted as λ max .

Original: This metric is widely accepted to be a good indicator of the sharpness/flatness of a solution [27].
Revised: This metric is widely accepted to be a good indicator of the sharpness/flatness of a solution [27].

Original: We calculate λ max over the full train data using the power iteration, as implemented by [8].
Revised: We calculate λ max over the full train data using the power iteration, as implemented by [8].

Original: We use the ResNet50 and
Revised: We use the ResNet50 and

Original: WRN-28-10 models trained on CIFAR100 (see Section 3.1) to calculate λ max .
Revised: WRN-28-10 models trained on CIFAR100 (see Section 3.1) to calculate λ max .

Original: The average results for these experiments are reported in Table 4.
Revised: The average results for these experiments are reported in Table 4.

Original: We see that mSAM leads to solutions with smaller λ max than SAM and vanilla SGD, confirming our conjecture
Revised: We see that mSAM leads to solutions with smaller λ max than SAM and vanilla SGD, confirming our conjecture.

Original: [Table]

Original: mSAM runtime:
Revised: mSAM runtime:

Original: A general misconception about mSAM is that it is computationally inefficient, as the total number of forward-backwards passes in the network is multiplied by m [2].
Revised: A general misconception about mSAM is that it is computationally inefficient, as the total number of forward-backward passes in the network is multiplied by m [2].

Original: However, note that these passes are performed on micro-batches, which are m times smaller than the actual minibatch.
Revised: However, note that these passes are performed on micro-batches, which are m times smaller than the actual minibatch.

Original: Hence, the overall computational cost gets amortized and is never as high as m times the cost of SAM.
Revised: Hence, the overall computational cost gets amortized and is never as high as m times the cost of SAM.

Original: In practice, on large networks, the runtime of mSAM is only 1 . 2 - 1 .
Revised: In practice, on large networks, the runtime of mSAM is only 1 . 2 - 1 .

Original: 3 times more compared to SAM.
Revised: 3 times more compared to SAM.

Original: In Appendix A, we present numerical evidence for mSAM efficiency and discuss a few hybrid algorithms to reduce the computational cost of mSAM even further.
Revised: In Appendix A, we present numerical evidence for mSAM efficiency and discuss a few hybrid algorithms to reduce the computational cost of mSAM even further.

Original: 5. Discussion
Revised: 5. Discussion

Original: The empirical study presented in this paper shows that mSAM outperforms SAM on different datasets (image classification and NLP tasks) and model architectures (CNNs and transformers); the exact performance gap is primarily a function of the data and the architecture.
Revised: The empirical study presented in this paper shows that mSAM outperforms SAM on different datasets (image classification and NLP tasks) and model architectures (CNNs and transformers); the exact performance gap is primarily a function of the data and the architecture.

Original: Moreover, our experiments suggest that mSAM does not necessarily incur a substantially higher computational cost than
Revised: Moreover, our experiments suggest that mSAM does not necessarily incur a substantially higher computational cost than

Original: SAM, making it amenable for large-scale problems.
Revised: SAM, making it amenable for large-scale problems.

Original: An exciting avenue of work is the theoretical justification of better generalization abilities of mSAM.
Revised: An exciting avenue of work is the theoretical justification of better generalization abilities of mSAM.

Original: Our analysis suggests that mSAM leads to flatter solutions than SAM, which may explain the better generalization performance.
Revised: Our analysis suggests that mSAM leads to flatter solutions than SAM, which may explain the better generalization performance.

Original: However, the formal proof regarding why mSAM promotes a flatter minimum is left for future work. 
Revised: However, the formal proof regarding why mSAM promotes a flatter minimum is left for future work. 6. Acknowledgements

Revised: Kayhan Behdin contributed to this work while he was an intern at LinkedIn during summer 2022.

Revised: This work is not a part of his MIT research.

Revised: Rahul Mazumder contributed to this work while he was a consultant for LinkedIn (in compliance with MIT’s outside professional activities policies).

Revised: This work is not a part of his MIT research.

Original: Appendix A. mSAM and Computational Efficiency
Revised: Appendix A. mSAM and Computational Efficiency

Original: Since SAM requires two forward-backwards passes for each batch of data, SAM is almost twice as slow as vanilla training.
Revised: Since SAM requires two forward-backwards passes for each batch of data, SAM is almost twice as slow as vanilla training.

Original: In our experiments, mSAM appears to be slower than SAM, although not m times slower, as suggested by, for example [2].
Revised: In our experiments, mSAM appears to be slower than SAM, although not m times slower, as suggested by, for example [2].

Original: To demonstrate that, we report the average runtime of our experiments from Section 3.1 on CIFAR100 data in Table A.1.
Revised: To demonstrate that, we report the average runtime of our experiments from Section 3.1 on CIFAR100 data in Table A.1.

Original: Expectedly, SAM is twice as slow as the vanilla method.
Revised: Expectedly, SAM is twice as slow as the vanilla method.

Original: Interestingly, in the worst case, mSAM is only twice as slow as SAM, and in the best case, the computational penalty is only about 10% of SAM.
Revised: Interestingly, in the worst case, mSAM is only twice as slow as SAM, and in the best case, the computational penalty is only about 10% of SAM.

Original: The primary reason is that the overall computation complexity of forward-backwards passes for mSAM and
Revised: The primary reason is that the overall computation complexity of forward-backwards passes for mSAM and

Original: SAM is not too different.
Revised: SAM is not too different.

Original: In each epoch, mSAM performs m times more passes on micro-batches that are m times smaller.
Revised: In each epoch, mSAM performs m times more passes on micro-batches that are m times smaller.

Original: Encouragingly, mSAM appears more efficient on larger architectures such as ViT-B/16 and WRN-28-10.
Revised: Encouragingly, mSAM appears more efficient on larger architectures such as ViT-B/16 and WRN-28-10.

Original: Although mSAM does not appear to be computationally prohibitive in our experiments, it is still not as efficient as vanilla training, leaving room for further improving its efficiency.
Revised: Although mSAM does not appear to be computationally prohibitive in our experiments, it is still not as efficient as vanilla training, leaving room for further improving its efficiency.

Original: To that end, we conduct the following set of experiments.
Revised: To that end, we conduct the following set of experiments.

Original: Building on our CIFAR100 experiments from
Revised: Building on our CIFAR100 experiments from

Original: Section 3.1, we start the training either with mSAM or vanilla training and then switch to the other training algorithm at some point.
Revised: Section 3.1, we start the training either with mSAM or vanilla training and then switch to the other training algorithm at some point.

Original: We keep all other training parameters fixed.
Revised: We keep all other training parameters fixed.

Original: The accuracy results for this setup for ResNet50 and WRN-28-10 are reported in Figure A.1.
Revised: The accuracy results for this setup for ResNet50 and WRN-28-10 are reported in Figure A.1.

Original: In this figure, the switch percent is the threshold in training when we transition from one algorithm to the other.
Revised: In this figure, the switch percent is the threshold in training when we transition from one algorithm to the other.

Original: For example, for the switch percent of 20 , if we start with mSAM, we use mSAM for the first 20% of epochs and vanilla updates for the rest.
Revised: For example, for the switch percent of 20 , if we start with mSAM, we use mSAM for the first 20% of epochs and vanilla updates for the rest.

Original: If mSAM is used for the initial and/or final part of training, the accuracy is always better than vanilla training.
Revised: If mSAM is used for the initial and/or final part of training, the accuracy is always better than vanilla training.

Original: In fact, in the WRN-28-10 case, as long as we partially use mSAM, the accuracy is almost the same as training with mSAM for the entire duration.
Revised: In fact, in the WRN-28-10 case, as long as we partially use mSAM, the accuracy is almost the same as training with mSAM for the entire duration.

Original: For
Revised: For

Original: ResNet50, not using mSAM for the whole training leads to a drop in performance; however, even in this case, the accuracy of the hybrid training is better than the SAM training.
Revised: ResNet50, not using mSAM for the whole training leads to a drop in performance; however, even in this case, the accuracy of the hybrid training is better than the SAM training.

Original: These observations suggest that it is possible to enjoy the superior performance of mSAM, at least to some degree, while not having to deal with the computational complexity of mSAM for the entire training.
Revised: These observations suggest that it is possible to enjoy the superior performance of mSAM, at least to some degree, while not having to deal with the computational complexity of mSAM for the entire training.

Original: A better theoretical and empirical understanding of the hybrid training method can be an exciting avenue for future work.
Revised: A better theoretical and empirical understanding of the hybrid training method can be an exciting avenue for future work.

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: [Table][Figure] [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: Appendix B. Hyper-parameters for Image Classification Experiments
Revised: Appendix B. Hyper-parameters for Image Classification Experiments

Original: As mentioned, our experiments on CIFAR data in this section are done on 4 Nvidia V100 GPUs, with effective batch size of 512.
Revised: As mentioned, our experiments on CIFAR data in this section are done on 4 Nvidia V100 GPUs, with effective batch size of 512.

Original: For mSAM, we have used the micro-batch size of 16 which corresponds to m = 32 .
Revised: For mSAM, we have used the micro-batch size of 16 which corresponds to m = 32 .

Original: Rest of the hyper-parameters are chosen as in Table B.1 for CNNs, and as in
Revised: Rest of the hyper-parameters are chosen as in Table B.1 for CNNs, and as in

Original: Table B.2 for ViT experiments.
Revised: Table B.2 for ViT experiments.

Original: Appendix C. Hyper-parameters for GLUE Experiments
Revised: Appendix C. Hyper-parameters for GLUE Experiments

Original: Similarly, our experiments in this section are done on 4 Nvidia V100 GPUs, with effective batch size of 32.
Revised: Similarly, our experiments in this section are done on 4 Nvidia V100 GPUs, with effective batch size of 32.

Original: For mSAM, we have used the microbatch size of 4 which corresponds to m = 8 .
Revised: For mSAM, we have used the microbatch size of 4 which corresponds to m = 8 .

Original: The rest of task-specific hyper-parameters can be found in Table C.1.
Revised: The rest of task-specific hyper-parameters can be found in Table C.1.

Original: [Table]
Revised: [Table]

Original: [Table] [Table]
Revised: [Table] [Table]