Low Rank Gradients and Where to Find Them
This work provides theoretical insights into gradient structure for neural networks, which could aid in optimization and regularization, but it is incremental as it builds on existing models with relaxed assumptions.
The paper investigates low-rank structure in gradients for two-layer neural networks under relaxed data and parameter assumptions, showing that the gradient is dominated by two rank-one terms influenced by data properties, scaling regimes, and activation functions, with experiments validating the findings.
This paper investigates low-rank structure in the gradients of the training loss for two-layer neural networks while relaxing the usual isotropy assumptions on the training data and parameters. We consider a spiked data model in which the bulk can be anisotropic and ill-conditioned, we do not require independent data and weight matrices and we also analyze both the mean-field and neural-tangent-kernel scalings. We show that the gradient with respect to the input weights is approximately low rank and is dominated by two rank-one terms: one aligned with the bulk data-residue , and another aligned with the rank one spike in the input data. We characterize how properties of the training data, the scaling regime and the activation function govern the balance between these two components. Additionally, we also demonstrate that standard regularizers, such as weight decay, input noise and Jacobian penalties, also selectively modulate these components. Experiments on synthetic and real data corroborate our theoretical predictions.