Understanding the Generalization Benefit of Normalization Layers: Sharpness Reduction

Kaifeng Lyu; Zhiyuan Li; Sanjeev Arora

Source-linked AI summary

Understanding the Generalization Benefit of Normalization Layers: Sharpness Reduction

Kaifeng Lyu, Zhiyuan Li, Sanjeev Arora

arXiv:2206.07085v3Jun 14, 2022cs.LGcs.AI

TL;DR

The paper asks how normalization can improve generalization when scale invariance makes ordinary sharpness measures vulnerable to rescaling. It analyzes GD with weight decay using spherical sharpness and characterizes its behavior near the Edge of Stability. The analysis and experiments report a sharpness-reduction bias associated with generalization benefits, within stated theoretical scope.

Problem
For normalized nets, Hessian sharpness can be changed by rescaling, while GD may increase that measure toward the Edge of Stability, complicating analysis of how normalization improves generalization.
Source: §1 · p. 2
Method
The paper analyzes GD with weight decay as projected updates with an adaptive effective learning rate and studies its sharpness-reduction dynamics near minimizers.
Source: §1.1 · p. 3
Results
Normalization and weight decay induce a sharpness-reduction bias, which experiments verify alongside generalization benefits on CIFAR-10 and matrix completion.
Source: §1 · p. 2, §7 · p. 10
Takeaways & Limitations
In the studied settings, decreasing spherical sharpness accompanies improved generalization, including after the training set is fitted.
Source: §1 · p. 2, §7 · p. 10
The theory concerns dynamics around a minimizer manifold at a small but finite learning rate; extending it to SGD and explaining behavior before that regime remain future work.
Source: §7 · p. 10

Abstract

from arXiv · show

Normalization layers (e.g., Batch Normalization, Layer Normalization) were introduced to help with optimization difficulties in very deep nets, but they clearly also help generalization, even in not-so-deep nets. Motivated by the long-held belief that flatter minima lead to better generalization, this paper gives mathematical analysis and supporting experiments suggesting that normalization (together with accompanying weight-decay) encourages GD to reduce the sharpness of loss surface. Here "sharpness" is carefully defined given that the loss is scale-invariant, a known consequence of normalization. Specifically, for a fairly broad class of neural nets with normalization, our theory explains how GD with a finite learning rate enters the so-called Edge of Stability (EoS) regime, and characterizes the trajectory of GD in this regime via a continuous sharpness-reduction flow.

1 Introduction

The paper studies how normalization and weight decay shape generalization in scale-invariant networks, where conventional sharpness measures can be distorted by rescaling. It analyzes GD+WD dynamics and reports experiments linking decreasing spherical sharpness with improved generalization.

1 Introduction

For BN matrix completion with U,V ∈ R50×50 and no explicit rank constraint, starting from step ∼2k, spherical sharpness drops significantly, encouraging low-rank solutions as all-entry test MSE decreases from 1.12 to 0.013.Source: §1 · p. 2
On CIFAR-10, VGG-11 test accuracy rises from 69.1% to 72.0% as spherical sharpness decreases during roughly 47k post-fit steps, then reaches 84.3% after destabilization.Training accuracy reaches 100% after about 680 steps; removing either BN or weight decay eliminates the reported phenomenon.Source: Figure 2 · p. 2
The paper attributes an implicit sharpness-reduction bias to normalization interacting with weight decay during full-batch GD.Prior ablations report better generalization with weight decay in normalized nets, while scale-invariant loss is unchanged by positive rescaling.Source: §1 · pp. 1–2

1.1 Our Contributions

Spherical sharpness is the largest Hessian eigenvalue after projecting parameters onto the unit sphere, avoiding the rescaling vulnerability of the ordinary Hessian measure.For scale-invariant loss, the unprojected Hessian measure can be made arbitrarily small by rescaling a minimizer.Source: §1 · p. 2, §1.1 · p. 3
Near a minimizer manifold, GD+WD’s effective learning rate grows toward the EoS regime, where GD approximately follows a gradient flow that reduces spherical sharpness.The paper also reports a BN linear-regression result: GD+WD finds the minimizer corresponding to the minimum-weight-norm linear model through this sharpness-reduction mechanism.Source: §1.1 · p. 3
The analysis connects finite-step-size oscillations around the minimizer manifold to motion along it, without stochastic gradient noise or additional regularization.The proof connects EoS dynamics to power methods and models key parameters with a one-dimensional Hamiltonian system.Source: §1.1 · p. 3

