1. Tensor-GaLore: Memory-Efficient Training via Gradient Tensor Decomposition
- Author
-
George, Robert Joseph, Pitt, David, Zhao, Jiawei, Kossaifi, Jean, Luo, Cheng, Tian, Yuandong, and Anandkumar, Anima
- Subjects
Computer Science - Machine Learning - Abstract
We present Tensor-GaLore, a novel method for efficient training of neural networks with higher-order tensor weights. Many models, particularly those used in scientific computing, employ tensor-parameterized layers to capture complex, multidimensional relationships. When scaling these methods to high-resolution problems makes memory usage grow intractably, and matrix based optimization methods lead to suboptimal performance and compression. We propose to work directly in the high-order space of the complex tensor parameter space using a tensor factorization of the gradients during optimization. We showcase its effectiveness on Fourier Neural Operators (FNOs), a class of models crucial for solving partial differential equations (PDE) and prove the theory of it. Across various PDE tasks like the Navier Stokes and Darcy Flow equations, Tensor-GaLore achieves substantial memory savings, reducing optimizer memory usage by up to 75%. These substantial memory savings across AI for science demonstrate Tensor-GaLore's potential.
- Published
- 2025