Original: When Do Flat Minima Optimizers Work?
Revised: When Do Flat Minima Optimizers Work?

Original: Abstract
Revised: Abstract

Original: Recently, flat-minima optimizers , which seek to find parameters in low-loss neighborhoods, have been shown to improve a neural network’s generalization performance over stochastic and adaptive gradient-based optimizers.
Revised: Recently, flat-minima optimizers , which seek to find parameters in low-loss neighborhoods, have been shown to improve a neural network’s generalization performance over stochastic and adaptive gradient-based optimizers.

Original: Two methods have received significant attention due to their scalability: 1.
Revised: Two methods have received significant attention due to their scalability: 1.

Original: Stochastic Weight Averaging (SWA), and 2.
Revised: Stochastic Weight Averaging (SWA), and 2.

Original: Sharpness-Aware Minimization (SAM).
Revised: Sharpness-Aware Minimization (SAM).

Original: However, there has been limited investigation into their properties and no systematic benchmarking of them across different domains.
Revised: However, there has been limited investigation into their properties and no systematic benchmarking of them across different domains.

Original: We fill this gap here by comparing the loss surfaces of the models trained with each method and through broad benchmarking across computer vision, natural language processing, and graph representation learning tasks.
Revised: We fill this gap here by comparing the loss surfaces of the models trained with each method and through broad benchmarking across computer vision, natural language processing, and graph representation learning tasks.

Original: We discover several surprising findings from these results, which we hope will help researchers further improve deep learning optimizers, and practitioners identify the right optimizer for their problem.
Revised: We discover several surprising findings from these results, which we hope will help researchers further improve deep learning optimizers, and practitioners identify the right optimizer for their problem.

Original:  Introduction
Revised:  Introduction

Original: Stochastic gradient descent (SGD) methods are central to neural network optimization [6].
Revised: Stochastic gradient descent (SGD) methods are central to neural network optimization [6].

Original: Recently, one class of algorithms has focused on biasing SGD methods towards so-called ‘ flat ’ minima, which are located in large weight space regions with very similar low loss values [41].
Revised: Recently, one class of algorithms has focused on biasing SGD methods towards so-called ‘ flat ’ minima, which are located in large weight space regions with very similar low loss values [42].

Original: Theoretical and empirical studies [21, 72, 9, 50, 46, 5, 12] postulate that such flatter regions generalize better than sharper minima, e g., due to the flat minimizer’s robustness against loss function shifts between train and test data, as illustrated in Fig.
Revised: Theoretical and empirical studies [20, 76, 9, 54, 48, 5, 12] postulate that such flatter regions generalize better than sharper minima, e g., due to the flat minimizer’s robustness against loss function shifts between train and test data, as illustrated in Fig.

Original: Two popular flat-minima optimization approaches are: 1.
Revised: Two popular flat-minima optimization approaches are: 1.

Original: Stochastic Weight Averaging (SWA) [45] and 2.
Revised: Stochastic Weight Averaging (SWA) [47] and 2.

Original: Sharpness-Aware Minimization (SAM) [22].
Revised: Sharpness-Aware Minimization (SAM) [21].

Original: While both strategies aim to find flatter minima they operate much differently.
Revised: While both strategies aim to find flatter minima they operate much differently.

Original: On the one hand
Revised: On the one hand

Original: SWA is based on the intuition that, near a flat minimum, gradients are smaller, leaving many iterates in that flat region.
Revised: SWA is based on the intuition that, near a flat minimum, gradients are smaller, leaving many iterates in that flat region.

Original: Therefore, averaging iterates will produce a solution that is pulled towards these flatter regions, see Fig. 1, top.
Revised: Therefore, averaging iterates will produce a solution that is pulled towards these flatter regions, see Fig. 1, top.

Original: On the other hand, SAM minimizes the maximum loss around a neighborhood of the current iterate.
Revised: On the other hand, SAM minimizes the maximum loss around a neighborhood of the current iterate.

Original: This way, a region around the iterate is designed to have uniformly low loss; see Fig.
Revised: This way, a region around the iterate is designed to have uniformly low loss; see Fig.

Original: 1, bottom.
Revised: 1, bottom.

Original: Crucially, SAM requires an additional forward/backward pass for each parameter update, making it more expensive than SWA.
Revised: Crucially, SAM requires an additional forward/backward pass for each parameter update, making it more expensive than SWA.

Original: Despite the successes [70, 3, 48, 12, 4] of SWA and SAM in some domains, we are unaware of a systematic comparison between them that would help practitioners to choose the right optimizer for their problem and researchers to develop better optimizers.
Revised: Despite the successes [75, 3, 50, 12, 4] of SWA and SAM in some domains, we are unaware of a systematic comparison between them that would help practitioners to choose the right optimizer for their problem and researchers to develop better optimizers.

Original: The SWA [45] paper was published in 2018, and the SAM [22] paper in 2021; however, the SAM paper, and its most noticeable followups [59, 12, 98], do not compare against SWA.
Revised: The SWA [47] paper was published in 2018, and the SAM [21] paper in 2021; however, the SAM paper, and its most noticeable follow-ups [64, 12, 102], do not compare against SWA.

Original: Further, there is very limited overlap in terms of the model architecture and dataset used in the experiments among both papers, which are likely further confounded by other differences in the training procedures (e
Revised: Further, there is very limited overlap in terms of the model architecture and dataset used in the experiments among both papers, which are likely further confounded by other differences in the training procedures (e

Original: [Table]
Revised: [Table]

Original: g. data augmentations, hyperparameters, etc.).
Revised: g. data augmentations, hyper-parameters, etc.).

Original: [Table]
Revised: [Table]

Original: [Equation]

Original: [Table]

Original: [Equation]

Original: [Table]

Original: //oHp41NEyVYS2ieRS9UKsKWeCtg0znPYSRXEcctoNJ7e5332iSjMpHs00oUGMR4JFjGBjJb8fYzMmmGf3s0G15tbdOdAq8QpSgwKtQfWrP5QkjakwhGOtfc9NTJBhZRjhdFbpp5ommEzwiPqWChxTHWTzyDN0ZpUhiqSyTxg0V39vZDjWehqHdjKPqJe9XPzP81MTXQcZE0lqqCCLj6KUIyNRfj8aMkWJ4VNLMFHMZkVkjBUmxrZUsSV4yyevks5F3busNx4ateZNUUcZTuAUzsGDK2jCHbSgDQQkPMMrvDnGeXHenY/FaMkpdo7hD5zPH4PzkWs=</latexit>
Revised: //oHp41NEyVYS2ieRS9UKsKWeCtg0znPYSRXEcctoNJ7e5332iSjMpHs00oUGMR4JFjGBjJb8fYzMmmGf3s0G15tbdOdAq8QpSgwKtQfWrP5QkjakwhGOtfc9NTJBhZRjhdFbpp5ommEzwiPqWChxTHWTzyDN0ZpUhiqSyTxg0V39vZDjWehqHdjKPqJe9XPzP81MTXQcZE0lqqCCLj6KUIyNRfj8aMkWJ4VNLMFHMZkVkjBUmxrZUsSV4yyevks5F3busNx4ateZNUUcZTuAUzsGDK2jCHbSgDQQkPMMrvDnGeXHenY/FaMkpdo7hD5zPH4PzkWs=</latexit>

Original: L +
Revised: L +

