Tensor-GaLore: Memory-Efficient Training via Gradient Tensor Decomposition

Authors: Robert Joseph George, David Pitt, Jiawei Zhao, Jean Kossaifi, Cheng Luo, Yuandong Tian, Anima Anandkumar
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
Keywords: neural operators, PDE, optimization, pre-training, Large scale training, AI4Science

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.

Decision: reject
Reviewer scores: 5, 5, 5, 6

Reviews and discussion:
<thread>
creation date:
1730061808051

thread_type:
official_review

<thread_object>
creator: reviewer_mexm

<content>
summary:
The authors present Tensor-GaLore as a method for compressing the weights using tensor compression. They show the use of their approach on Fourier Neural Operators (FNOs) used for solving PDEs.

strengths:
The focus on efficient learning techniques for PDEs is very timely and requires constant improvement to outperform traditional methods in terms of accuracy and computing time.

weaknesses:
- the authors discuss that TensorGalore is superior to Galore as it avoids SVDs. Nevertheless, the computation of tensor formats (Tucker, Tensor Train, etc.) relies on matricizations of the tensor and then typically singular value decomposition of those. So the SVD is still at the heart of TensorGalore.
- I think the idea of the manuscript is based on taking an FNO example and then using the tensor compression of the weights. I believe this to work well but it also seems a straightforward extension of previous work. There have been many applications of tensors for compression within neural networks.

questions:
- How do the authors implement their tensor format as they argue that the disadvantage of Galore is the need for the SVD, which typically used for efficient computations of popular tensor formats?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730691383806

thread_type:
official_review

<thread_object>
creator: reviewer_oxlj

<content>
summary:
This paper presents a modification of the GaLora method that allows to update the weights not directly, but in a low-parameter space. The authors present a modification of this method that uses a low rank tensor decomposition, namely the Tucker decomposition, instead of a low rank matrix decomposition. This approach is applied to neural operators for solving PDEs, where 4-way tensors arise naturally.

strengths:
- The paper uses a low-rank tensor decomposition, which preserves the original multidimensional structure
- Some experiments show the effectiveness of this technique
- Clearly presented paper

weaknesses:
- No code provided to reproduce the results
- No theoretical analysis (in the original GaLora paper there are theoretical justification of low-rank structure, convergence, etc.)
- The only significant change (other than technical changes) from the original GaLora method is the use of Tucker Decomposition for 4-way tensor instead of low-rank matrix decomposition
- The presented method shows good results only on the Darcy flow equation, in the other experiments the improvement is not strong
- Since this paper describes improvements to the GaLora method (which has been published previously), it would be fair to compare against it rather than baseline. For example, in Table 2 for the Darcy equation the presented method is 48.8% better than baseline, while the regular GaLora is 19% better than baseline. Thus, the improvement presented in this paper is a 25% improvement.

Overall, this paper is incremental to the original GaLora paper, without theoretical evaluations (which were in the original paper) and with inconclusive numerical results.

Minor 

- L458-459 "On Darcy flow (as shown in Table 6)" should be "Table 2"
- L397 word "Table" is missing

questions:
- Can one pick the efficient rank ratio in advance?
 - How is it that for Burgers Equation in Table 2 test loss is much (by an order of magnitude) less than train loss?
 - Have you tried using other low-rank tensor decompositions (CANDECOMP/PARAFAC, Tensor-train, etc.)?

In Algorithm 1

- is $r$ is rank of rank ratio?
- tensor $\mathcal{M}_0$ (with $\mathcal V_0$) has the same shape as $\mathcal W\in\mathbb{C}^{N_1\times N_2\times N_3\times N_4}$. Should it be $\mathcal M_0\in\mathbb{C}^{R_1\times R_2\times R_3\times R_4}$?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730700629834

thread_type:
official_review

<thread_object>
creator: reviewer_rfrz

<content>
summary:
This work presents Tensor-GaLore, an algorithm that leverages low-rank tensor decomposition on the gradients of tensorized weights. This work is built on top of the previous work (GaLore), which applies low-rank factorization (SVD) on the gradients. Experimental results show that applying it Fourier Neural Operators yield better memory usage and accuracy for numerical PDE problems.

strengths:
1. The quality of the presentation is good. The work is clear and easy to follow.
2. The idea of using Tucker decomposition to perform low-rank approximation makes sense for numerical PDE problems, and experimental results verify that.

weaknesses:
Despite being clear and effective, I believe the work has limited novelty. The tensor-GaLore approach has limited difference compared to GaLore. In addition, only empirical rather than theoretical results is provided to show the efficacy of the algorithm.

questions:
please see above
</content>
</thread_object>

</thread>
<thread>
creation date:
1730707002399