2 Related Works

Related work connects flatness with generalization and studies several benefits of normalization, but often treats different mechanisms or omits weight decay. The paper focuses on the joint normalization–weight-decay dynamics of full-batch GD.

2 Related Works

Prior studies report positive correlations between flatness and generalization, but sharpness measures can be sensitive to parameter rescaling and asymmetric valleys.Sharpness-based measures performed best overall in one extensive comparison of generalization measures.Source: §2 · p. 3
Normalization has been studied for effects on signal propagation, representations, optimization, and effective learning rates, while this paper focuses on its interaction with weight decay.The cited prior work either does not analyze training dynamics or assumes no weight decay.Source: §2 · p. 4
Prior work describes norm and update-size equilibria under normalization with weight decay, but why such equilibria generalize well remains unclear.This paper studies full-batch GD, the basic special case of SGD.Source: §2 · p. 4

3 Preliminaries

The preliminaries formalize scale-invariant losses and express GD with weight decay as projected updates on the unit sphere. They define the effective learning rate used in that representation.

3 Preliminaries

A scale-invariant loss is unchanged by positive rescaling, and its gradient is perpendicular to the parameter vector and scales inversely with it.The loss is defined away from the zero vector.Source: §3 · p. 4
Projecting parameters onto the unit sphere rewrites GD+WD as projected gradient descent with an effective learning rate that depends on parameter norm.The update rule uses θ_t as the normalized parameter direction and defines the effective rate from the original learning rate, weight decay, and parameter norm.Source: §3 · p. 4
The paper’s effective learning rate differs slightly from definitions in earlier work [113] and is chosen for analytical convenience..Source: §3 · p. 4

4 GD+WD on Scale-Invariant Loss Functions

For scale-invariant losses, GD with weight decay can leave the stable regime and enter the Edge of Stability, where its trajectory approximately follows a flow that reduces spherical sharpness. The analysis explains this behavior through oscillation around a manifold of local minimizers and states conditions under which the result holds.

4.1 GD+WD Eventually Leaves the Stable Regime

Near a local minimizer, weight decay lowers the parameter norm and raises the effective learning rate until GD+WD enters the Edge of Stability regime.The stated entry condition is that the effective learning rate is approximately 2 divided by the local sharpness.Source: §4.1 · p. 5
The stable-regime analysis uses a local PL condition around a minimizer to establish eventual EoS entry; PL is weaker than convexity.The theorem’s entry result also requires smoothness and sufficiently small weight-decay–learning-rate product.Source: §4.1 · pp. 5–6

4.2 Dynamics at the Edge of Stability

In the Edge of Stability regime, GD oscillates around the minimizer manifold while moving persistently in a direction that reduces spherical sharpness.The oscillation is not fully perpendicular to the manifold, producing the drift; the claim is for a manifold consisting of local minimizers.Source: §4.2 · p. 6

4.2.1 Assumptions

The analysis assumes a smooth scale-invariant loss and a smooth manifold of local minimizers with a unique largest Hessian eigenvalue.The authors note that sharpness reduction also appears empirically when the top eigenvalue has multiplicity greater than one, but leave that case unanalyzed.Source: §4.2.1 · p. 6

4.2.2 Main Theorem

With sufficiently small intrinsic learning rate and the stated initialization and manifold assumptions, GD+WD approximately tracks a flow that reduces spherical sharpness on the minimizer manifold.The theorem gives tracking error O(η_in^1/4−o(1)) with probability 1−O(η_in^1/2−o(1)).Source: §4.2.2 · p. 7

4.2.3 Proof Idea

The proof links period-two oscillation along the top Hessian eigenvector to drift on the minimizer manifold, while periodic effective learning rates yield the flow’s time scale.The authors also extend the analysis to projected-gradient methods with Quasi-RMSprop schedulers and describe an extension to RMSprop with a scalar learning rate.Source: §4.2.3 · pp. 8–9

5 Case Study: Linear Regression with Batch Normalization

The section applies the theory to over-parameterized linear regression with Batch Normalization. It shows that GD+WD selects a sharpness-minimizing solution, while gradient flow can converge to any zero-loss solution under the described setting.

Setup

