PaperlayerView exact revision on arXiv

Source-linked AI summary

A Theoretical Analysis of Generalization Dynamics in Neural Networks under Gradient Descent with Weight Decay

Yuqing Wang, Ioannis G. Kevrekidis, Mikhail Belkin

arXiv:2609.07755v1cs.LGmath.OC

TL;DR

Generalization theory must account for data, architecture, and training dynamics, including how performance changes across layers and over time. This paper analyzes gradient descent with weight decay using a partition-based error decomposition and bounds on prediction variation. It establishes convergence to a neighbourhood of empirical-loss minimizers and conditions for layerwise differences and delayed generalization, including a characterization of grokking.

  • Problem

    How data, architecture, and training dynamics jointly determine generalization over a broad neural-network trajectory, including layerwise behavior, remains unresolved by analyses tied to specialized settings.

    Source: §1 · pp. 3–4

  • Method

    The paper studies gradient descent with weight decay and decomposes population error by input partition into data, optimization, and prediction-variation terms, bounding prediction variation through local approximate homogeneity.

    Source: §1 · p. 4

  • Results

    Gradient descent enters a neighbourhood of global empirical-loss minimizers, and the derived conditions explain layerwise generalization differences and characterize delayed generalization, including grokking.

    Source: §1 · p. 4

  • Takeaways & Limitations

    The framework links population-error components to data, architecture, and optimization dynamics, while its conditions describe qualitative mechanisms rather than precise predictions for a given task or architecture.

    Source: §4 · p. 21

    The connections to loss reweighting, layerwise learning behavior, and grokking are qualitative mechanisms and conditions, not precise predictions for a given task or architecture.

    Source: §4 · p. 21

Abstract

from arXiv · show

Understanding generalization remains a central challenge in machine learning because it requires jointly considering data, architecture, and training dynamics. In this paper, we develop a theoretical framework that characterizes how these factors jointly shape generalization performance throughout training. More precisely, we study a broad class of neural networks trained under the $\ell^2$ loss by gradient descent (GD) with weight decay, and prove the convergence of GD to a neighbourhood of the global minimizers of the empirical loss. By partitioning the space based on the input data, we then decompose the population error into data error, optimization error, and prediction variation error, and bound them separately. In particular, for the prediction variation error, which measures the oscillations of the learned function, we propose (local) approximate homogeneity and derive explicit cellwise and layerwise bounds for its evolution along the training trajectory. These bounds yield two important implications: a necessary condition of improved generalization explains differences in layerwise generalization behavior; a sufficient condition describes delayed generalization and provides a theoretical characterization of grokking.

1 Introduction

The introduction asks how data, architecture, and training dynamics jointly shape generalization, especially across layers and over time. It presents a framework that decomposes population error and derives conditions for layerwise differences and delayed generalization, including grokking.

1 Introduction

  • The framework addresses the gap left by analyses tied to specialized architectures, data models, or limiting regimes: characterizing population-error dynamics along broad neural-network training trajectories.It also targets how generalization varies across layers.Source: §1 · p. 3
  • Gradient descent with weight decay enters a neighbourhood of global minimizers of empirical loss, with empirical loss and its gradient explicitly controlled by the learning rate.The regularized objective also decreases under mild assumptions.Source: §1 · p. 4
  • The population error is decomposed into data, optimization, and prediction-variation errors; the data bound reflects label noise and within-cell approximation, while the optimization bound depends on cell masses.The cell-mass dependence identifies when loss reweighting can remove the imbalance between empirical and population weights.Source: §1 · p. 4
  • Prediction variation captures architecture- and training-dynamics effects, with its evolution bounded using local approximate homogeneity, empirical-loss gradients, and the regularized objective.Cellwise and layerwise variations further characterize the dynamics of individual layers.Source: §1 · p. 4
  • A necessary condition for relative decay of the layerwise bound helps explain why layers may begin generalizing at different times and learn target components of different complexity.The motivation notes that shallower layers may stabilize earlier while deeper layers evolve toward more complex structures.Source: §1 · pp. 3–4
  • Delayed generalization occurs when optimization error becomes small before prediction variation has had time to decay; the sufficient condition also characterizes a grokking mechanism.That mechanism involves small approximate-homogeneity error, empirical gradient, and regularized objective near global minimizers, plus sufficient time in that region.Source: §1 · p. 4

2 Optimization Theory

The section adapts a gradient-descent convergence framework to weight decay, then uses geometric conditions to connect regularized-objective decay to a neighborhood of empirical-loss global minimizers.