thread_type:
official_review

<thread_object>
creator: reviewer_po1k

<content>
summary:
This paper extends GaLore [Zhao 2024] to neural networks with tensor weights by adding Tucker Decomposition and performing low-rank projection directly on tensor gradients. The experiment compares the proposed method to vanilla GaLora (with reshaping) on Fourier Neural Operators (FNOs), a class of tensor-weight models for solving partial differential equations.

strengths:
This paper appears to tackle a gap that has received little attention by the literature, the efficient training of tensor-weight models, with the only prior works being [Kossaifi 2024] and [George 2024]. 

The paper is generally well-written and easy to follow, with a clear story.

weaknesses:
Despite the novel application, the approach is a somewhat straight-forward extension of GaLore to tensor-weight models, replacing SVD decomposition with Tucker. 

There is a lack of discussion on the slowdown in training given the overhead.

questions:
See weaknesses
</content>
</thread_object>

</thread>
<thread>
creation date:
1731896808480

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response

comment:
### Thank you to all the reviewers for the thoughtful and insightful feedback on our paper. We greatly appreciate the time and effort you have put into reviewing our work. We have uploaded a revised version, which is more thorough and has more details to some of the reviewer's questions, along with over 20 pages of the theory of our method attached to it.

**In this paper, we have proposed Tensor-GaLore, a novel method for efficient training of Fourier Neural Operators (FNOs), a class of models crucial for solving partial differential equations (PDEs) in scientific computing. Our contributions include**

1. We propose **Tensor-GaLore, which leverages tensor decomposition techniques to project the high-dimensional tensor gradients onto low-rank subspaces, enabling substantial memory savings without compromising model performance**. Specifically, Tensor-GaLore utilizes the Tucker decomposition to decompose the gradient tensors into a core tensor and factor matrices. This allows us to project the gradients onto a low-rank subspace and perform the optimization in this compact representation, leading to significant reductions in memory usage for the optimizer states.

2. We now provide a **comprehensive theoretical analysis of Tensor-GaLore, proving both convergence guarantees and the emergence of the low-rank structure during training for FNOs**. Our theory shows that:

     2.1. The gradient tensors naturally develop low-rank structure in each mode during training

     2.2. Tensor-GaLore achieves convergence through mode-wise projections while preserving the multidimensional relationships

     2.3. We prove the convergence of Tensor-GaLore under mild conditions on the mode-k continuity of the operators

     2.4. Our theoretical results explain why the low-rank approximation works well in practice

3. We demonstrate that **Tensor-GaLore can achieve over 75% reduction in optimizer state** compared to the original FNO model while maintaining or even improving the model's performance across a diverse set of PDE tasks, including Navier-Stokes, Burgers' equation, and electromagnetic wave propagation. The ability to drastically reduce the memory footprint of FNOs is a crucial advancement that enables the application of these powerful models to high-resolution scientific computing problems. We carefully analyze the computational overhead introduced by the tensor decomposition operations. While Tensor-GaLore does introduce some computational overhead (5-20% slowdown depending on the configuration), we show that:

     3.1. The slowdown is modest compared to the significant memory savings achieved. The time cost for gradient projection remains constant regardless of input size, while the forward pass, backward pass, and gradient computation times scale linearly with input size. **This leads to decreasing relative slowdown as problem size increases**

     3.2. The overhead can be amortized through techniques like "warm-restart" initialization of tensor decomposition

     3.3. The trade-off between memory savings and computational cost can be controlled through the choice of rank parameters

      3.4. For many scientific computing applications, the memory reduction enables training of larger models that would otherwise be impossible, justifying the modest increase in computation time

We hope that these contributions together provide both theoretical understanding and practical benefits for efficiently training large-scale neural operators. The memory savings and theoretical guarantees make Tensor-GaLore a valuable tool for advancing AI in scientific computing.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732559659775

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response (Part 2)

comment:
We want to thank all the reviewers for their review of our paper. So far, we haven't gotten a single review back :(, and we have answered all the questions in detail and now included a detailed theoretical section of our work, which was missing before. **With the discussion period ending soon on Nov 26, 2024 (Anywhere on Earth), we kindly await your feedback or an updated assessment of our paper**. Please let us know if your concerns have been satisfactorily addressed; if so, we would greatly appreciate it if you could update your ratings. We are available and would be happy to answer any additional questions. Thank you once again!
</content>
</thread_object>

</thread>
<thread>
creation date:
1733269234276

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General summary of our work for the decision

comment:
Thank you to all the reviewers for the thoughtful and insightful feedback on our paper. We greatly appreciate the time and effort you have put into reviewing our work. We want to summarize the whole work and discussion since today is the last day we can respond. Although we are disappointed that 3 of the reviewers haven't responded to us since their comments, we are hopeful that the AC can take into the consideration of all the answers and revisions we made to make the final decision. To make it easier we include a summary of our revisions as well and the inclusion of the theoretical work.