The case study analyzes full-batch GD+WD on over-parameterized linear regression with Batch Normalization and targets generated by an unknown linear model.The BN output is a linear function of the input, with fixed output mean and variance matched to the targets.Source: §5 · p. 9

Case study result

GD+WD converges to a solution on the unit sphere that minimizes sharpness, regardless of initialization.The associated coefficients are also optimal for the constrained optimization problem M.Source: §5 · p. 9

Comparison with gradient flow

Unlike gradient flow, which may converge to any zero-loss solution, GD exhibits sharpness-reduction bias toward the constrained problem’s optimal solution, even with arbitrarily small learning rate.The authors present this as a concrete example where GD can generalize better than gradient flow under natural settings.Source: §5 · pp. 9–10

6 Discussion

Experiments show that test performance can keep improving after training fit when spherical sharpness decreases, while ablations implicate both normalization and weight decay. The analysis relates this behavior to progressive sharpening and progress in the Edge of Stability regime.

Experimental Verification of Sharpness Reduction

After fitting the training set, GD continues to improve test accuracy as spherical sharpness decreases in matrix-completion and CIFAR-10 experiments.The experiments include additional matrix-completion data sizes and ResNet-20 on CIFAR-10.Source: §6 · p. 10

Ablation Studies on Normalization and Weight Decay

Without normalization, CIFAR-10 test accuracy is 56.8%; without weight decay, it stalls at 66.4%, versus 84.3% with weight decay.The no-normalization run also stops improving after fitting the training set; without weight decay, spherical sharpness and test accuracy stop changing at low loss.Source: §6 · p. 10

Explaining the Progressive Sharpening and EoS Phenomena

The analysis attributes progressive sharpening to norm change and provides a theoretical account of training progress in the Edge of Stability regime.The related EoS observations describe the top Hessian eigenvalue hovering near or above 2/learning rate while training loss fluctuates short-term but decreases long-term.Source: §6 · p. 10

7 Conclusions and Future Work

The paper reports that GD with normalization and weight decay can implicitly reduce spherical sharpness, with experiments verifying the predicted bias and its generalization benefit on CIFAR-10. The theory is limited to dynamics near a minimizer manifold at a small but finite learning rate, and several extensions remain open.

Conclusions

The paper exhibits GD's implicit bias toward reducing spherical sharpness with normalization and weight decay, and verifies its predicted generalization benefit on CIFAR-10.The experiments also verify the sharpness-reduction bias predicted by the theorem.Source: §7 · p. 10

Limitations and Future Work

The theory covers dynamics around a minimizer manifold and requires a small but finite learning rate; explaining reductions observed earlier and at moderate rates remains future work.The paper also leaves extensions to SGD, task-specific effects, and architectures not scale-invariant on all parameters for future study.Source: §7 · p. 10
Future directions include extending the analysis to SGD, explaining task-specific effects such as low-rank matrix completion, and studying partially scale-invariant architectures.The paper notes that gradient noise may change the relevant sharpness measure for SGD.Source: §7 · p. 10

A Additional Related Works

The related-work discussion contrasts scale-aware sharpness measures and implicit-bias analyses, emphasizing that the paper studies discrete GD dynamics for scale-invariant loss with weight decay. It also positions its theory against prior work on EoS and minimizer-manifold flows.

Sharpness Measures and Parameter Rescaling

The top Hessian eigenvalue is sensitive to weight rescaling in normalized nets, whereas spherical sharpness is scale-invariant; whether GD or SGD reduces other symmetry-aware measures remains unclear.The paper identifies spherical sharpness as the only measure known to it that provably decreases during training normalized nets.Source: A Additional Related Works · p. 23

SGD Noise Helps to Escape Sharp Minima

Analyses of SGD escaping sharp minima include noise-based results under quadratic-loss and continuous-SDE assumptions [127, 121, 54], a one-point-convexity result, and stability analyses [117].These prior escaping analyses do not establish the paper's GD sharpness-reduction mechanism.Source: A Additional Related Works · p. 23

Implicit Bias of GD

Prior implicit-bias work often studies gradient flow or regimes where GD learns the same solutions as gradient flow, while this paper belongs to work relying fundamentally on discrete GD dynamics.Examples of discrete-dynamics results include gradient-norm regularization, multiscale losses, and balanced matrix norms.Source: A Additional Related Works · p. 23

Comparison with Arora et al.

