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.
翻译:摘要:贪婪低秩学习理论旨在解释深度学习令人瞩目的泛化能力。该理论证明,基于随机梯度的训练通过训练过程中秩的逐步增加,隐式地将神经网络正则化至低秩解。然而,该理论要求权重初始化为无穷小量,这在实践中难以实现(因为无穷小初始点为鞍点),因而理论与实际应用间存在差距。本文通过聚焦累积权重更新,消除了对无穷小初始化的假设。我们证明,在三层线性网络中,对于任意正交权重初始化,累积权重更新遵循增量低秩轨迹。实验表明,该理论在多种神经网络(如Transformer)及标准训练算法(如SGD、Adam)中均成立。然而,现有训练算法因未采用低秩参数化方式而无法利用低秩特性提升计算效率。为此,我们设计了一种新型训练算法——增量式低秩学习(InRank),该算法在训练过程中逐步增加秩的同时,将累积权重更新显式表示为低秩矩阵。我们在GPT-2模型上评估InRank,结果表明:在训练过程中,InRank的总秩不超过全秩模型的33%即可达到与之相当的预测性能。我们还提出一个高效版本的InRank,在WikiText-103数据集上从头训练GPT-Medium时,使总训练时间减少37%,模型规模缩减36%。