Emergent Low-Rank Training Dynamics in MLPs with Smooth Activations

Abstract: Recent empirical evidence has demonstrated that the training dynamics of large-scale deep neural networks occur within low-dimensional subspaces. While this has inspired new research into low-rank training, compression, and adaptation, theoretical justification for these dynamics in nonlinear networks remains limited. %compared to deep linear settings. To address this gap, this paper analyzes the learning dynamics of multi-layer perceptrons (MLPs) under gradient descent (GD). We demonstrate that the weight dynamics concentrate within invariant low-dimensional subspaces throughout training. Theoretically, we precisely characterize these invariant subspaces for two-layer networks with smooth nonlinear activations, providing insight into their emergence. Experimentally, we validate that this phenomenon extends beyond our theoretical assumptions. Leveraging these insights, we empirically show there exists a low-rank MLP parameterization that, when initialized within the appropriate subspaces, matches the classification performance of fully-parameterized counterparts on a variety of classification tasks.

• The paper demonstrates that MLP training with smooth activations confines weight updates to a fixed low-dimensional subspace, with an effective rank bounded by 2K.
• The study employs perturbation theory to show that most parameter update energy (>95%) remains aligned with the initial gradient's principal subspace, ensuring minimal drift.
• Practical experiments confirm that using initialization-derived low-rank projections achieves comparable performance to fully-parameterized models on benchmarks like Fashion MNIST and CIFAR-10.

Emergent Low-Rank Training Dynamics in MLPs with Smooth Activations

Overview and Motivation

The study systematically analyzes the emergence of low-rank structure in the weight dynamics of multi-layer perceptrons (MLPs) with smooth (i.e., twice differentiable) activation functions during gradient-based training. Motivated by empirical reports of parameter updates and optimization trajectories being highly concentrated in low-dimensional subspaces, the authors provide the first rigorous characterization of this phenomenon in nonlinear feedforward networks. The primary focus is on settings where the output dimensionality $K$ is significantly smaller than the input or hidden dimensions (i.e., $K \ll d, m$), a regime of practical importance for multiclass classification tasks.

Initial investigations reveal marked qualitative differences between networks with smooth activations (e.g., ELU, GELU, SiLU) versus nonsmooth variants (e.g., ReLU, Leaky-ReLU). For smooth activations, the evolution of most singular vectors of the weight matrices is slow and largely confined to a fixed, small-dimensional subspace over the course of training; in contrast, for nonsmooth activations, the parameter dynamics are spread over a larger dimensional space.

Figure 1: The middle singular subspace of the first-layer weight matrix in the ELU network evolves noticeably slower than that in the ReLU network, and the corresponding singular values remain closer to their initialization.

Theoretical Analysis: Two-Layer Networks with Smooth Activations

A central theoretical result comprehensively characterizes the low-rank nature of training-induced weight dynamics in the first-layer weights of two-layer MLPs with smooth activations and small output dimension. The authors prove that, under the assumptions of whitened input data, fixed output layer, and gradient descent (GD) training on squared-error loss, each GD update in the weight space is effectively contained within a fixed, low-dimensional subspace. This subspace, which is identified by the principal singular vectors of the gradient at initialization, is entirely determined by the data, initialization, activation smoothness, and the fixed output weights.

More formally:

• The effective rank of the parameter change over training is at most $2K$, with all significant updates occurring in a matrix subblock spanned by vectors derived from the initialization gradients and weights.
• The remainder of the parameter space exhibits only $\mathcal{O}(\epsilon)$-bounded perturbations for initialization scale $\epsilon$, and subsequent GD updates produce only exponentially decaying changes within those perturbation directions as a function of training time.

  Figure 2: Under the exact theoretical setting, $\widetilde{\bm W}_{1,1}(t)$ accounts for almost all of the change in $\widetilde{\bm W}_1(t)$, confirming that updates predominantly occur in the predicted low-dimensional subspace.

The analysis leverages perturbation theory (Wedin’s sin-theta theorem) to rigorously bound the drift of these crucial subspaces and establishes that low-rank weight evolution persists so long as the initialization scale and learning rate are not overly large.

Extension Beyond Theory: Deep Networks and Practical Optimizers

The study empirically verifies that the theoretical phenomenon is robust well beyond the narrow two-layer, squared-error setting:

