提出新型张量梯度分解方法,显著降低神经算子训练内存占用。
TensorGRaD: Tensor Gradient Robust Decomposition for Memory-Efficient Neural Operator Training
- 用低秩与稀疏张量之和分解梯度,高效压缩优化器状态
- 在雷诺数10⁵的湍流纳维-斯托克斯问题上,内存减少超50%
- 适合高分辨率科学计算场景,尤其需高效训练的神经算子模型
科学问题常涉及跨多尺度现象的求解及无限维函数空间中的解算子学习。神经算子为此提供强大框架,通过张量参数化层捕捉复杂多维关系。然而,将神经算子扩展至高分辨率问题时面临巨大计算开销,使工业级模型训练变得不可行。本文提出 extbf{TensorGRaD},一种直接应对大张量权重优化内存挑战的新方法。其基于鲁棒张量分解,将梯度分解为低秩与稀疏张量之和,以高效捕获优化器状态中的信息(包括异常值)。同时,我们提供混合精度训练方案,在不牺牲精度的前提下进一步节省内存。我们在傅里叶神经算子(Fourier Neural Operators)上验证了该方法的有效性,这是一类对求解偏微分方程至关重要的模型。理论证明表明,TensorGRaD 在根本上优于基于矩阵的梯度压缩方法。实证结果在多个偏微分方程任务中均显示显著提升,包括雷诺数达 $10^5$ 的挑战性湍流纳维-斯托克斯问题。相比基线,总内存使用降低超过 50%,且准确率保持或甚至提高。
原文摘要 · Abstract (English)
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 \textbf{TensorGRaD}, a novel method that directly addresses the memory challenges associated with optimizing large tensor-structured weights. Our approach, based on a \texit{robust 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 $10^5$. TensorGRaD reduces total memory usage by over $50\%$ while maintaining and sometimes even improving accuracy.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。