Convergence of Sharpness-Aware Minimization Algorithms using Increasing Batch Size and Decaying Learning Rate

arXiv:2409.09984 · cs.LG, math.OC · Submitted 2024-09-16 · Read on arXiv

Listen

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: Today's paper: "Convergence of Sharpness-Aware Minimization Algorithms using Increasing Batch Size and Decaying Learning Rate".

Jane: The gist The main contribution of this paper is to show an ϵ-approximation of the GSAM algorithm with an increasing batch size and constant learning rate ((7) in Table 1;…

Tom: First, who's behind it and why it matters.

Paper summary: Tom: So, wrapping up this discussion on "Convergence of Sharpness-Aware Minimization Algorithms using Increasing Batch Size and Decaying Learning Rate," what do we actually get from this work? It’s about showing that GSAM can converge to a good approximation under two different parameter schedules >

Jane: The authors are focused on verifying whether the SAM algorithm with either an increasing batch size or decaying learning rate performs well in training deep neural networks, and they do that by focusing on GSAM, which retains the core idea of sharpness-aware optimization >

Lu: The main contribution is showing this epsilon approximation of the GSAM algorithm with an increasing batch size and constant learning rate, as detailed in Theorem three > <ref:2409.09984#pg2,approximation of the GSAM algorithm with an increasing batch size and constant>

Meng: And they also show it for a constant batch size and decaying learning rate through Theorem four which proves convergence under those conditions > <ref:2409.09984#pg3>

Lalam: The implication is that we have a clearer theoretical understanding of how to tune the training process—specifically the batch size or the learning rate—to ensure we find those flatter minima that lead to better generalization >

Tom: It’s about giving practitioners concrete schedules they can follow, like increasing batch sizes, to get reliable convergence for this specific optimization problem >

Jane: And while they show these results, they also point out some limitations. The authors flag a theoretical limitation where Theorem three and Theorem four only hold if rho is proportional to epsilon raised to the power of four and the batch size b is proportional to epsilon raised to the power of negative four > <ref:2409.09984#pg3>

Lu: They also mentioned that estimating constants like G, sigma, C, and Pi in n is difficult mathematically >

Meng: Empirically, they only looked at a limited number of models and didn't test the combined setting where both an increasing batch size and a decaying learning rate are used at the same time >

Lalam: So what this means for the future is that we need to be careful when applying these theorems; they have conditions on rho and b that we still need to satisfy in real-world training scenarios >

Conclusion: Tom: So we’re wrapping up this look at the paper, "Convergence of Sharpness-Aware Minimization Algorithms using Increasing Batch Size and Decaying Learning Rate." Basically, they show you how to make sure that sharpness-aware optimization methods actually converge under two different training schedules.

Jane: Right. They’re looking at how the batch size or the learning rate affects whether an algorithm like GSAM settles down and finds a good solution for those DNN problems we talked about earlier.

Lu: The authors are proving that you can use increasing batch sizes with a constant learning rate, or you can use a constant batch size with a decaying learning rate, and in both cases, you get this epsilon approximation of the original GSAM algorithm.

Meng: So they’re giving us concrete schedules for training. It sounds like it moves this from just theory into something we can actually implement on real hardware.

Lalam: I see how that looks from my side too. If we can tune these parameters to guarantee convergence, it means the optimization process becomes much more stable and predictable when finding those flatter minima in the loss landscape.

Tom: Exactly. The main point is that you don't just pick one schedule; there are two ways to tune it, and both work under certain conditions they laid out.

Jane: But there are some fine print attached to those results, like the specific mathematical relationships between the error epsilon and those batch size or learning rate parameters they need to hit for the proof to hold.

Lu: Yeah, I saw a few limitations they mentioned too. They can't just assume anything; there are these constraints on how small rho has to be relative to epsilon, and how big the batch size needs to be depending on what schedule you choose.

Meng: From an engineering standpoint, those constraints are important. It tells us we can’t just throw in a massive batch size or some tiny learning rate without respecting those scaling rules for it to actually work as claimed.

Lalam: And they also noted that they haven't tested the scenario where you use *both* an increasing batch size and a decaying learning rate at the same time, so we don't know how that specific combination behaves yet.

Tom: So, while this is solid proof for two separate scenarios, it opens up a lot of questions about combining these strategies for maximum efficiency.

Jane: It really shows that the choice of training schedule isn't just an arbitrary tweak; it’s something you have to match with the right mathematical conditions to get a reliable result.

Lu: The real potential here is figuring out how to combine those schedules—like mixing increasing batch sizes with decaying learning rates—to get the best balance between finding good solutions and keeping training costs down.

Meng: That sounds like where the practical impact lies. If we can figure out that optimal combination, it could drastically cut down the time needed to train these complex models.

Tom: So, while they’ve given us these two solid paths forward, the next big question is figuring out which path—or maybe a hybrid of both—is actually the most efficient way to get those high-quality results.

Meiji University

cs.LG, math.OC

Submitted: 2024-09-16

Updated: 2026-10-12

Comments: Accepted at ACML2026

Code: https://github.com/iiduka-researches/GSAM_batch_inc

License: http://creativecommons.org/licenses/by/4.0/

Importance score: 83/100

The gist: The gist The main contribution of this paper is to show an ϵ-approximation of the GSAM algorithm with an increasing batch size and constant learning rate ((7) in Table 1; Theorem 3) and with a

Key concepts

Sharpness-Aware Minimization (SAM): SAM is an optimization problem designed to minimize a perturbed empirical loss. It seeks a parameter $x$ that minimizes the maximum empirical loss over a small neighborhood defined by $\epsilon$. This helps find flatter minima in the DNN's loss landscape, which often leads to better generalization.
GSAM Algorithm: GSAM is an optimization algorithm for SAM that simplifies the process by decomposing the gradient into orthogonal and parallel components. It retains the core idea of sharpness-aware optimization while having a simpler structure, making it easier to apply in practice.
Increasing Batch Size: This strategy involves gradually increasing the batch size during training while keeping the learning rate constant. The analysis shows that this approach helps GSAM achieve an $\epsilon$-approximation of the SAM problem, balancing computational cost reduction with maintaining effectiveness.
Decaying Learning Rate: This technique uses a learning rate that decreases over time during optimization. When combined with a constant batch size, it guarantees convergence to an $\epsilon$-approximation of the SAM problem, showing that SGD can be more effective than standard SAM setups.

Terminology

Generalization — 1× in this paper · explained in 13 episodes

Summary

The gist The main contribution of this paper is to show an ϵ-approximation of the GSAM algorithm with an increasing batch size and constant learning rate ((7) in Table 1; Theorem 3) and with a constant batch size and decaying learning rate ((8) in Table 1; Theorem 4).

SAM Problem and GSAM

The sharpness-aware minimization (SAM) problem is to minimize a perturbed empirical loss defined as the maximum empirical loss fS,ρ(x):= max∥ϵ∥≤ρ fS(x + ϵ) over a certain neighborhood of a parameter x ∈ R d of the DNN, where ρ ≥ 0 and ϵ ∈ R d The GSAM algorithm retains the core idea of sharpness-aware optimization in SAM while having a relatively simple structure, and it introduces a decomposition of the gradient into orthogonal and parallel components. The goal is to solve Problem 2.1, which is to minimize ˆfSAM S,ρ (x):= max∥ϵ∥2≤ρ fS(x) + ⟨∇fS(x), ϵ⟩2 subject to x ∈ R d.

Convergence Analysis of GSAM

The convergence analyses of GSAM are based on the search direction noise ηtωt (defined by (9)) between GSAM and gradient descent (GD). The norm of the noise is approximately Θ(√ ηt bt). Since this implies that GSAM using a large batch size b or a small learning rate η behaves approximately the same as GD in solving the SAM problem, GSAM eventually needs to use a large batch size or a small learning rate.

Increasing Batch Size and Constant Learning Rate

When using an increasing batch size and constant learning rate, Theorem 3 shows that there exists t0 ∈ N such that, for all T ≥ t0, min t∈[0:T −1] E[∥∇ˆfSAM S,ρ (xt)∥2] ≤ ϵ. This indicates that the parameters α and ρ in (13) become small and thereby achieve an ϵ–approximation of GSAM.

Constant Batch Size and Decaying Learning Rate

For a constant batch size and decaying learning rate, Theorem 4 guarantees that there exists t0 ∈ N such that, for all T ≥ t0, min t∈[0:T −1] E[∥∇ˆfSAM S,ρ (xt)∥2] ≤ ϵ. This result implies that SGD using α = ρ = 0 would minimize ∥∇fS(xt)∥2 more quickly than SAM/GSAM using α ≠ 0 or ρ ≠ 0.

Numerical Results and Conclusion

The numerical results showed that SAM/GSAM with an increasing batch size achieves lower test errors and lower measured sharpness than those of SAM/GSAM with a constant batch size. This is because using increasing batch sizes reduces computational costs to train DNNs while maintaining effectiveness in terms of generalization. The paper concludes that using increasing batch sizes is useful from the viewpoints of both generalization and computational cost.

Limitations

Theoretical limitations include the requirement that ρ = O(ϵ4) and b = Ω(ϵ−4) for Theorem 3 and Theorem 4 to hold, as well as the difficulty in estimating constants like G, σ, C, and P i∈[n]. Empirical limitations include examining only a limited number of models and not empirically examining the combined setting in which an increasing batch size and a decaying learning rate are used simultaneously.

The paper is 2409.09984.

Improvements for AI systems

Batch size scheduling for Sharpness-Aware Minimization (SAM) ensures better generalization by avoiding sharp local minima of empirical loss, as shown in the numerical results showed that SAM/GSAM with an increasing batch size achieves a lower test error and lower measured sharpness than those of SAM/GSAM with a constant batch size.
Learning rate scheduling, such as cosine annealing or linear decay, combined with the GSAM algorithm allows for an ϵ-approximation of GSAM under specific conditions, enabling the system to converge effectively even when using a constant batch size.
The increasing batch size strategy is particularly beneficial because it starts with small batches, which maintains the effectiveness of SAM compared to constant batch sizes, while simultaneously reducing computational costs.
The convergence analysis provides bounds on the search direction noise, showing that E[ηt∥ωt∥ 2] decreases as the batch size bt increases, meaning larger batches lead to a more accurate approximation of Gradient Descent in solving the SAM problem.
By tuning hyperparameters based on joint conditions (Equation 45), AI systems can achieve an ϵ–approximation of GSAM, where ρ = O(ϵ 4), ensuring that the achieved precision is controllable and linked to the required generalization capability.

Abstract

The sharpness-aware minimization (SAM) algorithm and its variants, including gap guided SAM (GSAM), have been successful at improving the generalization capability of deep neural network models by finding flat local minima of the empirical loss in training. Meanwhile, it has been shown theoretically and practically that increasing the batch size or decaying the learning rate avoids sharp local minima of the empirical loss. In this paper, we consider the GSAM algorithm with increasing batch sizes or decaying learning rates, such as cosine annealing or linear learning rate, and theoretically show its convergence. Moreover, we numerically compare SAM (GSAM) with and without an increasing batch size and conclude that using an increasing batch size achieves a lower worst-case infinity adaptive sharpness than compared with using a constant batch size and learning rate.