Gradient descent aligns the layers of deep linear networks

TL;DR

The paper investigates why gradient-based optimization biases deep linear networks toward simple, generalizable solutions when trained on linearly separable data. By analyzing both gradient flow and gradient descent, it proves that the empirical risk vanishes, layer weights become unbounded yet align to rank-1 forms, and adjacent layers converge to mutually aligned singular directions. For exponential and logistic losses with appropriate data assumptions, the first-layer singular vector and the overall network predictor converge in direction to the maximum-margin solution, revealing an implicit regularization mechanism that favors margin-maximizing configurations across layers. These results illuminate the role of depth in steering representations toward low-complexity, margin-optimal solutions and motivate extensions to nonlinear settings and rate analyses with practical step sizes.

Abstract

This paper establishes risk convergence and asymptotic weight matrix alignment --- a form of implicit regularization --- of gradient flow and gradient descent when applied to deep linear networks on linearly separable data. In more detail, for gradient flow applied to strictly decreasing loss functions (with similar results for gradient descent with particular decreasing step sizes): (i) the risk converges to 0; (ii) the normalized i-th weight matrix asymptotically equals its rank-1 approximation $u_iv_i^{\top}$; (iii) these rank-1 matrices are aligned across layers, meaning $|v_{i+1}^{\top}u_i|\to1$. In the case of the logistic loss (binary cross entropy), more can be said: the linear function induced by the network --- the product of its weight matrices --- converges to the same direction as the maximum margin solution. This last property was identified in prior work, but only under assumptions on gradient descent which here are implied by the alignment phenomenon.

Figures (3)

• Figure 1: Visualization of margin maximization and self-regularization of layers on synthetic data with a $4$-layer linear network compared to a $1$-layer network (a linear predictor). \ref{['fig:maxmarg']} shows the convergence of $1$-layer and $4$-layer networks to the same margin-maximizing linear predictor on positive (blue) and negative (red) separable data. \ref{['fig:align']} shows the convergence of $\|W_i\|_2/\|W_i\|_F$ to 1 on each layer, plotted against the risk.
• Figure 2: A visualization of inter-layer alignment on data consisting of two well-separated circles with a 3-layer linear network. \ref{['fig:il']} depicts, as in \ref{['fig:maxmarg']}, that optimizing 1- and 3-layer linear networks finds the same maximum margin solution. The other three plots show the data as it is mapped through progressively more and more layers. Due to alignment, the product $W_i\cdots{}W_1$ becomes $u_i\bar{u}^\top$, where $u_i$ is the top left singular vector of $W_i$, which means that asymptotically the mapped data will be well separated and lie along the span of $u_i$, as depicted by the flattening in \ref{['fig:il:0', 'fig:il:1', 'fig:il:2']}. Additionally, these three subfigures show that the top right singular vector $v_{i+1}$ of the subsequent layer is aligned with this $u_i$, which in these plots (with principal component axes) corresponds to following a horizontal line.
• Figure 3: Risk and alignment of dense layers (the ratio $\|W_i\|_2/\|W_i\|_F$) of (nonlinear!) AlexNet on CIFAR-10. \ref{['fig:cifar:vanilla']} uses default PyTorch initialization, while \ref{['fig:cifar:fancy']} forces initial Frobenius norms to be equal amongst dense layers.

Theorems & Definitions (28)

• theorem 2.2
• lemma 3
• lemma 4
• theorem 2.5
• corollary 6
• lemma 7
• lemma 8
• lemma 10
• lemma 11
• theorem 3.2