Constrained Factorization with Diagonal Scaling: Rank-Revealing Training and Pruning

Gradient descent for matrix factorization exhibits an implicit bias toward approximately low-rank solutions, even in regimes where the iterates may grow unbounded. We study this behavior through a constrained factorization with diagonal scaling, separating bounded outer factors from explicit diagonal scale parameters. For positive semidefinite matrix recovery, we use the model $X \approx UDU^\top$, where $U$ is constrained to a Frobenius norm ball and $D$ is a nonnegative diagonal factor. Although this reparameterization preserves the stationary points of classical Burer-Monteiro factorization, its projected dynamics consistently recover truly (rather than approximately) low-rank solutions across a wide range of step sizes and initializations. Motivated by this behavior, we extend the construction to neural networks by inserting diagonal layers between Frobenius-norm-constrained outer layers. The resulting UDV models exhibit strong rank-revealing structure during training, with many diagonal components and associated columns collapsing toward zero and becoming naturally prunable. We exploit this structure through an SVD-based pruning procedure (UDVPruning) that identifies and removes redundant units during training. Across ViT and MLP-Mixer benchmarks, UDVPruning substantially reduces parameters and FLOPs while maintaining competitive test accuracy and often improving measured training time, inference latency, and energy consumption.