Sparse Spectral Training and Inference on Euclidean and Hyperbolic Neural Networks

Jialin Zhao, Yingtao Zhang, Xinghang Li, Huaping Liu, Carlo Vittorio Cannistraci

Abstract: The growing demands on GPU memory posed by the increasing number of neural network parameters call for training approaches that are more memory-efficient. Previous memory reduction training techniques, such as Low-Rank Adaptation (LoRA) and ReLoRA, face challenges, with LoRA being constrained by its low-rank structure, particularly during intensive tasks like pre-training, and ReLoRA suffering from saddle point issues. In this paper, we propose **S**parse **S**pectral **T**raining **(SST)** to optimize memory usage for **pre-training**. SST **updates all singular values** and **selectively updates singular vectors** through a multinomial sampling method weighted by the magnitude of the singular values. Furthermore, SST employs **singular value decomposition to initialize and periodically reinitialize** low-rank parameters, reducing distortion relative to full-rank training compared to other low-rank methods. Through comprehensive testing on both Euclidean and hyperbolic neural networks across various tasks, SST demonstrates its ability to outperform existing memory reduction training methods and is comparable to full-rank training in various cases. On LLaMA-1.3B, with only 18.7\% of the parameters trainable compared to full-rank training (using a rank equivalent to 6\% of the embedding dimension), SST reduces the perplexity gap between other low-rank methods and full-rank training by **97.4\%**. This result highlights SST as an effective parameter-efficient technique for model pre-training.

Lay Summary:

Training large AI models usually requires a huge amount of computer memory and specialized hardware. This creates a high barrier for many researchers who want to contribute to AI development. Our work introduces a new method called **Sparse Spectral Training (SST)** that helps reduce the memory needed for training, without sacrificing performance. Think of it as a way to teach a model more efficiently by focusing only on the most important parts of its internal structure, kind of like learning the main points of a book instead of memorizing every word.

While existing methods like LoRA only update a small part of the model, SST updates a broader set of components in a smart, selective way. Our tests show that SST works better than previous methods across many tasks, from language understanding to graph analysis, and it even performs as well as full-scale training in some cases. By lowering the hardware barrier, we hope SST makes AI training more accessible to everyone in the research community.

Link To Code: https://github.com/biomedical-cybernetics/sparse-spectral-training

Primary Area: Deep Learning->Algorithms

Keywords: Parameter-efficient training; Pre-training; Hyperbolic Network

Submission Number: 6250