[Paper Review] Low-rank lottery tickets: finding efficient low-rank neural networks via matrix differential equations

Steffen Schotthöfer, Emanuele Zangrando+3 more|arXiv (Cornell University)|May 26, 2022

Model Reduction and Neural Networks5 citations

TL;DR

This paper proposes a novel training algorithm that dynamically maintains low-rank weight matrices in neural networks using matrix differential equations, enabling efficient training and inference with significantly reduced memory and computation. The method achieves performance comparable to full-rank networks while automatically adapting rank during training, demonstrating the existence of 'low-rank winning tickets'—high-performing subnetworks found directly during training without initialization dependence.

ABSTRACT

Neural networks have achieved tremendous success in a large variety of applications. However, their memory footprint and computational demand can render them impractical in application settings with limited hardware or energy resources. In this work, we propose a novel algorithm to find efficient low-rank subnetworks. Remarkably, these subnetworks are determined and adapted already during the training phase and the overall time and memory resources required by both training and evaluating them are significantly reduced. The main idea is to restrict the weight matrices to a low-rank manifold and to update the low-rank factors rather than the full matrix during training. To derive training updates that are restricted to the prescribed manifold, we employ techniques from dynamic model order reduction for matrix differential equations. This allows us to provide approximation, stability, and descent guarantees. Moreover, our method automatically and dynamically adapts the ranks during training to achieve the desired approximation accuracy. The efficiency of the proposed method is demonstrated through a variety of numerical experiments on fully-connected and convolutional networks.

Motivation & Objective

• To address the high memory and computational cost of training and deploying large neural networks in resource-constrained environments.
• To overcome limitations of existing low-rank and pruned network methods that require fixed ranks or post-training compression.
• To develop a training method that naturally discovers and maintains efficient low-rank subnetworks (low-rank winning tickets) during optimization.
• To ensure numerical stability and convergence guarantees when training on low-rank manifolds, especially near small singular values.
• To enable automatic, dynamic rank adaptation during training, eliminating the need for manual hyperparameter tuning of rank.

Proposed method

• The method restricts weight matrices to a low-rank manifold by representing them as $ W = U S V^ op $, where $ U $, $ S $, and $ V $ are low-rank factors updated during training.
• It formulates the training process as a continuous-time gradient flow on the low-rank manifold, using techniques from Dynamical Low-Rank Approximation (DLRA) for matrix ODEs.
• The algorithm employs low-rank numerical integrators that maintain the low-rank structure throughout training, ensuring stability even with small singular values.
• Backward pass gradients are computed via matrix chain rule, with efficient updates derived for $ U $, $ S $, and $ V $ using the adjoint method.
• For convolutional layers, the kernel tensor is reshaped into a matrix, and low-rank decomposition is applied to the flattened kernel, preserving the convolutional structure.
• The method provides theoretical guarantees on approximation accuracy, descent behavior, and stability during optimization on the low-rank manifold.

Experimental results

Research questions

• RQ1 Can low-rank subnetworks be discovered and trained efficiently directly during the training process, without post-hoc pruning or initialization-dependent search?
• RQ2 Can a training algorithm maintain low-rank structure while ensuring numerical stability and convergence, especially when singular values are small?
• RQ3 Does the proposed method automatically adapt the rank of the weight matrices during training to achieve desired accuracy with minimal parameters?
• RQ4 Do the resulting low-rank subnetworks—'low-rank winning tickets'—achieve performance comparable to full-rank models, even when trained from random initialization?
• RQ5 How does the proposed method compare in efficiency and accuracy to standard low-rank training and sparsity-based pruning approaches?

Key findings

• The proposed method achieves comparable test accuracy to full-rank networks on fully-connected and convolutional architectures, even with significantly reduced parameter counts.
• The method dynamically adapts the rank of weight matrices during training, eliminating the need for manual rank hyperparameter tuning.
• Numerical experiments show substantial reductions in memory usage and computational cost during both training and inference compared to full-rank models.
• The method provides theoretical guarantees on approximation quality, descent behavior, and stability, particularly in the presence of small singular values.
• The results suggest the existence of 'low-rank winning tickets'—high-performing low-rank subnetworks—that are discovered directly during training, independent of initialization.
• The approach outperforms vanilla low-rank training schemes (e.g., alternating SGD on $ U $ and $ V $) by maintaining stability and accuracy through proper manifold-constrained optimization.