An Augmented Backward-Corrected Projector Splitting Integrator for Dynamical Low-Rank Training

TL;DR

This paper introduces a new low-rank training method for neural networks that reduces computational costs by decreasing QR decompositions while maintaining convergence and robustness.

Contribution

We propose an augmented backward-corrected projector splitting integrator that improves efficiency and convergence in dynamical low-rank training of neural networks.

Findings

Reduces the number of QR decompositions needed during training.

Ensures convergence to a locally optimal solution.

Demonstrates effectiveness on multiple benchmark datasets.

Abstract

Layer factorization has emerged as a widely used technique for training memory-efficient neural networks. However, layer factorization methods face several challenges, particularly a lack of robustness during the training process. To overcome this limitation, dynamical low-rank training methods have been developed, utilizing robust time integration techniques for low-rank matrix differential equations. Although these approaches facilitate efficient training, they still depend on computationally intensive QR and singular value decompositions of matrices with small rank. In this work, we introduce a novel low-rank training method that reduces the number of required QR decompositions. Our approach integrates an augmentation step into a projector-splitting scheme, ensuring convergence to a locally optimal solution. We provide a rigorous theoretical analysis of the proposed method and demonstrate its effectiveness across multiple benchmarks.