Arora et al. derive a similar minimizer-manifold flow that decreases Hessian sharpness, but their setting and technique do not apply to scale-invariant loss with weight decay.Their analysis requires injected stochastic noise, whereas this paper analyzes GD directly and uses a one-dimensional Hamiltonian system in its proof.Source: A Additional Related Works · p. 23

Edge of Stability

Prior EoS work includes empirical observations, experimental and theoretical analyses [18], and a minimizer-manifold analysis for other loss settings.The paper presents its analysis as an explanation for EoS in the broader setting of scale-invariant loss with L2 regularization.Source: A Additional Related Works · pp. 23–24

B A General Theory for a Broader Class of Adaptive Gradient Methods

This appendix recasts GD with weight decay on scale-invariant loss as projected GD with an adaptive effective learning rate, then proves sharpness-reduction-flow results for a broader class of gradient methods. The theory covers spherical and full-space settings, including Scalar RMSprop as an example.

B A General Theory for a Broader Class of Adaptive Gradient Methods

The general theory treats GD with weight decay on scale-invariant loss as an instance of projected gradient descent with a quasi-RMSprop scheduler, enabling a broader theorem.The main theorem for GD+WD is a corollary of the scheduler-level result.Source: B A General Theory for a Broader Class… · p. 25

B.1 Scalar RMSprop and Quasi-RMSprop Scheduler

Quasi-RMSprop schedulers are gradient-based rules for setting effective learning rates with RMSprop-like updates, allowing negligible perturbations as β approaches 1.The appendix motivates the class through Scalar RMSprop and a single scalar learning rate.Source: §B.1 · pp. 25–26

B.2 Reformulation of GD+WD on Scale-Invariant Loss via Quasi-RMSprop Scheduler

For scale-invariant loss, GD+WD is reformulated as projected GD on the unit sphere, and its effective learning rate evolves according to the GWSI scheduler.The effective-rate recursion is shown to match GWSI exactly, which is itself a quasi-RMSprop scheduler.Source: §B.2 · p. 27

B.3 Main Results for GD/PGD with Quasi-RMSprop Scheduler

For small learning rates and an initial point near a local-minimizer manifold in EoS, projected GD with quasi-RMSprop rates approximately tracks the sharpness-reduction flow.The theorem also extends to GD with quasi-RMSprop scheduling.Source: §B.3 · p. 27

B.3.1 Spherical Optimization

On the sphere, initialized projected GD in EoS approximately tracks the sharpness-reduction flow from a local minimizer, under the stated smoothness and initialization conditions.The theorem gives a high-probability trajectory approximation for the stipulated initialization scheme.Source: §B.3.1 · p. 28

B.3.2 Full Space Optimization

In full space, GD with quasi-RMSprop scheduling likewise approximately tracks the flow; Scalar RMSprop is a corollary when its decay rate is set to β = 1 − Cbη^2 + O(η^2).The theorem assumes a local-minimizer manifold and specified regularity and initialization conditions.Source: §B.3.2 · p. 29

C PAC-Bayes Bounds Based on Spherical Sharpness

The appendix gives a PAC-Bayes generalization bound using spherical sharpness, then establishes properties of scale-invariant losses and local PL conditions that support its optimization analysis. It also gives convergence results for gradient flow and relates effective learning rates to scheduler formulations.

C PAC-Bayes Bounds Based on Spherical Sharpness

The PAC-Bayes result states that small spherical sharpness leads to small generalization error under its stated conditions.The setup assumes bounded loss on the sphere and includes a Gaussian perturbation condition for the stated implication.Source: C PAC-Bayes Bounds Based on Spherical… · p. 30

D.2 Scale-Invariant Functions

Scale-invariant losses have gradients perpendicular to the parameter and gradients and Hessians that scale with degrees −1 and −2, respectively.These properties are used in the appendix’s analysis of optimization on the sphere.Source: §D.2 · p. 32

D.3.1 Full Space Optimization

Near a compact subset of a smooth manifold of local minimizers, constant Hessian rank yields a neighborhood where the loss satisfies a PL condition.The theorem assumes a C3-smooth loss and a C1-smooth minimizer manifold.Source: §D.3.1 · p. 33

D.3 Polyak-Łojasiewicz Condition

Under a local PL condition, gradient flow from a sufficiently close starting point converges to a nearby point, both in full space and on the sphere.The two convergence results bound the final distance from the starting minimizer by a quantity of the same order as the initial distance.Source: §D.3.1 · p. 33, §D.3.2 · p. 34