Original: SWA <latexit sha1_base64="KIBUPXdFc9jbwGmlyaMi3aNFjHE=">AAAB8nicbVDLSgMxFL1TX7W+qi7dBIvgqsxIUZdFNy5cVLAPmA4lk2ba0EwyJBmhDP0MNy4UcevXuPNvzLSz0NYDgcM595JzT5hwpo3rfjultfWNza3ydmVnd2//oHp41NEyVYS2ieRS9UKsKWeCtg0znPYSRXEcctoNJ7e5332iSjMpHs00oUGMR4JFjGBjJb8fYzMmmGf3s0G15tbdOdAq8QpSgwKtQfWrP5QkjakwhGOtfc9NTJBhZRjhdFbpp5ommEzwiPqWChxTHWTzyDN0ZpUhiqSyTxg0V39vZDjWehqHdjKPqJe9XPzP81MTXQcZE0lqqCCLj6KUIyNRfj8aMkWJ4VNLMFHMZkVkjBUmxrZUsSV4yyevks5F3busNx4ateZNUUcZTuAUzsGDK2jCHbSgDQQkPMMrvDnGeXHenY/FaMkpdo7hD5zPH4PzkWs=</latexit>
Revised: SWA <latexit sha1_base64="KIBUPXdFc9jbwGmlyaMi3aNFjHE=">AAAB8nicbVDLSgMxFL1TX7W+qi7dBIvgqsxIUZdFNy5cVLAPmA4lk2ba0EwyJBmhDP0MNy4UcevXuPNvzLSz0NYDgcM595JzT5hwpo3rfjultfWNza3ydmVnd2//oHp41NEyVYS2ieRS9UKsKWeCtg0znPYSRXEcctoNJ7e5332iSjMpHs00oUGMR4JFjGBjJb8fYzMmmGf3s0G15tbdOdAq8QpSgwKtQfWrP5QkjakwhGOtfc9NTJBhZRjhdFbpp5ommEzwiPqWChxTHWTzyDN0ZpUhiqSyTxg0V39vZDjWehqHdjKPqJe9XPzP81MTXQcZE0lqqCCLj6KUIyNRfj8aMkWJ4VNLMFHMZkVkjBUmxrZUsSV4yyevks5F3busNx4ateZNUUcZTuAUzsGDK2jCHbSgDQQkPMMrvDnGeXHenY/FaMkpdo7hD5zPH4PzkWs=</latexit>

Original: L
Revised: L

Original: Contributions
Revised: Contributions

Original: In-depth comparison of minima found by SWA and SAM: We visualize linear interpolations between different models and quantify the minimizers’ flatnesses.
Revised: In-depth comparison of minima found by SWA and SAM: We visualize linear interpolations between different models and quantify the minimizers’ flatnesses.

Original: This analysis yields 4 insights, e g., despite SAM finding flatter solutions than SWA as quantified by Hessian eigenvalues, they can be close to sharp directions, a phenomenon that has been overlooked in the previous SAM literature.
Revised: This analysis yields 4 insights, e.g., despite SAM finding flatter solutions than SWA as quantified by Hessian eigenvalues, they can be close to sharp directions, a phenomenon that has been overlooked in the previous SAM literature.

Original: Averaging SAM iterates leads to the flattest among all minima.
Revised: Averaging SAM iterates leads to the flattest among all minima.

Original: Rigorous comparison of SWA and SAM’s performance
Revised: Rigorous comparison of SWA and SAM’s performance

Original: over 42 tasks: We empirically compare the optimizers with a rigorous model selection procedure on a broad range of tasks across different domains (CV, NLP, and GRL), model types (MLPs, CNNs, Transformers) and tasks (classification, self-supervised learning, open-domain question answering, natural language understanding, and node/graph/link property prediction).
Revised: over 42 tasks: We empirically compare the optimizers with a rigorous model selection procedure on a broad range of tasks across different domains (CV, NLP, and GRL), model types (MLPs, CNNs, Transformers) and tasks (classification, self-supervised learning, open-domain question answering, natural language understanding, and node/graph/link property prediction).

Original: We discuss 8 findings, e.g., that both dataset and architecture impact their effectiveness, that for
Revised: We discuss 9 findings, e.g., that both dataset and architecture impact their effectiveness, that for

Original: NLP tasks, SAM improves over SWA in most cases, and that the converse holds for GRL tasks.
Revised: NLP tasks, SAM improves over SWA in most cases, and that the converse holds for GRL tasks.

Original: When flat-minima optimizers do not help, we notice clear discrepancies between the shapes of loss and accuracy curves.
Revised: When flat-minima optimizers do not help, we notice clear discrepancies between the shapes of loss and accuracy curves.

Original: To assist future work, we open-source the code for all pipelines and hyper-parameters to reproduce the results.
Revised: To assist future work, we open-source the code for all pipelines and hyper-parameters to reproduce the results.

Original:  Background and Related Work 2.1 Stochastic Gradient Descent (SGD)
Revised:  Background and Related Work 2.1 Stochastic Gradient Descent (SGD)

Original: The classic optimization framework of machine learning is empirical risk minimization
Revised: The classic optimization framework of machine learning is empirical risk minimization

Original: [Equation]
Revised: [Equation]

Original: where θ ∈ R d is a vector of parameters, { x 1 , . .
Revised: where θ ∈ R d is a vector of parameters, { x 1 , . .

Original: , x N } is a training set of inputs x n ∈ R D , and ℓ ( x ; θ ) is a loss function quantifying the performance of parameters θ on x .
Revised: , x N } is a training set of inputs x n ∈ R D , and ℓ ( x ; θ ) is a loss function quantifying the performance of parameters θ on x .

Original: SGD samples a minibatch S ⊂ { 1 , . .
Revised: SGD samples a minibatch S ⊂ { 1 , . .

Original: . , N } of size |S| ≪
Revised: . , N } of size |S| ≪

Original: N from the training set and updates the parameters through
Revised: N from the training set and updates the parameters through

Original: [Equation]
Revised: [Equation]

Original: SGD

Original: [Equation]

Original: , where
Revised: where

Original: [Equation]
Revised: [Equation]

Original: for a length specified by η , the learning rate.
Revised: for a length specified by η , the learning rate.

Original: 2.2 Stochastic Weight Averaging (SWA)
Revised: 2.2 Stochastic Weight Averaging (SWA)

Original: The idea of averaging weights dates back to accelerating the convergence speed of SGD [73, 48].
Revised: The idea of averaging weights dates back to accelerating the convergence speed of SGD [77, 50].

Original: SWA’s motivation is based on the following observation about SGD’s behavior when training neural networks: it often traverses regions of the weight space that correspond to high-performing models, but rarely reaches the central points of this optimal set.
Revised: SWA’s motivation is based on the following observation about SGD’s behavior when training neural networks: it often traverses regions of the weight space that correspond to high-performing models, but rarely reaches the central points of this optimal set.

Original: Averaging the parameter values over iterations moves the solution closer to the centroid of this space of points.
Revised: Averaging the parameter values over iterations moves the solution closer to the centroid of this space of points.

Original: The SWA update rule is the cumulative moving average
Revised: The SWA update rule is the cumulative moving average

Original: [Equation]
Revised: [Equation]

Original: Algorithm
Revised: Algorithm

Original: Input: Loss function L , training budget in number of iterations b , training dataset D := ∪ ni =1 { x i } , mini-batch size |B| , averaging start epoch E , averaging frequency ν , (scheduled) learning rate η , initial weights θ 0 .
Revised: Input: Loss function L , training budget in number of iterations b , training dataset D := ∪ ni =1 { x i } , mini-batch size |B| , averaging start epoch E , averaging frequency ν , (scheduled) learning rate η , initial weights θ 0 .

Original: for k ← 1 , . . .
Revised: for k ← 1 , . . .

Original: , b do
Revised: , b do

Original: Sample a mini-batch B from D Compute gradient g ← ∇L ( θ t ) Update parameters θ t + ← θ t − η g if k ≥ E and mod( k, ν ) = 

Revised: Sample a mini-batch B from D Compute gradient g ← ∇L ( θ t ) Update parameters θ t + ← θ t − η g if k ≥ E and mod( k, ν ) =

Original: then θ SWA t +1 = (cid:0) θ SWA t · l + θ SWA t +1 (cid:1) /
Revised: then θ SWA t +1 = (cid:0) θ SWA t · l + θ SWA t +1 (cid:1) /

Original: ( l + 1) end if end for return
Revised: ( l + 1) end if end for return

Original: θ SWA
Revised: θ SWA

Original: Input: Loss function L , training budget in number of iterations b , training dataset D := ∪ ni =1 { x i } , mini-batch size |B| , neighborhood radius ρ , (scheduled) learning rate η , initial weights θ 0 .
Revised: Input: Loss function L , training budget in number of iterations b , training dataset D := ∪ ni =1 { x i } , mini-batch size |B| , neighborhood radius ρ , (scheduled) learning rate η , initial weights θ 0 .

Original: for k ← 1 , . . .
Revised: for k ← 1 , . . .

Original: , b do
Revised: , b do