2.1 Preliminaries from general GD convergence

  • The adapted convergence framework combines trajectory smoothness from dissipativity and poly-smoothness with an iterate-dependent PL-type inequality from analyticity and a generically nondegenerate network Jacobian.The underlying framework does not by itself ensure convergence to a global minimizer.Source: §2.1 · p. 6

2.2 General loss decay of the regularized objective

  • Theorem 1 establishes decay of the regularized objective under its stated assumptions, but that decay alone does not guarantee that the empirical loss decreases.The empirical-loss guarantee is established separately in Subsection 2.3.Source: §2.2 · pp. 7–8

2.2.1 Transferring dissipativity from L to Lλ

  • The analysis transfers dissipativity from empirical loss L to Lλ, using weight-decay drift to control the trajectory in the small-gradient region when ∥θ∥ ≥ 2ϵL/λ.Theorem 1 assumes dissipativity of L rather than Lλ, keeping the assumption tied to the data and model landscape.Source: §2.2.1 · pp. 8–9

2.3 Convergence to a neighbourhood of global minimizers of empirical loss

  • A separation condition on the global-minimizer set forces a path from initialization toward the small-norm region to enter a neighborhood of empirical-loss global minimizers.The condition connects decay of Lλ to control of L, since Lλ may decrease through its quadratic regularization term instead.Source: §2.3 · p. 9
  • Under Theorem 2’s assumptions, GD reaches a neighborhood of the empirical-loss global minimizers and, during a specified interval, has L(θk) = O(η²) and ∥∇L(θk)∥ = O(η).For fixed η, the result does not guarantee an arbitrarily small neighborhood without additional training strategies.Source: §2.3 · p. 10

3 Generalization Theory

This section analyzes generalization by decomposing population error into data, optimization, and prediction-variation terms over training-data-defined cells. It bounds prediction-variation dynamics and uses cellwise and layerwise analysis to characterize depth-dependent decay and delayed generalization.

3 Generalization Theory

  • Population error is decomposed into data error, optimization error, and within-cell prediction variation error using each training input as its cell representative.The cell masses weight local errors, so the same local error contributes more on high-mass cells and less on low-density cells.Source: §3 · pp. 10, 12

3.1 Data error

  • The data error separates Gaussian label-noise contributions from deterministic within-cell approximation of the ground truth, creating a tradeoff between balanced and geometry-adapted cells.Balanced cell masses minimize the noise contribution, while geometry-adapted cells can reduce deterministic error; the partition also affects optimization and prediction variation errors.Source: §3.1 · p. 13

3.2 Optimization error

  • The optimization-error bound converts the GD training guarantee into a population-weighted residual bound, with the largest cell mass controlling the conversion from empirical loss.The cell-mass factor depends on the partition and population distribution, not on training.Source: §3.2 · p. 13

3.2.1 Optimization error and loss reweighting

  • Reweighting empirical loss by population cell masses directly controls the weighted residuals and avoids the cell-imbalance factor when reduced to the standard empirical-loss scale.Balanced cells align empirical and population weights, but may increase deterministic data error, so balancing need not minimize total population error.Source: §3.2.1 · p. 14

3.3 Prediction variation error

  • Prediction variation is the key remaining generalization term, measuring within-cell changes in network predictions beyond data and optimization errors.The analysis uses (local) approximate homogeneity to bound its evolution and derives conditions for layerwise decay and delayed generalization, including grokking.Source: §3.3 · p. 14

3.3.1 Approximate homogeneity

  • Approximate homogeneity measures a nonlinear block’s deviation from Euler’s scaling relation, while its local form allows the effective degree to vary across bounded regions.The approximation error quantifies the residual; smaller parameter and input region diameters make the condition more local.Source: §3.3.1 · pp. 15–16

3.3.2 Dynamics of the prediction variation error

  • Under architectural and approximate-homogeneity assumptions, prediction variation contracts along GD up to residual contributions from homogeneity error, empirical gradients, and the regularized objective.A layerwise decomposition shows that controlling each layer’s variation suffices to control the total prediction variation error.Source: §3.3 · p. 14, §3.3.2 · p. 17
  • Weight decay contracts prediction variation, while approximate homogeneity, the regularized objective, and the empirical loss gradient contribute deviations from exact decay.The latter terms are small when approximate-homogeneity error is low, the regularized objective has decreased, and the iterates are near a stationary point; prediction variation and empirical loss need not decrease monotonically.Source: §3.3.2 · pp. 17–18