D.3.2 Spherical Optimization

The spherical counterpart establishes a local PL condition near compact subsets of a smooth manifold of spherical minimizers under a constant-rank assumption.The theorem places the neighborhood on the unit sphere.Source: §D.3.2 · p. 34

E.1 Proof for Theorem B.7

The appendix relates scale-invariant GD+WD effective learning rates to GWSI scheduling and shows that quasi-RMSprop can match them under specified settings.The matching uses a particular relation between the quasi-RMSprop hyperparameters and the GD+WD learning rate and weight decay.Source: §E.1 · p. 34, §E.2 · p. 35

F Details of the 3D Example

The 3D example constructs a scale-invariant loss with a manifold of minimizers and illustrates GD+WD moving along it toward a flatter point. The accompanying proof develops part of the argument that GD eventually enters the EoS regime.

F Details of the 3D Example

With learning rate 0.5 and weight decay 0.08, the normalized iterate continues moving after approaching its first minimizer and eventually oscillates around another point.The example starts from w0 = (0.3, 1.3, 1.2).Source: F Details of the 3D Example · p. 35
On the constructed minimizer manifold, spherical sharpness is controlled by |z|, and GD+WD moves toward the flattest point, where z is zero.The example identifies the minimizer manifold by x = y > 0 on the sphere and uses an orthogonal transformation for the plotted loss.Source: F Details of the 3D Example · pp. 35–36

G.2 Proof for Theorem 4.4: GD Eventually Enters the EoS Regime

In the EoS-regime proof, the loss decreases by a constant factor per step during a phase lasting at least Ω(1/(η̂λ̂)) steps.The passage then states that the loss reaches a value close to the minimum before the effective learning rate continues growing.Source: §G.2 · p. 37

G.3 Connection to the EoS Regime in Cohen et al.’s Definition

This section connects the EoS regime to sharpness reduction by extending Cohen et al.’s definition to PGD and carrying the spherical-case analysis through working-zone states, initialization, alignment, and drift. Its proof outlines relate long-term loss decrease to sharpness reduction through RMS-drift processes that track parameter projections along the flow, with alignment and drift statements extending to the full space under stated conditions.

G.3 Connection to the EoS Regime in Cohen et al.’s Definition

For PGD, EoS is defined by an effective learning rate near 2 divided by the local upper bound on spherical sharpness, capturing Cohen et al.’s first condition.The paper says the second condition—long-term loss decrease—also follows by repeating its argument for this setting.Source: §G.3 · p. 38
When effective learning rate is below the threshold, Theorem 4.4 justifies progressive sharpening; weight decay lowers the norm at small gradients, raising top-eigenvalue sharpness.The paper attributes this mechanism to the interplay between normalization and weight decay.Source: §G.3 · p. 38

H Proof Outlines of Our Theorems on Sharpness Reduction

The proof outline shows that two-step oscillations move parameter projections approximately along the manifold gradient of log top sharpness, and models those oscillations as RMS-drift processes.The process tracks displacement from the manifold and closeness to the stability edge to analyze drift speed.Source: H Proof Outlines of Our Theorems on… · p. 39

H.2 RMS-drift Process: Introduction

After warm-up, GD/PGD dynamics with a quasi-RMSprop scheduler can be represented by an RMS-drift process using a gradient-flow projection and two hidden variables.The projection evolves as gradient descent on log top sharpness with a changing learning rate; averaging that rate yields the flow approximation.Source: §H.2 · p. 40

H.3.1 Construction of Working Zones

For full-space optimization, nested working zones around the sharpness-reduction path keep parameters near the manifold and ensure their gradient-flow projections are well-defined.The working-zone construction also places those projections in a specified neighborhood on the manifold.Source: §H.3.1 · p. 41

H.3.3 Alignment Phase

After initialization, alignment to the top Hessian eigenvector reaches at most O(|h_t|η)-misalignment within η^-o(1) steps, while boundedness and proximity to the starting projection are maintained.The stated alignment theorem assumes the initial state satisfies boundedness, deviation, misalignment, and projection-distance conditions.Source: §H.3.3 · p. 42

H.3.4 Drifting Phase