# [New] Theoretical Contributions

## 1. Comprehensive Theory of Tensor-GaLore

We have developed a complete theoretical framework that explains why and how Tensor-GaLore works:

### Reversibility Analysis
- Proved that FNO is reversible when using reversible activations
- Showed that spectral layer, MLP layer, and their composition maintain reversibility
- This provides the foundation for analyzing gradient structure

### Low-Rank Structure
- Proved that gradients naturally develop mode-wise low-rank structure during training
- Showed that each tensor mode can have different rank behavior:
  - Fourier modes (3,4) exhibit natural spectral decay
  - Channel modes (1,2) maintain information structure
- Demonstrated that this emergent structure enables efficient compression

### Convergence Guarantees
- Proved convergence under mild mode-k continuity conditions
- Showed that fixed projection matrices are sufficient
- Established explicit bounds on convergence rate for each mode

## 2. Why Tensor-GaLore Outperforms GaLore

The final lemma (Lemma 11) provides a crucial theoretical justification for Tensor-GaLore's superiority:

1. **Mode-wise Independence**:
   - Tensor-GaLore achieves low rank in all modes simultaneously
   - GaLore cannot preserve low-rank structure across all modes due to matricization

2. **Structural Preservation**:
   - Tensor-GaLore maintains natural tensor structure of FNO weights
   - GaLore's matricization scrambles important multi-dimensional relationships

3. **Optimal Compression**:
   - Tensor-GaLore: All modes k satisfy sr_k(G_t) ≤ N_k/2 asymptotically
   - GaLore: At least one mode k must maintain sr_k(G_t) ≥ min(N_k/2, N')

# Practical Improvements

## 1. Memory Efficiency

- Achieved up to 75% reduction in optimizer memory usage
- Demonstrated scalability to high-resolution problems (1024×1024 Navier-Stokes)
- Enabled training of larger models that were previously infeasible

## 2. Performance Gains

- Maintained or improved model accuracy across all tested PDEs
- Showed implicit regularization benefits from tensor structure preservation
- Demonstrated better generalization in high-resolution settings

## 3. Implementation Optimizations

- Introduced "warm-restart" initialization for tensor decomposition
- Developed efficient mode-wise projection updates
- Carefully balanced computational overhead vs. memory savings

# Broader Impact

Our work has great implications across scientific computing, methodological advancement, and practical applications. In scientific computing, Tensor-GaLore enables the training of larger, more accurate FNOs while substantially reducing computational resource requirements, making advanced scientific modeling more accessible to researchers with limited resources (This is the key point we want to emphasize as academic labs/independent researchers do not have access to high-end GPU clusters). Our methodological contributions introduce novel techniques for tensor optimization and provide a comprehensive theoretical framework for analyzing tensor methods, effectively bridging the gap between matrix and tensor algorithms for the decomposition of gradients. These advances have direct practical applications in critical areas such as high-resolution climate modeling, precise fluid dynamics simulations, and electromagnetic wave propagation modeling.

# Summary of Changes after the discussions and the new revised version

1. **Theoretical Extensions**:
   - Added complete proofs of all theorems
   - Expanded mode-k continuity analysis
   - Enhanced tensor operation definitions and properties

2. **Technical Clarifications**:
   - Improved explanation of reversibility conditions
   - Added detailed analysis of computational complexity
   - Enhanced discussion of tensor rank properties

3. **Additional Results**:
   - Extended experimental validation
   - Added ablation studies
   - Included more detailed memory analysis

To conclude, we believe that this work represents a significant advancement in both theoretical understanding and practical implementation of memory-efficient training for tensor-based neural networks, especially neural operators (FNO), which are particularly important for scientific machine learning applications. We also provide code to reproduce all our experiments and detail where to get the datasets.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper proposes to use tensor factorization for gradients (instead of matrix factorization in the original GaLoRe paper) for applications for Fourier Neural Operators, showing memory savings.
The savings for the memory are not dramatic (i.e. 68 gigabytes -> 55 gigabytes) and for some cases there is a drop in the accuracy. 
Moreover, the usage of the memory can be reduced by using other techniques, such as checkpointing and quantization, and there is no strict need to use more complicated approaches for this particular task. There is also a slowdown effect for some of the parameters (again, should be compared to other memory footprint reduction techniques). 
The only modification is the generalization to the Tucker decomposition case (which is rather straightforward) and 
a specific application to FNO. The authors addressed some of the concerns in the rebuttal by adding new material and theoretical experiments, but I think it is not enough.