3.3.3 Layerwise and cellwise characterization of prediction variation error

  • Cellwise and layerwise variation contract at rates set by local homogeneity degrees and weight decay, with residual terms from approximate homogeneity, the empirical gradient, and the regularized objective.The cellwise quantities aggregate into layerwise variation using population cell masses.Source: §3.3.3 · pp. 18–19
  • Bias and variance components share a geometric contraction rate, but the approximate-homogeneity remainder for bias is twice the variance remainder, allowing greater residual bias.Inner weights contribute directly to approximate-homogeneity error, whereas outer weights do not; earlier outer weights can still influence later-layer variation through hidden nodes.Source: §3.3.3 · p. 19

3.3.4 Necessary condition for generalization: layerwise variation thresholds

  • A guaranteed relative decay requires initial layerwise variation to exceed the depth-scaled approximate-homogeneity error, so deeper layers may require larger initial variation and decay later.This necessary condition does not imply that every trained network has strictly ordered layerwise errors.Source: §3.3.4 · pp. 19–20

3.3.5 Sufficient condition for generalization: grokking

  • Layerwise variation decreases by a prescribed factor when approximate-homogeneity error, empirical-gradient and regularized-objective contributions leave a positive decay margin, and training remains in the global-minimizer neighbourhood long enough.When the condition holds for every layer and the pre-output variation, overall prediction variation decreases, tightening the generalization bound once data and optimization errors are controlled.Source: §3.3.5 · pp. 20–21
  • Grokking occurs when training and regularized objectives are already small but layerwise oscillations need a longer residence in that region to contract, delaying generalization.The sufficient condition identifies small approximate-homogeneity error, small empirical gradient and regularized objective, and sufficiently long residence time as mechanisms for delayed improvement.Source: §3.3.5 · p. 21

4 Discussions

The framework separates how data, architecture, and optimization shape generalization, but its conclusions are qualitative rather than precise predictions for a particular task or architecture.

4 Discussions

  • The framework's generality identifies how data partitioning, layerwise and blockwise architecture, and optimization dynamics enter population error, while limiting task-specific precision.The discussion proposes specializing the partition, architectural assumptions, and local approximate homogeneity bounds to obtain sharper estimates.Source: §4 · p. 21

A.1 Dissipative condition on the regularized objective

  • Under Assumption 2 and 1 − 2ρλ > 0, the regularized objective satisfies a dissipative condition.The appendix combines the preceding lemmas to establish dissipativity within a ball of radius Rλ.Source: §A.1 · pp. 22–23

A.2.2 Lower bound of gradient

  • Under Assumptions 1–3 and 0 < λ < ρ/2, Lemma 10’s lower-bound result covers θ ∈ B₀(Rλ) when either ∥∇L(θ)∥ > ϵL or (λ/2)∥θ∥² ≥ 2ϵL²/λ.Source: §A.2.2 · pp. 27–28

A.3 Convergence to a neighbourhood of global minimum

  • The convergence proof establishes that GD enters a neighbourhood of the global minimizers and remains there over an interval of iterates.During that interval, L(θk) = O(η^2) and ∥∇L(θk)∥ = O(η).Source: §A.3 · pp. 28, 30

A.3.1 Supplementary results

The supplementary results provide gradient bounds through separate cases, including a case where the dissipative condition is unavailable.

A.3.1 Supplementary results

  • The supplementary results include detailed upper and lower bounds for the regularized-objective gradient.Source: §A.3.1 · p. 30

Case 1

  • For points with ∥∇L∥ > ϵL, the proof invokes Lemma 10 to lower-bound the gradient.Source: §A.3.1 · p. 32

Case 2

  • When ∥∇L∥ ≤ ϵL, the dissipative condition is unavailable, but Lemma 10 still supplies a gradient lower bound.Source: §A.3.1 · p. 32

B Generalization

The generalization analysis uses a data-adapted partition and bounds prediction variation using cellwise and layerwise quantities.

B Generalization

  • With probability one, the training inputs are pairwise distinct, making the data-adapted partition well defined.Source: B Generalization · p. 34
  • The prediction variation error is bounded above by its layerwise counterpart.Source: B Generalization · p. 35
  • A refined decomposition separates within-cell variance of each layer contribution from the discrepancy between its cellwise mean and representative-input value.Source: B Generalization · p. 35

B.1 Data

The proof conditions on the training inputs and resulting partition, then gives a probability bound over the output noise.