• For deep MLPs trained with SGD (with momentum), Adam, and on cross-entropy loss (with or without input whitening), pronounced low-dimensional concentration of weight evolution is consistently observed for all layers.
• Quantitative measurement of subspace drift shows minimal deviation from the initialization-aligned low-rank structure, confirming the phenomenon’s generic character under both idealized and practical training regimes.

Figure 3: Training deep MLPs with SGD or Adam on unwhitened inputs with cross-entropy loss maintains the previously observed low-rank dynamics, especially for smooth activations.

Empirically, the key role of smoothness is further validated: when nonsmooth activations are used, a much larger fraction of the weight parameter space is involved in the learning dynamics, and the invariant subspaces are no longer present.

Low-Rank Parameterization: From Observations to Efficient Architectures

Building on these findings, the authors propose a constructive procedure for synthesizing low-rank MLP parameterizations that match the accuracy of their fully-parameterized counterparts on real-world tasks. The method involves:

• Projecting both the input and top-layer features into the aligned low-rank subspaces given by the singular vectors of the initialization gradient.
• Training only small internal weight matrices; the alignments are fixed at initialization.

Critically, the success of this approach depends sharply on using the correct, initialization-derived subspaces; random orthogonal choices for these projections consistently lead to poor convergence and local minima, affirming that the emergent low-rank structure is highly instance-specific and not random.

Experiments on Fashion MNIST and CIFAR-10 (using VGG-16 with a low-rank classifier MLP head) demonstrate that, when the construction is followed, test loss and accuracy learning curves for the low-rank MLP are virtually indistinguishable from those of the full model—not only under full fine-tuning but also under classifier-only settings.

Figure 4: The properly-initialized low-rank MLP achieves nearly-identical test loss and accuracy as the fully-parameterized MLP; with random subspace initialization, it is trapped in a suboptimal local minimum.

Figure 5: On CIFAR-10 with VGG-16, the classification performance of the low-rank MLP head with proper initialization closely matches that of the fully-parameterized counterpart. Random subspace initialization yields no training progress.

Technical Implications and Insights

Strong numerical outcomes:

• Across training runs, in multiple settings, greater than 95–99% of parameter update energy is confined to the identified low-rank subspaces.
• Low-rank MLPs constructed as described need only $r = 2K$ internal width to replicate full-model performance, with diminishing returns for $r > 4K$.

Sharp, nontrivial claim:

• The critical subspace is not data-universal nor random but is both data- and initialization-specific, governed by the top $K$ singular vectors of the initial gradient.

Contradictory to intuition: Despite the non-convexity and compositional nature of deep nonlinear architectures, most updates are not distributed throughout parameter space, but remain “trapped” in a small affine subspace determined at initialization.

Implications and Future Directions

This work both formalizes and validates the conjecture that overparameterized neural network training proceeds in (and can be restricted to) very low-dimensional subspaces, so long as smooth activation dynamics dominate—an observation with wide ramifications:

• Optimization and generalization: Explicit low-rank parameterization can yield massive reductions in training and storage costs, without impacting accuracy—especially in problems with small output dimension.
• Network compression and adaptation: Insights support and explain the empirical success of LoRA-type approaches, and suggest direct, theoretically-justified design of compressed architectures for warm-started training.
• Initialization protocols: The findings reinforce the criticality of initialization, as alignment with the emergent subspaces is necessary for optimal learning trajectories.
• Limits for nonsmooth activations: The lack of analogous low-rank dynamics for ReLU/Leaky-ReLU suggests differing implicit regularization and representation capacity properties with practical implications for architecture selection.

Future prospects include extending these findings to attention-based architectures (Transformers), structured data scenarios, or task transfer settings, where the low-rank phenomenon may underpin scalable adaptation to new problems.

Conclusion

The results decisively demonstrate that, under broad and realistic conditions, the evolution of weights in MLPs with smooth activations is strictly confined to invariant, low-dimensional, initialization-dependent subspaces. This provides both a new theoretical foundation for—and a practical prescription toward—directly exploiting low-rank structure in neural network training, with implementation-ready implications for model efficiency, transfer learning, and network design.

Reference:

"Emergent Low-Rank Training Dynamics in MLPs with Smooth Activations" (2602.06208)