Original: Sample a mini-batch B from D Compute worst-case perturbation (cid:98) ϵ ← ρ ∇L ( θ ) ∥∇L ( θ ) ∥ 2 Compute gradient g ← ∇L (cid:0) θ SAM t + (cid:98) ϵ (cid:1)
Revised: Sample a mini-batch B from D Compute worst-case perturbation (cid:98) ϵ ← ρ ∇L ( θ ) ∥∇L ( θ ) ∥ 2 Compute gradient g ← ∇L (cid:0) θ SAM t + (cid:98) ϵ (cid:1)

Original: Update parameters θ SAM t +1 ← θ SAM t − η g end for return θ SAM where l is the number of distinct parameters averaged so far and t is the SGD iteration number.
Revised: Update parameters θ SAM t +1 ← θ SAM t − η g end for return θ SAM where l is the number of distinct parameters averaged so far and t is the SGD iteration number.

Original: SWA has two hyper-parameters: the update frequency ν and starting epoch E .
Revised: SWA has two hyper-parameters: the update frequency ν and starting epoch E .

Original: When using a constant learning rate, Izmailov et al.
Revised: When using a constant learning rate, Izmailov et al.

Original: [45] suggests updating the parameters once after each epoch, i.e. ν ≈ N |B| , and starting at E ≈ 0 .
Revised: [47] suggests updating the parameters once after each epoch, i.e. ν ≈ N |B| , and starting at E ≈ 0 .

Original: 75 T , where T is the training budget required to train the model until convergence with conventional SGD training.
Revised: 75 T , where T is the training budget required to train the model until convergence with conventional SGD training.

Original: He et al.
Revised: He et al.

Original: [37] argue that SWA may always improve generalization, regardless of the loss function’s geometry.
Revised: [38] argue that SWA may always improve generalization, regardless of the loss function’s geometry.

Original: Kaddour [48] show that averaging a specific range of weights can speed up training convergence.
Revised: Kaddour [50] show that averaging a specific range of weights can speed up training convergence.

Original: Cha et al.
Revised: Cha et al.

Original: [8] argue that tuning ν and E carefully is necessary to make it work effectively in domain generalization (DG) tasks.
Revised: [8] argue that tuning ν and E carefully is necessary to make it work effectively in domain generalization (DG) tasks.

Original: Besides DG tasks, a list of tuned hyper-parameters based on a fair model selection procedure across different architectures and tasks has been missing in the literature.
Revised: Besides DG tasks, a list of tuned hyper-parameters based on a fair model selection procedure across different architectures and tasks has been missing in the literature.

Original: To the best of our knowledge, Cha et al.
Revised: To the best of our knowledge, Cha et al.

Original: [8] is the only study that compares SWA and
Revised: [8] is the only study that compares SWA and

Original: SAM over the same experiments, but it focuses on domain generalization tasks which we, therefore, leave out in this work.
Revised: SAM over the same experiments, but it focuses on domain generalization tasks which we, therefore, leave out in this work.

Original: 2.3 Sharpness-Aware Minimization (SAM)
Revised: 2.3 Sharpness-Aware Minimization (SAM)

Original: While SWA is implicitly biased towards flat minima, SAM explicitly approximates the flatness around parameters θ to guide the parameter update.
Revised: While SWA is implicitly biased towards flat minima, SAM explicitly approximates the flatness around parameters θ to guide the parameter update.

Original: It first computes the worst-case perturbation ϵ that maximizes the loss within a given neighborhood ρ , and then minimizes the loss w.r.t.
Revised: It first computes the worst-case perturbation ϵ that maximizes the loss within a given neighborhood ρ , and then minimizes the loss w.r.t.

Original: the perturbed weights θ + ϵ .
Revised: the perturbed weights θ + ϵ .

Original: Formally, SAM finds θ by solving the minimax problem:
Revised: Formally, SAM finds θ by solving the minimax problem:

Original: [Equation]
Revised: [Equation]

Original: where ρ ≥ 0 is a hyperparameter.
Revised: where ρ ≥ 0 is a hyperparameter.

Original: To find the worst-case perturbation ϵ ∗ efficiently in practice, Foret et al.
Revised: To find the worst-case perturbation ϵ ∗ efficiently in practice, Foret et al.

Original: [22] approximates Eq.
Revised: [21] approximates Eq.

Original: (4) via a first-order Taylor expansion of L ( θ + ϵ ) w.r.t.
Revised: (4) via a first-order Taylor expansion of L ( θ + ϵ ) w.r.t.

Original: ϵ around 0 , obtaining
Revised: ϵ around 0 , obtaining

Original: [Equation]
Revised: [Equation]

Original: In words, (cid:98) ϵ is simply the scaled gradient of the loss function w.r.t to the current parameters θ .
Revised: In words, (cid:98) ϵ is simply the scaled gradient of the loss function w.r.t to the current parameters θ .

Original: Given (cid:98) ϵ , the altered gradient used to update the current θ t (in place of g ( θ t ) ) is
Revised: Given (cid:98) ϵ , the altered gradient used to update the current θ t (in place of g ( θ t ) ) is

Original: [Equation]
Revised: [Equation]

Original: Due to Eq.
Revised: Due to Eq.

Original: (5), SAM’s computational overhead consists of an additional forward and backward pass per parameter update step compared to
Revised: (5), SAM’s computational overhead consists of an additional forward and backward pass per parameter update step compared to

Original: SAM’s performance strongly depends on the neighborhood radius ρ .
Revised: SAM’s performance strongly depends on the neighborhood radius ρ .

Original: For example, Chen et al.
Revised: For example, Chen et al.

Original: [12], Wu et al.
Revised: [12], Wu et al.

Original: [88] show that ρ should be set to values outside the originally considered ranges by
Revised: [92] show that ρ should be set to values outside the originally considered ranges by

Original: Foret et al.
Revised: Foret et al.

Original: Analogously to Sec.
Revised: Analogously to Sec.

Original: 2.2, this lack of coherence among hyper-parameter tuning protocols in the SAM literature makes it tricky to determine SAM’s comparative effectiveness.
Revised: 2.2, this lack of coherence among hyper-parameter tuning protocols in the SAM literature makes it tricky to determine SAM’s comparative effectiveness.

Original: 2.4 Other Flat-Minima Optimizers
Revised: 2.4 Other Flat-Minima Optimizers

Original: There are several extensions of SWA [34, 8] and SAM [59, 98, 96].
Revised: There are several extensions of SWA [35, 8] and SAM [64, 102, 100].

Original: For simplicity, we do not consider them in this work.
Revised: For simplicity, we do not consider them in this work.

Original: Besides SWA and SAM, other flat-minima optimizers include e.g. [9, 78].
Revised: Besides SWA and SAM, other flat-minima optimizers include e.g. [9, 83].

Original: However, due to their computational cost and/or lack of performance gains, we do not include them in this work.
Revised: However, due to their computational cost and/or lack of performance gains, we do not include them in this work.

Original: Chaudhari et al.
Revised: Chaudhari et al.

Original: [9] requires [5 , 20] forward and backward passes per parameter update.
Revised: [9] requires [5 , 20] forward and backward passes per parameter update.

Original: Sankar et al.
Revised: Sankar et al.

Original: [78] similarly requires [5 , 10] forward and backward passes to estimate the Hessian trace and 6 of 7 experiments yield minimal improvement of ≤ 0 .
Revised: [83] similarly requires [5 , 10] forward and backward passes to estimate the Hessian trace and 6 of 7 experiments yield minimal improvement of ≤ 0 .

Original: 27% , see Table 1 in Sankar et al.
Revised: 27% , see Table 1 in Sankar et al.

Original: In contrast, SWA and SAM have been shown to increase performance by multiple percentage points in some cases [8, 12], while requiringmany fewer computational resources.
Revised: In contrast, SWA and SAM have been shown to increase performance by multiple percentage points in some cases [8, 12], while requiring fewer computational resources.

Original:  How do minima found by SWA and SAM differ?
Revised:  How do minima found by SWA and SAM differ?

Original: In this section, we investigate SWA and SAM solutions for two prototypical deep learning tasks, where these optimizers improve over the baseline.
Revised: In this section, we investigate SWA and SAM solutions in two prototypical deep learning tasks, where these optimizers improve over the baseline.

Original: Our goal is to better understand their geometric properties (instead of their generalization performance, which is the focus of Sec. 4).
Revised: Our goal is to better understand their geometric properties (instead of their generalization performance, which is the focus of Sec. 4).

Original: First, we investigate the behavior of the loss landscape along the line between non-flat and flat solutions (Sec. 3.1).
Revised: First, we investigate the behavior of the loss landscape along the line between non-flat and flat solutions (Sec. 3.1).

