Pith. sign in

REVIEW 4 major objections 6 minor 2 cited by

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

T0 review · 4 major / 6 minor · reviewed 2026-08-10 · deepseek-v4-flash

Challenge this review Re-run · record.json Download PDF Read on arXiv ↗

Pith's one-line read TensorGRaD compresses gradients into a low-rank plus sparse tensor pair, cutting optimizer memory by up to 75% while matching or beating Adam's test error on turbulent Navier-Stokes.

desk verdict Useful empirical trick for compressing FNO gradient states, but the headline 50% memory claim doesn't survive the paper's own Table 1, and the theory rests on unverified assumptions. read the letter →

arxiv 2501.02379 v2 pith:CLPQ752G submitted 2025-01-04 cs.LG

Sebastian Loeschcke , David Pitt , Robert Joseph George , Jiawei Zhao , Cheng Luo , Yuandong Tian , Jean Kossaifi , Anima Anandkumar This is my paper · ORCID
classification cs.LG
keywords neuraloperatorsFourierOperatorrobusttensordecompositiongradientcompressionlow-rankplussparsememory-efficienttrainingmixedprecisionAdamoptimizer
verification ladder T0 review T1 audit T2 compute T3 formal

The pith

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The reading

TensorGRaD proposes that the memory bottleneck in training Fourier Neural Operators can be attacked at the optimizer-state level: decompose each gradient tensor into a low-rank part plus a sparse part before Adam stores its moments. The paper argues this robust tensor decomposition preserves the separate spatial, temporal, and channel structure of FNO weights, whereas flattening the gradient into a matrix destroys it. On a turbulent Navier-Stokes task at $1024\times1024$ resolution with Reynolds number $10^5$, a 25% optimizer-state budget (5% sparse plus 20% low-rank) achieves the lowest test loss in the paper, below full-precision Adam, while cutting total memory by over 55% when combined with a mixed-precision recipe. The paper also gives convergence and low-rank-emergence guarantees for the tensor-projected update.

What carries the argument

The machinery is a robust tensor decomposition of the gradient, $G=L+S$, combined with Adam updates computed in the compressed representation. The low-rank branch uses the Tucker (higher-order SVD) decomposition: factor matrices $U^{(n)}$ are computed once and reused for $T$ steps, the gradient is projected to a small core tensor, Adam moments live on the core, and the update is reconstructed by mode-wise products. The sparse branch is a COO-format tensor on a fixed index set $\Omega$ re-selected every $T$ steps by top-k magnitude or random sampling, so it preserves sharp outliers. The paper's theory (Theorem 1 and Lemma 10) rests on a parametric gradient form $G_t = (1/N) \sum_i (A_i - B_i \times_1 W_t \times_2 C_i)$ with mode-k PSD tensors $B_i, C_i$ and mode-k Lipschitz continuity; under that form, the mode-wise stable rank decays and the projected residual contracts, giving a convergence bound for TensorGRaD.

What would settle it

Re-run the Navier-Stokes $1024\times1024$ benchmark with the sparse branch disabled and the same 25% low-rank budget; if the resulting test $L^2$ no longer beats full-precision Adam, the specific claim that low-rank plus sparse is needed at this budget falls. Separately, record FNO gradient tensors over training and measure mode-wise stable rank; if stable rank does not decay in every mode, the low-rank-emergence mechanism that the convergence proof depends on is absent.

Watch

Extended reading notes

Core claim

The paper's central claim is that FNO gradients can be written, to good effect, as $G=L+S$, with $L$ captured by a mode-wise low-rank Tucker decomposition and $S$ captured by an unstructured sparse COO tensor, and that Adam moments can live entirely in those compressed spaces without hurting test accuracy. Because the sparse branch is extracted first, the low-rank basis is built on the residual, which removes outliers and leaves the Tucker factors a smoother signal. In the headline experiment, TensorGRaD with a 25% optimizer state (5% unstructured top-k sparsity plus 20% low-rank) reaches $16.82\times10^{-2}$ test $L^2$, beating full-precision Adam's $17.02\times10^{-2}$, and mixed-precision TensorGRaD matches Adam at $16.87\times10^{-2}$ while cutting total memory by 55%; pure low-rank and pure structured-sparse baselines at the same budget land at $17.19\times10^{-2}$ and $18.73\times10^{-2}$. The paper further claims, and proves under a parametric gradient model, that mode-wise tensor projections converge and that gradients of such models become low-rank in every mode simultaneously---something a matrix-based GaLore extension cannot do.

Load-bearing premise

The theory assumes every FNO gradient has the exact parametric form $G_t = (1/N) \sum_i (A_i - B_i \times_1 W_t \times_2 C_i)$ with mode-k positive-semidefinite tensors $B_i$ and $C_i$ and mode-k Lipschitz continuity, a form the paper asserts rather than verifies empirically.

Editorial extensions

If this is right

  • A fixed 25% optimizer-state budget (5% unstructured sparse plus 20% low-rank) beats both pure low-rank and pure structured-sparse at the same budget on the turbulent Navier-Stokes benchmark, so the two compression mechanisms are complementary rather than redundant.
  • Because optimizer states are the dominant memory category as FNO mode counts grow, compressing them to low-rank-plus-sparse directly enlarges the resolution or model size that fits on a given GPU.
  • The mixed-precision recipe---half-precision weights, activations, and gradients with full-precision optimizer states---removes another roughly 30% of peak memory without a measurable accuracy penalty, and half-precision optimizer states are explicitly shown to degrade performance.
  • A direct tensor extension of GaLore by matricizing the gradient performs substantially worse than the tensor Tucker projection, supporting the paper's claim that mode-wise structure is load-bearing for FNO training.
  • The theoretical results imply that, under the assumed gradient form, TensorGRaD converges with fixed projections, so the update can be run without frequent recomputation of the low-rank basis.

Reading between the lines

Editorial extensions of the paper, not claims the author makes directly.

  • We infer that the same low-rank-plus-sparse decomposition should transfer to other tensorized architectures---tensorized attention, tensor-train layers, or tensorized LLM adapters---because the memory category it attacks (optimizer states on high-order weights) is architecture-independent; this is a testable extension the paper does not run.
  • A direct way to test the core assumption: record gradient tensors during an FNO training run and fit them to the parametric form $G_t = (1/N) \sum_i (A_i - B_i \times_1 W_t \times_2 C_i)$; if mode-wise stable rank does not decay or the residual fit fails, the theoretical guarantee would be off the table even though the empirical method might still work.
  • The observation that compressed runs often show higher training loss but lower test loss than Adam suggests the compression acts as a regularizer; if so, tuning the sparse/low-rank budget could be used deliberately to control overfitting, not just memory.
  • Reversing the decomposition order (sparse first, low-rank on the residual) is reported as best; we infer this is because outliers otherwise perturb the Tucker subspace, and a natural follow-up would be to adapt the order per layer based on measured gradient outlierness.
Share X Bluesky LinkedIn Reddit HN

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, and a circularity audit.

Referee Report

4 major / 6 minor