In the full-space drifting phase, an α-bounded state with at most O(|h_t|η)-misalignment yields an O(1)-RMS-drift transition over two steps.This is the condition under which the hidden variables’ transition is treated as an RMS-drift transition.Source: §H.3.4 · p. 42

H.4.1 Construction of Working Zones

For spherical optimization, the reduction follows the same basic logic as the full-space case, with a corresponding working-zone lemma.The section presents only the new lemma and theorem statements for this case.Source: §H.4 · p. 42, §H.4.1 · p. 43
Working-zone analysis defines α-bounded and α-deviated states using the full-space definitions with the working zone changed.An α-bounded state has θ_t in Z_ϵ0 ∩ S_D−1, ∥x_t∥₂ ≤ αη, and |u_t| ≤ α; an α-deviated state violates the stated distance or zone condition.Source: §H.4.1 · p. 43

H.4.2 Good Initialization

Lemma H.17 states that the good-initialization result of Lemma H.8 also holds for the spherical case.Source: §H.4.2 · p. 43

H.4.3 Alignment Phase

For sufficiently small base learning rate η, α-bounded states with 1 ≤ α ≤ η^−o(1) remain in the working zone, and relations (17)–(21) continue to hold in the spherical case.Theorem H.19 also states that Theorem H.10’s alignment result holds for the spherical case.Source: §H.4.3 · p. 43

H.4.4 Drifting Phase

For sufficiently small base learning rate η, α-bounded states that are at most O(|h_t|η)-misaligned remain in the working zone after two steps, and relations (22)–(25) continue to hold in the spherical case.Source: §H.4.4 · p. 43

H.5 RMS-drift Process: Analysis

This section analyzes RMS-drift processes by establishing energy, initialization, alignment, and flow-approximation bounds under stated conditions. It develops two-step and segment-wise trajectory estimates, including extensions to spherical optimization.

H.5 RMS-drift Process: Analysis

For a process in the working zone with E(S0) ≤ α², energy remains O(α²) at every even step over M = O(1/η²) steps.The energy conservation theorem applies when 1 ≤ α ≤ η^-o(1).Source: §H.5 · p. 44
If initialized within O(α²η^1/2) of the sharpness-reduction flow and O(α)-bounded, the process stays within O(α²η^1/2) at every even step.The tracking result assumes the process is in the working zone and compares φt with ζ(tη²).Source: §H.5 · p. 44

H.6 Finalizing Proofs

The final proof uses induction over even and odd steps to maintain bounded energy, limited misalignment, and flow tracking; the spherical proof invokes corresponding spherical results.The induction proceeds to T/η².Source: §H.6 · p. 44

I.1 Construction of Working Zones

Working zones are chosen so Φ is well-defined and C3-smooth, while nearby gradient flows move O(ε′0) and converge within Zε1.The construction also uses compactness to establish the final stated property of the zone.Source: §I.1 · p. 44, §I.2 · p. 44

I.2 Gradient Flow Projection

Projection lemmas show vanishing second derivatives in specified normal directions and bound the projected displacement by O(∥θ − φ∥²) in full-space and spherical settings.The spherical second-derivative result additionally requires ⟨x, φ⟩ = 0.Source: §I.2.1 · p. 44, §I.2.2 · p. 45

J.1 Full Space Optimization; J.2 Spherical Optimization

Under near-minimizer and effective-learning-rate conditions, full-space and spherical lemmas derive zeroth-, first-, and second-order approximations for one-step updates.Both settings also provide two-step update lemmas.Source: §J.1 · pp. 45–46, §J.2 · pp. 47–48

K Lemmas for Quasi-RMSprop Schedulers

The quasi-RMSprop lemmas analyze one- and two-step moment estimates when gradients and the initial moment deviation are small.They use a scheduler with β = 1 − Cbη² + O(η⁴), where Cb = Θ(1).Source: K Lemmas for Quasi-RMSprop Schedulers · p. 50

L.1 Good Initialization

The good-initialization proof uses Gaussian concentration and anti-concentration to establish its initial high-probability bounds.The supplied passage gives probability 1 − δ but does not expose the resulting bounds.Source: §L.1 · p. 51
Good initialization yields |h0|η ≥ Ω(δr) and an initial state that is O(1/δ)-misaligned.The initialization proof chooses log(1/η) with a sufficiently large constant C0.Source: §L.1 · p. 51

L.2 Alignment Phase