Original: Previous studies successfully used such linear interpolations to gain novel insights, e.g., for training dynamics [30, 24], regularization [63, 27], and network pruning [25].
Revised: Previous studies successfully used such linear interpolations to gain novel insights, e.g., for training dynamics [31, 24], regularization [68, 27], and network pruning [25].

Original: Second, motivated by findings in Sec. 3.1, we average SAM iterates and visualize interpolations between averaged and non-averaged solutions (Sec. 3.2).
Revised: Second, motivated by findings in Sec. 3.1, we average SAM iterates and visualize interpolations between averaged and non-averaged solutions (Sec. 3.2).

Original: Interestingly, the averaged SAM solution is less susceptible to asymmetric directions.
Revised: Interestingly, the averaged SAM solution is less susceptible to asymmetric directions.

Original: Third, we compare quantitative measurements of all solutions’ flatnesses (Sec. 3.3).
Revised: Third, we compare quantitative measurements of all solutions’ flatnesses (Sec. 3.3).

Original: Here, we compute dominant Hessian eigenvalues, as commonly used in the flat minima literature [9, 93, 12, 22].
Revised: Here, we compute dominant Hessian eigenvalues, as commonly used in the flat minima literature [9, 97, 12, 21].

Revised: Lastly, in Appendix A.1, we further compute CKA [60] and cosine similarities between SWA/SAM’s network output logits.

Original: We choose the following two disparate learning settings: (i) a well-known image classification task, widely used for evaluation in flat-minima optimizer papers, and (ii) a novel, challenging Python code summarization task, on which state-of-the-art models achieve only around 16% F1 score on the test set (which is higher than its commonly achieved accuracy on the more challenging training set), and that has not been explored yet in the flat-minima literature.
Revised: We choose the following two disparate learning settings: (i) a well-known image classification task, widely used for evaluation in flat-minima optimizer papers, and (ii) a novel, challenging Python code summarization task, on which state-of-the-art models achieve only around 16% F1 score on the test set (which is higher than its commonly achieved accuracy on the more challenging training set), and that has not been explored yet in the flat-minima literature.

Original: Specifically, for (i), we investigate the loss/accuracy surfaces of a WideResNet28-10 [94] model on CIFAR-100 [57] (baseline non-flat optimizer: SGD with momentum (SGD-M)) [77].
Revised: Specifically, for (i), we investigate the loss/accuracy surfaces of a WideResNet28-10 [98] model on CIFAR-100 [62] (baseline non-flat optimizer: SGD with momentum (SGD-M)) [82].

Original: For (ii), we use the theoretically-grounded Graph
Revised: For (ii), we use the theoretically-grounded Graph

Original: Isomorphism Network [90] model on OGB-Code2 [42] (baseline optimizer: Adam [51]).
Revised: Isomorphism Network [94] model on OGB-Code2 [44] (baseline optimizer: Adam [55]).

Original: All optimizers start from the same initialization.
Revised: All optimizers start from the same initialization.

Original: We denote the minimizer produced by the non-flat methods (SGD-M and Adam) by θ NF and the flat ones by θ SWA and θ SAM .
Revised: We denote the minimizer produced by the non-flat methods (SGD-M and Adam) by θ NF and the flat ones by θ SWA and θ SAM .

Original: 3.1 What is between non-flat and flat solutions?
Revised: 3.1 What is between non-flat and flat solutions?

Original: We start by comparing the similarity of flat and non-flat minimizers through linear interpolations.
Revised: We start by comparing the similarity of flat and non-flat minimizers through linear interpolations.

Original: This analysis allows us to understand if they are in the same basin, and how close they are to a region of sharply-increasing loss, where we expect loss/accuracy to differ widely between train and test.
Revised: This analysis allows us to understand if they are in the same basin, and how close they are to a region of sharply-increasing loss, where we expect loss/accuracy to differ widely between train and test.

Original: Further, for each of our 4 observations, we recommend a future work direction.
Revised: Further, for each of our 4 observations, we recommend a future work direction.

Original: To linearly interpolate between two sets of parameters θ and θ ′ , we parameterize the line connecting these two by choosing a scalar parameter α , and defining the weighted average θ ( α ) =
Revised: To linearly interpolate between two sets of parameters θ and θ ′ , we parameterize the line connecting these two by choosing a scalar parameter α , and defining the weighted average θ ( α ) =

Original: (1 − α )
Revised: (1 − α )

Original: θ + α θ ′ .
Revised: θ + α θ ′ .

Original: If there exists no high-loss barrier between two networks θ , θ ′ along the linear interpolation, we say that they are located in the same basin , i.e., { θ , θ ′ } ∈ Ω .
Revised: If there exists no high-loss barrier between two networks θ , θ ′ along the linear interpolation, we say that they are located in the same basin , i.e., { θ , θ ′ } ∈ Ω .

Original: A basin is an area in the parameter space where the loss function has relatively low values.
Revised: A basin is an area in the parameter space where the loss function has relatively low values.

Original: Due to NN non-linearities, the linear combination of the weights of two accurate models does not necessarily define an accurate model.
Revised: Due to NN non-linearities, the linear combination of the weights of two accurate models does not necessarily define an accurate model.

Original: Hence, we generally expect high-loss barriers along the linear interpolation path.
Revised: Hence, we generally expect high-loss barriers along the linear interpolation path.

Original: While there are alternative distance measures that could be used to compare two networks, they typically either (a) do not offer clear interpretations, as pointed out by Frankle et al.
Revised: While there are alternative distance measures that could be used to compare two networks, they typically either (a) do not offer clear interpretations, as pointed out by Frankle et al.

Original: [25], or (b) yield trivial network connectivity results, such as non-linear low-loss paths, which can be found for any two network minimizers [20, 26, 31, 23].
Revised: [25], or (b) yield trivial network connectivity results, such as non-linear low-loss paths, which can be found for any two network minimizers [19, 26, 32, 22].

Original: [Figure]
Revised: [Figure]

Original: θ ( α ) =
Revised: θ ( α ) =

Original: (1 − α )
Revised: (1 − α )

Original: θ
Revised: θ

Original: + α θ ′ (for α ∈ [ − 1 , 1 . 5] ) between SWA ( + ) and SAM ( × ) solutions ( α = 0 . 0 ) and non-flat baseline solutions ( • , α = 1 . 0 ).
Revised: + α θ ′ (for α ∈ [ − 1 , 1 . 5] ) between SWA ( + ) and SAM ( × ) solutions ( α = 0 . 0 ) and non-flat baseline solutions ( • , α = 1 . 0 ).

Original: Obs.
Revised: Obs.

Original: 1: { θ SWA , θ NF } ∈ Ω NF .
Revised: 1: { θ SWA , θ NF } ∈ Ω NF .

Original: θ SWA and θ NF are in the same basin, as can be seen in Figures 2a and 2e.
Revised: θ SWA and θ NF are in the same basin, as can be seen in Figures 2a and 2e.

Original: Additionally, θ NF is near the periphery of a sharp increase in loss, as can be seen when moving in the direction from θ SWA to θ NF (i.e., α > 1 ).
Revised: Additionally, θ NF is near the periphery of a sharp increase in loss, as can be seen when moving in the direction from θ SWA to θ NF (i.e., α > 1 ).

Original: Conversely, θ SWA finds flat regions that change slowly in the loss.
Revised: Conversely, θ SWA finds flat regions that change slowly in the loss.

Original: This bias of SWA to flatter loss beneficially transfers to the accuracy landscape too: Figures 2b and 2f show the accuracy/F1 score rapidly dropping off approaching and beyond θ NF .
Revised: This bias of SWA to flatter loss beneficially transfers to the accuracy landscape too: Figures 2b and 2f show the accuracy/F1 score rapidly dropping off approaching and beyond θ NF .

Original: Interestingly, in Figures 2e and 2f, we see that for Code2, for α < 0 , there exist solutions with even better training loss/accuracy but worse test loss/accuracy.
Revised: Interestingly, in Figures 2e and 2f, we see that for Code2, for α < 0 , there exist solutions with even better training loss/accuracy but worse test loss/accuracy.

Original: However, θ SWAGIN is close to the test accuracy maximizer along this interpolation.
Revised: However, θ SWAGIN is close to the test accuracy maximizer along this interpolation.

