# TensorGRaD: Tensor Gradient Robust Decomposition for Memory-Efficient Neural Operator Training

arXiv: https://arxiv.org/abs/2501.02379 · HTML: https://mlsanity.com/paper/2501.02379/

## Key facts

- Cited by 6 papers in the map (data through 14 Oct 2026).
- Builds on 24 of its 32 references that are in the map.
- First posted on arXiv 4 Jan 2025.

## Builds on (most cited first, 24 of 24)

- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [LoRA: Low-Rank Adaptation of Large Language Models](https://mlsanity.com/paper/2106.09685/) (2021)
- [Segment Anything](https://mlsanity.com/paper/2304.02643/) (2023)
- [Fourier Neural Operator for Parametric Partial Differential Equations](https://mlsanity.com/paper/2010.08895/) (2020)
- [Training Deep Nets with Sublinear Memory Cost](https://mlsanity.com/paper/1604.06174/) (2016)
- [GaLore: Memory-Efficient LLM Training by Gradient Low-Rank Projection](https://mlsanity.com/paper/2403.03507/) (2024)
- [Neural Operator: Learning Maps Between Function Spaces](https://mlsanity.com/paper/2108.08481/) (2021)
- [Geometry-Informed Neural Operator for Large-Scale 3D PDEs](https://mlsanity.com/paper/2309.00583/) (2023)
- [Tensorizing Neural Networks](https://mlsanity.com/paper/1509.06569/) (2015)
- [Spherical Fourier Neural Operators: Learning Stable Dynamics on the Sphere](https://mlsanity.com/paper/2306.03838/) (2023)
- [Neural Operators for Accelerating Scientific Simulations and Design](https://mlsanity.com/paper/2309.15325/) (2023)
- [Speeding-up Convolutional Neural Networks Using Fine-tuned CP-Decomposition](https://mlsanity.com/paper/1412.6553/) (2014)
- [Compression of Deep Convolutional Neural Networks for Fast and Low Power Mobile Applications](https://mlsanity.com/paper/1511.06530/) (2015)
- [TensorLy: Tensor Learning in Python](https://mlsanity.com/paper/1610.09555/) (2016)
- [Multi-Grid Tensorized Fourier Neural Operator for High-Resolution PDEs](https://mlsanity.com/paper/2310.00120/) (2023)
- [Tensor Methods in Computer Vision and Deep Learning](https://mlsanity.com/paper/2107.03436/) (2021)
- [SLTrain: a sparse plus low-rank approach for parameter and memory efficient pretraining](https://mlsanity.com/paper/2406.02214/) (2024)
- [Incremental Spatial and Spectral Learning of Neural Operators for Solving Large-Scale PDEs](https://mlsanity.com/paper/2211.15188/) (2022)
- [Rethinking Conventional Wisdom in Machine Learning: From Generalization to Scaling](https://mlsanity.com/paper/2409.15156/) (2024)
- [Grass: Compute Efficient Low-Memory LLM Training with Structured Sparse Gradients](https://mlsanity.com/paper/2406.17660/) (2024)
- [LoQT: Low-Rank Adapters for Quantized Pretraining](https://mlsanity.com/paper/2405.16528/) (2024)
- [Optimization Landscape of Tucker Decomposition](https://mlsanity.com/paper/2006.16297/) (2020)
- [Guaranteed Approximation Bounds for Mixed-Precision Neural Operators](https://mlsanity.com/paper/2307.15034/) (2023)
- [Coarse Graining with Neural Operators for Simulating Chaotic Systems](https://mlsanity.com/paper/2408.05177/) (2024)

## Built on by (most cited first, 6 of 6)

- [GaLore 2: Large-Scale LLM Pre-Training by Gradient Low-Rank Projection](https://mlsanity.com/paper/2504.20437/) (2025)
- [Scalable Parameter and Memory Efficient Pretraining for LLM: Recent Algorithmic Advances and Benchmarking](https://mlsanity.com/paper/2505.22922/) (2025)
- [F-Adapter: Frequency-Adaptive Parameter-Efficient Fine-Tuning in Scientific Machine Learning](https://mlsanity.com/paper/2509.23173/) (2025)
- [SirenFNO: Efficient and Full Frequency Learning of Fourier Neural Operators](https://mlsanity.com/paper/2606.11518/) (2026)
- [M+Adam: Low-Precision Training via Additive-Multiplicative Optimization](https://mlsanity.com/paper/2607.10611/) (2026)
- [FTD-GNO: Memory-Efficient Graph Neural Operators through Functional Tensor Decomposition of the Kernel](https://mlsanity.com/paper/2610.04212/) (2026)

## Abstract (arXiv)

Scientific problems require resolving multi-scale phenomena across different resolutions and learning solution operators in infinite-dimensional function spaces. Neural operators provide a powerful framework for this, using tensor-parameterized layers to capture complex, multi-dimensional relationships. However, scaling neural operators to high-resolution problems leads to significant computational demands, making the training of industrial-scale models prohibitive. In this work, we introduce TensorGRaD, a novel method that directly addresses the memory challenges associated with optimizing large tensor-structured weights. Our approach, based on a texitrobust tensor decomposition, factorizes gradients as the sum of a low-rank tensor and a sparse one to efficiently capture information within optimizer states, including outliers. Additionally, we provide a recipe for mixed precision training of TensorGRaD, achieving further memory savings without sacrificing accuracy. We showcase the effectiveness of TensorGRaD on Fourier Neural Operators, a class of models crucial for solving partial differential equations (PDE). We provide theoretical guarantees for TensorGRaD, demonstrating its fundamental advantage over matrix-based gradient compression methods. We empirically demonstrate large improvements across various PDE tasks, including the challenging turbulent Navier-Stokes case at a Reynolds number of 105. TensorGRaD reduces total memory usage by over 50% while maintaining and sometimes even improving accuracy.