During alignment, the iterates remain O(αmax)-bounded, xt is O(h0η), and ht stays within |h0|/4 of its alternating-sign reference.The analysis also bounds the change in φt by O(αmaxη).Source: §L.2 · pp. 52–53

L.3 Drifting Phase

Under η^-o(1)-boundedness and O(|ht|η)-misalignment, the drifting-phase analysis establishes two-step approximations for the state, ht, scheduler, and φt.These approximations support the analysis of the RMS-drift trajectory.Source: §L.3 · pp. 53, 55, 56

M Reduction to RMS-drift Process: The Case of Spherical Optimization

For spherical optimization, projection changes the random-initialization argument by only O(r²), while the alignment and drifting lemmas retain the corresponding full-space statements.The spherical alignment proof uses the unchanged linear approximation of the update rule; the drifting lemmas use the spherical versions of the supporting results.Source: §M.1 · p. 57, §M.2 · p. 57, §M.3 · p. 57

N.1 Conservation of Energy

For an O(1)-RMS-drift process in the working zone, E(St) = E(S0) + O(α²η^0.5) through M = O(1/η^1.5) transitions when E(S0) ≤ α².The bound applies at every even step and assumes 1 ≤ α ≤ η^-o(1).Source: §N.1 · p. 59

N.2 Flow Approximation

For M = Θ(1/η^1.5), the flow approximation has error O(α²η^1/2) when the even-step states are α-bounded.The result is stated for the setting of Theorem H.23.Source: §N.2 · pp. 59–60
Grouping M transitions into O(1/η^1/2) segments of length O(1/η^1.5), Lemma N.4 and induction yield ∥φ_t − ζ(tη^2)∥_2 ≤ O(α^2η^1/2) for even 0 ≤ t ≤ 2M.Source: §N.2 · p. 60

O Proofs for Linear Regression with Batch Normalization

The linear-regression proof characterizes the limit of the sharpness-reduction flow, while experiments examine its predicted dynamics and generalization patterns across matrix completion and CIFAR-10, including normalization and weight-decay ablations.

O Proofs for Linear Regression with Batch Normalization

The sharpness-reduction flow converges to the least-squares solution on mean-shifted data; the proof identifies this as the flow's optimal solution at convergence.The theorem states that the flow converges to a point minimizing spherical sharpness.Source: O Proofs for Linear Regression with… · p. 62

P Experiments

The experiments test sharpness reduction and generalization in matrix completion and CIFAR-10, then use ablations to examine normalization, weight decay, and initial effective learning rate.The experiments also examine periodic evolution of oscillation magnitude and effective learning rate in EoS.Source: P Experiments · p. 63

P.1 Validation of Sharpness Reduction on Matrix Completion

In overparameterized matrix completion, entering EoS coincides with decreasing spherical sharpness, lower test loss, and an emerging low-rank solution despite no explicit rank constraint.The gap between the second- and third-largest singular values widens as sharpness and test loss decrease.Source: Figure 7 · p. 63, §P.1 · pp. 64–65

P.2 Validation of Sharpness Reduction on CIFAR-10

Spherical sharpness falls as test accuracy rises in smooth VGG-11 and non-smooth scale-invariant ResNet-20; the same pattern appears in partially scale-invariant ResNet-20.For the standard ResNet-20, sharpness is evaluated only on its scale-invariant parameters.Source: §P.2 · p. 65, Figure 9 · p. 66, Figure 10 · p. 66

P.3 Periodic Behaviors

In EoS, loss fluctuations validate periodic behavior in matrix completion, while CIFAR-10 retains periodic fluctuations but often lacks the threshold switching seen in the theory's alternating sub-regimes.The authors attribute the CIFAR-10 difference to less-smooth loss and practical learning-rate and weight-decay settings.Source: §P.3 · p. 67

P.4 Ablation Study: Normalization

With normalization, test accuracy can keep improving after training accuracy reaches 100%; removing normalization eliminates that post-fit increase in the VGG-11 experiment.The paper says spherical sharpness is meaningful only for normalized nets.Source: §P.4 · p. 68

P.5 Ablation Study: Weight Decay

With ˆλ=5 × 10−4, test accuracy rises from 69.1% to 84.3%; without weight decay, sharpness stops falling and test accuracy is 66.4%.Smaller non-zero weight decay still produces sharpness reduction, but more slowly.Source: Figure 15 · p. 69, §P.5 · p. 70