Original: Future work may inspect why the cross entropy loss function used for GIN/Code2 seems less well correlated with its accuracy compared to WRN/CIFAR100.
Revised: Future work may inspect why the cross entropy loss function used for GIN/Code2 seems less well correlated with its accuracy compared to WRN/CIFAR100.

Original: Obs. 2: θ SAM ∈ Ω SAM ̸ = Ω NF .
Revised: Obs. 2: θ SAM ∈ Ω SAM ̸ = Ω NF .

Original: θ SAM and θ NF are not in the same basin: Figures 2c and 2g show that there is a high loss barrier between them, respectively.
Revised: θ SAM and θ NF are not in the same basin: Figures 2c and 2g show that there is a high loss barrier between them, respectively.

Original: Figures 2d and 2h show that θ SAM and even nearby points in parameter space achieve higher accuracies/F1 scores (i.e. generalize better) than θ NF and points around it.
Revised: Figures 2d and 2h show that θ SAM and even nearby points in parameter space achieve higher accuracies/F1 scores (i.e. generalize better) than θ NF and points around it.

Original: This is an interesting result because we expect different basins to produce qualitatively different predictions, one of the motivations behind combining models, even if they exhibit different performances [43, 61].
Revised: This is an interesting result because we expect different basins to produce qualitatively different predictions, one of the motivations behind combining models, even if they exhibit different performances [45, 66].

Original: Grewal & Bui [32] successfully combine models yielded by different optimizers, and we think future work should study ensembling SAM and non-SAM solutions.
Revised: Grewal & Bui [33] successfully combine models yielded by different optimizers, and we think future work should study ensembling SAM and non-SAM solutions.

Original: Obs. 3: SAM finds a saddle point.
Revised: Obs. 3: SAM finds a saddle point.

Original: Figure 2g shows θ SAMGIN being located in a sharp training loss minimum whose loss is much higher than θ NF .
Revised: Figure 2g shows θ SAMGIN being located in a sharp training loss minimum whose loss is much higher than θ NF .

Original: Yet, its test loss is only slightly higher, and its F score is better.
Revised: Yet, its test loss is only slightly higher, and its F1 score is better.

Original: We visualize 2D plots moving along random directions (not shown here due to space) to confirm that θ SAMGIN is a saddle point.
Revised: We visualize 2D plots moving along random directions (not shown here due to space) to confirm that θ SAMGIN is a saddle point (Appendix A.2).

Original: A common pathology among curvature-based methods is that they attract saddle points [15].
Revised: A common pathology among curvature-based methods is that they attract saddle points [15].

Original: Since SAM takes some form of curvature into account, too, we believe that future work should investigate SAM’s propensity to find saddle points and potential remedies.
Revised: Since SAM takes some form of curvature into account, too, we believe that future work should investigate SAM’s propensity to find saddle points and potential remedies.

Original: Obs. 4: θ SAM is closer to sharper directions than θ SWA , as can be seen by L tr/te ( θ SAM (0 . 1))
Revised: Obs. 4: θ SAM is closer to sharper directions than θ SWA , as can be seen by L tr/te ( θ SAM (0 . 1))

Original: ≈ 2 · L tr/te ( θ SAM ( − 0 . 1)) , while L tr/te ( θ SWA (0 . 1))
Revised: ≈ 2 · L tr/te ( θ SAM ( − 0 . 1)) , while L tr/te ( θ SWA (0 . 1))

Original: ≈ L tr/te ( θ SWA ( − 0 . 1)) , where L ( · ) tr/te refers to both training and test loss functions.
Revised: ≈ L tr/te ( θ SWA ( − 0 . 1)) , where L ( · ) tr/te refers to both training and test loss functions.

Original: A possible explanation for SAM being closer to sharp sides is that while it finds different basins than SGD/SWA by smoothing the loss surface (as illustrated in Fig. 1), within a local basin, it may oscillate around the minimizer similarly as SGD.
Revised: A possible explanation for SAM being closer to sharp sides is that while it finds different basins than SGD/SWA by smoothing the loss surface (as illustrated in Fig. 1), within a local basin, it may oscillate around the minimizer similarly as SGD.

Original: One cause for this can be that Ω SAM ’s hypersphere is larger than SAM’s radius ρ .
Revised: One cause for this can be that Ω SAM ’s hypersphere is larger than SAM’s radius ρ .

Original: If that holds, then given a small enough learning rate, we expect it to oscillate around θ ∗ ∈ Ω SAM (the smaller the learning rate, the less likely it escapes the basin due to that stochasticity).
Revised: If that holds, then given a small enough learning rate, we expect it to oscillate around θ ∗ ∈ Ω SAM (the smaller the learning rate, the less likely it escapes the basin due to that stochasticity).

Original: Two possible remedies for that are: (1) adapt/schedule ρ , or (2) average SAM iterates to bias its solution towards the flatter side.
Revised: Two possible remedies for that are: (1) adapt/schedule ρ , or (2) average SAM iterates to bias its solution towards the flatter side.

Original: (1) has been explored by [98, 96].
Revised: (1) has been explored by [102, 100].

Original: We try (2) in the next subsection.
Revised: We try (2) in the next subsection.

Original: Future work may study SAM’s basin escape time, e.g., using convolutions [53] or stochastic differential equations [97].
Revised: Future work may study SAM’s basin escape time, e.g., using convolutions [57] or stochastic differential equations [101].

Original: 3.2 What happens if we average SAM iterates?
Revised: 3.2 What happens if we average SAM iterates?

Original: Based on observation 4: “ θ SAM is closer to sharper directions than θ SWA ”, averaging SAM iterates may further improve generalization.
Revised: Based on observation 4: “ θ SAM is closer to sharper directions than θ SWA ”, averaging SAM iterates may further improve generalization.

Original: The reason is that while SAM finds better-performing basins, within the basin, its final iterate may still be near a side that increases sharply in the loss.
Revised: The reason is that while SAM finds better-performing basins, within the basin, its final iterate may still be near a side that increases sharply in the loss.

Original: GIN on Code2 WRN on CIFAR
Revised: GIN on Code2 WRN on CIFAR

Original: [Figure]
Revised: [Figure]

Original: [Figure]
Revised: [Figure]

Original: Starting with the first of the two previously analyzed settings (WRN/CIFAR100), Figures 3a, and 3b show that θ SWA+SAMWRN (marker: ⋆ ) achieves the lowest test loss and highest test accuracy, respectively.
Revised: Starting with the first of the two previously analyzed settings (WRN/CIFAR100), Figures 3a, and 3b show that θ SWA+SAMWRN (marker: ⋆ ) achieves the lowest test loss and highest test accuracy, respectively.

Original: What stands out in comparison to the previous plots is θ SAM WRN ’s ( × ) proximity to sharp sides, surprisingly similar to θ NFWRN ( • ) here and in Figures 2c and 2e.
Revised: What stands out in comparison to the previous plots is θ SAM WRN ’s ( × ) proximity to sharp sides, surprisingly similar to θ NFWRN ( • ) here and in Figures 2c and 2e.

