{"id":"8561f3b4-a92a-431f-9a4c-12bf044a6481","arxiv_id":"2501.02379","paper_version":2,"verdict":"CONDITIONAL","confidence":"MODERATE","novelty_score":6.0,"correctness_risk":"medium","formal_verification":"none","parameter_count":5,"one_line_summary":"TensorGRaD compresses tensor gradients into low-rank plus sparse pieces and shows this cuts optimizer memory by up to 75% for Fourier neural operators without hurting test error.","lead":"This paper introduces TensorGRaD, a way to compress the gradients used in training neural operators by splitting each gradient into a low-rank part and a sparse part, so the optimizer stores far less memory. It reports 50-75% memory savings on high-resolution PDE benchmarks like turbulent Navier-Stokes, while keeping test accuracy on par with full Adam.","discovery_kind":"new_method","skeptic_critique":{"model":"deepseek-v4-flash","headline":"The >50% total-memory claim is not supported by the paper's own Table 1: recomputed reductions are ~11% (full), ~22% (mixed vs mixed), or ~44% (mixed vs full Adam), so either the headline uses a different unreported configuration or the central claim is overstated.","rationale":"The reader identified the unverified parametric gradient assumption behind Theorem 1 as the weakest assumption; that is a real concern and I partially agree. But the most load-bearing issue for the central claim is more direct: the paper's own memory table contradicts the headline. A method's theoretical guarantees can be wrong while empirical claims hold; here, the quantitative central claim is not derivable from the reported evidence. I would still keep the disposition conditional rather than reject outright, because the over-50% figure could refer to the TensorGRaD-H+ configuration shown only in Fig. 2, and the released code could settle it. However, the authors should be required to supply a corrected memory table, report which baseline supports each percentage, and add error bars before the memory claim is accepted. The theory section should also be separated or repaired, but that is secondary to the memory accounting.","tokens_in":32025,"tokens_out":8519,"duration_ms":85897,"concrete_test":"Profile peak CUDA memory on the NS1024 Re=1e5 FNO with the released code (or the appendix profiler) for four configurations under identical settings: AdamW, AdamW-H+, TensorGRaD, and TensorGRaD-H+. Verify which pair yields the claimed >50% total reduction and report the same rows in a corrected Table 1. If the only reproducible reductions are the 21.6% (29 vs 37 GB) and 44.2% (29 vs 52 GB) values in the current table, the abstract and intro claims must be revised downward.","verdict_should_be":"CONDITIONAL","load_bearing_attack":"Abstract and intro state TensorGRaD reduces total memory usage by over 50% and that total memory cost is cut by more than 55%. Section 3, Table 1 is the only place where total peak memory is reported for the headline NS1024 experiment, and those rows do not support the claim. Full precision: Adam 52 GB vs TensorGRaD 25% 46 GB = 11.5% total reduction. Mixed precision: Adam 37 GB vs TensorGRaD 29 GB = 21.6% total reduction. Comparing mixed TensorGRaD (29 GB) to full-precision Adam (52 GB) gives 44.2%, still below 50%; the text's statement after Table 1 that reducing the total memory by 55% is consistent with none of these. The claim may have been computed for TensorGRaD-H+ (half precision plus activation checkpointing, Fig. 2 right), but that configuration's memory numbers are not reported in Table 1 or the text. Because the abstract's central claim is precisely a >50% total-memory reduction at matched or better accuracy, this internal inconsistency is the load-bearing weak point: it must be resolved before the empirical contribution can be assessed.","agreement_with_reader":"partial"},"referee_report":{"model":"deepseek-v4-flash","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.","tokens_in":32260,"tokens_out":4659,"duration_ms":49607,"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":[{"comment":"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.","section":"Abstract, Section 1, Table 1"},{"comment":"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.","section":"Appendix M, Lemma 10, Remark 1; Theorem 1"},{"comment":"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.","section":"Section 2.5, Theorem 1; Appendix M, Proof 8"},{"comment":"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.","section":"Section 3, Appendix G, Tables 5-9"}],"minor_comments":[{"comment":"The text refers to 'Appendix 6' when listing extended benchmark results; this should be 'Appendix G' (or the appropriate appendix label).","section":"Section 3"},{"comment":"Table 2's caption does not state the dataset and Reynolds number; the body text implies NS128, but the caption should say so explicitly.","section":"Section 3, Table 2"},{"comment":"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.","section":"Algorithm 1 and Theorem 1"},{"comment":"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.","section":"Throughout"},{"comment":"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.","section":"Section 3, Introduction"},{"comment":"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.","section":"Figure 2 and Figure 3"}],"recommendation":"major_revision","confidential_remarks":"The empirical idea is interesting and the NS1024 result is worth reporting, but the abstract's headline memory claim is contradicted by Table 1 and must be corrected or properly qualified. The theory section needs substantial reworking: as written, Theorem 1 does not analyze the actual algorithm and depends on an unverified gradient model. I would not reject, since the empirical contribution may stand once the memory accounting is fixed and the claims are aligned with the evidence."},"author_rebuttal":null,"desk_editor":{"model":"deepseek-v4-flash","letter":"First, the positive: the empirical core is genuinely useful. TensorGRaD combines top-k unstructured sparsity with Tucker low-rank projection on gradient tensors, applied sequentially with the residual, and the paper shows on Navier-Stokes 1024x1024 that a 25% optimizer budget matches full Adam in test L2. The ablation against matricized GaLore on Darcy is informative—tensor low-rank wins by 48%. The precision ablation (full, mixed-1, mixed-2) is well designed and shows that half-precision optimizer states are the problem. That is a solid contribution to a real scaling bottleneck.\n\nThe soft spots are real. First, the headline memory claim: the abstract and intro say total memory is cut by over 50% or more than 55%, but Table 1 does not support those numbers. Adam full is 52 GB, TensorGRaD full is 46 GB—an 11.5% reduction. Mixed Adam is 37 GB, mixed TensorGRaD is 29 GB—21.6%. Even the generous mixed-vs-full comparison is 44%. The 55% figure might come from a configuration like TensorGRaD-H+ with activation checkpointing, but that configuration's memory is not reported in the table or the text. So the central empirical claim is overstated, and this needs to be fixed before the empirical contribution can be trusted.\n\nSecond, the theory. The convergence proof (Theorem 1 in the main text, Theorem 2 in Appendix M) assumes gradients have the exact parametric 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 conditions. That form is asserted, not verified. The proof also contains a hand-wavy step citing a favorable landscape result to get from 'gradients are low-rank' to convergence—a non-sequitur. If the authors want to keep the theory, they need to either verify the gradient structure empirically on FNOs or significantly weaken the claims.\n\nMinor but worth noting: the intro says code is available on GitHub, while Section 3 says code will be provided in the supplementary. Also, the full sparse+low-rank method is only reported on NS1024 in the main tables; the other PDE results are for the low-rank-only variant, which is less than the abstract suggests.\n\nWho this is for: people working on memory-efficient training for neural operators and scientific ML. The method is worth looking at, but the memory accounting and theory need work. I'd send it to a serious referee, with instructions to reconcile the memory numbers and either repair or temper the theoretical section before acceptance.","headline":"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.","tokens_in":32852,"tokens_out":3502,"would_cite":false,"duration_ms":31845,"reading_group":"yes","serious_thinker":"yes","would_accept_peer_review":true},"rs_alignment":null,"lean_confirmation":null,"pith_extraction":{"msc":[],"pacs":[],"model":"deepseek-v4-flash","headline":"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.","keywords":["neural operators","Fourier Neural Operator","robust tensor decomposition","gradient compression","low-rank plus sparse","memory-efficient training","mixed precision","Adam optimizer"],"falsifier":"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.","tokens_in":1891,"feed_emoji":"🌊","tokens_out":2469,"duration_ms":82121,"temperature":0.7,"pith_summary":"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.","feed_headline":"One gradient trick halves memory and beats full-precision Adam on PDEs","feed_subtitle":"On 1024×1024 turbulent Navier-Stokes, a 25% optimizer state gives the lowest test error reported","key_machinery":"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.","core_discovery":"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.","pith_inferences":["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."],"forward_implications":["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."],"supporting_citations":[{"why":"Supplies the matrix-based gradient low-rank projection method that TensorGRaD extends to tensors and against which it proves a structural advantage.","marker":"[5]"},{"why":"Introduces structured sparse gradient selection strategies (top-k, rand-k) that TensorGRaD adapts for its unstructured sparse branch.","marker":"[6]"},{"why":"Provides the robust tensor decomposition framework that separates low-rank components from sparse corruptions, the conceptual basis of the method.","marker":"[12]"},{"why":"Establishes the low-rank plus sparse separation principle that robust tensor decomposition builds on.","marker":"[13]"},{"why":"Gives convergence guarantees for low-rank plus sparse recovery, which the paper invokes for the robustness of the decomposition.","marker":"[14]"},{"why":"Defines the Tucker decomposition used as the low-rank branch of TensorGRaD.","marker":"[15]"},{"why":"Provides approximation guarantees for mixed-precision FNO training, which TensorGRaD extends to full half-precision weight storage.","marker":"[17]"},{"why":"Supplies the efficient Tucker decomposition implementation used to compute the low-rank factors in practice.","marker":"[18]"},{"why":"Supports the convergence argument in Theorem 1 by describing the optimization landscape of Tucker decomposition.","marker":"[36]"}],"fun_headline_variants":["TensorGRaD: 50% less memory, beats Adam on 1024²","Split gradients into low-rank + sparse: 55% memory cut","25% optimizer state beats full-precision Adam on PDEs","FNO training: tensor decomp halves memory, outdoes Adam","Robust tensor decomp: low-rank + sparse gradients, 55% savings"],"cache_read_input_tokens":34944,"weakest_assumption_plain":"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.","fun_headline_variants_meta":{"raw":{"variants":["TensorGRaD: 50% less memory, beats Adam on 1024²","Split gradients into low-rank + sparse: 55% memory cut","25% optimizer state beats full-precision Adam on PDEs","FNO training: tensor decomp halves memory, outdoes Adam","Robust tensor decomp: low-rank + sparse gradients, 55% savings"]},"model":"deepseek-v4-flash","effort":"low","cost_usd":0.000608,"raw_usage":{"total_tokens":2883,"prompt_tokens":1050,"completion_tokens":1833,"prompt_tokens_details":{"cached_tokens":384},"prompt_cache_hit_tokens":384,"prompt_cache_miss_tokens":666,"completion_tokens_details":{"reasoning_tokens":1751}},"tokens_in":666,"tokens_out":1833,"duration_ms":12079,"temperature":1.0,"reasoning_tokens":1751,"cache_read_input_tokens":384,"cache_creation_input_tokens":0},"cache_creation_input_tokens":0},"created_at":"2026-08-10T22:13:33.594636+00:00","model_set":{"reader":"deepseek-v4-flash"},"falsifier":"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.","supporting_citations":[{"cited_title":"Robust Principal Component Analysis?","cited_arxiv_id":"0912.3599","evidence_quote":"Establishes the low-rank plus sparse separation principle that robust tensor decomposition builds on."},{"cited_title":"Optimization Landscape of Tucker Decomposition","cited_arxiv_id":"2006.16297","evidence_quote":"Supports the convergence argument in Theorem 1 by describing the optimization landscape of Tucker decomposition."}],"review_version":1}