InRank: Incremental Low-Rank Learning

TL;DR

This paper introduces InRank, a new training algorithm that incrementally increases low-rank representations of neural network weights, improving efficiency while maintaining performance, based on theoretical insights into low-rank learning.

Contribution

It removes the impractical initialization assumption in low-rank learning theory and develops InRank, an algorithm that explicitly enforces low-rank weight updates during training.

Findings

InRank achieves comparable accuracy to full-rank models with significantly lower rank.

InRank reduces training time by up to 37% and model size by 36%.

Theoretical results hold across various neural network architectures and training algorithms.

Abstract

The theory of greedy low-rank learning (GLRL) aims to explain the impressive generalization capabilities of deep learning. It proves that stochastic gradient-based training implicitly regularizes neural networks towards low-rank solutions through a gradual increase of the rank during training. However, there is a gap between theory and practice since GLRL requires an infinitesimal initialization of the weights, which is not practical due to the fact that it is a saddle point. In this work, we remove the assumption of infinitesimal initialization by focusing on cumulative weight updates. We prove the cumulative weight updates follow an incremental low-rank trajectory for arbitrary orthogonal initialization of weights in a three-layer linear network. Empirically, we demonstrate that our theory holds on a broad range of neural networks (e.g., transformers) and standard training algorithms (e.g., SGD, Adam). However, existing training algorithms do not exploit the low-rank property to improve computational efficiency as the networks are not parameterized in low-rank. To remedy this, we design a new training algorithm Incremental Low-Rank Learning (InRank), which explicitly expresses cumulative weight updates as low-rank matrices while incrementally augmenting their ranks during training. We evaluate InRank on GPT-2, and our results indicate that InRank achieves comparable prediction performance as the full-rank counterpart while requiring at most 33% of the total ranks throughout training. We also propose an efficient version of InRank that achieves a reduction of 37% in total training time and 36% in model size when training GPT-medium on WikiText-103 from scratch.