Original: As we hoped, θ SWA+SAMWRN is indeed closer to a flatter region, as can be seen by L tr/te ( θ SWA+SAM
Revised: As we hoped, θ SWA+SAMWRN is indeed closer to a flatter region, as can be seen by L tr/te ( θ SWA+SAM

Original: [Equation]
Revised: [Equation]

Original: In GIN/OGB-Code2, one unanticipated finding is that θ SWA+SAMGIN escapes the (previously discussed) saddle point of θ SAMGIN , appearing here as a maximum in Figure 3c.
Revised: In GIN/OGB-Code2, one unanticipated finding is that θ SWA+SAMGIN escapes the (previously discussed) saddle point of θ SAMGIN , appearing here as a maximum in Figure 3c.

Original: A likely reason for that is that
Revised: A likely reason for that is that

Original: SAM traversed nearby flatter regions before arriving at the saddle point, especially if it is a non-strict saddle.
Revised: SAM traversed nearby flatter regions before arriving at the saddle point, especially if it is a non-strict saddle.

Original: In terms of F1 score, Figure 3d shows that while θ SWAGIN ( + ) and θ SAMGIN perform about equally well, the flatter region found by θ SWA+SAMGIN improves over both.
Revised: In terms of F1 score, Figure 3d shows that while θ SWAGIN ( + ) and θ SAMGIN perform about equally well, the flatter region found by θ SWA+SAMGIN improves over both.

Original: 3.3 How “flat” are the found minima?
Revised: 3.3 How “flat” are the found minima?

Original: We now quantify the flatnesses of all four optimizers over both tasks by computing the median of the dominant Hessian eigenvalue across all training set batches using the Power Iteration algorithm [68, 93].
Revised: We now quantify the flatnesses of all four optimizers over both tasks by computing the median of the dominant Hessian eigenvalue across all training set batches using the Power Iteration algorithm [73, 97].

Original: This metric measures the worst-case loss landscape curvature.
Revised: This metric measures the worst-case loss landscape curvature.

Original: We choose this metric as it is very commonly used in the minima flatness literature, e.g., [9, 12, 22, 92, 18, 56, 79].
Revised: We choose this metric as it is very commonly used in the minima flatness literature, e.g., [9, 12, 21, 96, 17, 61, 84].

Original: Table 1 shows that SAM leads to flatter minima than SWA in both cases.
Revised: Table 1 shows that SAM leads to flatter minima than SWA in both cases.

Original: Interestingly λ max ( θ NFWRN ) ≈ 2 . 5 · λ max ( { θ SWA , θ SAM } ) ,
Revised: Interestingly λ max ( θ NFWRN ) ≈ 2 . 5 · λ max ( { θ SWA , θ SAM } ) ,

Original: while λ max ( θ NFWRN ) ≈ 5 .
Revised: while λ max ( θ NFWRN ) ≈ 5 .

Original: 75 λ max ( θ SWA+SAMWRN ) , indicating room for improvement in terms of flatness for both SWA and SAM.
Revised: 75 λ max ( θ SWA+SAMWRN ) , indicating room for improvement in terms of flatness for both SWA and SAM.

Original: The relative differences are less dramatic for GIN/Code2, although surprisingly λ max ( θ NFGIN )
Revised: The relative differences are less dramatic for GIN/Code2, although surprisingly λ max ( θ NFGIN )

Original: ≈ λ max ( θ SWAGIN ) .
Revised: ≈ λ max ( θ SWAGIN ) .

Original: In sum, averaging SAM iterates leads to the flattest minima and best-performing minima in both cases (see Sec. 4).
Revised: In sum, averaging SAM iterates leads to the flattest minima and best-performing minima in both cases (see Sec. 4).

Original: [Table]
Revised: [Table]

Original:  How do SWA and SAM perform on a broad set of experiments?
Revised:  How do SWA and SAM perform on a broad set of experiments?

Original: As we point out in the introduction, there is almost no overlap and consistency regarding reported SWA and SAM results in the literature.
Revised: As we point out in the introduction, there is almost no overlap and consistency regarding reported SWA and SAM results in the literature.

Original: This section addresses this gap.
Revised: This section addresses this gap.

Original: For example, Bahri et al.
Revised: For example, Bahri et al.

Original: [4], Chen et al.
Revised: [4], Chen et al.

Original: [12] illustrate that the flat minima found by SAM improve generalization on Transformer [85] architectures compared to non-flat optimizers, but they do not compare against SWA.
Revised: [12] illustrate that the flat minima found by SAM improve generalization on Transformer [89] architectures compared to non-flat optimizers, but they do not compare against SWA.

Original: Hence, it is unclear if the computationally cheaper SWA may provide better or similar performance.
Revised: Hence, it is unclear if the computationally cheaper SWA may provide better or similar performance.

Original: We compare flat minimizers SWA, SAM, and averaged SWA iterates (SWA+SAM) over the non-flat minimizers across a range of different tasks in the domains of computer vision, natural language processing, and graph representation learning.
Revised: We compare flat minimizers SWA, SAM, and averaged SWA iterates (SWA+SAM) over the non-flat minimizers across a range of different tasks in the domains of computer vision, natural language processing, and graph representation learning.

Original: We average all runs at least three times across random seeds (more often for experiments with higher variability, see details in Appendix B), and we report
Revised: We average all runs at least three times across random seeds (more often for experiments with higher variability, see details in Appendix B), and we report

Original: [Table]
Revised: [Table]

Original: the corresponding standard error.
Revised: the corresponding standard error.

Original: We bold the best-performing approach and any approach whose average performance plus standard error overlaps it.
Revised: We bold the best-performing approach and any approach whose average performance plus standard error overlaps it.

Original: Hyper-parameters.
Revised: Hyper-parameters.

Original: For all architectures and datasets, we set hyperparameters shared by all methods (e.g., learning rate) mostly to values cited in prior work 3 As explained in Secs.
Revised: For all architectures and datasets, we set hyperparameters shared by all methods (e.g., learning rate) mostly to values cited in prior work 3 As explained in Secs.

Original: and 2.3, the effectiveness of flat-minima optimizers is highly sensitive to their additional hyperparameters.
Revised: and 2.3, the effectiveness of flat-minima optimizers is highly sensitive to their additional hyperparameters.

Original: We select hyper-parameters using a grid search over a held-out validation set.
Revised: We select hyper-parameters using a grid search over a held-out validation set.

Original: Specifically, for SWA we follow Izmailov et al.
Revised: Specifically, for SWA we follow Izmailov et al.

Original: [45] and hold the update frequency ν constant to once per epoch and tune the start time E ∈ { 0 .
Revised: [47] and hold the update frequency ν constant to once per epoch and tune the start time E ∈ { 0 .

Original: 5 T, 0 .
Revised: 5 T, 0 .

Original: 6 T, 0 .
Revised: 6 T, 0 .

Original: 75 T, 0 .
Revised: 75 T, 0 .

Original: 9 T } ( T is the number of baseline training epochs).
Revised: 9 T } ( T is the number of baseline training epochs).

Original: Izmailov et al.
Revised: Izmailov et al.

Original: [45] argue that a cyclical learning rate starting from E helps to encourage exploration of the basin.
Revised: [47] argue that a cyclical learning rate starting from E helps to encourage exploration of the basin.

Original: For the sake of simplicity, we average the iterates of the baseline directly but include even earlier starting times (i.e., 0 .
Revised: For the sake of simplicity, we average the iterates of the baseline directly but include even earlier starting times (i.e., 0 .

Original: 5 T, 0 .
Revised: 5 T, 0 .

Original: 6 T ).
Revised: 6 T ).

Original: For SAM, we tune its neighborhood size ρ ∈ { 0 . 01 , 0 . 02 , 0 .
Revised: For SAM, we tune its neighborhood size ρ ∈ { 0 . 01 , 0 . 02 , 0 .

Original: 05 , 0 . 1 , 0 . 2 } , as in previous work [22, 4].
Revised: 05 , 0 . 1 , 0 . 2 } , as in previous work [21, 4].

Original: Appendix B contains the values of all hyper-parameters and additional training details (including public model checkpoints, hardware infrastructure, software libraries, etc.) to ensure full reproducibility alongside open-sourcing our code.
Revised: Appendix B contains the values of all hyper-parameters and additional training details (including public model checkpoints, hardware infrastructure, software libraries, etc.) to ensure full reproducibility alongside open-sourcing our code.

Original: 4.1 Computer Vision
Revised: 4.1 Computer Vision

Original: Supervised Classification (SC).
Revised: Supervised Classification (SC).

Original: We evaluate the CNN architectures WideResNets [94] withlayers and width 10, and PyramidNet (PN) with 110 layers and widening factor 272 [36] as well as
Revised: We evaluate the CNN architectures WideResNets [98] withlayers and width 10, and PyramidNet (PN) with 110 layers and widening factor 272 [37] as well as

Original: Vision Transformer (ViT) [19] and MLP-Mixer [81] on CIFAR{10, 100} [57].
Revised: Vision Transformer (ViT) [18] and MLP-Mixer [86] on CIFAR{10, 100} [62].

Original: All experiments use basic data augmentations: horizontal flip, padding by four pixels, random crop, and cutout [16].
Revised: All experiments use basic data augmentations: horizontal flip, padding by four pixels, random crop, and cutout [16].

Original: Self-Supervised Learning (SSL).
Revised: Self-Supervised Learning (SSL).

Original: We consider the following methods on CIFAR10 and ImageNette 4 : Momentum Contrast [39], a Simple framework for Contrastive Learning (SimCLR) [10],
Revised: We consider the following methods on CIFAR10 and ImageNette 4 : Momentum Contrast [40], a Simple framework for Contrastive Learning (SimCLR) [10],

Original: Simple Siamese representation learning (SimSiam)
Revised: Simple Siamese representation learning (SimSiam)

Original: [11], Barlow Twins [95], Bootstrap your own
Revised: [11], Barlow Twins [99], Bootstrap your own

Original: Latent (BYOL) [33], and Swapping Assignments between multiple Views of the same image (SwAV) [7].
Revised: Latent (BYOL) [34], and Swapping Assignments between multiple Views of the same image (SwAV) [7].

Original: All SSL methods use a ResNet-18 [38] as backbone network.
Revised: All SSL methods use a ResNet-18 [39] as backbone network.

Original: To test the frozen representations, we use k -nearest-neighbor classification with a memory bank [89].
Revised: To test the frozen representations, we use k -nearest-neighbor classification with a memory bank [93].

Original: We choose k =200 and temperature τ = 0 .
Revised: We choose k =200 and temperature τ = 0 .

Original: 1 to reweight similarities.
Revised: 1 to reweight similarities.

Original: Compared to learning a linear model on top of the representations, this evaluation procedure is more robust to hyperparameter changes [54].
Revised: Compared to learning a linear model on top of the representations, this evaluation procedure is more robust to hyperparameter changes [58].

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: 4.2 Natural Language Processing
Revised: 4.2 Natural Language Processing

Original: We consider the task of open domain question answering (ODQA) using a T5-based model FusionIn-Decoder (FiD) [44].
Revised: We consider the task of open domain question answering (ODQA) using a T5-based model FusionIn-Decoder (FiD) [46].

Original: We evaluate FiD-base on the test sets of Natural Questions (NQ) [58] and TriviaQA [47].
Revised: We evaluate FiD-base on the test sets of Natural Questions (NQ) [63] and TriviaQA [49].

Original: We also consider natural language understanding tasks included in the GLUE benchmark [86], which cover acceptability, sentiment, paraphrase, similarity
Revised: We also consider natural language understanding tasks included in the GLUE benchmark [90], which cover acceptability, sentiment, paraphrase, similarity

Original: [Table]
Revised: [Table]

Original: and inference.
Revised: and inference.

Original: We fine-tune RoBERTa-base [66] for each task individually and report the results on the GLUE dev set.
Revised: We fine

Original: 4.3 Graph Representation Learning
Revised: 4.3 Graph Representation Learning

Original: We use a subset of the Open Graph Benchmark (OGB) datasets [42].
Revised: We use a subset of the Open Graph Benchmark (OGB) datasets [44].

Original: The tasks are node property prediction (NPP), graph property prediction (GPP), and link property prediction (LPP).
Revised: The tasks are node property prediction (NPP), graph property prediction (GPP), and link property prediction (LPP).

Original: For each task, we use two of the following GNN architectures and matrix factorization methods: GCN [52],
Revised: For each task, we use two of the following GNN architectures and matrix factorization methods: GCN [56],

Original: DeeperGCN (DGCN) [62], SAGE [35], GIN [90], ComplEx [83], and CP [60].
Revised: DeeperGCN (DGCN) [67], SAGE [36], GIN [94], ComplEx [88], and CP [65].

Original: We use popular training schemes, such as virtual nodes, cluster sampling [14], or relation prediction as auxiliary training objective [13].
Revised: We use popular training schemes, such as virtual nodes, cluster sampling [14], or relation prediction as auxiliary training objective [13].

Original: The reported metrics are ROC-AUC for Proteins, Accuracy for Products, F score for Code2, Average precision for Molpcba, and Mean Reciprocal Rank for Biokg/Citation2.
Revised: The reported metrics are ROC-AUC for Proteins, Accuracy for Products, F score for Code2, Average precision for Molpcba, and Mean Reciprocal Rank for Biokg/Citation2.

Original: 4.4 8 Findings
Revised: 4.4 9 Findings

Original: We use G ( · ) to describe the generalization accuracy/F1/ROCAUC/AP/MRR of all optimizers { Non-flat baseline ( NF ) , SWA , SAM , SWA+SAM } .
Revised: We use G ( · ) to describe the generalization accuracy/F1/ROCAUC/AP/MRR of all optimizers { Non-flat baseline ( NF ) , SWA , SAM , SWA+SAM } .

Original: Dataset and architecture interact. To explain the relationship between flatness and generalization, some works focused on the network’s architecture [17, 12, 71], while others focused on the data distribution [84, 72, 21]. Our results show that both matter.Datasets: e
Revised: Datasets matter.

Original: g., for node property prediction (Proteins), we see that no flat optimizer improves over the baseline optimizer, however for (Products), flat-minima optimizers on the same architectures significantly improve over the baseline.
Revised: For example, for node property prediction (Proteins), we see that no flat optimizer improves over the baseline optimizer; however, for (Products), flat-minima optimizers on the same architectures significantly improve over the baseline.

Original: Architectures: e.g., there is a vast difference across model architectures for link property prediction on the Citation2 dataset: using a GNN with GCN layers, SAM achieves a statistically significant boost of > 1 . 30% and SWA slightly hurts the performance.
Revised: Architectures matter, e.g., there is a vast difference across model architectures for link property prediction on the Citation2 dataset: using a GNN with GCN layers, SAM achieves a statistically significant boost of > 1 . 30% and SWA slightly hurts the performance.

Original: When we replace the GCN layers with SAGE layers (and fix everything else), we see a boost of > 1 .
Revised: When we replace the GCN layers with SAGE layers (and fix everything else), we see a boost of > 1 .

Original: 15% for SWA, while SAM hurts the performance by − 0 .
Revised: 15% for SWA, while SAM hurts the performance by − 0 .

Original: SWA underperforms on NLP tasks.
Revised: SWA underperforms on NLP tasks.

Original: SAM achieves the best performance in 7/10 experiments on NLP tasks, consistent with the findings of Bahri et al.
Revised: SAM achieves the best performance in 7/10 experiments on NLP tasks, consistent with the findings of Bahri et al.

Original: [4], which show that SAM can boost performance across a wide range of NLP tasks.
Revised: [4], which show that SAM can boost performance across a wide range of NLP tasks.

Original: However, SWA never performs best, only improving the results in 1/10 cases, and even hurting the performance on 4 tasks.
Revised: However, SWA never performs best, only improving the results in 1/10 cases, and even hurting the performance on 4 tasks.

Original: Surprisingly, G ( SWA+SAM ) > G ( SAM ) for {SST, QNLI} while SWA decreases performance in these cases. 3.
Revised: Surprisingly, G ( SWA+SAM ) > G ( SAM ) for {SST, QNLI} while SWA decreases performance in these cases. 4.

Original: SWA beats SAM on GRL tasks , G ( SWA ) >
Revised: SWA beats SAM on GRL tasks.

Original: G (  NF ) in 10/12 experiments, while G ( SAM ) > G ( NF ) only in 4.
Revised: G ( SWA ) > G ( NF ) in 10/12 experiments, while G ( SAM ) > G ( NF ) only in 4.

Original: We examine why SAM under-performs in the Appendix. 4.
Revised: We examine why SAM under-performs in Appendix A.4.

Original: SWA does not work well with Transformers.
Revised: SWA does not work well with Transformers.

Original: SWA often does not improve and sometimes hurts performances, as can be seen in the ViT results in Table 2, and NLP results in Fig. 4a.
Revised: SWA often does not improve and sometimes hurts performances, as can be seen in the ViT results in Table 2, and NLP results in Fig. 4a.

Original: In contrast, SAM has some positive effects in these settings.
Revised: In contrast, SAM has some positive effects in these settings.

Original: We explore this further in the Appendix.
Revised: We explore this further in Appendix A.3.

Original: SWA and SAM improve SSL task performance :
Revised: SWA and SAM improve SSL task performance.

Original: This is non-trivial as the theoretical motivation behind finding flat minima is linked to supervised learning losses [41, 72, 21].
Revised: This is non-trivial as the theoretical motivation behind finding flat minima is linked to supervised learning losses [42, 76, 20].

Original: Concurrently to our work, Ramesh et al.
Revised: Concurrently to our work, Ramesh et al.

Original: 76] report that SAM helps for contrastive CLIP [74] models too.
Revised: [81] report that SAM helps for contrastive CLIP [78] models too.

Original: Flat optimizers do not strictly improve over non-flat optimizers.
Revised: Flat optimizers do not strictly improve over non-flat optimizers.

Original: The non-flat optimizer is nearly always the best for NPP Proteins and SSL methods on CIFAR10.
Revised: The non-flat optimizer is nearly always the best for NPP Proteins and SSL methods on CIFAR10.

Original: We investigate the NPP Proteins solutions in the next subsection and recommend a more thorough investigation of the landscapes of SSL objectives for future work.
Revised: We investigate the NPP Proteins solutions in the next subsection and recommend a more thorough investigation of the landscapes of SSL objectives for future work.

Original: Flat-minima optimizers offer asymmetric payoffs: at worst, they decreased performance by − 0 .
Revised: Flat-minima optimizers offer asymmetric payoffs: at worst, they decreased performance by − 0 .

Original: 30% , at best, they increased it by 2 .
Revised: 30% , at best, they increased it by 2 .

Original: Averaging SAM iterates often improves over SWA or SAM alone .
Revised: Averaging SAM iterates often improves over SWA or SAM alone .

Original: G ( SWA+SAM ) > min( G ( SWA ) , G ( SAM )) in 39/42 cases.
Revised: G ( SWA+SAM ) > min( G ( SWA ) , G ( SAM )) in 39/42 cases.

Original: We hypothesize that asymmetric payoffs are the reason: when either SWA or SAM does not improve over the baseline (as discussed above), it does not hurt (much) either, hence SWA+SAM is more robust across all tasks.
Revised: We hypothesize that asymmetric payoffs are the reason: when either SWA or SAM does not improve over the baseline (as discussed above), it does not hurt (much) either, hence SWA+SAM is more robust across all tasks.

Original: 4.5 Why do flat-minima optimizers fail?
Revised: 4.5 Why do flat-minima optimizers fail?

Original: Here, we audit one of the cases, where neither θ SWA nor θ SAM improves over θ NF (this happens inout of 42 cases): training a GraphSAGE [35] model on OGB-Proteins: a protein-protein interaction graph where the goal is to predict the presence of protein functions (multi-label binary classification) [42].
Revised: Here, we audit one of the cases, where neither θ SWA nor θ SAM improves over θ NF (this happens inout of 42 cases): training a GraphSAGE [36] model on OGB-Proteins: a protein-protein interaction graph where the goal is to predict the presence of protein functions (multi-label binary classification) [44].

Original: θ SWA performs noticeably worse; θ SAM performs about equally well.
Revised: θ SWA performs noticeably worse; θ SAM performs about equally well.

Original: [Figure]
Revised: [Figure]

Original: [Figure]

Original: [Figure]

Original: Fig.
Revised: Fig.

Original: 5 shows two linear interpolations: between θ NF (ADAM) and (1) θ
Revised: 5 shows two linear interpolations: between θ NF (ADAM) and (1) θ

Original: SWA (Figures 5a and 5b), and (2) θ SAM (Figures 5c and 5d).
Revised: SWA (Figures 5a and 5b), and (2) θ SAM (Figures 5c and 5d).

Original: In contrast to success cases in Fig. 2, here: (a) for both SWA and
Revised: In contrast to success cases in Fig. 2, here: (a) for both SWA and

Original: SAM, the training loss minimizer is very uncorrelated with the test loss minimizer; (b) SAM and
Revised: SAM, the training loss minimizer is very uncorrelated with the test loss minimizer; (b) SAM and

Original: ADAM seem to be contained in the same test loss/accuracy basin.
Revised: ADAM seem to be contained in the same test loss/accuracy basin.

Original: More analyses can be found in the
Revised: More analyses can be found in the

Original: Appendix.
Revised: Appendix.

Original:  Limitations and Future Work
Revised:  Limitations and Future Work

Original: First, some of the fixed, shared hyperparameter values we used from previous works may harm the effect of flat optimizers.
Revised: First, some of the fixed, shared hyperparameter values we used from previous works may harm the effect of flat optimizers.

Original: The ideal experimental design includes tuning all hyperparameters independently for the non-flat optimizer, SWA, SAM, and SWA+SAM.
Revised: The ideal experimental design includes tuning all hyperparameters independently for the non-flat optimizer, SWA, SAM, and SWA+SAM.

Original: However, this forces the number of required runs to grow exponentially in unique hyperparameters and quickly renders this benchmark infeasible.
Revised: However, this forces the number of required runs to grow exponentially in unique hyperparameters and quickly renders this benchmark infeasible.

Original: Second, despite our best efforts to evaluate the optimizers on a broad range of benchmark tasks, there are still plenty of unexplored domains; especially some of which are known to be sensitive to careful optimization, such as generative modeling [40], deep reinforcement learning [2], or causal machine learning [49].
Revised: Second, despite our best efforts to evaluate the optimizers on a broad range of benchmark tasks, there are still plenty of unexplored domains; especially some of which are known to be sensitive to careful optimization. For example, bi-level optimization problems [23] are common in generative modeling [30, 41], deep reinforcement learning [59, 43], meta-learning [80, 51], or causal machine learning [52, 53].

Original:  Third, in general, we believe fruitful directions of research include (a) optimizers that explicitly find basins where training loss flatness more directly corresponds to higher hold-out accuracy, (b) post-processing methods for existing optimization runs to move into flatter regions of these basins, (c) loss functions whose contoursmore tightly align with accuracy contours, (d) the study of flat-minima hyperparameter interactions (e.g., learning rate and neighborhood radius in SAM), (e) analyses of flat
Revised: We are unaware of an investigation of flat minima optimization for such problems. Third, in general, we believe fruitful directions of research include (a) optimizers that explicitly find basins where training loss flatness more directly corresponds to higher hold-out accuracy, (b) post-processing methods for existing optimization runs to move into flatter regions of these basins [2], (c) loss functions whose contours

Original: Our benchmark results point to which tasks would most benefit from improving these future work directions: graph learning tasks would clearly benefit from improvements in (a), as SAM is never among the best performing method, and language tasks would benefit if (b) is improved, as SWA is never among the best performing method).
Revised: Our benchmark results point to which tasks would most benefit from improving these future work directions: graph learning tasks would clearly benefit from improvements in (a), as SAM is never among the best performing method, and language tasks would benefit if (b) is improved, as SWA is never among the best performing method).

