论文TensorGRaD: Tensor Gradient Robust Decomposition for Memory-Efficient Neural Operator Training提出了一种张量级的梯度鲁棒分解方法——TensorGRaD。TensorGRaD为神经算子的训练引入了结构感知的张量级梯度压缩机制，在理论与实践上均展现出卓越的表现。它不仅保留了张量结构信息，也降低了对显存的要求，特别适用于高分辨率PDE类科学任务的高效训练。该方法为今后更广泛采用张量神经网络和低资源场景中的训练提供了技术基础。
论文作者为Sebastian Loeschcke, David Pitt, Robert Joseph George, Jiawei Zhao, Cheng Luo, Yuandong Tian, Jean Kossaifi, Anima Anandkumar，来自University of Copenhagen, California Institute of Technology, Meta FAIR, NVIDIA AI。
一、研究背景与动机
随着深度学习模型不断扩展，尤其是在科学计算等对高维数据建模要求极高的领域，神经算子（Neural Operators, NOs）成为解决偏微分方程（PDEs）问题的有力工具。然而，NO模型的权重和梯度通常为高阶张量结构，这在训练过程中带来了巨大的内存消耗，尤其是在使用诸如AdamW等自适应优化器时，需要存储每个参数的梯度一阶和二阶动量，进一步加剧了问题。现有如GaLore、GRASS等方法多聚焦于低秩或稀疏结构的矩阵梯度压缩，难以保留张量多模态结构，因此对NO类模型效果不佳。为解决这一难题，作者提出了一种张量级的梯度鲁棒分解方法——TensorGRaD，兼具低秩与稀疏结构建模能力，在不损失精度的前提下可显著节省内存。
二、方法概述与核心思想
TensorGRaD的核心思想是将梯度张量 GGG 分解为低秩张量 L 与稀疏张量 S 的和，即：G=L+S
其中：
- 低秩部分 L 采用 Tucker 分解获取核心张量与正交矩阵的乘积，保持梯度张量的多模结构，类似矩阵的SVD。
- 稀疏部分 S 采用非结构化的COO格式，仅保留Top-k梯度元素的信息，可高效表示尖锐的、局部的变化。
该分解结构具有如下特性：
- 保留了张量结构中的模态信息，避免信息丢失；
- 稀疏部分可容纳异常值、离群点，提升梯度表示鲁棒性；
- 低秩部分可捕捉张量中全局的平滑模式；
- 在训练中分别对 L 与 S 进行AdamW更新，然后重构梯度。
三、方法实现细节
- 梯度压缩流程：
- 首先抽取稀疏部分 GS，剩余作为残差；
- 对残差进行Tucker低秩分解，得出低秩部分 GL；
- 两部分分别进行Adam更新，最后合并得到重构梯度；
- 为节省内存，低秩部分重构后直接将稀疏部分以scatter方式叠加，避免两份完整张量同时存在。
- 混合精度训练：
- 权重、激活、梯度使用半精度（FP16）存储和运算；
- 优化器状态（即Adam的动量）仍保留为全精度（FP32），防止精度丢失导致模型性能下降；
- 验证了在保留精度的同时，混合精度策略可额外节省显存。
- 算法伪代码：
- 提供了完整的TensorGRaD训练流程，包括梯度压缩、低秩分解、稀疏抽样与优化器更新等；
- 每T步重新选择稀疏索引与低秩子空间，缓解计算开销。
四、理论分析
作者在理论上证明了TensorGRaD在各个模态上的梯度残差会逐步收敛至零，条件是每个模态的投影算子具有足够的“最小奇异值”保障。这一结果意味着即使使用压缩梯度表示，也能保证训练过程收敛。该理论拓展了GaLore的收敛性分析，从矩阵空间推广至张量空间。
五、实验设置与结果分析
作者在多个具有挑战性的PDE任务上测试了TensorGRaD方法，包括：
- Navier-Stokes 方程：二维Kolmogorov流，Reynolds数达 105，极具湍流特性；
- Darcy Flow：非线性椭圆型PDE，测试孔隙介质中流体传播；
- Burgers Equation：一维非线性粘性守恒律；
- 电磁波传播：非线性Schrödinger方程，具有复数特征和二次谐波生成。
结果显示：
- 在Navier-Stokes 1024×1024任务中，TensorGRaD在使用仅25%内存（5%稀疏 + 20%低秩）的情况下，达到比全精度Adam更低的测试L2误差（16.82e-2 vs 17.02e-2）；
- 相较单独使用稀疏或低秩方法，TensorGRaD组合结构显著提升精度，稀疏优先分解优于低秩优先；
- 在混合精度下，保持优化器状态为全精度至关重要，否则模型性能严重下降；
- 与GaLore方法的直接张量化拓展相比，TensorGRaD显著优于其在性能与内存消耗两方面。
六、相关工作比较
TensorGRaD的优势在于：
- 将低秩与稀疏压缩同时用于张量梯度空间，而不是仅限于参数权重或矩阵结构；
- 弥补了GaLore在张量结构中的信息破坏问题；
- 相比GRASS方法中的结构化稀疏策略，非结构化稀疏方法能更精准捕捉梯度中的高值异常项；
- 可以与其他技术如iFNO、FLoRA等组合以进一步节省内存。
七、局限性与未来方向
尽管TensorGRaD在内存与性能之间达成了良好平衡，仍存在若干局限：
- Tucker分解带来的计算开销在高频更新下可能较大；
- 如何自动选择最佳低秩秩值与稀疏比例仍未解决；
- 当前只在科学问题中评估（PDEs），未来可推广至视觉、语言等高阶张量场景；
- 当前方法仍需保留优化器状态为全精度，量化版本的优化器状态尚未探索。