Emergent Low-Rank Training Dynamics in MLPs with Smooth Activations

TL;DR

This paper investigates the low-dimensional subspace dynamics of MLP training with smooth activations, providing theoretical insights and demonstrating that low-rank parameterizations can match full models' performance.

Contribution

It offers a theoretical characterization of invariant low-dimensional subspaces in nonlinear MLPs and shows low-rank models can achieve comparable accuracy.

Findings

Training dynamics concentrate in low-dimensional subspaces

Invariant subspaces are precisely characterized for two-layer networks

Low-rank parameterizations can match full model performance

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.