Original:  Conclusion
Revised:  Conclusion

Original: We investigated when flat minimia optimizers work by conducting a fair comparison of two popular flat-minima optimizers.
Revised: We investigated when flat minima optimizers work by conducting a fair comparison of two popular flatminima optimizers.

Original: We examined the behavior of SWA/SAM by analyzing their loss landscapes on two representative deep learning tasks.
Revised: We examined the behavior of SWA/SAM by analyzing their loss landscapes on two representative deep learning tasks.

Original: Our next step was to evaluate their generalization performance on a broad and diverse set of tasks (in data, learning settings, and model architectures).
Revised: Our next step was to evaluate their generalization performance on a broad and diverse set of tasks (in data, learning settings, and model architectures).

Original: Based on this benchmarking, we identified 8 findings, of which some directly guide future work directions.
Revised: Based on this benchmarking, we identified 9 findings, of which some directly guide future work directions.

Original: Finally, when SWA/SAM did not improve over baselines, common assumptions seemed broken (i e., train-to-test loss minimizers were not correlated).
Revised: Finally, when SWA/SAM did not improve over baselines, common assumptions seemed broken (i.e., train-to-test loss minimizers were not correlated).

Original: Acknowledgements
Revised: Acknowledgements

Original: We are very grateful to Kilian Q. Weinberger and Gao Huang for initial discussions and intuition,
Revised: We are very grateful to Kilian Q. Weinberger and Gao Huang for initial discussions and intuition,

Original: Pontus Stenetorp for NLP experimental design advice and Oscar Key for feedback on the draft.
Revised: Pontus Stenetorp for NLP experimental design advice, and Oscar Key for feedback on the draft.

Original: JK and LL acknowledge support by the Engineering and Physical Sciences Research Council with grant number EP/S021566/1.
Revised: JK and LL acknowledge support by the Engineering and Physical Sciences Research Council with grant number EP/S021566/1.

Original: This research was supported through Azure resources provided by The Alan
Revised: This research was supported through Azure resources provided by The Alan

Original: Turing Institute and credits awarded by Google Cloud.
Revised: Turing Institute and credits awarded by Google Cloud.