REVIEW 4 major objections 4 minor 2 references
Learning, Solving and Optimizing PDEs with TensorGalerkin: an efficient high-performance Galerkin assembly algorithm
T0 review · 4 major / 4 minor · reviewed 2026-08-03 · deepseek-v4-flash
Pith's one-line read The paper claims that Galerkin assembly of finite element matrices can be reformulated as a deterministic Map-Reduce: a single dense tensor contraction for all elements followed by one sparse matrix multiply, so the computation graph has O(
desk verdict Useful engineering contribution; the central O(1)-graph claim is asserted not shown, and headline numbers are overstated — deserves a referee. read the letter →
The pith
A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.
The reading
What carries the argument
The central object is the Map–Reduce assembly itself. In Stage I (Batch-Map), basis gradients are pushed forward onto physical elements and contracted over the quadrature dimension in a single batched einsum, fusing all element loops into one GPU kernel. In Stage II (Sparse-Reduce), the local contributions are flattened and projected into the global sparse stiffness matrix and load vector through precomputed binary routing matrices S_mat and S_vec, implemented as sparse matrix multiplications. This carries the argument because it converts an O(E × k^2) graph-fragmenting loop into a constant number of autograd nodes, which is what delivers both the speedup and the clean gradient flow.
What would settle it
Take a diffusion coefficient defined by a neural network with a hard if/else branch inside the integrand, assemble on a sequence of meshes from 10^2 to 10^6 elements, and count the autograd nodes in the forward graph and the backward time per step. If the node count grows with the number of elements instead of staying constant, or backward time scales super-linearly while forward time stays flat, the O(1)-graph claim is falsified for that class of forms; a negative result on any single variational form of this type is enough to bound the method's generality.
Extended reading notes
Core claim
On its own terms, the paper's central discovery is that Galerkin assembly, traditionally a scatter-add loop over elements and basis pairs, factorizes cleanly into two monolithic operations. The Map stage tensorizes physics: with the element index promoted to a batch dimension, all local stiffness matrices and load vectors are computed as a single einsum-style contraction over quadrature points, basis indices, Jacobians and coefficients (Eq. 7). The Reduce stage tensorizes topology: once binary routing matrices encoding local-to-global DoF maps are precomputed, global assembly is a sparse matrix multiplication (Eq. 8). Because each stage is one operation in an autodiff library, autograd sees
Load-bearing premise
The approach assumes that the local physics of any supported PDE can be written as a single dense tensor contraction (the einsum function F in Eq. 7); if a variational form needs per-element control flow, data-dependent branching, or non-polynomial adaptive quadrature that won't batch, the monolithic O(1)-graph property and the speedups that follow from it can weaken or disappear.
Editorial extensions
If this is right
- A GPU-native finite element solver built on this assembly reaches an order of magnitude speedup over leading CPU-based FEM stacks on 3D Poisson and nearly two orders of magnitude over CPU baselines on 3D linear elasticity, at residual levels equal to or below the baselines.
- Physics-informed operator learning needs no autodiff for spatial derivatives; it uses analytical shape-function gradients, so loss evaluation scales nearly flat with degrees of freedom where strong-form PINN loss evaluation grows sharply.
- On wave and Allen–Cahn operator learning with 16 training samples, the data-free Galerkin-trained model beats a supervised data-driven model on out-of-distribution rollout and outperforms a physics-informed DeepONet-style baseline by at least an order of magnitude.
- PDE-constrained optimization becomes end-to-end differentiable through assembly and sparse solve; the topology optimization benchmark completes with a 3.7× faster optimization loop than an XLA-based differentiable FEM framework while converging to the same design.
- Batch data generation amortizes fixed overhead on GPU: runtime stays nearly constant for batch sizes up to 10^2 before scaling sub-linearly, making massive physics-dataset construction cheap.
Reading between the lines
- An implied but untested consequence is that the O(1)-graph property should transfer beyond linear elasticity to any variational form whose physics evaluation can be written as an einsum-compatible contraction; forms with data-dependent control flow or adaptive quadrature may not factor this way and would break the monolithic-node claim.
- The analytical-gradient trick suggests a general recipe: replace autodiff spatial derivatives in any variational, mesh-based learning method with precomputed shape-function gradients, which should also extend to higher-order elements where the contraction structure is richer.
- A testable extension is to treat the routing-matrix construction itself as differentiable or learned, which could turn the method into a differentiable mesh-coarsening or remeshing tool.
- The framework's dependence on a specific autodiff engine's kernel fusion means the O(1) claim is an empirical property of the backend; on a different engine that splits einsum or SpMM into sub-graphs, the constant-node property would need re-verification.
Editorial analysis
A structured set of objections, weighed in public.
Referee Report
Summary. The paper proposes TensorGalerkin, a two-stage Map–Reduce algorithm for Galerkin finite-element assembly in PyTorch. Stage I evaluates all element-local bilinear/linear forms as batched dense tensor contractions (Eq. 7), and Stage II aggregates local contributions into the global sparse stiffness matrix/load vector through precomputed sparse routing matrices and one SpMM (Eq. 8). The authors claim that this yields a computational graph with O(1) nodes regardless of the number of elements E and local degrees of freedom k, eliminating the Python-loop and autograd fragmentation of standard FEM assembly. They deploy the framework in three downstream settings: TensorMesh (GPU FEM solver), TensorPils (physics-informed operator learning / neural PDE solver), and TensorOpt (PDE-constrained optimization). Benchmarks cover 3D Poisson and elasticity, 2D checkerboard Poisson, wave and Allen–Cahn operator learning, and a 2D topology-optimization problem, comparing against FEniCS, scikit-fem, JAX-FEM, PINNs, VPINNs, Deep Ritz, PI-DeepONet, and data-driven GNNs.
Significance. If the central claims hold, the contribution is significant: an efficient, GPU-compatible, differentiable FEM assembly primitive would provide a practical building block for numerical PDE solving, physics-informed learning, and PDE-constrained optimization. The paper is commendably broad in scope, with a reproducible project page, extensive ablations, and a serious effort to compare against multiple strong baselines under shared training schedules. The analytical shape-gradient approach in TensorPils, which avoids autograd for spatial derivatives, is a principled and potentially impactful design. However, the manuscript's headline claims—O(1) graph complexity and order-of-magnitude accuracy/efficiency gains—are currently supported only in part. The O(1) claim is presented as a consequence of the formulation but is not measured and is not obviously valid for the nonlinear/case-dependent forms the paper claims to support. Several numerical overclaims are directly contradicted by the paper's own tables. These issues are load-bearing for the paper's stated significance and need to be resolved before the contribution can be fully assessed.
major comments (4)
- [§2, 'Analysis of the Computational Graph'; Eq. (7); SM B.3.1, Eq. (B.19)] The paper claims that TensorGalerkin consolidates assembly into 'two monolithic graph nodes (O(1) complexity)' independent of E and k. This is not established. For the semi-linear Allen–Cahn residual in SM B.3.1, evaluating F_nonlin(U^{k+1}) requires gathering U to quadrature points, applying a cubic nonlinearity, multiplying by test-function values and quadrature weights, and reducing—a sequence of gather/pow/mul/reduce operations that cannot be expressed as a single einsum of the form in Eq. (7). It may still be O(1) in E and k if implemented as a fixed number of batched ops, but the paper never states this, and whether torch.einsum/torch.sparse.mm are recorded as one autograd node is an empirical property of the PyTorch version, not a mathematical consequence. I ask the authors to (i) state precisely the class of variational forms for which the Map stage is a single dense contraction,
- [Table 1, §3, 'Neural PDE Solver'] The text states that TensorPils achieves '50% less error than the nearest baseline (Deep Ritz).' This holds at K=2 and K=4, but at K=8 the reported relative L2 errors are 10.05 (TensorPils) vs 10.60 (Deep Ritz), i.e., about 5% less error, not 50%. As K=8 is the most challenging case and the one where the claimed advantage matters most, the blanket statement is inaccurate. Please qualify the claim by frequency or remove it.
- [Table 2, §3, 'Physics-informed Operator Learning'] The abstract and discussion state that TensorPils is 'at least an order of magnitude more accurate than PI-DeepONet.' The reported wave ID error is 0.085 vs 0.626, a factor of 7.4; wave OOD is 0.090 vs 0.863, a factor of 9.6. Neither reaches an order of magnitude. The AC ID case (0.110 vs 0.743) is a factor of 6.8. Only the AC OOD case (0.083 vs 8.536) exceeds an order of magnitude. The claim should be revised to match the data.
- [SM B.2.2, Figure B.6; §3, Table 1] The neural-PDE-solver comparison is confounded: TensorPils imposes Dirichlet boundary conditions as hard constraints by reducing the linear system, while PINN, VPINN, and Deep Ritz use soft penalty BCs. Since the same SIREN backbone is used, the accuracy differences at higher K may reflect the well-known difficulty of enforcing BCs via soft penalties rather than the TensorGalerkin loss itself. This does not invalidate the efficiency comparison, but the accuracy claims should be interpreted with this confound acknowledged, or a hard-BC variant of a baseline should be included.
minor comments (4)
- [Eq. (7)] The notation F(G_eq a, G_eq b, C_eq) is introduced abstractly but the general conditions on F (smoothness, polynomial/quadrature structure) are never stated. A brief definition or reference would help. SM A.2 gives the linear diffusion example, but the nonlinear case is only implicitly covered.
- [§2, Algorithm 1] Step 1 says 'Compute Jacobians J(X) and determinants |detJ| in batch.' For affine elements this is a single contraction, but for isoparametric/higher-order elements the Jacobian depends on quadrature points and may require additional gather operations. Please clarify whether the paper currently assumes affine elements only.
- [Table 2] The AC OOD PI-DeepONet error has a very large standard deviation (8.536±6.306). Reporting the median or per-seed values would make the comparison more robust, especially since the 'order of magnitude' claim rests heavily on this entry.
- [Throughout] Several references are duplicated or inconsistently formatted (e.g., Svanberg appears as both [1987a] and [1987b] with overlapping content; Pfaff et al. appears as 2020 and 2021; the citation for DOLFINx is listed as 'preprint' without a venue). A careful reference cleanup is needed.
Circularity Check
No significant circularity: the core claim is an engineering reformulation of FEM assembly, validated against external baselines, and no prediction reduces by construction to a fitted input.
full rationale
The paper's central derivation is a Map-Reduce factorization of standard Galerkin assembly: local element stiffness/load tensors are computed by a tensorized contraction (Eq. 7) and aggregated by precomputed routing matrices via SpMM (Eq. 8). This is a re-expression of the classical scatter-add loop, and the O(1)-graph claim is an implementation-level assertion about PyTorch autograd, not a mathematically forced consequence of the inputs; it may be empirically unverified, but that is an evidence concern, not circularity. TensorPils minimizes the discrete variational residual ||KU-F||^2, which is the standard Galerkin condition, and compares against external FEM references (FEniCS, scikit-fem, JAX-FEM, fine-mesh FEM ground truth, PI-DeepONet); no parameter is fitted to the reported error and then reported as a prediction. Self-citations (e.g., Poseidon, RIGNO, wPINNs, the Acta Numerica survey) are background literature and are not load-bearing assumptions, and the paper invokes no uniqueness theorem or ansatz from the authors' prior work to force its choices. The identified weaknesses (nonlinear terms may not collapse to a single einsum; graph-size measurements are absent) are correctness/robustness risks, not circular reductions.
Assumptions & free parameters
assumptions (5)
- domain assumption The PDE has a variational structure a_rho(u,v) = l_rho(v) on suitable Banach spaces (Eq 1).
- ad hoc to paper Every element-local physics computation can be expressed as a single dense tensor contraction over the batch of elements (Eq 7).
- standard math Sparse routing matrices S_mat and S_vec exactly reproduce scatter-add assembly (Eq 8).
- domain assumption PyTorch autograd represents the einsum and SpMM as monolithic O(1) graph nodes.
- domain assumption Time-dependent PDEs are accurately captured by method-of-lines Galerkin discretization with the given time integrators (SM A.1, B.3.1).
Cite this review
Pith. "Pith review of Learning, Solving and Optimizing PDEs with TensorGalerkin: an efficient high-performance Galerkin assembly algorithm." pith.science (2026). https://pith.science/paper/VS2W7RON
@misc{pith2026260205052,
author = {Pith},
title = {Pith review of: Learning, Solving and Optimizing PDEs with TensorGalerkin: an efficient high-performance Galerkin assembly algorithm},
year = {2026},
howpublished = {\url{https://pith.science/paper/VS2W7RON}},
note = {Machine review of arXiv:2602.05052}
}
abstract
We present a unified algorithmic framework for the numerical solution, constrained optimization, and physics-informed learning of PDEs with a variational structure. Our framework is based on a Galerkin discretization of the underlying variational forms, and its high efficiency stems from a novel highly-optimized and GPU-compliant TensorGalerkin framework for linear system assembly (stiffness matrices and load vectors). TensorGalerkin operates by tensorizing element-wise operations within a Python-level Map stage and then performs global reduction with a sparse matrix multiplication that performs message passing on the mesh-induced sparsity graph. The Map and Reduce stages are co-designed inside PyTorch's autograd so that the assembly graph contains $O(1)$ nodes regardless of how the number of elements and local DoFs scale. We validate this $O(1)$-graph property by deploying TensorGalerkin downstream as i) a highly-efficient numerical PDEs solver, ii) an end-to-end differentiable framework for PDE-constrained optimization, and iii) a physics-informed operator learning algorithm for PDEs. With multiple benchmarks, including 2D and 3D elliptic, parabolic, and hyperbolic PDEs on unstructured meshes, we demonstrate that the proposed framework provides significant computational efficiency and accuracy gains over a variety of baselines in all the targeted downstream applications.
Figures
Reference graph
Works this paper leans on
-
[1992]
URLhttps://doi.org/10.1137/0913035
doi: 10.1137/0913035. URLhttps://doi.org/10.1137/0913035. Sifan Wang and Paris Perdikaris. Long-time integration of parametric evolution equations with physics- informed DeepONets.arXiv preprint arXiv:2106.05384, 2021. Shizheng Wen, Arsh Kumbhat, Levi Lingsch, Sepehr Mousavi, Yizhou Zhao, Praveen Chandrashekar, and Siddhartha Mishra. Geometry aware operat...
arXiv 2021
-
[2023]
URLhttps://openreview.net/forum?id=MtekhXRP4h. Y. Saad.Iterative methods for sparse linear systems. SIAM, 2003. Vincent Sitzmann, Julien Martel, Alexander Bergman, David Lindell, and Gordon Wetzstein. Implicit neural representations with periodic activation functions.Advances in neural information processing systems, 33:7462–7473, 2020. Krister Svanberg. ...
2003
Reviewed August 3, 2026 · model on record in the stance chip above.
Discussion (0). Sign in to comment.