Summary. TensorGRaD proposes to compress gradient tensors during neural-operator training as the sum of an unstructured sparse component and a low-rank (Tucker) component, keeping Adam moments in the compressed spaces and optionally using mixed precision. The paper reports state-of-the-art test L2 loss on a 1024x1024 turbulent Navier-Stokes task at 25% optimizer-state budget, claims total-memory reductions above 50%, and provides convergence theorems for a projected low-rank formulation. Extended experiments on Burgers, Darcy, and electromagnetic wave datasets are reported in the appendix, along with ablations of sparsification order, sparsity type, and update frequency.

Significance. If the claims hold, TensorGRaD would be a useful memory-efficient training method for tensor-structured scientific models, and the idea of a sparse-plus-low-rank decomposition of gradients is a natural and interesting extension of GaLore and GRASS. The NS1024 result, where 5% unstructured sparsity plus 20% low-rank matches or slightly beats Adam at lower memory, is plausible and worth reporting. The paper also provides code and machine-checkable pseudocode, and it attempts a mode-wise theoretical comparison against matrix-based GaLore. However, the central quantitative memory claim is not supported by the manuscript's own table, and the convergence theory rests on an unverified parametric model of FNO gradients, with the main theorem actually analyzing a regime that does not correspond to the proposed sparse-plus-low-rank algorithm. The empirical contribution is therefore real but currently overstated.

major comments (4)
  1. [Abstract, Section 1, Table 1] The abstract and introduction claim that TensorGRaD 'reduces total memory usage by over 50%' and 'cut[s] total memory cost of more than 55%', but Table 1, the only place where total peak memory is reported for the headline NS1024 experiment, does not support these numbers. From Table 1, TensorGRaD 25% uses 46 GB full / 29 GB mixed, while Adam uses 52 GB full / 37 GB mixed. The reductions are 11.5% (46 vs 52), 21.6% (29 vs 37), and 44.2% (29 vs 52). The 55% figure may refer to the TensorGRaD-H+ configuration in Figure 2, but that configuration's numerical memory values are not reported. The authors must either report the exact configuration and measured peak memory that yields >50% total-memory reduction, or revise the abstract and intro claims to match Table 1.
  2. [Appendix M, Lemma 10, Remark 1; Theorem 1] All convergence and low-rank-emergence results assume the parametric gradient form G_t = (1/N) sum_i (A_i - B_i x1 W_t x2 C_i) with mode-k PSD tensors B_i, C_i and mode-k Lipschitz continuity. This form is asserted, not verified empirically, and Remark 1's claim that mode-k continuity 'holds generically' is not proved. Moreover, the derivation in Appendix L is for reversible networks with reversible activations and specific losses, whereas the actual FNO experiments use standard FNO layers and complex-valued weights; no argument is given that the trained FNO gradients satisfy Lemma 8's form. Since the theoretical claims are the paper's stated 'fundamental advantage', the authors should either verify the model on real FNO gradients or explicitly characterize the theory as conditional on this unverified assumption.
  3. [Section 2.5, Theorem 1; Appendix M, Proof 8] Theorem 1 states convergence for 'TensorGRaD with rho_t = 1'. In Algorithm 1, rho is the sparsity density, so rho_t = 1 means the 'sparse' component is the full dense gradient and the low-rank component is empty; this is not the proposed 5%+20% method. The theorem therefore does not analyze TensorGRaD as defined, and no theorem in the paper covers the combined sparse-plus-low-rank decomposition with subspace updates. The proof also contains a non sequitur: it cites the optimization landscape of Tucker decomposition [36] to assert that local minima are globally optimal, without proof, and it ignores the Adam dynamics and the sparse component. The convergence claim needs to be either proved for the actual algorithm or removed/qualified.
  4. [Section 3, Appendix G, Tables 5-9] The paper's broad empirical claims across Burgers, Darcy, and electromagnetic wave propagation are supported only by experiments with the low-rank component alone ('Tensor low-rank'), not by the full TensorGRaD sparse-plus-low-rank method. Table 1 and Table 4 are the only evaluations of the full method, both on NS1024. The abstract's 'large improvements across various PDE tasks' should either be backed by full-method results on those datasets or explicitly restricted to the low-rank variant.
minor comments (6)
  1. [Section 3] The text refers to 'Appendix 6' when listing extended benchmark results; this should be 'Appendix G' (or the appropriate appendix label).
  2. [Section 3, Table 2] Table 2's caption does not state the dataset and Reynolds number; the body text implies NS128, but the caption should say so explicitly.
  3. [Algorithm 1 and Theorem 1] The symbol rho is used for sparsity density, but the theorem's condition 'rho_t = 1' is never defined in context; please clarify the notation and state which rho the theorem refers to.
  4. [Throughout] The naming is inconsistent between 'TENSORGRAD' in the main text and 'Tensor low-rank' in the appendix; the appendix results appear to describe a different method, and the reader cannot tell whether 'Tensor low-rank' is a component of TENSORGRAD or a separate algorithm.
  5. [Section 3, Introduction] The sentence 'cutting and total memory cost of more than 55%' contains a typo ('and' should be removed), and the phrase 'reduces memory usage by up to 75%' should distinguish optimizer-state memory from total memory.
  6. [Figure 2 and Figure 3] The right panel of Figure 3 appears to duplicate the left panel of Figure 2; if these are separate experiments, please make the difference explicit, otherwise remove the duplicate.

Circularity Check

0 steps flagged · score 0.0 of 10

No significant circularity: theoretical guarantees are conditional on an explicit (if unverified) gradient model and empirical claims are benchmark-based.

full rationale

The paper's derivation chain is not circular. The central theoretical results (Section 2.5, Appendix M) are conditional theorems: Lemma 10 and Theorem 1 explicitly suppose that G_t = (1/N) sum_i (A_i - B_i x1 W_t x2 C_i) with PSD tensors B_i and C_i, and then derive stable-rank decay and contraction under a stated eigenvalue condition. The weak point is real but of a different kind: this parametric gradient model is asserted rather than verified on trained FNOs, and no evidence is given that the convergence condition min_{t,k} kappa_t^(k) > L_A^(k) + L_B^(k) L_C^(k) D_k^2 actually holds. A theorem whose assumptions are unverified is fragile, but it is not circular: the conclusion is not identical to the premise, and the algebra from the assumption to the bound is a genuine derivation. The empirical claims are benchmark measurements, not predictions fitted from the theory. The >50% and >55% total-memory statements in the abstract and introduction are not supported by Table 1, where 29 GB mixed-precision TensorGRaD versus 52 GB full-precision Adam is about 44% and versus 37 GB mixed-precision Adam is about 22%; this is an internal reporting inconsistency, not a circular reduction. The use of GaLore's Lemma B.3 (Appendix M, Lemma 11) is a self-citation by overlapping authors, but the lemma is concrete, parameter-free, and its assumptions do not include the tensor mode-wise target being proved, so it functions as a stated mathematical input rather than an appeal to authority. No fitted parameter is renamed as a prediction, and no known result is merely re-labeled. The paper's own limitations section honestly admits the computational overhead, rank-selection difficulty, and full-precision optimizer-state requirement, all of which are scope concerns rather than evidence of circularity.

Assumptions & free parameters 5 free parameters · 6 assumptions · 0 invented entities

The central empirical claim depends on hyperparameters r, rho, lambda, T, and the decomposition order, all chosen by hand or ablation. The theoretical claims depend on strong unverified structural assumptions about FNO gradients, including a parametric form with mode-k PSD tensors and a landscape condition that is asserted rather than proven. No new physical or model entities are introduced.

