GaLore icon indicating copy to clipboard operation
GaLore copied to clipboard

Extend GaLore Algorithm for General Tensor Decomposition

Open Robertboy18 opened this issue 9 months ago • 0 comments

The GaLore algorithm was originally designed to perform lower-order gradient approximation for matrices using Singular Value Decomposition (SVD). This pull request extends the algorithm to support general tensor decomposition, allowing it to handle tensors of dimension greater than 2. A particular example is given in the usage of Neural Operators which are used to solve Partial Differential Equations (The tensors here are 5 dimensional).

Changes:

  • Implemented the GaLoreProjectorTensor class in Python to perform projection using orthogonal matrix decomposition for low-rank approximation of gradients for general tensors.
  • Utilized the TensorLy library (https://tensorly.org/stable/index.html) for tensor decomposition operations.
  • Modified the get_orthogonal_matrix method to compute the orthogonal matrix using Tucker decomposition instead of SVD.
  • Updated the transform and inverse_transform methods to handle general tensors using the tenalg.multi_mode_dot function from TensorLy.
  • Adjusted the input and output types of the project and project_back methods to work with general tensors.

Benefits:

  • Extends the applicability of the GaLore algorithm to handle tensors of higher dimensions.
  • Leverages the TensorLy library for efficient and optimized tensor decomposition operations.
  • Enhances the flexibility and versatility of the algorithm for various tensor-based applications.

Please review the changes and provide any feedback or suggestions for improvement. Let me know if you have any questions or concerns.

Robertboy18 avatar May 14 '24 21:05 Robertboy18