Convergence of SGD in deep linear networks
A proof that training on random mini-batches, as done in practice, keeps a deep linear network's weights bounded and drives its gradient to zero, with step-size rules you can compute.
The short version
- Stochastic gradient descent (SGD) trains most modern deep networks, yet its convergence on nonconvex losses is only partly understood. Most guarantees assume a globally Lipschitz gradient, and the square loss of a deep linear network doesn't have one.
- This preprint, with Bubacarr Bah, Holger Rauhut, and Ulrich Terstiege, extends the convergence results of our gradient descent paper to mini-batch SGD on a regularized square loss. With probability one, the iterates stay bounded, the loss converges, the gradients go to zero, and the distance to the set of critical points goes to zero. It holds for deterministic and adaptive step sizes, at any batch size down to 1.
- The key step is proving boundedness directly, with no Lipschitz assumption. Bounded iterates let the SGD path track gradient flow as an asymptotic pseudotrajectory, and the regularized loss is a Lyapunov function that fixes where SGD ends up. To our knowledge, it's the first almost sure convergence result for SGD in deep linear networks in this setting.
- What that buys you in practice: step-size schedules with explicit constants, computed from the data, the batch size, the regularization, and the starting weights, under which convergence is guaranteed.
Problem statement and research objective
Context and significance
SGD and its variants are how highly parameterized models get trained. Each step uses the gradient on a random mini-batch, which is cheaper than a full-batch gradient on large datasets. The cost is noise, and in nonconvex settings, where the loss landscape is riddled with saddle points and local minima, the theory lags behind the practice.
Most SGD convergence results assume bounded gradients or a Lipschitz gradient. That property gives the descent lemma, is key to bounding the second moment of the gradient, and is how existing proofs show the iterates stay bounded. The gradient of the square loss of a deep linear network is not globally Lipschitz, so those general results don't apply unless boundedness is proved some other way.
Mathematical formulation
A deep linear network of depth N ≥ 2 multiplies its input by N weight matrices in turn. It computes a linear map, but the factored parameters make the loss nonconvex, so it keeps key traits of deep learning while staying tractable. The training loss is the square loss plus a regularizer λ(½‖W‖²)^β on all the weights, with λ > 0 and β ≥ N. At each step, SGD draws a mini-batch of size T uniformly at random and updates every weight matrix with the gradient of the regularized mini-batch loss, whose expectation is the full gradient. The paper adds the regularizer for a stated reason: for the unregularized square loss, proving that SGD stays bounded turned out to be very hard and is still open.
Scope
In scope: mini-batch SGD on deep linear networks of any depth N ≥ 2 with the regularized square loss, batch sizes from 1 up to the full dataset, deterministic and adaptive step-size schedules, and numerical experiments on synthetic data, including a check on networks with a tanh activation on the last layer. Out of scope: the unregularized square loss, which is still open for SGD, and a proof for nonlinear networks, which the experiments probe but the theory doesn't cover.
Success criteria
Prove almost sure boundedness of the SGD iterates without a Lipschitz-gradient assumption, then almost sure convergence to the critical set under step sizes with explicit constants, for both deterministic and adaptive schedules. Then test numerically how sharp the step-size conditions are.
Methodology and approach
Boundedness of the SGD iterates, proved directly
Everything rests on the iterates staying bounded, which earlier SGD analyses got from a Lipschitz gradient. This paper gets it from the regularizer instead. For β ≥ N and step sizes scaled by explicit constants, the norm of every iterate stays below max{‖W(0)‖, M + B}, where M and B depend on the data, the batch size, λ, β, N, and the starting weights. When β = N, λ must also exceed a threshold set by the data. The adaptive schedule gets an explicit bound of the same kind, set by M, the cap D, and the data.
From SGD to gradient flow: asymptotic pseudotrajectories
With boundedness in hand, the analysis treats SGD as a noisy discretization of gradient flow on the regularized loss. Joining the iterates into a continuous path, it shows that this path is an asymptotic pseudotrajectory, a tool from stochastic approximation theory: over any fixed time window, the path follows the gradient flow ever more closely as training goes on. That needs the mini-batch noise to average out, which comes from step sizes whose sum diverges and whose squares have a finite sum.
Landing on the critical set
The regularized loss is a Lyapunov function for the gradient flow, so it decreases along the flow. A result of Benaïm on asymptotic pseudotrajectories then turns a bounded pseudotrajectory plus this Lyapunov function into convergence of the loss values and of the iterates to the set of critical points, with probability one.
Two step-size schedules
Both schedules multiply a sequence b_k in (0, 1] by a scale. The theorem needs the b_k to have an infinite sum and a finite sum of squares, as b_k = 1/(k + 1)^γ does for γ above ½.
- Deterministic. The scale is fixed before training, from the data, the batch size, the starting weights, and the regularization constants.
- Adaptive. The step size reads the current weight norm and shrinks as the weights grow, capped by a constant D. The paper suggests a large D in practice, so the adaptive term applies most of the time.
What the batch size does to the step size
A section the paper labels heuristic, not a theorem, compares SGD at batch size 1 with full-batch gradient descent through these step-size bounds. On data where one point carries all the signal, gradient descent may come out ahead. On data made of copies of one point, SGD may. The paper reads this as a hint that SGD is likely more efficient when it can see the whole dataset in fewer steps.
Results and impact
Key theoretical results
| Result | Step sizes | Guarantee |
|---|---|---|
| Boundedness | Deterministic | Every iterate stays within the explicit radius max{‖W(0)‖, M + B} |
| Boundedness | Adaptive | Every iterate stays within an explicit radius set by M, D, and the data |
| Loss convergence | Both | The regularized loss along the SGD iterates converges, with probability one |
| Vanishing gradient | Both | The full gradient at the SGD iterates goes to zero, with probability one |
| Critical-set convergence | Both | The distance from the iterates to the set of critical points goes to zero, with probability one |
Experiments on synthetic data
The experiments train linear networks of depth 3 and 4 on 300 Gaussian data points in dimension 70, with targets produced by a random linear network and mini-batches of 30. Under adaptive step sizes that meet the theorem, the regularized loss falls in stages and then stays flat for the rest of a 500,000-epoch run: from about 5,000 epochs on at depth 3 (Figure 1), and from about 100,000 epochs on at depth 4 (Figure 2). The paper reports the same converging loss and bounded weight norms under the deterministic schedule. Step sizes that break the theorem's conditions still kept the iterates bounded, which suggests the bound isn't tight.
Figure 3 compares the two schedules on the same depth-3 network, with the loss divided by its starting value. Under deterministic step sizes it declines slowly and levels off near 0.91 after about 20,000 epochs. Under adaptive step sizes with D = 1 it barely moves for the first few hundred epochs, then drops sharply around epoch 1,000 and levels off near 0.68. The adaptive schedule converged faster, and to a lower loss.
A last experiment puts a tanh activation on the last layer, a case the theory doesn't cover. Under the adaptive step-size rule, the regularized loss of a depth-3 network falls from about 0.048 to about 0.013, most steeply between 100 and 1,000 epochs, and is flat from about 10,000 epochs on (Figure 4), much like the linear network in Figure 1. Its weight norms dip from about 1.76 to 1.32 within the first 50 epochs, climb back, and hold steady near 1.78 from about 2,000 epochs on (Figure 5). The paper reads this as a hint that some of the results may extend to nonlinear networks.
Limitations and what comes next
The guarantee is convergence to the set of critical points of the regularized loss. It doesn't say which critical point SGD reaches, or whether that point is a global minimum. The analysis needs the regularizer with β ≥ N: SGD on the unregularized square loss, the setting of the earlier gradient descent work, is still open, and the authors believe this analysis gives a blueprint for many of its steps. Nonlinear networks are tested numerically only.