free parameters (5)
  • low-rank ratio r (25% setting: 20%) = 0.2
    Chosen by hand or ablation as the best 25% setting; directly determines memory and accuracy of the low-rank component.
  • sparsity density rho (25% setting: 5%) = 0.05
    Chosen by ablation; controls the number of top-k unstructured sparse indices retained.
  • sparse scale factor lambda = not reported
    Scales the sparse update in Algorithm 1; no value or tuning procedure is given in the paper.
  • subspace update frequency T = 500 (NS1024), 1000 (NS128)
    Controls how often Tucker factors and sparse index sets are recomputed; ablations show it affects generalization.
  • decomposition order (sparse then low-rank) = US -> LR
    The order is chosen by ablation; Table 2 shows reversing the order increases test loss.
assumptions (6)
  • domain assumption FNO gradients have parametric form G_t = (1/N) * sum_i (A_i - B_i x1 W_t x2 C_i)
    Introduced in Appendix M, Lemma 10; derived for a specific reversible-network setup and used to prove low-rank emergence. Not verified for real FNO training.
  • domain assumption B_i and C_i are mode-k PSD tensors and mode-k Lipschitz
    Assumed in Theorem 1 and Lemma 10 to define kappa_t and ensure contraction; the paper claims this holds generically in Remark 1 without proof.
  • domain assumption Mode-k spectral norm bound ||W_t||_(k) <= D_k holds
    Used in the convergence proof of Theorem 1; no verification is provided.
  • ad hoc to paper Optimization landscape of low-rank tensor problems contains only easily escapable saddle points
    Invoked in Proof 8 to conclude convergence from low-rank gradient structure; this is a non-sequitur and is not derived in the paper.
  • standard math Tucker/HOSVD properties and tensor mode products
    Used throughout for compression, projection, and reconstruction of gradient tensors.
  • domain assumption Small logits assumption for the logsoftmax gradient form
    Used in Lemma 9 to derive a closed-form gradient; this approximation is not validated for FNO training.

how reviews work

0 comments
Cite this review

Pith. "Pith review of TensorGRaD: Tensor Gradient Robust Decomposition for Memory-Efficient Neural Operator Training." pith.science (2026). https://pith.science/paper/CLPQ752G

