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 [43].
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 [43].

Original: Theoretical and empirical studies [21, 78, 9, 55, 49, 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 [21, 78, 9, 55, 49, 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) [48] and 2.
Revised: Stochastic Weight Averaging (SWA) [48] and 2.

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

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 [76, 3, 51, 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 [76, 3, 51, 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 [48] paper was published in 2018, and the SAM [22] paper in 2021; however, the SAM paper, and its most noticeable follow-ups [65, 12, 105], do not compare against SWA.
Revised: The SWA [48] paper was published in 2018, and the SAM [22] paper in 2021; however, the SAM paper, and its most noticeable follow-ups [65, 12, 105], 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, hyper-parameters, etc.).
Revised: g. data augmentations, hyper-parameters, etc.).

Original: [Table]
Revised: [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
Revised: Rigorous comparison of

Original: SWA and SAM’s performance
Revised: 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 8 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: 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 [79, 51].
Revised: The idea of averaging weights dates back to accelerating the convergence speed of SGD [79, 51].

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: [48] suggests updating the parameters once after each epoch, i.e. ν ≈ N |B| , and starting at E ≈ 0 .
Revised: [48] 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: [39] argue that SWA may always improve generalization, regardless of the loss function’s geometry.
Revised: [39] argue that SWA may always improve generalization, regardless of the loss function’s geometry.

Original: Kaddour [51] show that averaging a specific range of weights can speed up training convergence.
Revised: Kaddour [51] 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: [22] 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: [95] show that ρ should be set to values outside the originally considered ranges by
Revised: [95] 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 [36, 8] and SAM [65, 105, 103].
Revised: There are several extensions of SWA [36, 8] and SAM [65, 105, 103].

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, 85].
Revised: Besides SWA and SAM, other flat-minima optimizers include e.g. [9, 85].

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: [85] similarly requires [5 , 10] forward and backward passes to estimate the Hessian trace and 6 of 7 experiments yield minimal improvement of ≤ 0 .
Revised: [85] 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 requiring 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 in 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 [32, 25], regularization [69, 28], and network pruning [26].
Revised: Previous studies successfully used such linear interpolations to gain novel insights, e.g., for training dynamics [32, 25], regularization [69, 28], and network pruning [26].

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, 100, 12, 22].
Revised: Here, we compute dominant Hessian eigenvalues, as commonly used in the flat minima literature [9, 100, 12, 22].

Original: Lastly, in Appendix A.1, we further compute CKA [61] and cosine similarities between SWA/SAM’s network output logits.
Revised: Lastly, in Appendix A.1, we further compute CKA [61] 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 [101] model on CIFAR-100 [63] (baseline non-flat optimizer: SGD with momentum (SGD-M)) [84].
Revised: Specifically, for (i), we investigate the loss/accuracy surfaces of a WideResNet28-10 [101] model on CIFAR-100 [63] (baseline non-flat optimizer: SGD with momentum (SGD-M)) [84].

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

Original: Isomorphism Network [97] model on OGB-Code2 [45] (baseline optimizer: Adam [56]).
Revised: Isomorphism Network [97] model on OGB-Code2 [45] (baseline optimizer: Adam [56]).

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: [26], or (b) yield trivial network connectivity results, such as non-linear low-loss paths, which can be found for any two network minimizers [20, 27, 33, 23].
Revised: [26], or (b) yield trivial network connectivity results, such as non-linear low-loss paths, which can be found for any two network minimizers [20, 27, 33, 23].

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 [46, 67].
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 [46, 67].