P.6 Ablation Study: Initial Effective Learning Rate

Across the tested initial effective learning rates, spherical sharpness decreases during the Edge of Stability (EoS) regime, and the runs finish at nearly the same sharpness and test accuracy. The ablation varies the learning rate and weight decay inversely while holding the intrinsic learning rate fixed.

P.6 Ablation Study: Initial Effective Learning Rate

With ˆλ = 5 × 10−4, initial spherical sharpness is much larger than 2/˜η0; after a few steps, it decreases to the same level as 2/˜ηt, bringing the dynamics into the EoS regime.Source: §P.6 · p. 70
The ablation changes the learning rate and inversely rescales weight decay, keeping the intrinsic learning rate fixed while testing different initial effective learning rates.Experiments use scale-invariant VGG-11 on CIFAR-10-2k.Source: §P.6 · p. 70, Figure 17 · p. 71
All runs finish at nearly the same spherical sharpness and test accuracy, although their EoS entry times differ with initial effective learning rate.For the smallest initial effective learning rate, 10^-5, EoS begins only after training accuracy reaches 100%; larger values enter quickly.Source: §P.6 · p. 71
During EoS, spherical sharpness persistently decreases while test accuracy increases across the tested runs.Source: §P.6 · p. 71

Q Experiment Details

The longest full-batch CIFAR-10 training experiments each took seven days on NVIDIA RTX A5000 GPUs.Source: Q Experiment Details · p. 71

Q.1 Additional Details of CIFAR-10 Experiments

The experimental appendix specifies the CIFAR-10 architectures and training setup, linear-regression data and hyperparameters, and procedures for computing spherical sharpness. It also explains how the final Batch Normalization affine parameters are chosen for regression and classification.

Q.1 Additional Details of CIFAR-10 Experiments

CIFAR-10 experiments use no data augmentation and ghost batch normalization in 50 batches of size 1000; CIFAR-10-2k takes the first 200 images per class and uses full-batch BN.The tested architectures are scale-invariant VGG-11, standard ResNet-20, and scale-invariant ResNet-20; the appendix describes their architectural changes.Source: §Q.1 · pp. 71–72, Figure 18 · p. 72

Q.2 Additional Details of Linear Regression Experiments

Linear-regression data use dimension 40, 20 Gaussian samples, and a uniformly sampled unit-norm ground-truth weight; Figure 5 uses learning rate 0.5 and weight decay 2 × 10^-4 to visualize periodic behavior.Source: §Q.2 · p. 74

Q.3 Computing Spherical Sharpness

For scale-invariant loss, spherical sharpness is computed as ||w||^2 times the largest Hessian eigenvalue at w, found using Lanczos with Hessian-vector products.For full-batch CIFAR-10, Hessian-vector products use only the first 5000 training examples; CIFAR-10 and matrix-completion sharpness is measured at current parameters, while linear-regression sharpness is measured at computed minimizers.Source: §Q.3 · p. 74

R Discussion on the Affine Parameters of the Final BN

The experiments add a final Batch Normalization layer with fixed affine parameters; for squared loss, its output mean and variance are set by β1 and γ1^2, which should match the regression targets’ mean and variance.For multiple output units, the corresponding βk and γk^2 are matched to each target coordinate’s mean and variance.Source: R Discussion on the Affine Parameters… · p. 74, §R.1 · p. 74
The final Batch Normalization layer normalizes pre-layer outputs using training-set means and variances before applying its affine parameters.Source: R Discussion on the Affine Parameters… · p. 74

R.2 Crossentropy Loss

The paper relates fixed final Batch Normalization parameters in classification to label smoothing and examines how changing γ affects training dynamics. The sharpness-reduction bias persists across tested γ values, while larger γ makes macroscopic periodic behavior more likely.

R.2 Crossentropy Loss

Fixing final-layer γ and β to γ and zero, respectively, encourages output probabilities matching soft labels equivalent to label smoothing with ϵ = 0.1 when γ ≈ 1.353.Source: §R.2 · p. 75
With fixed affine parameters, the minimum loss is non-zero; for γ = 3/10 ln 91, the stated minimum is approximately 0.0943.Source: §R.2 · p. 75
Across tested final-BN γ values, sharpness-reduction bias remains, but larger γ makes periodic behavior more likely at a macroscopic scale.Source: §R.2 · p. 75, Figure 20 · p. 76