@misc{pith2026250102379,
  author       = {Pith},
  title        = {Pith review of: TensorGRaD: Tensor Gradient Robust Decomposition for Memory-Efficient Neural Operator Training},
  year         = {2026},
  howpublished = {\url{https://pith.science/paper/CLPQ752G}},
  note         = {Machine review of arXiv:2501.02379}
}
abstract

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.

Figures

Figures reproduced from arXiv: 2501.02379 by the authors.

Figure 1
Figure 1. Overview of TENSORGRAD. Low-rank plus sparse decomposition Unlike neural networks, neural operators learn a mapping between function spaces, making them naturally suited for capturing the multi-scale structure of scientific data. To capture these multi￾scale relationships, Neural Operators leverage the inherent (multi-dimensional) structure in the data, which requires maintaining high-order tensor weights and gradie… view at source ↗
Figure 2
Figure 2. Memory and performance. Left: Comparison of low-rank, structured sparse, and TEN￾SORGRAD (mixed precision) vs. Adam on Navier–Stokes 1024 × 1024. TENSORGRAD offers the best memory–accuracy trade-off. Right: Peak CUDA memory for FNO models with 256 channels. H: uses half precision for weights/gradients and full for optimizer states, and + includes activation checkpointing. We demonstrate that our robust tensor gradie… view at source ↗
Figure 3
Figure 3. Left: Gradient reconstruction error. Box plot of gradient reconstruction error for different compression strategies on a complex FNO layer. TENSORGRAD variants reduce high￾magnitude outliers. Right: Accuracy–memory trade-off. Performance of low-rank and structured sparse across compression ratio vs. TENSORGRAD at 25% and Adam baselines. TENSORGRAD achieves the best trade-off at. Datasets. We report results on the se… view at source ↗
Figures from the paper (5 more)
Figure 4
Figure 4. Figure 4: Navier-Stokes 1024 × 1024 with Reynolds number at 2 × 105 Appendix A Dataset A.1 Navier-Stokes Datasets Navier-Stokes 1024: We use the 2D Kolmogorov flow of Wang et al. [19], a periodically forced, incompressible variant of the Navier–Stokes equations. The velocity fie…
Figure 5
Figure 5. Figure 5: Memory usage in FNO as a function of the number of modes [PITH_FULL_IMAGE:figures/full_fig_p016_5.png]
Figure 6
Figure 6. Figure 6: Memory usage of FNO and GINO on various datasets on an NVIDIA A100. On top of the [PITH_FULL_IMAGE:figures/full_fig_p017_6.png]
Figure 7
Figure 7. Figure 7: Left: Low-rank, structured, at various rank percentages vs T [PITH_FULL_IMAGE:figures/full_fig_p022_7.png]
Figure 8
Figure 8. Figure 8: Comparison of the higher-order low-rank decomposition Tensor Low-Rank with standard [PITH_FULL_IMAGE:figures/full_fig_p028_8.png]

Discussion (0). Continue with ORCID to comment.

Forward citations

Cited by 2 Pith papers

Reviewed papers in the Pith corpus that reference this work. Sorted by Pith novelty score. Full citation record

  1. M+Adam: Low-Precision Training via Additive-Multiplicative Optimization

    cs.LG 2026-07 conditional novelty 6.0 of 10

    M+Adam combines additive and multiplicative update branches, avoiding low-precision rounding stalls and improving LLaMA-style pretraining perplexity compared with AdamW at BF16, FP8, and FP4 master-weight storage.

  2. Scalable Parameter and Memory Efficient Pretraining for LLM: Recent Algorithmic Advances and Benchmarking

    cs.LG 2025-05 conditional novelty 4.0 of 10

    A benchmark and two low-cost tricks (weight refactorization and momentum reset) that make low-rank LLM pre-training competitive with GaLore and Fira at about 25% lower memory.

Reference graph

Works this paper leans on

129 extracted references · 71 canonical work pages · cited by 2 Pith papers

  1. [36]

    Optimization Landscape of Tucker Decomposition

    Abraham Frandsen and Rong Ge. Optimization landscape of tucker decomposition.CoRR, abs/2006.16297, 2020. 13 T=0 T=256 T=512 T=768 T=1024 T=1280 T=1536 T=1792 T=2048 80 60 40 20 0 20 40 60 80 Figure 4: Navier-Stokes1024×1024with Reynolds number at2×10 5 Appendix A Dataset A.1 Navier-Stokes Datasets Navier-Stokes 1024:We use the 2D Kolmogorov flow of Wang e...

    work page Pith review arXiv 2006

  2. [1]

    Language mod- els are few-shot learners.Advances in neural information processing systems, 33:1877–1901, 2020

    Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language mod- els are few-shot learners.Advances in neural information processing systems, 33:1877–1901, 2020

    1901

  3. [2]

    Segment anything

    Alexander Kirillov, Eric Mintun, Nikhila Ravi, Hanzi Mao, Chloe Rolland, Laura Gustafson, Tete Xiao, Spencer Whitehead, Alexander C Berg, Wan-Yen Lo, et al. Segment anything. InProceedings of the IEEE/CVF International Conference on Computer Vision, pages 4015– 4026, 2023

    2023

  4. [3]

    Rethinking conventional wisdom in machine learning: From generalization to scaling, 2025

    Lechao Xiao. Rethinking conventional wisdom in machine learning: From generalization to scaling, 2025

    2025

  5. [4]

    Adam: A method for stochastic optimization.arXiv preprint arXiv:1412.6980, 2014

    Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization.arXiv preprint arXiv:1412.6980, 2014

    arXiv 2014

  6. [5]

    Galore: Memory-efficient llm training by gradient low-rank projection.arXiv preprint arXiv:2403.03507, 2024

    Jiawei Zhao, Zhenyu Zhang, Beidi Chen, Zhangyang Wang, Anima Anandkumar, and Yuan- dong Tian. Galore: Memory-efficient llm training by gradient low-rank projection.arXiv preprint arXiv:2403.03507, 2024

    arXiv 2024

  7. [6]

    Grass: Compute efficient low-memory llm training with structured sparse gradients.arXiv preprint arXiv:2406.17660, 2024

    Aashiq Muhamed, Oscar Li, David Woodruff, Mona Diab, and Virginia Smith. Grass: Compute efficient low-memory llm training with structured sparse gradients.arXiv preprint arXiv:2406.17660, 2024

    arXiv 2024

  8. [7]

    Kastoryano, Serge Belongie, and V ´esteinn Snæ bjarnarson

    Sebastian Loeschcke, Mads Toftrup, Michael J. Kastoryano, Serge Belongie, and V ´esteinn Snæ bjarnarson. Loqt: Low-rank adapters for quantized pretraining. In A. Globerson, L. Mackey, D. Belgrave, A. Fan, U. Paquet, J. Tomczak, and C. Zhang, editors,Advances in Neural Information Processing Systems, volume 37, pages 115282–115308. Curran Asso- ciates, Inc., 2024

    2024

Show all 129 references
  1. [8]

    Neural operators for accelerating scientific simulations and design

    Kamyar Azizzadenesheli, Nikola Kovachki, Zongyi Li, Miguel Liu-Schiaffini, Jean Kossaifi, and Anima Anandkumar. Neural operators for accelerating scientific simulations and design. Nat Rev Phys, 6:320–328, 2024

    2024

  2. [9]

    Kolda and Brett W

    Tamara G. Kolda and Brett W. Bader. Tensor decompositions and applications.SIAM Review, 51(3):455–500, 2009

    2009

  3. [10]

    Spherical fourier neural operators: Learning stable dy- namics on the sphere, 2023

    Boris Bonev, Thorsten Kurth, Christian Hundt, Jaideep Pathak, Maximilian Baust, Karthik Kashinath, and Anima Anandkumar. Spherical fourier neural operators: Learning stable dy- namics on the sphere, 2023. 11

    2023

  4. [11]

    Geometry-informed neural operator for large-scale 3d pdes, 2023

    Zongyi Li, Nikola Borislavov Kovachki, Chris Choy, Boyi Li, Jean Kossaifi, Shourya Prakash Otta, Mohammad Amin Nabian, Maximilian Stadler, Christian Hundt, Kamyar Azizzade- nesheli, and Anima Anandkumar. Geometry-informed neural operator for large-scale 3d pdes, 2023

    2023

  5. [12]

    Robust tensor decomposition with gross corruption

    Quanquan Gu, Huan Gui, and Jiawei Han. Robust tensor decomposition with gross corruption. In Z. Ghahramani, M. Welling, C. Cortes, N. Lawrence, and K.Q. Weinberger, editors,Ad- vances in Neural Information Processing Systems, volume 27. Curran Associates, Inc., 2014

    2014

  6. [13]

    Cand `es, Xiaodong Li, Yi Ma, and John Wright

    Emmanuel J. Cand `es, Xiaodong Li, Yi Ma, and John Wright. Robust principal component analysis?CoRR, abs/0912.3599, 2009

    2009 arXiv

  7. [14]

    A unified framework for nonconvex low- rank plus sparse matrix recovery

    Xiao Zhang, Lingxiao Wang, and Quanquan Gu. A unified framework for nonconvex low- rank plus sparse matrix recovery. In Amos Storkey and Fernando Perez-Cruz, editors,Pro- ceedings of the Twenty-First International Conference on Artificial Intelligence and Statistics, volume 84...

    2018

  8. [15]

    Ledyard R. Tucker. Some mathematical notes on three-mode factor analysis.Psychometrika, 31:279–311, 1966

    1966

  9. [16]

    Tensor decompositions and applications.SIAM Rev., 51(3):455–500, 2009

    Tamara G Kolda and Brett W Bader. Tensor decompositions and applications.SIAM Rev., 51(3):455–500, 2009

    2009

  10. [17]

    Guaranteed approximation bounds for mixed-precision neural operators, 2024

    Renbo Tu, Colin White, Jean Kossaifi, Boris Bonev, Nikola Kovachki, Gennady Pekhimenko, Kamyar Azizzadenesheli, and Anima Anandkumar. Guaranteed approximation bounds for mixed-precision neural operators, 2024

    2024

  11. [18]

    Tensorly: Tensor learning in python.Journal of Machine Learning Research, 20(26):1–6, 2019

    Jean Kossaifi, Yannis Panagakis, Anima Anandkumar, and Maja Pantic. Tensorly: Tensor learning in python.Journal of Machine Learning Research, 20(26):1–6, 2019

    2019

  12. [19]

    Beyond closure models: Learning chaotic-systems via physics-informed neural opera- tors.arXiv preprint arXiv:2408.05177, 2024

    Chuwei Wang, Julius Berner, Zongyi Li, Di Zhou, Jiayun Wang, Jane Bae, and Anima Anand- kumar. Beyond closure models: Learning chaotic-systems via physics-informed neural opera- tors.arXiv preprint arXiv:2408.05177, 2024

    2024 arXiv

  13. [20]

    Training deep nets with sublinear memory cost, 2016

    Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin. Training deep nets with sublinear memory cost, 2016

    2016

  14. [21]

    Tensorizing neu- ral networks.arXiv preprint arXiv:1509.06569, 2015

    Alexander Novikov, Dmitry Podoprikhin, Anton Osokin, and Dmitry Vetrov. Tensorizing neu- ral networks.arXiv preprint arXiv:1509.06569, 2015

    2015 arXiv

  15. [22]

    Speeding-up convolutional neural networks using fine-tuned cp-decomposition

    Vadim Lebedev, Yaroslav Ganin, Maksim Rakhuba, Ivan Oseledets, and Victor Lempitsky. Speeding-up convolutional neural networks using fine-tuned cp-decomposition. InInterna- tional Conference on Learning Representations, 2015

    2015

  16. [23]

    Compression of deep convolutional neural networks for fast and low power mobile applica- tions

    Yong-Deok Kim, Eunhyeok Park, Sungjoo Yoo, Taelim Choi, Lu Yang, and Dongjun Shin. Compression of deep convolutional neural networks for fast and low power mobile applica- tions. InInternational Conference on Learning Representations, 2016

    2016

  17. [24]

    Chrysos, James Oldfield, Mihalis A

    Yannis Panagakis, Jean Kossaifi, Grigorios G. Chrysos, James Oldfield, Mihalis A. Nicolaou, Anima Anandkumar, and Stefanos Zafeiriou. Tensor methods in computer vision and deep learning.Proceedings of the IEEE, 109(5):863–890, 2021

    2021

  18. [25]

    Chrysos, James Oldfield, Taylor Patti, Mi- halis A

    Yannis Panagakis, Jean Kossaifi, Grigorios G. Chrysos, James Oldfield, Taylor Patti, Mi- halis A. Nicolaou, Anima Anandkumar, and Stefanos Zafeiriou. Chapter 15 - tensor methods in deep learning. In Paulo S.R. Diniz, editor,Signal Processing and Machine Learning Theory, pages ...

    2024

  19. [27]

    Neural operator: Learning maps between function spaces.arXiv preprint arXiv:2108.08481, 2021

    Nikola Kovachki, Zongyi Li, Burigede Liu, Kamyar Azizzadenesheli, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Neural operator: Learning maps between function spaces.arXiv preprint arXiv:2108.08481, 2021

    2021 arXiv

  20. [28]

    Fourier neural operator for parametric partial differ- ential equations.arXiv preprint arXiv:2010.08895, 2021

    Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Fourier neural operator for parametric partial differ- ential equations.arXiv preprint arXiv:2010.08895, 2021

    2010 arXiv

  21. [29]

    Lora: Low-rank adaptation of large language models.arXiv preprint arXiv:2106.09685, 2022

    Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, and Weizhu Chen. Lora: Low-rank adaptation of large language models.arXiv preprint arXiv:2106.09685, 2022

    2022 arXiv

  22. [30]

    Flora: Low-rank core space for n-dimension, 2024

    Chongjie Si, Xuehui Wang, Xue Yang, Zhengqin Xu, Qingyun Li, Jifeng Dai, Yu Qiao, Xi- aokang Yang, and Wei Shen. Flora: Low-rank core space for n-dimension, 2024

    2024

  23. [31]

    Multi-grid tensorized fourier neural operator for high-resolution pdes.arXiv preprint arXiv:2403.00071, 2024

    Jean Kossaifi, Nikola Kovachki, Kamyar Azizzadenesheli, and Anima Anandkumar. Multi-grid tensorized fourier neural operator for high-resolution pdes.arXiv preprint arXiv:2403.00071, 2024

    2024 arXiv

  24. [32]

    Incre- mental spatial and spectral learning of neural operators for solving large-scale pdes, 2024

    Robert Joseph George, Jiawei Zhao, Jean Kossaifi, Zongyi Li, and Anima Anandkumar. Incre- mental spatial and spectral learning of neural operators for solving large-scale pdes, 2024

    2024

  25. [33]

    Candes, Xiaodong Li, Yi Ma, and John Wright

    Emmanuel J. Candes, Xiaodong Li, Yi Ma, and John Wright. Robust principal component analysis?, 2009

    2009

  26. [34]

    Sltrain: a sparse plus low rank approach for parameter and memory effi- cient pretraining

    Andi Han, Jiaxiang Li, Wei Huang, Mingyi Hong, Akiko Takeda, Pratik Jawanpuria, and Bamdev Mishra. Sltrain: a sparse plus low rank approach for parameter and memory effi- cient pretraining. InThe Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024

    2024

  27. [35]

    Scatterbrain: Unifying sparse and low-rank attention.Advances in Neural Information Processing Systems, 34:17413–17426, 2021

    Beidi Chen, Tri Dao, Eric Winsor, Zhao Song, Atri Rudra, and Christopher R ´e. Scatterbrain: Unifying sparse and low-rank attention.Advances in Neural Information Processing Systems, 34:17413–17426, 2021

    2021

  28. [37]

    Trained with H1 loss. NS128 FNO • 4 layers, 64 x 64 modes • 64 hidden channels, 256 projection channels • Skip: ’linear’ • Use channel MLP: 1 • Channel MLP expansion: 0.5, dropout: 0 Complex Adam with step LR 3e-4, weight decay 1e-4, 500 epochs, batch size 8. Update decomposit...

  29. [39]

    Darcy Flow FNO • 4 layers, 64 modes • 128 hidden channels, 128 projection channels • Skip: ’linear’ Adam with step LR 1e−3, weight de- cay1e−4, 250 epochs, batch size

    Trained with L2loss. Darcy Flow FNO • 4 layers, 64 modes • 128 hidden channels, 128 projection channels • Skip: ’linear’ Adam with step LR 1e−3, weight de- cay1e−4, 250 epochs, batch size

  30. [40]

    Trained withL 2 loss. EM Wave Complex- FNO • 8 layers, 128 modes • 128 hidden channels, 128 projection channels • Skip: ’linear’ • Complex data: True • Complex activation function: True Complex Adam with step LR 1e-4, weight decay 2e-6, batch size 32, 1000 epochs. Trained with...

  31. [41]

    The savings become even more pronounced as the spatial dimensions (M) increase, which is crucial for high-resolution problems

    Matrix GaLore (best case): PM atrix=r 2 max(N 2 +M 2)(9) 2.Tensor low-rank: PT ensor=r 4 max + 2rmaxN+ 2r maxM(10) In typical neural operator architectures: •N≫r max (number of channels much larger than rank) •M≫r max (spatial dimensions much larger than rank) Therefore: • Mat...

  32. [42]

    The Tucker decomposition preserves the natural tensor structure of the operator

  33. [43]

    Each mode has its own rank parameter, allowing for more flexible approximation

  34. [44]

    30 K Tensor Operations and Notation Definition 1 (Tensor)An order-d tensorA ∈R I1×I2×···×Id is a d-dimensional array with entries ai1,i2,...,id, where1≤i k ≤I k fork= 1,

    The core tensor captures higher-order interactions between modes This explains whyTensor low-rankcan achieve comparable or better performance while using significantly less memory than matrix-based approaches. 30 K Tensor Operations and Notation Definition 1 (Tensor)An order-d...

  35. [45]

    Frobenius norm:∥A∥ F = p ⟨A,A⟩

  36. [46]

    Mode-k spectral norm:∥A∥ (k) =∥A (k)∥2

  37. [47]

    Spectral norm:∥A∥= max ∥x(k)∥=1 ∥A ×1 x(1) ×2 · · · ×d x(d)∥ Definition 6 (Tensor Outer Product)The outer product of vectorsu (k) ∈R Ik fork= 1, . . . , dis a tensorA=u (1) ◦u (2) ◦ · · · ◦u(d) with entries: ai1,...,id =u (1) i1 u(2) i2 · · ·u(d) id Definition 7 (Tensor Contra...

  38. [48]

    Mode-wise trace: trk(A) = IkX ik=1 Ai1,...,ik,...,id |ik=ik

  39. [49]

    Using mode-k unfolding: tr(A(k)) = IkX i=1 (A(k))i,i

  40. [50]

    Inner product interpretation: When used in expressions like tr(dW ⊤ l ×1 X× 2 Y), this is actually computing: ⟨dWl, X⊗Y⟩ Proposition 2 (Key Properties)For the trace operation in tensor gradients:

  41. [51]

    Inner Product Form: tr(dW ⊤ ×1 X× 2 Y) =⟨dW, X⊗Y⟩

  42. [52]

    Differential Form: For scalar functionϕand tensorW: dϕ=tr(dW ⊤ ×1 X× 2 Y) =⇒ ∂ϕ ∂W =X⊗Y

  43. [53]

    Mode-wise Consistency: tr(dW ⊤ ×1 X× 2 Y) =tr(X ⊤dW(1)Y) wheredW (1) is the mode-1 unfolding. Example 1In the logsoftmax gradient computation: −dϕ=tr(dW ⊤ l ×1 (P ⊥ 1 y)⊤Jl ×2 f ⊤ l−1) =⟨dW l,J ⊤ l P ⊥ 1 y⊗f l−1⟩ This leads to the gradient term: Gl =J ⊤ l P ⊥ 1 y⊗f l−1 Remark ...

  44. [54]

    sr k(T)is invariant under orthogonal transformations in modek

  45. [55]

    For any matrixM, we know∥M∥ 2 F ≥ ∥M∥2

    For a rank-1 tensor, sr k(T) = 1for allk Proof 11. For any matrixM, we know∥M∥ 2 F ≥ ∥M∥2

  46. [56]

    Therefore: srk(T) = ∥T ∥2 F ∥T(k)∥2 2 = ∥T(k)∥2 F ∥T(k)∥2 2 ≥1 where we used the tensor-matrix norm relation lemma

  47. [57]

    For any matrixMof rankr: ∥M∥ 2 2 ≥ ∥M∥ 2 F r Applying this toT (k): srk(T) = ∥T(k)∥2 F ∥T(k)∥2 2 ≤rank(T (k))

  48. [58]

    For any orthogonal transformationUin modek: ∥UT (k)∥F =∥T (k)∥F and∥UT (k)∥2 =∥T (k)∥2

  49. [59]

    Remark 4 (Connection to Low-Rank Approximation)The stable rank of a tensor in each mode provides insight into how well it can be approximated by a low-rank decomposition:

    For a rank-1 tensorT=a 1 ⊗...⊗a d: 33 • Each mode-k unfolding is rank-1 • For rank-1 matrices,∥M∥ 2 F =∥M∥ 2 2 • Therefore sr k(T) = 1 Definition 12 (Multilinear Stable Rank)For a tensorT, the multilinear stable rank is: msr(T) := min k srk(T) This provides a lower bound on th...

  50. [60]

    If srk(T)is close to 1 in modek, thenTis nearly low-rank in that mode

  51. [61]

    For a Tucker decomposition: T ≈ G ×1 U (1) ×2 U (2)...× d U (d) The stable rank helps determine appropriate ranks for each mode Remark 5 (Application to FNO)For FNO weight tensorsR ∈R N1×N2×N3×N4:

  52. [62]

    Mode-3 and Mode-4 cor- respond to Fourier modes 3

    Mode-1 and Mode-2 typically correspond to input/output channels 2. Mode-3 and Mode-4 cor- respond to Fourier modes 3. Stable rank in Fourier modes often naturally decreases due to spectral decay K.3 Positive Semi-Definiteness for Tensors Definition 13 (Mode-k PSD Tensor)A tens...

  53. [63]

    Mode-specific PSD condition: Bi andC i are mode-k PSD for relevant modes k

  54. [64]

    This means for each mode k: •(B i)(k) is a PSD matrix 34 •(C i)(k) is a PSD matrix •The tensor operatorS k = 1 N PN i=1 Ci ⊗k Bi is well-defined

  55. [65]

    This ensures: •The mode-k eigenvaluesλ (k) 1 , λ(k) 2 are real and non-negative •The projection onto minimal eigenspace is well-defined for each mode •The stable rank bounds make sense mode-wise Proposition 4 (For FNO)In FNO, the tensorsB i andC i naturally satisfy mode-k PSD ...

  56. [66]

    For channel modes (1,2): • Unfoldings correspond to standard channel operations • PSD property follows from network structure

  57. [67]

    For Fourier modes (3,4): • Unfoldings correspond to frequency domain operations • PSD property follows from spectral properties Corollary 1 (Implications for Gradient Analysis)The mode-k PSD property ensures:

  58. [68]

    Each mode has real, non-negative eigenvalues: 0≤λ (k) 1 < λ(k) 2 ≤ · · ·

  59. [69]

    Mode-wise stable rank bounds are well-defined: srk(Gt)≤sr k(G∥ t0 ) +decay term

  60. [70]

    The gradient naturally becomes low-rank in each mode independently. Definition 16 (Lipschitz Continuity)A functionh:X → Ybetween normed spaces hasL- continuity (isL-Lipschitz) if for anyx 1, x2 ∈ X: ∥h(x1)−h(x 2)∥Y ≤L∥x 1 −x 2∥X For tensors, this generalizes to mode-wise conti...

  61. [71]

    Backward:dx=J(x) ⊤dy whereJ(x)can be a function of both input and weights. 35 L.2 Spectral Layer Lemma 3 (Spectral Layer Reversibility)The FNO spectral convolution layer(Kv)(x) =F −1(R· Fv)(x)is reversible, whereRis the learnable weight tensor in Fourier space. The spectral la...

  62. [72]

    Fourier transform:F:v7→ˆv

  63. [73]

    Linear transform in Fourier space:R·: ˆv7→Rˆv

  64. [74]

    L.3 MLP Layer Lemma 4 (MLP Layer Reversibility)The MLP layer with weight matrixWmappingv7→W vis reversible

    Inverse Fourier:F −1 :Rˆv7→ F−1(Rˆv) We can express the complete operation as: Kv=J K(x)vwhereJ K(x) =F −1RF For the backward pass: dv=J K(x)⊤dy=F ⊤R⊤(F −1)⊤dy SinceFis unitary:F ⊤ =F −1 and(F −1)⊤ =F, we have: dv=F −1R⊤Fdy Therefore: • Forward pass:y=J K(x)x • Backward pass:d...

  65. [75]

    L.4 Activation Function Lemma 5 (Activation Reversibility)If the activation functionσis reversible (e.g., LeakyReLU), then its application is reversible

    Backward pass:dv=W ⊤dy=J W (x)⊤dy The linear layer satisfies reversibility conditions directly, even whenWis rank-deficient. L.4 Activation Function Lemma 5 (Activation Reversibility)If the activation functionσis reversible (e.g., LeakyReLU), then its application is reversible...

  66. [76]

    Forward:y= max(ax, x)

  67. [77]

    SetJ σ(x) =diag(1[x >0] +a·1[x≤0])

  68. [78]

    36 L.5 Full FNO Analysis Lemma 6 (FNO Block Reversibility)An FNO block consisting of spectral layer(K), MLP layer (W), and reversible activation(σ)is reversible

    Backward:dx=J σ(x)⊤dy This matches the required reversibility form. 36 L.5 Full FNO Analysis Lemma 6 (FNO Block Reversibility)An FNO block consisting of spectral layer(K), MLP layer (W), and reversible activation(σ)is reversible. LetN= (σ◦W◦K)be an FNO block. From previous the...

  69. [79]

    Each blockN i has itsJ i(x)from previous lemma

  70. [80]

    By sequential composition: y=J FNO(x)v whereJ FNO(x) =J L(x)JL−1(x)...J1(x)

  71. [81]

    The backward pass follows from composition: dv=J 1(x)⊤...JL−1(x)⊤JL(x)⊤dy=J FNO(x)⊤dy Therefore, the full FNO with reversible activations satisfies the reversibility conditions. Lemma 8 (Gradient Form for Tensor Reversible Models)Consider a chained reversible neural networkN(x...

  72. [82]

    For linear layerl: •f l−1 ∈R Nl−1 is input •W l ∈R N1×N2×N3×N4 is weight tensor •W l ×1 fl−1 maps toR Nl 38 •J l ∈R K×N l is Jacobian

  73. [83]

    Then, like before we do the differential computation: dϕ= (y− N(x))⊤dN(x)[R K ×R K →R] = (y− N(x))⊤Jl(dWl ×1 fl−1)[R K ×R K×N l ×R Nl →R] = (y− Jl ×1 (Wl ×1 fl−1))⊤Jl ×1 (dWl ×1 fl−1)

  74. [84]

    Mode-wise analysis for gradient: • First term:J ⊤ l y⊗f l−1 -J ⊤ l y∈R Nl -f l−1 ∈R Nl−1 - Outer product gives tensor in RN1×N2×N3×N4 • Second term:(J ⊤ l JlWl ×1 fl−1)⊗f l−1 -J ⊤ l Jl ∈R Nl×Nl -W l ×1 fl−1 ∈R Nl - Result is tensor inR N1×N2×N3×N4

  75. [85]

    Remark 8We only wanted to show an example of checking all the dimensions to ensure they match the generalized version for tensors

    Therefore final gradient: Gl =J ⊤ l y⊗f l−1 −(J ⊤ l JlWl ×1 fl−1)⊗f l−1 ∈R N1×N2×N3×N4 We finally have a gradient tensor of the same shape asW l. Remark 8We only wanted to show an example of checking all the dimensions to ensure they match the generalized version for tensors. ...

  76. [86]

    For reversible network,dN(x) =J l ×1 (dWl ×1 fl−1)

  77. [87]

    The zero-mean projection in the tensor form: d ˆf=P ⊥ 1 dN(x) =P ⊥ 1 Jl ×1 (dWl ×1 fl−1)

  78. [88]

    Substituting into the logsoftmax differential: −dϕ=y ⊤P ⊥ 1 Jl ×1 (dWl ×1 fl−1) −γK −1 ˆf ⊤P ⊥ 1 Jl ×1 (dWl ×1 fl−1) +O( ˆf 2/K)terms

  79. [89]

    Under small logits assumption, theO( ˆf 2/K)terms become negligible

  80. [90]

    Express in tensor form: −dϕ=tr(dW ⊤ l ×1 (P ⊥ 1 y)⊤Jl ×2 f ⊤ l−1) −γK −1tr(dW ⊤ l ×1 (P ⊥ 1 Jl ×1 (Wl ×1 fl−1))⊤Jl ×2 f ⊤ l−1)

  81. [91]

    We study vanilla SGD weight update: Wt =W t−1 +ηG t−1

    Therefore, the gradient is: Gl = (Jl ×1 P ⊥ 1 y−γK −1J ⊤ l ×1 P ⊥ 1 ×2 Jl ×1 (Wl ×1 fl−1))⊗f l−1 39 M Theoretical Results of Tensor low-rank for Neural Operators Lemma 10 (Tensor Gradient becomes low-rank during training)Suppose the gradient tensor fol- lows the parametric for...

  82. [92]

    Define the tensor operator: Sk := 1 N NX i=1 Ci ⊗k Bi 40 Then for the mode-k unfolding(G t)(k): (Gt)(k) = (Gt−1)(k) −ηS k(Gt−1)(k) (13)

    For each mode k, let’s consider the mode-k unfolding. Define the tensor operator: Sk := 1 N NX i=1 Ci ⊗k Bi 40 Then for the mode-k unfolding(G t)(k): (Gt)(k) = (Gt−1)(k) −ηS k(Gt−1)(k) (13)

  83. [93]

    Letλ (k) 1 < λ(k) 2 be its two smallest distinct eigenvalues

    SinceB i andC i are mode-k PSD,S k is a PSD operator. Letλ (k) 1 < λ(k) 2 be its two smallest distinct eigenvalues. LetV (k) 1 be the eigenspace corresponding toλ (k) 1

  84. [94]

    For any mode k, we can decompose(G t0 )(k) into parallel and perpendicular components: (Gt0 )(k) = (G∥ t0 )(k) + (G⊥ t0 )(k) where(G ∥ t0 )(k) is the projection ontoV (k) 1

  85. [95]

    The mode-k unfolded gradient follows: (Gt)(k) = (I−ηS k)t−t0 (Gt0 )(k)

  86. [96]

    Using the spectral properties ofS k: ∥(Gt)(k)∥2 F ≤(1−ηλ (k) 2 )2(t−t0)∥(G⊥ t0 )(k)∥2 F + (1−ηλ (k) 1 )2(t−t0)∥(G∥ t0 )(k)∥2 F

  87. [97]

    For the mode-k stable rank: srk(Gt) = ∥(Gt)(k)∥2 F ∥(Gt)(k)∥2 2 ≤sr k(G∥ t0 ) + 1−ηλ (k) 2 1−ηλ (k) 1 !2(t−t0) ∥G0 − G∥ t0 ∥2 F ∥G∥ t0 ∥2 2

  88. [98]

    Remark 9For FNO specifically:

    Finally, for the multilinear stable rank: msr(Gt) = min k srk(Gt) Therefore, the bound holds for each mode independently. Remark 9For FNO specifically:

  89. [99]

    Fourier modes (3,4) may have different stable rank behavior than channel modes (1,2)

  90. [100]

    Natural frequency decay affects eigenvalue structure in Fourier modes

  91. [101]

    Channel modes might maintain a higher stable rank due to information preservation needs

  92. [102]

    41 Proof 6Similar to the GaLore paper, it’s easy to analyze mode by mode

    Overall low-rank structure emerges from combined effect across all modes Corollary 2 (Low-rank Tensor Gradient)If the gradient takes the parametric form Gt = 1 N NX i=1 (Ai − Bi ×1 Wt ×2 fi)⊗f i with allB i mode-k full-rank, andN ′ :=rank({f i})< n, then for each mode k: srk(G...

  93. [103]

    SinceN ′ :=rank({f i}N i=1)< nandf i ∈R n, the collections of vectors {fi}N i=1 cannot span the entire spaceR n

    LetC i =f i ⊗f ⊤ i . SinceN ′ :=rank({f i}N i=1)< nandf i ∈R n, the collections of vectors {fi}N i=1 cannot span the entire spaceR n

  94. [104]

    For each mode k: • Let{u j}n−N ′ j=1 be orthonormal bases for the null space of{f i}N i=1 • Let{e k}nk k=1 be orthonormal bases forR nk • The product bases{u j ⊗e k}form a set of bases for the minimal eigenspaceV (k) 1 ofS k with minimal eigenvalue 0 • SinceB i are mode-k full...

  95. [105]

    For the mode-k projection ofG t0 ontoV (k) 1 : (G∥ t0 )(k) = n−N ′ X j=1 nkX l=1 cjluje⊤ l = n−N ′ X j=1 uj nkX l=1 cjlel !⊤

  96. [106]

    Therefore: srk(G∥ t0 )≤rank((G ∥ t0 )(k))≤n k −N ′ since stable rank is a lower-bound of the rank in each mode

  97. [107]

    On the other hand,G t can be written as a summation ofN ′ rank-1 tensors by representing each fi = PN ′ j=1 bijf ′ j as a linear combination of{f ′ j}N ′ j=1: Gt = 1 N NX i=1 (Ai − Bi ×1 Wt ×2 fi)⊗   N ′ X j=1 bijf ′ j   = 1 N N ′ X j=1 " NX i=1 bij(Ai − Bi ×1 Wt ×2 fi) # ⊗f ′ j

  98. [108]

    Thus each mode-k unfolding has rank at mostN ′. When t is sufficiently large so that the second term in the mode-k stable rank bound is negligible, by the tensor version of Lemma 3.3: srk(Gt)≤min(n k −N ′, N′)≤n k/2 sinceN ′ < nk. Corollary 3 (Tensor Low-rank with Special Stru...

  99. [109]

    The mode-k unfolding of the projected gradient is: (G∥ t0 )(k) =v kv⊤ k g0 ∝v k

  100. [110]

    Sincev k =y k ⊗z k is decomposable: • The resulting(G ∥ t0 )(k) is a rank-1 matrix 42 • Thus sr k(G∥ t0 ) = 1

  101. [111]

    From the main lemma, when t is large: srk(Gt)≈sr k(G∥ t0 ) = 1

  102. [112]

    Theorem 2 (Tensor low-rank Convergence)For a gradient tensorG t ∈R I1×I2×···×Id, let{P k ∈ RIk×rk }d k=1 be fixed orthonormal factor matrices for each mode k with ranks{r k}d k=1

    This meansG t becomes effectively rank-1 in mode k. Theorem 2 (Tensor low-rank Convergence)For a gradient tensorG t ∈R I1×I2×···×Id, let{P k ∈ RIk×rk }d k=1 be fixed orthonormal factor matrices for each mode k with ranks{r k}d k=1. TheTensor low-rankupdate consists of:

  103. [113]

    Project the gradient: Rt =G t ×1 P ⊤ 1 ×2 P ⊤ 2 ×3 · · · ×d P ⊤ d

  104. [114]

    Update optimizer states usingR t

  105. [115]

    Project back for weight update: ˜Gt =R t ×1 P1 ×2 P2 ×3 · · · ×d Pd Suppose for each mode k: •A i,B i,C i haveL (k) A ,L (k) B ,L (k) C mode-k continuity •∥W t∥(k) ≤D k (mode-k spectral norm bound) • ˆB(k) it :=P ⊤ k B(k) i (Wt)Pk • ˆC(k) it :=P ⊤ k C(k) i (Wt)Pk •κ (k) t := 1...

  106. [116]

    Using the assumption that gradient follows the parametric form: Gt = 1 N NX i=1 (Ai − Bi ×1 Wt ×2 Ci)

    First, we establish the mode-k unfolding of the gradient tensor update. Using the assumption that gradient follows the parametric form: Gt = 1 N NX i=1 (Ai − Bi ×1 Wt ×2 Ci)

  107. [117]

    For any mode k, the mode-k unfolding gives: (Gt)(k) = 1 N NX i=1 (Ai)(k) −(B i)(k)Wt(k)(Ci)⊤ (k) 43 whereW t(k) is the mode-k unfolding ofW t

  108. [118]

    The projected gradient in mode-k has unfolding: (Rt)(k) =P ⊤ k (Gt)(k) = 1 N NX i=1 P ⊤ k (Ai)(k) −P ⊤ k (Bi)(k)Wt(k)(Ci)⊤ (k)

  109. [119]

    Using the SGD updateW t =W t−1 +η ˜Gt−1, we can write: Wt(k) =W t−1(k) +ηP k(Rt−1)(k)

  110. [120]

    Substituting this into the gradient expression: (Rt)(k) = (Rt−1)(k) −η 1 N NX i=1 P ⊤ k (Bi)(k)Pk(Rt−1)(k)(Ci)⊤ (k) +E (k) t whereE (k) t captures the differences inA i andB i,C i terms

  111. [121]

    Define the mode-k operator: S (k) t := 1 N NX i=1 P ⊤ k (Bi)(k)Pk ⊗P ⊤ k (Ci)(k)Pk

  112. [122]

    Then the update can be written compactly as: (Rt)(k) = (I−ηS (k) t−1)(Rt−1)(k) +E (k) t

  113. [123]

    For the error term, using mode-k continuity: ∥E (k) t ∥F ≤L (k) A ∥Wt − Wt−1∥F +L (k) B L(k) C D2 k∥Wt − Wt−1∥F =η(L (k) A +L (k) B L(k) C D2 k)∥Rt−1∥F

  114. [124]

    Using properties of projection matricesP k: •P ⊤ k Pk =I rk (orthonormal) •∥P k∥2 = 1(projection)

  115. [125]

    The minimal eigenvalue ofS (k) t−1 satisfies: λmin(S (k) t−1)≥κ (k) t−1 due to mode-k PSD properties ofB i andC i

  116. [126]

    Therefore: ∥(Rt)(k)∥F ≤ ∥I−ηS(k) t−1∥2∥(Rt−1)(k)∥F +∥E (k) t ∥F ≤[1−η(κ (k) t−1 −L (k) A −L (k) B L(k) C D2 k)]∥(Rt−1)(k)∥F

  117. [127]

    Let: (a) GaLore with matricization along dimension d unfoldG t toG (d) t ∈R Nd×(N1N2N3N4/Nd) (b)Tensor low-rankpreserve the tensor structure and apply mode-wise projections Then:

    Whenmin t,k κ(k) t > L(k) A +L (k) B L(k) C D2 k for all modes k: 44 • Each mode-k unfolding converges:(R t)(k) →0 • Thus the full tensor converges:R t →0 Lemma 11 (Tensor low-rank vs GaLore Rank Structure)Consider a gradient tensorG t ∈R N1×N2×N3×N4 following the parametric f...

  118. [128]

    Under GaLore with any dimension d: ∃k̸=d: lim t→∞ srk(Gt)≥min(N k/2, N′) whereN ′ is the rank of the training data

  119. [129]

    UnderTensor low-rank: ∀k: lim t→∞ srk(Gt)≤N k/2 That is, GaLore cannot achieve low rank in all modes simultaneously, whileTensor low-rank achieves low rank across all modes. Proof 91) First, let’s analyze GaLore’s behavior: a) When GaLore matricizes along dimension d, it resha...

  120. [130]

    M.1 Sparsity types 46

    Now forTensor low-rank: a) Each mode k is handled independently with its own projection: Rt =G t ×1 P T 1 ×2 P T 2 ×3 · · · ×d P T d 45 b) From Theorem 2 (proven earlier), under SGD: ∥(Rt)(k)∥F ≤ h 1−η(κ (k) t−1 −L (k) A −L (k) B L(k) C D2 k) i ∥(Rt−1)(k)∥F c) From Corollary 2...

  121. [500]

    Trained with L2loss. NS1024 Re= 10 5 FNO • 4 layers, 128 modes • 128 hidden channels, 256 projection channels • Skip: ’linear’ • Channel MLP expansion: 0.5, dropout: 0 Complex Adam with step LR 5e-3, weight decay 1e-4, 100 epochs in total: batch size 8 for 50 iterations and re...

Pith tools

Reviewed August 10, 2026 · model on record in the stance chip above.