Original: Grewal & Bui [34] successfully combine models yielded by different optimizers, and we think future work should study ensembling SAM and non-SAM solutions.
Revised: Grewal & Bui [34] 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 F1 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 (Appendix A.2).
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 [105, 103].
Revised: (1) has been explored by [105, 103].

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 [58] or stochastic differential equations [104].
Revised: Future work may study SAM’s basin escape time, e.g., using convolutions [58] or stochastic differential equations [104].

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 [74, 100].
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 [74, 100].

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, 99, 18, 62, 86].
Revised: We choose this metric as it is very commonly used in the minima flatness literature, e.g., [9, 12, 22, 99, 18, 62, 86].

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 [92] 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 [92] 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: [48] and hold the update frequency ν constant to once per epoch and tune the start time E ∈ { 0 .
Revised: [48] 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: [48] argue that a cyclical learning rate starting from E helps to encourage exploration of the basin.
Revised: [48] 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 [22, 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 [101] withlayers and width 10, and PyramidNet (PN) with 110 layers and widening factor 272 [38] as well as
Revised: We evaluate the CNN architectures WideResNets [101] withlayers and width 10, and PyramidNet (PN) with 110 layers and widening factor 272 [38] as well as

Original: Vision Transformer (ViT) [19] and MLP-Mixer [88] on CIFAR{10, 100} [63].
Revised: Vision Transformer (ViT) [19] and MLP-Mixer [88] on CIFAR{10, 100} [63].

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 [41], a Simple framework for Contrastive Learning (SimCLR) [10],
Revised: We consider the following methods on CIFAR10 and ImageNette 4 : Momentum Contrast [41], a Simple framework for Contrastive Learning (SimCLR) [10],

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

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

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

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

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

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 [59].
Revised: Compared to learning a linear model on top of the representations, this evaluation procedure is more robust to hyperparameter changes [59].

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) [47].
Revised: We consider the task of open domain question answering (ODQA) using a T5-based model FusionIn-Decoder (FiD) [47].

Original: We evaluate FiD-base on the test sets of Natural Questions (NQ) [64] and TriviaQA [50].
Revised: We evaluate FiD-base on the test sets of Natural Questions (NQ) [64] and TriviaQA [50].

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

Original: [Table]
Revised: [Table]

Original: and inference.
Revised: and inference.

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

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

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

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 [57],
Revised: For each task, we use two of the following GNN architectures and matrix factorization methods: GCN [57],

Original: DeeperGCN (DGCN) [68], SAGE [37], GIN [97], ComplEx [90], and CP [66].
Revised: DeeperGCN (DGCN) [68], SAGE [37], GIN [97], ComplEx [90], and CP [66].

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 8 Findings

Original: We use G ( · ) to describe the generalization accuracy/F1/ROCAUC/AP/MRR of all optimizers { Non-flat
Revised: We use G ( · ) to describe the generalization accuracy/F1/ROCAUC/AP/MRR of all optimizers { Non-flat

Original: baseline ( NF ) ,
Revised: baseline ( NF ) ,

Original: SWA
Revised: SWA

Original: , SAM ,
Revised: , SAM ,

Original: SWA+
Revised: SWA+

Original: SAM }
Revised: SAM }

Original: Both dataset and architecture matter.
Revised: Both dataset and architecture matter.

Original: To explain the relationship between flatness and generalization, some works focused on the network’s architecture [17, 12, 77], while others focused on the data distribution [91, 78, 21].
Revised: To explain the relationship between flatness and generalization, some works focused on the network’s architecture [17, 12, 77], while others focused on the data distribution [91, 78, 21].

Original: Our results show that both matter.
Revised: Our results show that both matter.

Original: Datasets: e
Revised: Datasets: e

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: 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.

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: 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
Revised: SWA

Original: underperforms on NLP tasks.
Revised: 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,
Revised: Surprisingly,

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

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

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

Original: G ( NF ) only in 4.
Revised: G ( NF ) only in 4.

Original: We examine why SAM under-performs in Appendix A.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
Revised: SWA often does not improve and sometimes hurts performances, as can be seen in the ViT results in Table

Original: and NLP results in Fig.
Revised: and NLP results in Fig.

Original: 4a
Revised: 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 Appendix A.
Revised: We explore this further in Appendix A.

Original: SWA and SAM can work on SSL tasks too : This is non-obvious because the theoretical motivation behind finding flat minima is linked to supervised learning
Revised: SWA and SAM can work on SSL tasks too : This is non-obvious because the theoretical motivation behind finding flat minima is linked to supervised learning

Original: losses [
Revised: losses [

Original: Concurrently
Revised: Concurrently

Original: to our work
Revised: to our work

Original: Ramesh et al.
Revised: Ramesh et al.

Original: [83] report that SAM helps for contrastive CLIP [80] models
Revised: [83] report that SAM helps for contrastive CLIP [80] models

Original: too.
Revised: too.

Original: Flat optimizers can underperform non-flat optimizers.
Revised: Flat optimizers can underperform 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 [37] model on OGB-Proteins: a protein-protein interaction graph where the goal is to predict the presence of protein functions (multi-label binary classification) [45].
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 [37] model on OGB-Proteins: a protein-protein interaction graph where the goal is to predict the presence of protein functions (multi-label binary classification) [45].

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: °
Revised: °

Original: [Figure]
Revised: [Figure]

Original: [Figure] [Figure]
Revised: [Figure] [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.
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.

Original: For example, bi-level optimization problems [24] are common in generative modeling [31, 42], deep reinforcement learning [60, 44], meta-learning [82, 52], or causal machine learning [53, 54].
Revised: For example, bi-level optimization problems [24] are common in generative modeling [31, 42], deep reinforcement learning [60, 44], meta-learning [82, 52], or causal machine learning [53, 54].

Original: We are not aware of an investigation of flat minima optimization for such problems.
Revised: We are not aware of an investigation of flat minima optimization for such problems.

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 [2], (
Revised: 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], (

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 minimia optimizers work by conducting a fair comparison of two popular flat-minima 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 8 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.

Original: Checklist
Revised: Checklist

Original: For all authors... (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?
Revised: For all authors... (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?

Original: [Yes] (b) Did you describe the limitations of your work?
Revised: [Yes] (b) Did you describe the limitations of your work?

Original: [Yes] See Sec.(c)
Revised: [Yes] See Sec.(c)

Original: Did you discuss any potential negative societal impacts of your work?
Revised: Did you discuss any potential negative societal impacts of your work?

Original: [No] We consider this work to be an investigation into optimization algorithms that are central to modern ML systems.
Revised: [No] We consider this work to be an investigation into optimization algorithms that are central to modern ML systems.

Original: As these systems can be used in vastly different ways, we believe that it is intractable to isolate the societal impact of novel insights into optimization.
Revised: As these systems can be used in vastly different ways, we believe that it is intractable to isolate the societal impact of novel insights into optimization.

Original: (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
Revised: (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?

Original: [Yes] 2.
Revised: [Yes] 2.

Original: If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?
Revised: If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?

Original: [N/A] (b) Did you include complete proofs of all theoretical results?
Revised: [N/A] (b) Did you include complete proofs of all theoretical results?

Original: [N/A] 3.
Revised: [N/A] 3.

Original: If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
Revised: If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?

Original: [Yes] Supplemental material (b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
Revised: [Yes] Supplemental material (b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?

Original: [Yes] See Appendix.
Revised: [Yes] See Appendix.

Original: (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
Revised: (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?

Original: [Yes] See Sec.(d)
Revised: [Yes] See Sec.

Original:  Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
Revised: (d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?

Original: [Yes] See Appendix. 4.
Revised: [Yes] See Appendix. 4.

Original: If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?
Revised: If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?

Original: [Yes] (b) Did you mention the license of the assets?
Revised: [Yes] (b) Did you mention the license of the assets?

Original: [Yes] See Appendix.
Revised: [Yes] See Appendix.

Original: (c) Did you include any new assets either in the supplemental material or as a URL?
Revised: (c) Did you include any new assets either in the supplemental material or as a URL?

Original: [N/A] (d) Did you discuss whether
Revised: [N/A] (d) Did you discuss whether

Original: and how consent was obtained from people whose data you’re using/curating?
Revised: and how consent was obtained from people whose data you’re using/curating?

Original: [N/A] (e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content?
Revised: [N/A] (e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content?

Original: [N/A] 5.
Revised: [N/A] 5.

Original: If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?
Revised: If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?

Original: [N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
Revised: [N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?

Original: [N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
Revised: [N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?

Original: [N/A]
Revised: [N/A]

Original: A Additional Analyses A.1 CKA Similarities
Revised: A Additional Analyses A.1 CKA Similarities

Original: We compute the CKA [61, 98] and cosine similarities of network output logits on train and test set, respectively.
Revised: We compute the CKA [61, 98] and cosine similarities of network output logits on train and test set, respectively.

Original: Table 3 shows the results.
Revised: Table 3 shows the results.

Original: [Table]
Revised: [Table]

Original: The results show that the SAM solutions produce predictions that are less similar to the non-flat baseline than SWA solutions, as indicated by lower CKA and cosine similarities.
Revised: The results show that the SAM solutions produce predictions that are less similar to the non-flat baseline than SWA solutions, as indicated by lower CKA and cosine similarities.

Original: This result is in line with Observation 1 and 2 from Sec. 3.1.
Revised: This result is in line with Observation 1 and 2 from Sec. 3.1.

Original: A.2 Saddle point of SAM’s GIN solution
Revised: A.2 Saddle point of SAM’s GIN solution

Original: − 0 .0 .0 .1 .
Revised: − 0 .0 .0 .1 .

Original: 0 − 1 .− 0 .0 .0 .1 .0.
Revised: 0 − 1 .− 0 .0 .0 .1 .0.

Original: (a) SAM solution.
Revised: (a) SAM solution.

Original: − 1 .− 0 .0 .0 .1 .
Revised: − 1 .− 0 .0 .0 .1 .

Original: 0 − 1 .− 0 .0 .0 .1 .0.
Revised: 0 − 1 .− 0 .0 .0 .1 .0.

Original: (b) SWA solution (for comparison).
Revised: (b) SWA solution (for comparison).

Original: To gather further evidence on whether the SAM solution is a saddle point we analyze its Hessian eigenvalue density, following Ghorbani et al.
Revised: To gather further evidence on whether the SAM solution is a saddle point we analyze its Hessian eigenvalue density, following Ghorbani et al.

Original: [29] and using the Stochastic Lanczos Quadrature algorithm [30].
Revised: [29] and using the Stochastic Lanczos Quadrature algorithm [30].

Original: Fig.
Revised: Fig.

Original: 7 shows the density, including significant probability mass for both positive and negative eigenvalues, indicating that the solution is a saddle point.
Revised: 7 shows the density, including significant probability mass for both positive and negative eigenvalues, indicating that the solution is a saddle point.

Original: A.3 Why does SWA not work well on NLP tasks?
Revised: A.3 Why does SWA not work well on NLP tasks?

Original: In Fig.
Revised: In Fig.

Original: 4a, we saw that SWA had only a mild effect on the generalization performance of NLP tasks, and sometimes even decreases it.
Revised: 4a, we saw that SWA had only a mild effect on the generalization performance of NLP tasks, and sometimes even decreases it.

Original: Here, we seek to investigate why that is.
Revised: Here, we seek to investigate why that is.

Original: [Figure]
Revised: [Figure]

Original: We consider two tasks: (i) the RTE task, for which SWA decreases the performance by around − 0 .
Revised: We consider two tasks: (i) the RTE task, for which SWA decreases the performance by around − 0 .

Original: 23 ± 0 .
Revised: 23 ± 0 .

Original: 20 compared to Adam, (ii) the QNLI task, for which SWA decreases the performance by − 0 .
Revised: 20 compared to Adam, (ii) the QNLI task, for which SWA decreases the performance by − 0 .

Original: 08 ± 0 .
Revised: 08 ± 0 .

Original: In both cases, SAM improved the performance statistically significantly over Adam.
Revised: In both cases, SAM improved the performance statistically significantly over Adam.

Original: For the QNLI task in Fig.
Revised: For the QNLI task in Fig.

Original: 8a, we observe that SWA finds a lower/higher training loss/accuracy than
Revised: 8a, we observe that SWA finds a lower/higher training loss/accuracy than

Original: Adam, respectively.
Revised: Adam, respectively.

Original: However, the test loss/accuracy is higher/lower at the SWA solutions and the loss functions seem less well correlated in between both solutions (i.e, for α ∈ [0 , 1] ).
Revised: However, the test loss/accuracy is higher/lower at the SWA solutions and the loss functions seem less well correlated in between both solutions (i.e, for α ∈ [0 , 1] ).

Original: For the RTE task in Fig.
Revised: For the RTE task in Fig.

Original: 8b, we note that SWA finds a solution that is closer to a sharply increasing side.
Revised: 8b, we note that SWA finds a solution that is closer to a sharply increasing side.

Original: This may happen if the baseline optimizer skips or goes around sharper solutions (e.g., due to large step sizes) and the average pulls it towards these suboptimal regions.
Revised: This may happen if the baseline optimizer skips or goes around sharper solutions (e.g., due to large step sizes) and the average pulls it towards these suboptimal regions.

Original: Further, in Table 3, we notice very high values of s CKA ( θ NF , θ SWA ) , s cosine ( θ NF , θ SWA ) for both training and test sets, indicating that the predictions are indeed very similar.
Revised: Further, in Table 3, we notice very high values of s CKA ( θ NF , θ SWA ) , s cosine ( θ NF , θ SWA ) for both training and test sets, indicating that the predictions are indeed very similar.

Original: In contrast, s CKA ( θ NF , θ SAM ) , s cosine ( θ NF , θ SAM ) is lower, especially for the test set.
Revised: In contrast, s CKA ( θ NF , θ SAM ) , s cosine ( θ NF , θ SAM ) is lower, especially for the test set.

Original: A.4 Why does SAM not work well on GRL tasks?
Revised: A.4 Why does SAM not work well on GRL tasks?

Original: Fig.
Revised: Fig.

Original: 10 shows the interpolations between the Adam and SAM solution.
Revised: 10 shows the interpolations between the Adam and SAM solution.

Original: We do not observe a significant loss/accuracy difference between the two different basins.
Revised: We do not observe a significant loss/accuracy difference between the two different basins.

Original: One possible explanation for this phenomenon is that the loss surface for this task is “globally well-connected” [98], yielding many basins with very similar geometric properties.
Revised: One possible explanation for this phenomenon is that the loss surface for this task is “globally well-connected” [98], yielding many basins with very similar geometric properties.

Original: [Figure]
Revised: [Figure]

Original: Æ
Revised: Æ

Original: [Figure]
Revised: [Figure]

Original: RoBERTa-QNLI:
Revised: RoBERTa-QNLI:

Original: Adam to SWA
Revised: Adam to SWA

Original: Train
Revised: Train

Original: Test
Revised: Test

Original: Loss
Revised: Loss

Original: Accuracy
Revised: Accuracy

Original: [Figure]
Revised: [Figure]

Original: RoBERTa-RTE: Adam to SWA
Revised: RoBERTa-RTE: Adam to SWA

Original: Train
Revised: Train

Original: Test
Revised: Test

Original: Loss
Revised: Loss

Original: Accuracy
Revised: Accuracy

Original: [Figure]
Revised: [Figure]

Original: GIN-Molpcba: Adam to SAM
Revised: GIN-Molpcba: Adam to SAM

Original: Train
Revised: Train

Original: Test
Revised: Test

Original: Loss
Revised: Loss

Original: Accuracy
Revised: Accuracy

Original: [Figure]
Revised: [Figure]

Original: °     Æ 0 .0 .0 .0 .0 .
Revised: °     Æ 0 .0 .0 .0 .0 .

Original: SAM
Revised: SAM

Original: ADAM °
Revised: ADAM °

Original: °°     Æ 0 .0 .0 .
Revised: °°     Æ 0 .0 .0 .

Original: [Figure]
Revised: [Figure]

Original: A.5 Does changing SAM’s ρ result in different basins?
Revised: A.5 Does changing SAM’s ρ result in different basins?

Original: Here, we plot linear interpolations of solutions obtained by smaller and larger
Revised: Here, we plot linear interpolations of solutions obtained by smaller and larger

Original: ρ values.
Revised: ρ values.

Original: Overall, we find that all solutions seem to lie in different basins indicated by high loss barriers in between ( α = 0 . 5 ) them.
Revised: Overall, we find that all solutions seem to lie in different basins indicated by high loss barriers in between ( α = 0 . 5 ) them.

Original: A.5.1 WRN-28- 
Revised: A.5.1 WRN-28-

Original: For the WRN-28-10 model investigated in Sec.
Revised: For the WRN-28-10 model investigated in Sec.

Original: 3.1 and Sec. 4, we set ρ = 0 .
Revised: 3.1 and Sec. 4, we set ρ = 0 .

Original: 1 (as determined by hyper-parameter tuning on validation loss).
Revised: 1 (as determined by hyper-parameter tuning on validation loss).

Original: Accuracy
Revised: Accuracy

Original: Loss
Revised: Loss

Original: TrainTest
Revised: TrainTest

Original: [Figure]
Revised: [Figure]

Original: [Equation]
Revised: [Equation]

Original: latexit> From ⇢ = 0 .
Revised: latexit> From ⇢ = 0 .

Original: 1 ( ↵ = 0) to ⇢ = 0 .
Revised: 1 ( ↵ = 0) to ⇢ = 0 .

Original: 2 ( ↵ = 1)
Revised: 2 ( ↵ = 1)

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: A.5.
Revised: A.5.

Original: -Code
Revised: -Code

Original: For the GIN model investigated in Sec.
Revised: For the GIN model investigated in Sec.

Original: 3.1 and Sec. 4, we set ρ = 0 .
Revised: 3.1 and Sec. 4, we set ρ = 0 .

Original: 15 (as determined by hyperparameter tuning on validation loss).
Revised: 15 (as determined by hyperparameter tuning on validation loss).

Original: Accuracy
Revised: Accuracy

Original: Loss
Revised: Loss

Original: TrainTest
Revised: TrainTest

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

Original: =
Revised: =

Original: [Table]
Revised: [Table]

Original: to
Revised: to

Original: A.6 Does changing the base optimizer impact SWA and SAM’s effectiveness?
Revised: A.6 Does changing the base optimizer impact SWA and SAM’s effectiveness?

Original: The influence of the base optimizer on SWA/SAM’s effectiveness is under-explored.
Revised: The influence of the base optimizer on SWA/SAM’s effectiveness is under-explored.

Original: In the previous experiments, we use the default base optimizers from existing code repositories or, as reported in previous works.
Revised: In the previous experiments, we use the default base optimizers from existing code repositories or, as reported in previous works.

Original: Here, we want to conduct an initial investigation into its effect on SWA and SAM.
Revised: Here, we want to conduct an initial investigation into its effect on SWA and SAM.

Original: In the following experiments, we only switch the base optimizer and keep everything else fixed (including hyper-parameters such as learning rate, etc.).
Revised: In the following experiments, we only switch the base optimizer and keep everything else fixed (including hyper-parameters such as learning rate, etc.).

Original: We train (i) a ResNet-34 (similar to WRN-2810 but smaller) on CIFAR100, once per SGD with momentum and AdamW [73], and (ii) a GIN model on Code2, as in Sec. 3, but using RMSprop instead of Adam.
Revised: We train (i) a ResNet-34 (similar to WRN-2810 but smaller) on CIFAR100, once per SGD with momentum and AdamW [73], and (ii) a GIN model on Code2, as in Sec. 3, but using RMSprop instead of Adam.

Original: We choose AdamW and RMSProp, since they are commonly used in image classification [73, 19, 89] and graph representation learning (GRL) [71, 70, 81], respectively.
Revised: We choose AdamW and RMSProp, since they are commonly used in image classification [73, 19, 89] and graph representation learning (GRL) [71, 70, 81], respectively.

Original: Due to time constraints, we only report results obtained with one random seed (except for GIN-Code with Adam, which we already evaluated across three random seeds in our initial submission, see
Revised: Due to time constraints, we only report results obtained with one random seed (except for GIN-Code with Adam, which we already evaluated across three random seeds in our initial submission, see

Original: Fig. 4b).
Revised: Fig. 4b).

Original: Further, again due to time constraints, for the ResNet-34 task, we do not conduct a hyperparameter search of SAM’s ρ but set it to ρ = 0 .
Revised: Further, again due to time constraints, for the ResNet-34 task, we do not conduct a hyperparameter search of SAM’s ρ but set it to ρ = 0 .

Original: 05 , as this value has been reported to be a good default value [22].
Revised: 05 , as this value has been reported to be a good default value [22].

Original: Table 4 show the test performances with switched base optimizers.
Revised: Table 4 show the test performances with switched base optimizers.

Original: First, discussing task (i), we note that AdamW under-fits the model and generalizes poorly compared to SGD.
Revised: First, discussing task (i), we note that AdamW under-fits the model and generalizes poorly compared to SGD.

Original: Here, SAM exacerbates the performance even further, performing even worse than the AdamW baseline.
Revised: Here, SAM exacerbates the performance even further, performing even worse than the AdamW baseline.

Original: The reasons for that are unclear, and we leave an investigation into them for future work.
Revised: The reasons for that are unclear, and we leave an investigation into them for future work.

Original: For task (ii) using RMSprop, we observe that both flat-minima optimizers improve over the baseline performance.
Revised: For task (ii) using RMSprop, we observe that both flat-minima optimizers improve over the baseline performance.

Original: However, compared to when using Adam, they perform even more similarly, with SAM only being 0 .
Revised: However, compared to when using Adam, they perform even more similarly, with SAM only being 0 .

Original: 03% better.
Revised: 03% better.

Original: Interestingly, the combination SWA+SAM again performs best.
Revised: Interestingly, the combination SWA+SAM again performs best.

Original: [Table]
Revised: [Table]

Original: B Experimental details B.1 Computer Vision
Revised: B Experimental details B.1 Computer Vision

Original: We mostly adapt the hyper-parameter values from Foret et al.
Revised: We mostly adapt the hyper-parameter values from Foret et al.

Original: [22] for WRN-28-10 and PyramidNet272, from Dosovitskiy et al.
Revised: [22] for WRN-28-10 and PyramidNet272, from Dosovitskiy et al.

Original: [19] for ViT, and from [88] for MLP-Mixer models.
Revised: [19] for ViT, and from [88] for MLP-Mixer models.

Original: We average all results across three random seeds.
Revised: We average all results across three random seeds.

Original: B.1.1 Supervised Classification
Revised: B.1.1 Supervised Classification

Original: We train WideResNets [101] with 28 layers and width 10 (WRN28-10) and PyramidNet [38] with 110 layers and widening factor α = 272 (PyramidNet-272) from scratch.
Revised: We train WideResNets [101] with 28 layers and width 10 (WRN28-10) and PyramidNet [38] with 110 layers and widening factor α = 272 (PyramidNet-272) from scratch.

Original: The Vision
Revised: The Vision

Original: Transformer (ViT) base model with input patch size 16 (ViT-B/16) and MLP-Mixer base model with input patch size 16 (MLP-Mixer-B/16) start from pre-trained checkpoints available at https://console.cloud.google.com/storage/vit_models/.
Revised: Transformer (ViT) base model with input patch size 16 (ViT-B/16) and MLP-Mixer base model with input patch size 16 (MLP-Mixer-B/16) start from pre-trained checkpoints available at https://console.cloud.google.com/storage/vit_models/.

Original: The reason for using pre-trained checkpoints for the ViT and MLP-Mixer models is that, due to their lack of some inductive biases inherent to
Revised: The reason for using pre-trained checkpoints for the ViT and MLP-Mixer models is that, due to their lack of some inductive biases inherent to

Original: CNNs, such as translation equivariance and locality, they do not generalize well when trained on insufficient amounts of data [19].
Revised: CNNs, such as translation equivariance and locality, they do not generalize well when trained on insufficient amounts of data [19].

Original: Table 5 shows the hyper-parameters for each architecture.
Revised: Table 5 shows the hyper-parameters for each architecture.

Original: B.1.2 Self-Supervised Learning
Revised: B.1.2 Self-Supervised Learning

Original: Table 6 shows the hyper-parameters for each SSL method.
Revised: Table 6 shows the hyper-parameters for each SSL method.

Original: We use implementations from the lightly package, available at https://github.com/lightly-ai/lightly [87].
Revised: We use implementations from the lightly package, available at https://github.com/lightly-ai/lightly [87].

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

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: B.2 Natural Language Processing
Revised: B.2 Natural Language Processing

Original: For the task of Open Domain Question Answering, we adapt the hyper-parameter values and theretrieved passages for each question from 47.
Revised: For the task of Open Domain Question Answering, we adapt the hyper-parameter values and theretrieved passages for each question from 47.

Original: We report the Exact Match score of FiD-base model on
Revised: We report the Exact Match score of FiD-base model on

Original: Natural Questions (NQ) and TriviaQA test sets.
Revised: Natural Questions (NQ) and TriviaQA test sets.

Original: For GLUE benchmark, we report Matthew’s Corr for
Revised: For GLUE benchmark, we report Matthew’s Corr for

Original: CoLA, Pearson correlation coefficient for STSB, and accuracy for the the rest of the datasets.
Revised: CoLA, Pearson correlation coefficient for STSB, and accuracy for the the rest of the datasets.

Original: Results are all evaluated on the dev set of GLUE benchmark.
Revised: Results are all evaluated on the dev set of GLUE benchmark.

Original: We use the RoBERTa-base as our backbone language model, implemented with Huggingface Transformers [94].
Revised: We use the RoBERTa-base as our backbone language model, implemented with Huggingface Transformers [94].

Original: Most of the task-specific hyper-parameter values are adapted from 1.
Revised: Most of the task-specific hyper-parameter values are adapted from 1.

Original: B.3 Graph Representation Learning
Revised: B.3 Graph Representation Learning

Original: We mostly adapt the hyper-parameter values from Hu et al.
Revised: We mostly adapt the hyper-parameter values from Hu et al.

Original: [45] for GCN [57], SAGE [37], and GIN [97], from Chen et al.
Revised: [45] for GCN [57], SAGE [37], and GIN [97], from Chen et al.

Original: [13] for [66] and ComplEx [90], and from Li et al.
Revised: [13] for [66] and ComplEx [90], and from Li et al.

Original: [68] for DGCN.
Revised: [68] for DGCN.

Original: Due to high standard errors, we averaged the results of a few tasks more than three times, as mentioned in the following tables.
Revised: Due to high standard errors, we averaged the results of a few tasks more than three times, as mentioned in the following tables.

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

Original: [Figure]
Revised: [Figure]

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

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

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]