B.1 Data

  • Conditioned on the probability-one input event, the remaining randomness is the independent Gaussian output noise.Source: §B.1 · p. 36
  • With probability at least 1 − δ over output noise, the stated bound holds conditional on the probability-one input event.Source: §B.1 · p. 36

B.2 Optimization error

The proof of Corollary 1 works on a probability-one event where training inputs are pairwise distinct and the adapted partition is well defined.

B.2 Optimization error

  • For K1 ≤ K ≤ K2, the proof applies Theorem 2 and the definition of Eopt on the probability-one event with distinct training inputs and a well-defined adapted partition.The event is absorbed into the probability-one event underlying Theorem 2.Source: §B.2 · p. 36

B.3 Prediction variation error

Theorem 6 is stated under Theorem 1 and Assumptions 6 and 7; the proof bounds initial layerwise prediction-variation terms by geometric contraction.

B.3 Prediction variation error

  • Theorem 6 is stated for k ≥ 0 under the assumptions of Theorem 1 and Assumptions 6 and 7.The constants used in the following consequence are defined in Lemma 15.Source: §B.3 · p. 37
  • The first terms are bounded by (1 − 2ηλ)^k E_PV,layer(θ_K0).The proof notes that the relevant quantities and contraction factors are nonnegative.Source: §B.3 · p. 38

B.4 Necessary/sufficient conditions of generalization

The supplied proof fragments give limited detail about the conditions and consequences in these subsections.

B.4 Necessary/sufficient conditions of generalization

  • The proof of Corollary 3 uses a lower bound on K2 − K0 to bound its first right-hand-side term by ϵK0.The excerpt does not identify that term or state the resulting conclusion.Source: §B.4 · p. 39

B.5.1 Different decay scales of variance and bias

Layerwise prediction variation separates into within-cell variance and a squared bias-type term, which share a contraction factor but have different approximate-homogeneity remainders. When other contributions are asymptotically smaller, the bias bound has one additional leading-order copy of the homogeneity error, while the centered variance benefits from cancellation.

B.5.1 Different decay scales of variance and bias

  • The within-cell variance and squared bias-type terms share a geometric contraction factor but have different approximate-homogeneity remainders.These are the two components of each layerwise prediction variation.Source: §B.5.1 · p. 39
  • When remaining contributions are asymptotically smaller than the approximate-homogeneity scale, geometrically decaying contraction terms vanish as k → ∞.Under that regime, the estimates yield the stated leading-order asymmetry between bias and variance.Source: §B.5.1 · p. 40
  • At leading order in ϵah, the bias-type upper bound exceeds the variance bound by one additional copy of Cah,Vℓ,var(ℓ + 2)ϵah.The bias compares the cellwise mean with the representative value, so the approximation remainder enters through both; centered variance benefits from cancellation around the cellwise mean.Source: §B.5.1 · p. 40

B.5.2 Weights inside and outside the nonlinear nodes

The analysis distinguishes weights inside nonlinear maps from weights outside them: inner weights contribute directly to approximate-homogeneity error, while outer weights do not directly affect their own layer’s error but can affect later layers.

B.5.2 Weights inside and outside the nonlinear nodes

  • Inner weights appear inside the nonlinear map and contribute directly to approximate-homogeneity error.The passage contrasts them with outer weights, which enter their own block linearly.Source: §B.5.2 · p. 40
  • An outer weight creates no direct approximate-homogeneity error for variation in its own layer.The outer weight enters its own block linearly, yielding an exact identity without an approximate-homogeneity remainder.Source: §B.5.2 · p. 40
  • Through Jacobian composition, outer weights can contribute to error bounds for every later layer, so depth propagates both weight types’ effects forward.The direct contribution is asymmetric between inner and outer weights.Source: §B.5.2 · p. 40

B.6 Supplementary lemmas for generalization

The supplementary lemmas establish a smoothness-preserving integration result and bounds involving layer-indexed quantities and weight-specific constants.

B.6 Supplementary lemmas for generalization

  • Integrating a function that is poly-smooth with respect to its parameter preserves poly-smoothness.Proposition 2 states this result, and a later passage restates its conclusion.Source: §B.6 · pp. 40–41
  • Lemma 14 states that two constants have at most quadratic dependence on the layer index ℓ.The excerpt names the constants as Cah,Vℓ and Cah,Vℓ,var.Source: §B.6 · p. 48
  • For each preceding layer index r, the supplementary results provide bounds for both inner- and outer-weight constants.The excerpt identifies these constants as Ci,ℓ,r,in and Ci,ℓ,r,out.Source: §B.6 · p. 49
Loading 2609.07755v1…