Hierarchical Transformer Preconditioning for Interactive Physics Simulation
arXiv:2605.13343v1 [cs.GR] 13 May 2026
CARL OSBORNE, MIT CSAIL, USA MINGHAO GUO, MIT CSAIL, USA CRYSTAL OWENS, MIT CSAIL, USA WOJCIECH MATUSIK, MIT CSAIL, USA
Fig. 1. Teaser. Hierarchical neural preconditioning for interactive multiphase Poisson solves: (a) weak-admissibility H-matrix prior, (b) learned preconditioner 𝑀 ≈𝐴 −1 , (c) multiscale residual transport induced by 𝑀, and (d) 3D multiphase fluid simulation application context. Neural preconditioners for real-time physics simulation offer promising datadriven priors, but they often fail to capture long-range couplings efficiently because they inherit local message passing or sparse-operator access patterns. We introduce the Hierarchical Transformer Preconditioner, a neural preconditioner anchored to a weak-admissibility H-matrix partition. The partition provides a multiscale structural prior (dense diagonal leaves plus coarsening off-diagonal tiles) that enables full-graph approximate-inverse computation with 𝑂 (𝑁 ) scaling at fixed block sizes. The network models the inverse through low-rank far-field factors and uses highway connections (axial buffers plus a global summary token) to propagate context across transformer depth. At each PCG iteration, preconditioner application reduces to batched dense GEMMs with regular memory access. The key training contribution is a cosine-Hutchinson probe objective that learns the action of 𝑀𝐴 on convergence-critical spectral subspaces, optimizing angular alignment of 𝑀𝐴z with z rather than forcing eigenvalue clusters to a prescribed location. This removes unnecessary spectral-placement constraints from SAI-style objectives and improves conditioning on irregular spectra. Because both inference and apply are dense, dependency-free tensor programs, the full solve loop is captured as a single CUDA Graph. On stiff multiphase Poisson systems (up to 100 : 1 density contrast, 𝑁 = 1,024–16,384), the solver runs from ∼ 143 to ∼ 21 fps. At 𝑁 =8,192, it reaches 17.9 ms/frame, with 2.2× speedup over GPU Jacobi, ∼ 28× over GPU IC/DILU (AMGX multicolor_dilu), and 2.7× over neural SPAI retrained per scale on the same benchmark. CCS Concepts: • Computing methodologies → Physical simulation; Neural networks. Authors’ Contact Information: Carl Osborne, MIT CSAIL, Cambridge, USA, osbo@mit. edu; Minghao Guo, MIT CSAIL, Cambridge, USA, [email protected]; Crystal Owens, MIT CSAIL, Cambridge, USA, [email protected]; Wojciech Matusik, MIT CSAIL, Cambridge, USA, [email protected]. 2026. ACM 1557-7368/2026/12-ART1 https://doi.org/10.1145/nnnnnnn.nnnnnnn
Additional Key Words and Phrases: preconditioning, iterative solvers, hierarchical matrices, transformers, neural operators, physics simulation, real-time graphics ACM Reference Format: Carl Osborne, Minghao Guo, Crystal Owens, and Wojciech Matusik. 2026. Hierarchical Transformer Preconditioning for Interactive Physics Simulation. ACM Trans. Graph. 45, 6, Article 1 (December 2026), 10 pages. https://doi. org/10.1145/nnnnnnn.nnnnnnn
1
Introduction
Real-time simulation of fluids, soft bodies, and coupled multiphysics often reduces every time step to a large sparse SPD solve attacked with preconditioned conjugate gradient (PCG) [Saad 2003]. The preconditioner 𝑀 ≈𝐴 −1 is the dominant design variable. Classical choices — Jacobi, incomplete Cholesky/LU (IC/ILU), sparse approximate inverses (SPAI), and algebraic multigrid (AMG) [Benzi et al. 1996; Kolotilina and Yeremin 1993; Lin and Moré 1999; Ruge and Stüben 1987] — excel when the matrix is reused across many solves so a heavy setup amortizes. In an interactive setting 𝐴 changes every frame, so any preconditioner whose setup approaches the per-frame budget cannot amortize; IC and AMG suffer worst [Liu et al. 2016; Naumov et al. 2015; Yamazaki et al. 2020]. Machine learning offers an alternative: train a neural preconditioner once and amortize the cost across frames. Graph neural networks have been the dominant choice [Chen 2024; Häusner et al. 2023; Li et al. 2023; Yang et al. 2025] because a sparse matrix has a natural graph interpretation, but the choice forces two limitations that become acute in real-time simulation. (i) GNN preconditioners typically inherit the sparsity of 𝐴, so non-adjacent interactions ACM Trans. Graph., Vol. 45, No. 6, Article 1. Publication date: December 2026.
1:2
•
Osborne et al.
are either dropped or only reached through repeated graph convolutions that oversmooth local detail [Trifonov et al. 2025; Wu et al. 2019]. (ii) The dominant operations at apply time — triangular solves on a learned IC, or two SpMVs on a learned SPAI — are bandwidth-bound and either serialize on data hazards or scatter through irregular memory. From PCG geometry to hierarchical structure. The ideal preconditioner is 𝐴 −1 , which is dense even when 𝐴 is sparse, so the question is what structure makes 𝐴 −1 cheap to store and apply. In modern simulation codes the degrees of freedom are already laid out along a space-filling curve (Morton, Hilbert) or a bandwidth-reducing permutation, because the surrounding pipeline (spatial hashing, BVH, neighbor search) wants that ordering [Ihmsen et al. 2011; Karras 2012; Teschner et al. 2003]. Under such orderings the nonzeros of 𝐴 cluster near the diagonal and the corresponding off-diagonal blocks of 𝐴 −1 describe long-range interactions whose numerical rank decays with separation. Hierarchical matrices (H -matrices) [Börm et al. 2003; Hackbusch 1999; Hackbusch and Khoromskij 2002] exploit exactly this — a recursive block partition with dense diagonal blocks and low-rank off-diagonal blocks — and the fast multipole method [Greengard and Rokhlin 1987] factors the same prior for translation-invariant kernels. An H -matrix preconditioner is architecturally an FMM-shaped operator on the assembled system, and that is the prior we hand the network. Approach. We anchor the preconditioner to a weak-admissibility H -matrix partition [Hackbusch and Khoromskij 2002], computed analytically from the leaf-index geometry and reused for every frame: 𝐾 =𝑁 /𝐿 dense diagonal blocks of size 𝐿×𝐿 plus off-diagonal tiles spanning 𝑆 ×𝑆 leaves with 𝑆 = 2 𝑗 . A two-stream transformer operates on this layout. A diagonal stream emits one factor per leaf at full rank; an off-diagonal stream emits one factor per tile at a constant coarse token count 𝐿𝑠 ≪ 𝐿 regardless of 𝑆, so the rank fraction 𝐿𝑠 /(𝑆𝐿) shrinks automatically with separation — exactly the bias H -matrix compression calls for. To let the network route information between blocks without breaking the Θ(𝑁 ) budget, we add per-layer highway buffers (axial row/column bands plus a single global token, §4.3). A loss with no preferred eigenvalue location. We train self-supervised with a cosine-similarity Hutchinson probe loss on the preconditioned image 𝑀𝐴z. The conceptual move over the Frobenius/SAI loss used by [Yang et al. 2025] is to replace a distance between vec1 tors (∥ ∥𝐴∥ 𝐴𝑀z − z∥ 2 ) with a distance between subspaces (the angle between the lines spanned by z and 𝑀𝐴z). This is invariant under positive rescaling of 𝑀, exactly matching PCG, which cares about the direction 𝑀𝐴 sends each probe and not the magnitude. The network is then free to cluster the spectrum of 𝑀𝐴 wherever angular alignment is easiest, with no implicit demand that the cluster live near any particular eigenvalue (Prop. 1; spectral evidence in Fig. 2). Empirically, this is not a minor objective tweak: with architecture and training split fixed, the loss swap alone accounts for a major quality gap in Table 6. One CUDA Graph for the whole iteration. The choices above together remove the data dependencies that normally fragment a CG iteration into separately dispatched kernels: the partition is fixed, ACM Trans. Graph., Vol. 45, No. 6, Article 1. Publication date: December 2026.
Table 1. Per-iteration preconditioner application: arithmetic cost, dominant memory-access pattern, GPU throughput limit, and whether the apply admits single-CUDA-Graph capture of the PCG inner loop. nnz is the nonzero count of the relevant factor. “𝐿𝑠 ” is our fixed coarse-token count per offdiagonal tile. Method
FLOPs
Memory access
GPU bottleneck
Graph
Jacobi
𝑂 (𝑁 )
stride-1 read
✓
IC (tri. solve)
𝑂 (nnz)
wavefront-dep. SpTRSV
Neural SPAI [Yang et al. 2025] AMG (V-cycle)
𝑂 (nnz)
random gather (SpMV)
memory bandwidth latency / data hazards memory bandwidth
𝑂 (𝑁 log 𝑁 ) multilevel, irregular
×
Ours
𝑂 (𝑁 𝐿𝑠 )
level synchronization compute (Tensor Cores)
batched GEMM, stride-1
× ✓
✓
every tile shape is the same, the apply is batched GEMMs over preallocated tensors with no triangular solves and no allocations. The whole PCG inner loop — preconditioner apply, SpMV, SAXPYs, dot reductions — records into one CUDA Graph that replays per iteration with zero CPU dispatch overhead. In our measurements this is a substantial share of the gap to classical GPU preconditioners. Contributions. In short: [c1] a neural preconditioner that bakes the H -matrix block partition in as a structural prior with a fixed coarse rank on every off-diagonal tile, keeping total work at Θ(𝑁 ) (intro paragraph c1, §4.2); [c2] highway buffers that route information between blocks without growing the budget (§4.3); [c3] a cosine-similarity Hutchinson loss that minimizes a distance between subspaces rather than between vectors, exactly matching PCG’s scale-invariance and freeing the network to put the spectrum of 𝑀𝐴 wherever convergence is easiest (intro paragraph c3, Prop. 1 in §6.1); [c4] an allocation-free apply with no triangular solves, no data-dependent branching, no per-kernel CPU dispatch — the whole PCG inner loop captured as one CUDA Graph (intro paragraph c4, §5). Table 1 summarizes the per-iteration access pattern and throughput limit of each preconditioner family at the one step where every fixed-pattern alternative loses time on modern GPUs.
2
Related Work
Classical preconditioners. Jacobi, incomplete Cholesky/LU, sparse approximate inverses (SPAI), and algebraic multigrid [Benzi et al. 1996; Kolotilina and Yeremin 1993; Lin and Moré 1999; Naumov et al. 2015; Ruge and Stüben 1987; Saad 2003] are mature and remain default choices when the system is reused across many solves so a careful setup amortizes. The known GPU costs of triangular solves [Liu et al. 2016; Yamazaki et al. 2020] and of V-cycle communication [Naumov et al. 2015] bite hardest in the regime we target (𝐴 changing every frame at interactive rates), which motivates the amortized neural alternative pursued here. Learned PDE solvers and preconditioners. Neural operators (Fourier [Li et al. 2021], DeepONet [Lu et al. 2021], graph-based [Li et al. 2020a,b]) and physics-informed networks [Raissi et al. 2019] learn functionspace mappings between coefficient and solution fields, and have
Hierarchical Transformer Preconditioning for Interactive Physics Simulation
been used both as surrogates and as inner preconditioners for flexible Krylov solvers [Rudikov et al. 2024]; we operate one rung lower, on the assembled algebraic system. A productive line of work trains GNNs to emit a learned IC factor [Häusner et al. 2023; Li et al. 2023; Trifonov et al. 2024] or a sparse approximate inverse in factored 𝐺𝐺 ⊤ form using a scale-invariant Frobenius (SAI) loss [Yang et al. 2025]; we compare to the latter directly in Table 4. Chen [Chen 2024] reports strong performance across SuiteSparse problems where classical IC and AMG struggle, and Trifonov et al. [Trifonov et al. 2025] argue that message-passing GNNs cannot approximate the non-local elimination structure of sparse triangular factors — a motivation for architectures (like ours) that route global information through explicit channels. Hierarchical neural and matrix machinery. A parallel line embeds FMM/H -matrix structure into neural architectures [Fan et al. 2019; Fognini et al. 2025; Li et al. 2020a; Luz et al. 2020; Lyu et al. 2026; Sittoni et al. 2026; Xu et al. 2025]. Classical H - and H 2 -matrix arithmetic [Börm et al. 2003; Hackbusch 1999; Hackbusch and Khoromskij 2002] provides near-linear-cost machinery for storing and applying operators with low-rank well-separated sub-blocks; HODLR variants [Hartland et al. 2023] apply the same prior to dense Hessians. The partitioning machinery is what we instantiate; the contribution lies in how the off-diagonal factors are produced and applied (§4).
3
Background
CG converges through the polynomial bound ∥e𝑘 ∥𝐴 ≤ ∥e0 ∥𝐴
min
max |𝑝 (𝜆)|,
𝑝 ∈ P𝑘 , 𝑝 (0)=1 𝜆∈𝜎 (𝐴)
(1)
so CG only sees 𝐴 through how small a polynomial fixed to 1 at the origin can be made on 𝜎 (𝐴). The practical consequence we lean on is that when 𝜎 (𝑀𝐴) collapses to 𝑐 tight clusters away from zero, CG converges in ∼𝑐 iterations regardless of where on the real line those clusters sit. The preconditioning objective is therefore “cluster 𝜎 (𝑀𝐴),” not “make 𝑀𝐴 close to 𝐼 ,” with cluster location essentially free — the property our cosine loss exploits. An H -matrix [Börm et al. 2003; Hackbusch 1999; Hackbusch and Khoromskij 2002] represents a dense matrix as a recursive block partition: dense diagonal blocks plus off-diagonal blocks stored as low-rank factors 𝑈𝑉 ⊤ , with admissible rank shrinking off the diagonal because the underlying Green’s function is smooth there. We use the weak-admissibility variant [Hackbusch and Khoromskij 2002; Hartland et al. 2023], whose fewer, larger off-diagonal tiles batch well on a GPU; classical H -matrix arithmetic costs 𝑂 (𝑁 log 𝑁 ), our learned realization runs at Θ(𝑁 ) thanks to the fixed coarse-token count in c1. The training loss uses Hutchinson-style probes z with E[zz⊤ ] = 𝐼 [Hutchinson 1989], but not as a trace estimator. The property we use is spectral whiteness: isotropy gives E[(u𝑖⊤ z) 2 ] = 1 for every eigenvector u𝑖 of 𝑀𝐴, so a scalar built from (z, 𝑀𝐴z) weighs every eigenmode equally rather than favoring smooth or oscillatory ones — the property that makes the cosine loss of c3 a faithful global indicator of preconditioner quality.
4
•
1:3
Method
Notation and pipeline. 𝐴 ∈ R𝑁 ×𝑁 is the assembled sparse SPD system, indexed along a space-filling-curve or bandwidth-reducing ordering. The 𝑁 indices are partitioned into 𝐾 = 𝑁 /𝐿 contiguous leaves of size 𝐿, inducing an H -matrix partition of 𝐾 diagonal blocks and 𝑀 H unique weakly-admissible off-diagonal tiles (tile 𝑚 spans 𝑆𝑚×𝑆𝑚 leaves). The network is a four-stage pipeline: a physics-aware encoder, a diagonal attention stack (full per-node resolution 𝐿), an off-diagonal attention stack (coarse token count 𝐿𝑠 ≪ 𝐿, the same for every tile), and linear decoder heads. The whole network dispatches exactly two attention kernels regardless of 𝑁 — one batched over leaves, one batched over tiles — which, together with the static partition, is what later lets the entire PCG inner loop run inside a single CUDA Graph. We use 𝐿 =128, 𝐿𝑠 =32, 𝑑 =128, 𝑛𝑙 =3 in every reported result.
4.1
Encoder
A two-layer MLP lifts whichever per-node features the simulator already exposes (position, velocity, density, pressure, material parameters) plus the geometric and coupling features (Δx𝑖 𝑗 , 𝐴𝑖 𝑗 ) for every nonzero of 𝐴 to width-𝑑 node embeddings; 𝑛 gcn =2 residual graph-convolutional layers [Kipf and Welling 2017] mix in neighborhood information using 𝐴 as the message-passing weight. This is the only stage that observes individual graph edges; everything downstream operates on the block layout of the partition, decoupling cost from nnz(𝐴). Because 𝑀 is neural rather than algebraic, the network can look at (𝜌𝑖 , Δx𝑖 𝑗 ) directly — physical structure a fixed analytical preconditioner cannot see (full feature list in Supplementary ??).
4.2
Diagonal and off-diagonal attention stacks
Diagonal stack. For each of the 𝐾 leaves, 𝑛𝑙 transformer layers [Liu et al. 2021; Vaswani et al. 2017] with within-leaf (Swin-style) attention and edge-bias logits emit a dense factor 𝐹𝑘 ∈ R𝐿×𝐿 . The corresponding diagonal block of 𝑀 is the PSD outer product 𝐹𝑘 𝐹𝑘⊤ . Off-diagonal stack. For each off-diagonal tile 𝑚 of size 𝑆𝑚 𝐿×𝑆𝑚 𝐿, node embeddings on its row- and column-strips are pooled to a fixed coarse token count 𝐿𝑠 =𝐿/𝑝 off ≪ 𝐿 regardless of 𝑆𝑚 . 𝑛𝑙 transformer layers over these 𝐿𝑠 tokens emit a factor pair 𝑈𝑚 , 𝑉𝑚 ∈ R𝐿𝑠 ×𝐿𝑠 ; the tile’s coarse representation is 𝐵𝑚 =𝑈𝑚𝑉𝑚⊤ . Because the token count is the same on every tile, the implied rank fraction 𝐿𝑠 /(𝑆𝑚 𝐿) shrinks automatically with separation — exactly the H -matrix prior that distant blocks need less rank. A rank-fraction audit on real frames confirms the architecture-provided 𝐿𝑠 /(𝑆𝐿) stays above what truncated SVD needs at 𝜀 ∈ {10−3, 10−6, 10−9 } in every distance class (Fig. 3; assembled-𝑀 visualization in Supplementary, Fig. ??). With 𝑆𝑚 doubling geometrically, 𝑀 H = 𝑂 (𝐾) unique off-diagonal tiles, and a tile-batched attention of fixed shape, the total work over the off-diagonal stack stays linear in 𝑁 .
4.3
Highway connections
Within-block attention is local by construction. To let the network route information across the matrix without giving up the Θ(𝑁 ) budget, after every attention sublayer we scatter-add block-token embeddings into three buffers per transformer layer: a row-band axial buffer rhw , a column-band axial buffer chw , and a single global ACM Trans. Graph., Vol. 45, No. 6, Article 1. Publication date: December 2026.
1:4
•
Osborne et al.
Table 2. Subsystem breakdown of one forward pass at 𝑁 =8 192. Attention layers scale with 𝑛𝑙 ; everything else is independent of depth. Subsystem Diagonal attention stack Off-diagonal attention stack Decoder heads Encoder (MLP + GCN) Layout helpers / scatter
% device time
36% 23% 18% 16% < 4%
Scales with
𝑛𝑙 𝑛𝑙 — — —
Table 3. Kernel-level breakdown at 𝑁 = 8 192. Dominant kernels are compute-bound Tensor-Core GEMMs and fused attention; no SpMV or indexed-gather kernel appears among the top kernels. Kernel family
% device time
CUTLASS Tensor-Op GEMMs (s1688gemm, sm90_xmma, TF32) CUTLASS fused attention (fmha_cutlassF) Elementwise / normalization Copies / reshapes / layout Other
31% 14% 17% 17% 21%
summary token ghw . Each token then concatenates its row, column, and global context with its own embedding before the FFN sublayer. The four channels per layer (2D intra-block, 1D row, 1D column, 0D global) cost 𝑂 (𝑁𝑑) scatter-gather and preserve Θ(𝑁 ) scaling. We ablate the highways in §7.3 and illustrate the per-layer connectivity in Supplementary, Fig. ??b.
4.4
Decoder, geometric picture, and complexity
Three linear decoder heads project tokens to the final factors: a leaf head emits 𝐹𝑘 ∈ R𝐿×𝐿 (diagonal block 𝐹𝑘 𝐹𝑘⊤ , PSD by construction); two off-diagonal heads emit 𝑈𝑚 , 𝑉𝑚 ∈ R𝐿𝑠 ×𝐿𝑠 ; two node heads emit per-leaf bridges 𝑈˜𝑘 , 𝑉˜𝑘 ∈ R𝐿×𝐿𝑠 between node resolution and coarse tile resolution. A small learned per-node Jacobi gate 𝜆𝑖 adds a diagonal correction diag(𝝀)diag(𝐴) −1 that absorbs gradient early in training and decays toward zero as the structured branches take over (Supplementary ??). In reported runs we do not explicitly enforce strict SPD and did not observe instability; if a strict guarantee is required, one can add a tiny positive diagonal shift (e.g., softplus-parameterized, as in prior learned preconditioners [Yang et al. 2025]). All factors pack into a single tensor of width 𝑃 = 𝐾𝐿 2 + 𝑀 H 𝐿𝑠2 + 2𝑁 𝐿𝑠 + 𝑁 that the apply consumes without ever materializing 𝑀. Complexity and measured breakdown. Encoder cost is 𝑂 (𝑁𝑑 2 ), diagonal stack 𝑂 (𝑛𝑙 𝑁 𝐿𝑑), off-diagonal stack 𝑂 (𝑛𝑙 𝑁 𝐿𝑠2𝑑/𝐿), highways 𝑂 (𝑛𝑙 𝑁𝑑); total inference is Θ(𝑁 ). At 𝑁 =8 192, the subsystem profile in Table 2 shows that the two attention stacks account for 59% of device time, while encoder + decoder remain secondary. Kernellevel profiling in Table 3 then explains why: CUTLASS Tensor-Op GEMMs (31%) and fused attention (14%) dominate, so the workload is primarily compute-bound on Tensor Cores rather than bandwidthbound on sparse gathers. This directly supports the access-pattern argument in Table 1. ACM Trans. Graph., Vol. 45, No. 6, Article 1. Publication date: December 2026.
5
Preconditioner Application
𝑀 is never assembled. The apply consumes the packed factor tensor and produces 𝑀r through three stages on preallocated buffers, with shapes fixed at solve setup by the static partition. (i) A diagonal stage applies 𝐹𝑘 to each leaf residual r𝑘 in one batched GEMM of shape (𝐾, 𝐿, 𝐿). (ii) An off-diagonal FMM-style chain makes the data movement explicit: û𝑘 = 𝑈˜𝑘⊤ r𝑘 , v̂𝑘 = 𝑉˜𝑘⊤ r𝑘 (restriction) (2) ∑︁ ∑︁ r c s𝑚 = û𝑘 , s𝑚 = v̂𝑘 (strip aggregation) (3) 𝑘 ∈ R𝑚 c ⊤ r t𝑚 = 𝐵𝑚 s𝑚 ,
𝑘 ∈ C𝑚 r c t𝑚 = 𝐵𝑚 s𝑚
Í r + 𝑉˜ Í c Δy𝑘 = 𝑈˜𝑘 𝑚:𝑘 ∈ R𝑚 t𝑚 𝑘 𝑚:𝑘 ∈ C𝑚 t𝑚
(coarse coupling)
(4)
(prolongation).
(5)
Stages 1/3/4 are batched GEMMs of shapes (𝐾, 𝐿, 𝐿𝑠 ) and (𝑀 H , 𝐿𝑠 , 𝐿𝑠 ) (with (𝐾, 𝐿, 𝐿𝑠 ) reused in prolongation); stage 2 is two partitionindexed scatter-adds. (iii) The CSR SpMV for 𝐴r is the inner loop’s only sparse operation, and the space-filling-curve ordering keeps gathers nearly banded. Total apply cost is 𝑂 (𝑁 (𝐿 + 𝐿𝑠 )) FLOPs. Supplementary ?? contains the same equations with full shape progression notes. Single-graph capture. Every kernel in the inner loop has fixed launch shape, no host-side allocation, and no data-dependent control flow, so the entire iteration — preconditioner apply, SpMV, SAXPYs, dot-product reductions — records into one CUDA Graph that subsequent iterations replay with a single cudaGraphLaunch. CPU dispatch between kernels disappears, which in a real-time engine usually is the floor on per-frame latency. The graph-capturable property falls out of the design rather than being engineered: IC/DILU triangular solves serialize through data-dependent wavefronts and AMG V-cycles need level-by-level barriers, so neither is single-graph capturable in the same way. Among learned alternatives, neural SPAI [Yang et al. 2025] is graph-capturable but applies y =𝐺 ⊤ (𝐺r) as two SpMVs against a sparse 𝐺 whose nonzero pattern follows 𝐴, so its apply remains bandwidth-bound on irregular gathers; in contrast, our factorized approximate-inverse preconditioner apply reduces to dense block GEMMs on the partition and is compute-bound on Tensor Cores.
6 Training 6.1 Cosine Hutchinson Probe Loss We train the network self-supervised. Given a batch of probe vectors Z ∈ R𝑁 ×𝐾𝑧 drawn from the spectrally-balanced distribution described below, we compute 𝐴Z via SpMV (treated as a fixed input, with no gradient through 𝐴), apply the preconditioner to obtain 𝑀𝐴Z, and minimize the angle between 𝑀𝐴Z and Z as flattened tensors: ⟨Z, 𝑀𝐴Z⟩𝐹 Lcos = 1 − , (6) ∥Z∥ 𝐹 ∥𝑀𝐴Z∥ 𝐹 where ⟨𝑋, 𝑌 ⟩𝐹 = tr(𝑋 ⊤𝑌 ) is the Frobenius inner product and ∥ · ∥ 𝐹 the Frobenius norm. Equivalently, this is the cosine similarity between vec(Z) and vec(𝑀𝐴Z) — a single global angle over all 𝑁 ·𝐾𝑧 entries rather than an average of 𝐾𝑧 per-probe cosines. The singledenominator form ties all probes together through the same normalization, which we found to give noticeably more stable gradients at
Hierarchical Transformer Preconditioning for Interactive Physics Simulation
small 𝐾𝑧 than the per-probe mean. A perfect preconditioner gives Lcos =0; the worst case is Lcos =2 (anti-aligned). Why cosine: from vector distance to subspace distance. The conceptual upgrade Lcos makes over Frobenius- and SAI-style objectives is to replace a distance between vectors with a distance between sub2 1 spaces. Frobenius-type losses — including the SAI loss ∥𝐴∥ 𝐴𝑀 −𝐼 𝐹 used by [Yang et al. 2025] — penalize the pointwise Euclidean de1 viation of the specific vector ∥𝐴∥ 𝐴𝑀z from the specific vector z. That distance is sensitive to the magnitude of 𝐴𝑀z, so it implicitly demands that 𝐴𝑀 act as identity at a particular absolute scale — in the SAI case, at scale ∥𝐴∥. Lcos , by contrast, sees only the direction of 𝑀𝐴Z. Two preconditioners that send a probe to the same onedimensional subspace incur the same loss, no matter how they scale the vector inside that subspace. Geometrically, Lcos is a distance on the projective space P(R𝑁 𝐾𝑧 ) rather than a Euclidean distance on R𝑁 𝐾𝑧 : Proposition 1 (Cosine Hutchinson loss is a subspace distance). Let 𝑀 : R𝑁 → R𝑁 be a linear preconditioner, 𝐴 a fixed SPD matrix, and Z ∈ R𝑁 ×𝐾𝑧 a probe matrix with Z ≠ 0 and 𝑀𝐴Z ≠ 0. Write b Z =vec(Z) and b Y =vec(𝑀𝐴Z), both elements of R𝑁 𝐾𝑧 . (1) (Positive-scale invariance.) For every 𝛼 > 0, Lcos (𝛼𝑀) =Lcos (𝑀). The loss therefore descends to a well-defined function on the quotient of preconditioners modulo positive rescaling. (2) (Subspace interpretation.) Let 𝜃 ∈ [0, 𝜋/2] be the principal angle between the lines Rb Z and Rb Y in R𝑁 𝐾𝑧 , and let Π Rv = ⊤ 2 vv /∥v∥ 2 denote the orthogonal projector onto the line through v. Then Lcos (𝑀) = 1 − cos 𝜃,
1 2
2
Π RbZ − Π RbY 𝐹 = 1 − cos2 𝜃,
(7) so Lcos depends on b Z, b Y only through the unsigned angle between the lines they span in R𝑁 𝐾𝑧 . Both functionals vanish exactly when those two lines coincide. (3) (SAI loss is a vector distance.) The SAI-style loss LSAI (𝑀) = 2 1 ∥𝐴∥ 𝐴𝑀Z − Z 𝐹 is not invariant under 𝑀 ↦→ 𝛼𝑀 for 𝛼 ≠ 1, and is minimized uniquely (over scalings of 𝑀) at the choice 1 that places the vector ∥𝐴∥ 𝐴𝑀Z as close as possible as a vector 𝑁 𝐾 in R 𝑧 to the specific target Z. Sketch. (1) Scaling 𝑀 by 𝛼 > 0 scales the numerator ⟨Z, 𝑀𝐴Z⟩𝐹 by 𝛼 and the denominator ∥Z∥ 𝐹 ∥𝑀𝐴Z∥ 𝐹 by the same 𝛼, leaving the cosine unchanged. (2) The first equality is by definition of 𝜃 ; the second follows from the rank-one projector identity 12 ∥ΠRu − Π Rv ∥ 2𝐹 = 1 − ⟨u, v⟩ 2 /(∥u∥ 2 ∥v∥ 2 ) = 1 − cos2 𝜃 . Both 1 − cos 𝜃 and 1−cos2 𝜃 are valid notions of squared chordal distance on P(R𝑁 𝐾𝑧 ) near 𝜃 =0; we use 1 − cos 𝜃 in (6) for better-conditioned gradients near the minimum. (3) LSAI has the form ∥𝛽𝑀u − u∥ 2 for fixed 𝛽 =1/∥𝐴∥, strictly convex in 𝛽𝑀, hence not rescaling-invariant. □ The proposition matches the PCG geometry one-to-one. PCG’s convergence depends only on the relative spread of the eigenvalues of 𝑀𝐴, not on their absolute location (Sec. 3); accordingly, the correct space to optimize 𝑀 over is the projective space 𝑀/R>0 , and the natural loss on that space is a distance between the subspaces Rb Z
•
1:5
and Rb Y they span, exactly what Lcos provides. Frobenius- and SAIstyle losses live on the wrong space — they pin down the absolute scale of 𝑀𝐴 even though PCG does not care — and as a side effect implicitly demand that the eigenvalues of 𝑀𝐴 cluster near a chosen value (∥𝐴∥ in the SAI case), wasting capacity on a constraint with no algorithmic payoff. By dropping that constraint, Lcos frees the network to cluster the spectrum wherever the current preconditioner makes angular alignment easiest — a behavior we observe directly in §7.3 (Figs. 4–2) and ablate against SAI in §7.3.
6.2
Spectrally Biased Probe Vectors
An isotropic Gaussian probe places equal expected energy on every eigenmode of 𝑀𝐴 (Sec. 3), so the resulting gradient signal is also white in the probe’s eigenbasis. The blocks of our preconditioner, however, are not all at the same spatial scale: leaf-diagonal blocks resolve fine, high-frequency structure over 𝐿 adjacent nodes, while an off-diagonal tile of size 𝑆𝐿 resolves much lower-frequency, largerscale structure. A spectrally white probe distribution therefore distributes the gradient signal unevenly across these block scales — the high-frequency components, which the fine-scale (diagonal) blocks are tuned to, dominate the signal; the coarse-scale (off-diagonal, large-𝑆) blocks receive proportionally weaker gradients. The downstream effect is that blocks at different scales saturate at different times during training, with the coarse blocks plateauing late and limiting overall convergence. We rebalance the gradient signal across block scales by shifting probe energy toward lower spatial frequencies. A small number of damped-Jacobi smoothing steps acts as a spectral low-pass on the probe: z (𝑡 +1) = z (𝑡 ) − 𝜔 𝐷 −1𝐴 z (𝑡 ) ,
𝐷 = diag(𝐴),
(8)
with 𝜔 =0.6 and two steps in every reported run. The high-frequency components of the probe are damped more strongly than the lowfrequency components, redistributing probe energy toward the eigenmodes the coarse blocks are responsible for. The result is more even gradient magnitudes across block scales and substantially more synchronous training of fine and coarse blocks. Probes are detached, so gradients do not flow back through the smoothing.
6.3
Training Setup
We train with AdamW under a reduce-on-plateau schedule and global gradient clipping. Training contexts (graph, 𝐴, masks, padded sizes, smoothed probes) are precomputed once and cached on disk; at each step a mini-batch is drawn at random and padded to a common node count. We set the number of probe vectors to 𝐾𝑧 = √ max(64, ⌈ 𝑁 ⌉), balancing gradient noise against compute. The model is compiled with torch.compile. Preconditioner weights and the apply path use float32; only the PCG scalar accumulators (dot products, residual norms, step sizes) are computed in float64, which we found is sufficient to prevent residual drift on stiff systems without paying for full mixed-precision GEMMs. ACM Trans. Graph., Vol. 45, No. 6, Article 1. Publication date: December 2026.
1:6
•
Osborne et al.
7 Experiments 7.1 Setup All GPU experiments run on a single NVIDIA H200. Our model uses 𝑑 =128, 𝑛𝑙 =3, 𝐿 =128, 𝑝 off =4, 𝑛 gcn =2, highways on, trained once for ∼ 15 min per scale and reused for every test frame at that scale. All reported solve times are to relative residual ∥r𝑘 ∥ 2 /∥r0 ∥ 2 ≤ 10−8 — two to three orders of magnitude tighter than the ∼ 10−3 tolerance typical of graphics-grade pressure projection, chosen so the ranking reflects preconditioner quality, not early termination. PCG timing uses single-graph CUDA Graph replay for every method that admits it (unpreconditioned CG, Jacobi, AMGX SPAI, neural SPAI, ours) and per-kernel launches for IC/AMG-class methods. AMGX runs with vendor defaults (we swept neighboring configurations at 𝑁 =8 192 and saw no improvement). The neural SPAI baseline of Yang et al. [Yang et al. 2025] is re-trained per scale on the same multiphase distribution using their SAI loss and applied as a CUDA-Graphcaptured pair of SpMVs. Full hardware/software, precision, and dataset details are in Supplementary §??. Benchmark. We instantiate the target regime — stiff, every-framedifferent SPD systems with a hard real-time budget — as 2D multiÍ phase pressure-Poisson: the 5-point Laplacian 𝐴𝑖𝑖 = 𝑗 𝑤𝑖 𝑗 , 𝐴𝑖 𝑗 = −𝑤𝑖 𝑗 with harmonic-mean face conductances 𝑤𝑖 𝑗 =2𝜌𝑖 𝜌 𝑗 /(𝜌𝑖 +𝜌 𝑗 ), on a per-cell density field 𝜌 randomized per frame across three axes: contrast (𝜌 heavy ∼ loguniform[5, 100], so 𝜅 (𝐴) ∈ [103, 105 ]), barrier topology (1–3 rectangular barriers with gap configurations including closed, creating near-disconnected sub-domains that force long-range coupling), and orientation. We pick 2D because it is a harder setting for local preconditioners (only 4 neighbors per node) and because it covers the pressure-projection workload that dominates graphics-grade specialized simulators (FLIP/PIC, MPM, fractional-step Navier–Stokes). The architecture is not specific to structured grids — it needs only a sparse graph and a loose spatial ordering — and extends directly to 3D (Fig. 1); we evaluate quantitatively in 2D for fair, fully-tuned comparison against classical baselines. We do not target 𝑁 ≫ 106 regular voxel grids where structured multigrid amortizes [Lyu et al. 2026], nor batched offline PDE workloads. Per-frame randomization and discretization details are in Supplementary §??; representative frame in Supplementary, Fig. ??.
7.2
Main Performance
Main result. Table 4 (also plotted on the figure pages as Fig. 5; per-method runtime breakdown in Supplementary, Table ??) reports per-frame mean solve time and iteration counts. Our method runs at interactive framerates across the full size range: 17.9 ms (∼ 56 fps, 168 iters) at 𝑁 =8 192 and 47.6 ms (∼ 21 fps, 394 iters) at 𝑁 =16 384. The closest GPU baseline that converges in the same regime is Jacobi, at 39.5/65.7 ms (968/1543 iters) respectively — ∼ 6× fewer iterations at 𝑁 =8 192 (968 vs. 168) and ∼ 4× at 𝑁 =16 384 (1543 vs. 394) on identical hardware, and a 1.4–2.2× wall-clock gap. Neural SPAI [Yang et al. 2025], re-trained per scale with its CUDA apply path, lands at 48.1/70.9 ms (338/496 iters), trailing our method by 2.7× at 𝑁 =8 192 and 1.5× at 𝑁 =16 384 despite an iteration count ACM Trans. Graph., Vol. 45, No. 6, Article 1. Publication date: December 2026.
Table 4. Per-frame end-to-end PCG solve time (ms) and PCG iteration count (in parentheses) on the multiphase Poisson benchmark, 20 frames per scale, relative residual tolerance 10 −8 . GPU methods that admit single-graph capture (unpreconditioned CG, Jacobi, AMGX SPAI, neural SPAI, ours) run the inner loop as a single CUDA Graph replay (§5); IC-class GPU methods and AMGX V-cycles launch per kernel. Our Jacobi/CG numbers already include the CUDA-Graph speedup — an upper bound on what graph capture alone buys without our preconditioner. CPU rows are reference only. Bold marks the fastest GPU time at each scale. 1 024
Method Unprec. CG (GPU) Jacobi (GPU) AMGX SPAI (GPU) IC / DILU (GPU) Neural SPAI (GPU) [Yang et al. 2025] Ours (GPU)
18.5 (497) 12.7 (325) 53.2 (1) 139.5 (11) 18.4 (118) 7.0 (47)
2 048
4 096
24.7 (650) 43.6 (1 153) 16.1 (429) 31.7 (839) 82.0 (1) 76.0 (1) 208.4 (15) 312.8 (22) 26.0 (167) 38.3 (246) 8.8 (66) 9.2 (80)
8 192
16 384
77.6 (1 765) 39.5 (968) 134.5 (2) 503.0 (30) 48.1 (338) 17.9 (168)
89.2 (2 103) 65.7 (1 543) 188.3 (3) 579.7 (40) 70.9 (496) 47.6 (394)
CPU reference (not in competition with GPU rows) IC (CPU) 128.9 (59) 51.6 (96) 159.9 (113) 1 478.9 (170) 13 968.4 (254) AMG (CPU) 22.2 (5) 24.8 (9) 24.5 (5) 109.8 (7) 674.2 (10) Neural SPAI (CPU) [Yang et al. 2025] 6.7 (118) 10.9 (167) 25.6 (246) 62.3 (338) 202.4 (496)
within ∼ 2× of ours — the gap is dominated by the two randomgather SpMVs its apply dispatches per iteration, in contrast to our single batched-GEMM apply on contiguous tensors (§5). The ICand AMG-class baselines reach very low iteration counts (1–40) but pay for it in sequential triangular solves or V-cycle synchronization, falling below 2–5 fps at 𝑁 =16 384 — confirming the architectural argument of Table 1. The measured curve is not perfectly linear in 𝑁 because kernels cross warp/tile-quantization and cache-transition thresholds as the working set grows, even though the algorithmic order remains Θ(𝑁 ). Fig. 6 traces how convergence actually looks on a challenging frame — a closed cross-shaped barrier with stiff density contrast chosen to maximize long-range coupling — across unpreconditioned CG, Jacobi, IC, AMG, and ours, with the right-hand side supported only on thin density interfaces. By 𝑘 =1, Jacobi and IC have damped the residual only locally around supp(𝑏), while AMG and ours have already attenuated it across the whole domain — the visual signature of multiscale transport that the highway buffers implement (§4.3). AMG matches the spread but pays in V-cycle synchronization per iteration; ours runs the inner loop as a single CUDA-Graph capture. Iteration counts on this frame (288/1097/334/19/384 for unprec/Jacobi/IC/AMG/ours) span nearly two orders of magnitude, but per-iteration cost reverses the ranking for AMG and IC. Standard deviations across the 20 test frames track condition number rather than 𝑁 : Jacobi and unprec. CG show 𝜎/𝜇 ≈ 50–100% at 𝑁 ≥ 8 192 across the [5, 100] contrast range, while ours stays at ≤ 21% even at 𝑁 =16 384.
7.3
Training dynamics and ablations
Training dynamics and loss. Fig. 4 tracks Lcos , the SAI loss on the same checkpoints, and PCG iteration count across training. The three curves move together until step ∼ 8 000, after which the SAI loss rises from ∼ 10−3 to ∼ 0.5 while Lcos keeps falling in lockstep with PCG iterations — direct evidence of Prop. 1: the model is moving the eigenvalues of 𝑀𝐴 away from ∥𝐴∥ in pointwise terms while tightening the relative cluster wherever angular alignment
Hierarchical Transformer Preconditioning for Interactive Physics Simulation
Table 5. Core architecture ablations (subset of Supplementary Table ??). Group
Infer. (ms)
Iters
Total (ms)
Width (avg. 𝑁 ∈ {2048, 4096, 8192} ) 𝑑=64, 𝑛𝑙 =3, hw Width (avg. 𝑁 ∈ {2048, 4096, 8192} ) 𝑑=128, 𝑛𝑙 =3, hw Width (avg. 𝑁 ∈ {2048, 4096, 8192} ) 𝑑=256, 𝑛𝑙 =3, hw
Configuration
3.0 3.1 3.2
147 105 191
15.4 12.0 19.3
Depth (𝑁 =8 192) Depth (𝑁 =8 192)
𝑑=128, 𝑛𝑙 =1, hw 𝑑=128, 𝑛𝑙 =3, hw
1.3 3.4
421 168
36.9 17.9
Highways (𝑁 =2 048) Highways (𝑁 =2 048)
𝑑=128, 𝑛𝑙 =3, hw 𝑑=128, 𝑛𝑙 =3, no-hw
3.3 2.2
66 149
8.8 13.8
Table 6. Generalization at 𝑁 = 4 096 across the four eval distributions, in PCG iterations and speedup vs. Jacobi (mean over 20 frames, rtol =10−8 ). “Jacobi ms” is the baseline denominator used to compute speedup. Our model and the SAI-loss ablation are trained on the complement of each eval cell (held-out OOD setup); neural SPAI [Yang et al. 2025] is trained with the same split protocol (train on the corresponding restricted distribution for each eval row).
Jacobi
Ours (cosine loss)
Ours arch + SAI loss Neural SPAI [Yang et al. 2025] [Yang et al. 2025]
Eval distribution
ms
iters
speedup
iters
speedup
iters
speedup
Full / in-distribution Closed barriers only High contrast Closed + high contrast
29.4 24.7 24.4 25.1
82 68 142 147
3.4× 3.3× 1.9× 1.9×
405 208 175 174
0.9× 1.3× 1.5× 1.6×
236 258 264 396
0.9× 0.7× 0.7× 0.5×
is easiest. Fig. 2 confirms this on the spectrum at 𝑁 = 1 024: SAI delivers a 16× 𝜅 reduction with the cluster anchored near ∥𝐴∥, the same architecture trained with Lcos delivers 68× with the cluster wherever it pleased — a 4.3× gap attributable to the loss alone. Total wall-clock training is 16.2 min on a single H200 (24 300 steps); the model overtakes GPU Jacobi by step ∼ 2 000 (∼ 1.3 min) and drops below 200 PCG iterations by step ∼ 12 000 (Fig. 7 corroborates the link from probe-space alignment to spectral clustering). Architecture ablations. Width 𝑑, depth 𝑛𝑙 , and highway connections each move total solve time non-trivially. Default (𝑑 = 128, 𝑛𝑙 =3, hw on, ∼ 2.6M parameters): (i) shrinking depth 𝑛𝑙 =3 → 1 cuts inference 2.6× but doubles total solve time because iterations rise 2.5× (additional layers are needed to compose information routed through the highway tokens); (ii) removing highways raises PCG iterations 2.3× at 𝑁 = 2 048 and the penalty grows with 𝑁 ; (iii) 𝑑 = 64 is 28% slower overall, 𝑑 = 256 is competitive at small 𝑁 but raises mean PCG iterations from 105 to 191 when averaged over 𝑁 ∈ {2048, 4096, 8192}. Table 5 summarizes the core sweep results; full rows remain in Supplementary, Table ??.
7.4
Generalization and what the loss buys
We probe within-family deployment robustness at 𝑁 =4 096 across three shifts — topology (closed barriers withheld from training), contrast (train on 𝜌 heavy ∈ [5, 25], evaluate on (25, 100]), and their composition — and compare three systems on identical eval sets: ours, the same architecture trained instead with the SAI loss of [Yang et al. 2025], and neural SPAI [Yang et al. 2025] trained/evaluated under the same split protocol (Table 6). The same-architecture row is a clean loss ablation; the neural SPAI row is a matched-split learned baseline rather than a full-train upper bound.
•
1:7
Three observations carry the section. (i) Topology generalization is essentially free: withholding closed barriers from training leaves iteration counts unchanged (68 vs. 82), because the H -matrix partition is keyed to spatial indexing, not barrier geometry. (ii) The loss, not the architecture or the data, is what unlocks the quality. Replacing only the loss — same network, same training distribution — raises iteration counts ∼ 5× at the in-distribution eval cell (82 → 405); the SAI gradient pins eigenvalues near ∥𝐴∥, wasting capacity on a constraint PCG does not care about (Prop. 1). Under matched-split training, neural SPAI sits at 236–264 iterations on the first three rows and degrades to 396 on the compositional row. Relative to Jacobi, our row stays at 1.9–3.4× speedup across all eval cells, while neural SPAI is 0.9×, 0.7×, 0.7×, and 0.5× (slower than Jacobi in three of four rows, and substantially slower in the compositional case); the same-architecture SAI ablation reaches only 0.9–1.6×. (iii) Contrast is the dominant remaining OOD axis for our method. Pure amplitude growth is absorbed for free by Lcos ’s scale invariance, but far-field interactions at high contrast push the spectrum past the cluster the network has seen — iteration count roughly doubles (82 → 142– 147) and 𝜎/𝜇 grows fivefold. Full robustness grid in Supplementary, Table ??.
8
Discussion and Future Work
The recipe — a weak-admissibility H -matrix prior, a scale-invariant cosine-Hutchinson objective, and a single-graph apply path — is most useful where it targets: stiff, every-frame-different SPD systems with a hard real-time budget. The architecture itself depends on nothing fluid-specific, only a loose spatial ordering of degrees of freedom, and extends directly to 3D (Fig. 1). We expect the largest gains to persist on other SPD families with (i) geometric locality, (ii) frame-to-frame coefficient changes, and (iii) hard real-time budgets (implicit viscosity/diffusion, soft-body and contact dynamics) and smaller gains where a single matrix is reused long enough for heavy classical setup to amortize. Two known limitations are worth flagging: (a) training pre-sizes the partition to a maximum 𝑁 , so pushing past it currently requires retraining — a dynamic-partition variant (constant leaf count at 𝑂 (𝑁 log 𝑁 ) rather than constant leaf size at Θ(𝑁 ), §4.2) removes this ceiling at the cost of one extra pooling pass per layer; (b) the H -matrix prior assumes some spatial locality of 𝐴 under its indexing, and degrades on operators without a natural spatial coordinate (power-grid Laplacians, social-network matrices) or where the far-field rank does not decay with separation.
References Michele Benzi, Carl D. Meyer, and Miroslav Tŭma. 1996. A Sparse Approximate Inverse Preconditioner for the Conjugate Gradient Method. SIAM Journal on Scientific Computing 17, 5 (1996), 1135–1149. Steffen Börm, Lars Grasedyck, and Wolfgang Hackbusch. 2003. Introduction to Hierarchical Matrices with Applications. Engineering Analysis with Boundary Elements. Jie Chen. 2024. Graph Neural Preconditioners for Iterative Solutions of Sparse Linear Systems. arXiv preprint arXiv:2406.00809 (2024). https://arxiv.org/abs/2406.00809 Yuwei Fan, Lin Lin, Lexing Ying, and Leonardo Zepeda-Núñez. 2019. A Multiscale Neural Network Based on Hierarchical Matrices. arXiv preprint arXiv:1807.01883 (2019). https://arxiv.org/abs/1807.01883 Emilio McAllister Fognini, Marta M. Betcke, and Ben T. Cox. 2025. Learning Green’s Operators through Hierarchical Neural Networks Inspired by the Fast Multipole Method. arXiv preprint arXiv:2509.20591 (2025). https://arxiv.org/abs/2509.20591 Leslie Greengard and Vladimir Rokhlin. 1987. A Fast Algorithm for Particle Simulations. J. Comput. Phys. 73, 2 (1987), 325–348. doi:10.1016/0021-9991(87)90140-9 ACM Trans. Graph., Vol. 45, No. 6, Article 1. Publication date: December 2026.
1:8
•
Osborne et al.
Wolfgang Hackbusch. 1999. A Sparse Matrix Arithmetic Based on H -Matrices. Part I: Introduction to H -Matrices. Springer, Berlin. Wolfgang Hackbusch and Boris N. Khoromskij. 2002. Adaptive H -Matrix Approximation on General Domains. Springer. Tucker Hartland, Georg Stadler, Mauro Perego, Kim Liegeois, and Noémi Petra. 2023. Hierarchical Off-Diagonal Low-Rank Approximation of Hessians in Inverse Problems, with Application to Ice Sheet Model Initialization. arXiv preprint arXiv:2301.03644 (2023). https://arxiv.org/abs/2301.03644 Paul Häusner, Ozan Öktem, and Jens Sjölund. 2023. Neural Incomplete Factorization: Learning Preconditioners for the Conjugate Gradient Method. arXiv preprint arXiv:2305.16368 (2023). https://arxiv.org/abs/2305.16368 Michael F. Hutchinson. 1989. A Stochastic Estimator of the Trace of the Influence Matrix for Laplacian Smoothing Splines. Communications in Statistics—Simulation and Computation 18, 3 (1989), 1059–1076. Markus Ihmsen, Nadir Akinci, Markus Becker, and Matthias Teschner. 2011. A Parallel SPH Implementation on Multi-Core CPUs. Computer Graphics Forum 30, 1 (2011), 99–112. Tero Karras. 2012. Maximizing Parallelism in the Construction of BVHs, Octrees, and 𝑘 d Trees. In Proc. ACM SIGGRAPH/Eurographics Conf. on High-Performance Graphics (HPG). 33–37. Thomas N. Kipf and Max Welling. 2017. Semi-Supervised Classification with Graph Convolutional Networks. In International Conference on Learning Representations (ICLR). L. Yu. Kolotilina and A. Yu. Yeremin. 1993. Factorized Sparse Approximate Inverse Preconditionings I: Theory. SIAM J. Matrix Anal. Appl. 14, 1 (1993), 45–58. Yichen Li, Peter Yichen Chen, Tao Du, and Wojciech Matusik. 2023. Learning Preconditioners for Conjugate Gradient PDE Solvers. arXiv preprint arXiv:2305.16432 (2023). https://arxiv.org/abs/2305.16432 Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. 2020a. Multipole Graph Neural Operator for Parametric Partial Differential Equations. Advances in Neural Information Processing Systems 33 (2020). Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. 2020b. Neural Operator: Graph Kernel Network for Partial Differential Equations. arXiv preprint arXiv:2003.03485 (2020). https://arxiv.org/abs/2003.03485 Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. 2021. Fourier Neural Operator for Parametric Partial Differential Equations. In International Conference on Learning Representations. https://openreview.net/forum?id=c8P9NQVtmnO Chih-Jen Lin and Jorge J. Moré. 1999. Incomplete Cholesky Factorizations with Limited Memory. SIAM Journal on Scientific Computing 21, 1 (1999), 24–45. Lijun Liu, Shengguo Li, Xiangke Hu, Yutong Wang, Xuejun Liu, and Jingling Xue. 2016. Exploring Data Level Parallelism in Incomplete LU Factorization on GPUs. IEEE Transactions on Parallel and Distributed Systems 27, 12 (2016), 3397–3410. Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. 2021. Swin Transformer: Hierarchical Vision Transformer using Shifted Windows. arXiv preprint arXiv:2103.14030 (2021). https://arxiv.org/abs/2103.14030 Lu Lu, Pengzhan Jin, Guofei Pang, Zhongqiang Zhang, and George Em Karniadakis. 2021. Learning Nonlinear Operators via DeepONet Based on the Universal Approximation Theorem of Operators. Nature Machine Intelligence 3, 3 (2021), 218–229. Ilay Luz, Meirav Galun, Haggai Maron, Ronen Basri, and Irad Yavneh. 2020. Learning Algebraic Multigrid Using Graph Neural Networks. arXiv preprint arXiv:2003.05744 (2020). https://arxiv.org/abs/2003.05744 Kangbo Lyu, Ruihong Cen, Yushen Wu, and Tao Du. 2026. A Multigrid-Inspired Neural Iterative Solver for Poisson Equations on Large Voxel Grids. https://openreview. net/forum?id=lNcbGSWhJo. Maxim Naumov, Michael Chien, Paul Vandermersch, Ujval Kapasi, Boris Catanzaro, and Michael Garland. 2015. cuSPARSE Library. In GPU Technology Conference (GTC). Maziar Raissi, Paris Perdikaris, and George Em Karniadakis. 2019. Physics-Informed Neural Networks: A Deep Learning Framework for Solving Forward and Inverse Problems Involving Nonlinear Partial Differential Equations. J. Comput. Phys. 378 (2019), 686–707. Kirill Rudikov, Anastasia Markeeva, Vasily Bulatov, and Dmitry Vetrov. 2024. Neural Functional Operator for Parametric PDEs. arXiv preprint arXiv:2402.01030 (2024). https://arxiv.org/abs/2402.01030 John W. Ruge and Klaus Stüben. 1987. Algebraic Multigrid (AMG). Multigrid Methods 3 (1987), 73–130. Yousef Saad. 2003. Iterative Methods for Sparse Linear Systems (2 ed.). SIAM, Philadelphia, PA. Pietro Sittoni, Emanuele Zangrando, Angelo Alberto Casulli, Nicola Guglielmi, and Francesco Tudisco. 2026. Neural-HSS: Hierarchical Semi-Separable Neural PDE Solver. arXiv preprint arXiv:2602.18248 (2026). https://arxiv.org/abs/2602.18248 Matthias Teschner, Bruno Heidelberger, Matthias Müller, Danat Pomerantes, and Markus H. Gross. 2003. Optimized Spatial Hashing for Collision Detection of Deformable Objects. In Vision, Modeling, and Visualization (VMV). 47–54.
ACM Trans. Graph., Vol. 45, No. 6, Article 1. Publication date: December 2026.
Vladislav Trifonov, Ekaterina Muravleva, and Ivan Oseledets. 2025. Message-Passing GNNs Fail to Approximate Sparse Triangular Factorizations. arXiv preprint arXiv:2502.01397 (2025). https://arxiv.org/abs/2502.01397 Vladislav Trifonov, Alexander Rudikov, Oleg Iliev, Yuri M. Laevsky, Ivan Oseledets, and Ekaterina Muravleva. 2024. Learning from Linear Algebra: A Graph Neural Network Approach to Preconditioner Design for Conjugate Gradient Solvers. arXiv preprint arXiv:2405.15557 (2024). https://arxiv.org/abs/2405.15557 Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Łukasz Kaiser, and Illia Polosukhin. 2017. Attention Is All You Need. In Advances in Neural Information Processing Systems 30 (NIPS 2017). 5998–6008. Felix Wu, Amauri Souza, Tianyi Zhang, Christopher Fifty, Tao Yu, and Kilian Weinberger. 2019. Simplifying Graph Convolutional Networks. In Proceedings of the 36th International Conference on Machine Learning (Proceedings of Machine Learning Research, Vol. 97). PMLR, 6861–6871. http://proceedings.mlr.press/v97/wu19e.html Tianshi Xu, Rui Peng Li, and Yuanzhe Xi. 2025. Neural Approximate Inverse Preconditioners. arXiv preprint arXiv:2510.13034 (2025). https://arxiv.org/abs/2510.13034 Ichitaro Yamazaki, Stanimire Tomov, and Jack Dongarra. 2020. Mixed-Precision Cholesky QR Factorization and Its Parallelization on GPUs. Parallel Comput. 99 (2020), 102693. Zherui Yang, Zhehao Li, Kangbo Lyu, Yixuan Li, Tao Du, and Ligang Liu. 2025. Learning Sparse Approximate Inverse Preconditioners for Conjugate Gradient Solvers on GPUs. arXiv preprint arXiv:2510.27517 (2025). https://arxiv.org/abs/2510.27517
Hierarchical Transformer Preconditioning for Interactive Physics Simulation
Fig. 2. Spectra of 𝑀𝐴 at 𝑁 =1 024 on a representative multiphase Poisson frame (full eigendecomposition; 𝑦 jittered). Top: unpreconditioned (left) and Jacobi (right). Bottom: same architecture trained with SAI (left, 16× 𝜅 reduction, cluster anchored near 𝜆 = ∥𝐴∥ as the Frobenius objective demands) versus cosine-Hutchinson (right, 68× reduction, cluster wherever angular alignment is easiest). The 4.3× gap is attributable entirely to the loss. The bimodal structure in both neural preconditioners reflects the two-phase density contrast of the benchmark.
Fig. 3. Off-diagonal rank audit: provided vs. required. Red: the architectureprovided rank fraction 𝐿𝑠 /(𝑆𝐿) at fixed 𝐿𝑠 =32, plotted against off-diagonal distance class 𝑆. Dashed: the mean numerical rank fraction required for truncated-SVD approximations of 𝐴 −1 tiles to satisfy ∥𝑋 − 𝑋𝑟 ∥ 𝐹 /∥𝑋 ∥ 𝐹 ≤ 𝜀 at 𝜀 ∈ {10−3 , 10−6 , 10 −9 }, with bands giving ±1𝜎 over tiles in the same distance class. Required rank drops monotonically with 𝑆 (far interactions are more compressible) while the architecture’s fixed rank stays above it everywhere, with ∼ 2.6× headroom in the largest-distance class even at 𝜀 =10−9 — the structural evidence that one shared off-diagonal rank covers every far-field tile (justifying c1, §4.2).
•
1:9
Fig. 4. Training profile at 𝑁 =8 192 (default config, 𝑑 =128, 𝑛𝑙 =3, highways on). Blue: cosine-Hutchinson loss Lcos (mean ±1𝜎 over last 100 steps, linear axis). Green dashed: SAI loss [Yang et al. 2025] on a held-out frame at the same checkpoints (log axis); the down-then-up trajectory is the empirical signature of the loss-ablation argument in §7.3 — the model keeps reducing the cosine objective and the PCG iteration count after the SAI surrogate has bottomed out and started to climb. Orange: PCG iterations to rtol =10−8 on that frame (log). Bottom axis: wall-clock minutes; top axis: optimizer steps. Grey: unpreconditioned CG and Jacobi baselines.
Fig. 5. End-to-end PCG solve time on the multiphase Poisson benchmark across five problem sizes, 20 frames per scale (mean ±1𝜎 band; relative residual tolerance 10 −8 ). Dashed grey lines mark the 24 and 60 fps interactive budgets. Ours covers ∼ 143 down to ∼ 21 fps; the closest classical GPU baseline (Jacobi) covers ∼ 79 down to ∼ 15 fps; the closest learned baseline (neural SPAI [Yang et al. 2025], re-trained per scale on our benchmark with its CUDA apply path) covers ∼ 54 down to ∼ 14 fps. AMGX SPAI and GPU IC sit above 40 ms across all scales despite much lower iteration counts, dominated by unparallelizable triangular solves or V-cycle synchronization. Full per-method times and iteration counts appear in Table 4.
ACM Trans. Graph., Vol. 45, No. 6, Article 1. Publication date: December 2026.
1:10
•
Osborne et al.
Fig. 7. Probe-alignment dynamics during training at 𝑁 = 8 192 (𝑑 = 128, 𝑛𝑙 =3, highways on), over checkpoints in the 8k–24k step range. Each probe (point) is plotted by the angle its image 𝑀𝐴z makes with z in R𝑁 , early in training (left) vs. late (right). As Lcos drives the per-probe angle to zero, the point cloud collapses inward — the spatial signature of the spectral clustering of 𝑀𝐴 that Fig. 2 shows at 𝑁 = 1 024 and that Fig. 4 shows in PCG iteration counts.
Fig. 6. Multiscale error transport across preconditioner families on a challenging multiphase Poisson frame (closed cross-barrier topology, stiff density contrast), distinct from the teaser frame. Top row: density 𝜌, right-hand side 𝑏, pressure 𝑝, |𝑣 |, and a top-left view of the learned 𝑀. Rows below, top to bottom: unpreconditioned CG, Jacobi, IC, AMG, ours. Each row shows |𝑟𝑘 | on a shared log color scale at five iteration snapshots; row labels give each method’s iteration count to rtol =10 −8 . At 𝑘 =1, Jacobi and IC damp error only locally around supp(𝑏 ), while AMG and ours have already attenuated the residual across the whole domain — direct evidence that the learned 𝑀 routes correction signals across graph-distant regions in a single apply, the global routing pattern the highway buffers implement (§4.3). AMG matches the spread but pays in V-cycle synchronization per iteration; ours runs the inner loop as a single CUDA-Graph-captured sequence of batched GEMMs (§5).
ACM Trans. Graph., Vol. 45, No. 6, Article 1. Publication date: December 2026.
Supplementary Material: Hierarchical Transformer Preconditioning for Interactive Physics Simulation CARL OSBORNE, MIT CSAIL, USA MINGHAO GUO, MIT CSAIL, USA CRYSTAL OWENS, MIT CSAIL, USA WOJCIECH MATUSIK, MIT CSAIL, USA ACM Reference Format: Carl Osborne, Minghao Guo, Crystal Owens, and Wojciech Matusik. 2026. Supplementary Material: Hierarchical Transformer Preconditioning for Interactive Physics Simulation. ACM Trans. Graph. 45, 6, Article 1 (December 2026), 5 pages. https://doi.org/10.1145/nnnnnnn.nnnnnnn
This document collects the implementation details, dataset construction, per-scale performance numbers, full train × eval generalization grid behind Table 6, and the precision/apply-path discussion that did not fit the main-paper page budget. Section numbering is independent of the main paper.
1
Experimental Setup — Full Details
Hardware. All GPU experiments run under a batch scheduler on a shared institutional GPU cluster. Each job allocates one NVIDIA H200 (140 GB HBM3), eight Intel Xeon Platinum 8580 host cores, and 32 GiB host DRAM alongside the device (partial-node slice on dual-socket nodes with eight GPUs per node). CPU baselines use the same eight host cores. Software stack. GPU code is built with the CUDA 12 toolchain and a PyTorch 2.6.x build from a fixed Conda environment, using torch.compile for graph capture. CPU baselines use Eigen, SciPy, and PyAMG. Sparse matrix-vector products on the GPU use the CSR layout produced by the dataset generator (Sec. 4). PCG timing uses single-graph CUDA Graph replay for methods that admit it (unpreconditioned CG, Jacobi, AMGX SPAI, neural SPAI, ours); ICand AMG-class methods are timed per launch since their sequential triangular solves and V-cycle synchronization preclude single-graph capture. Default model configuration. Unless noted, our model uses 𝑑 =128, 𝑛𝑙 = 3, 𝐿 = 128, 𝑝 diag = 1, 𝑝 off = 4 (so 𝐿𝑠 = 32), 𝑛 gcn = 2, ℎ = 8 attention heads, highways enabled (d128_L3_hw in figures). This single checkpoint is reused at every 𝑁 in the main performance table. Sec. 2 expands the per-stage shapes and Sec. 3 the optimizer and probe details. Baseline tuning. GPU IC (multicolor_dilu) and GPU AMG settings follow AMGX vendor defaults. At 𝑁 =8 192 we additionally swept the nearest neighboring AMGX hyperparameter configurations (smoother type, coarsening strategy, maximum levels) and Authors’ Contact Information: Carl Osborne, MIT CSAIL, Cambridge, USA, osbo@mit. edu; Minghao Guo, MIT CSAIL, Cambridge, USA, [email protected]; Crystal Owens, MIT CSAIL, Cambridge, USA, [email protected]; Wojciech Matusik, MIT CSAIL, Cambridge, USA, [email protected]. 2026. ACM 1557-7368/2026/12-ART1 https://doi.org/10.1145/nnnnnnn.nnnnnnn
verified that none meaningfully reduces total time relative to the default. CPU PyAMG uses its default classical AMG cycle. AMGX SPAI is included as the sparse-approximate-inverse data point; we keep its vendor defaults because the neural SPAI baseline whose public implementation does not currently support our problem class (main paper, footnote in §7.1) would be the natural learned comparison, not AMGX SPAI’s hand-tuned variants. Convergence criterion. All reported solve times and iteration counts are to a relative residual tolerance of ∥r𝑘 ∥ 2 /∥r0 ∥ 2 ≤ 10−8 . This is two to three orders of magnitude tighter than the ∼ 10−3 accuracy typical of graphics-grade pressure projection. We chose the tighter tolerance so per-method comparisons reflect preconditioner quality rather than early termination; at 10−3 every method finishes in fewer iterations, but the relative ordering and ratios we report are preserved. Precision. PCG scalar accumulators (dot products, residual norms, step sizes) use float64; preconditioner weights and the apply path use float32. Sec. 7 discusses why this split is sufficient on stiff systems at our target tolerance and where each component sits within the apply-path memory budget.
2
Network Architecture Details
This section expands the four-stage pipeline of main-paper §4 with the exact tensor shapes and module-level choices needed to reproduce the network. Encoder (Sec. 4.1). The per-node feature vector concatenates the simulator-exposed scalar fields (density, pressure, geometric position, boundary indicators) with a broadcast global context of dimension 𝑑 glob = 12 summarizing per-frame statistics of the assembled system. A two-layer MLP with GELU activations lifts this to width 𝑑. The encoder then applies 𝑛 gcn =2 residual graph-convolutional layers [Kipf and Welling 2017] that use 𝐴 itself (normalized by diag(𝐴)) as the message-passing weight. The encoder is the only stage that touches individual edges; downstream stages consume only the per-node embedding tensor of shape (𝑁 , 𝑑). Diagonal attention stack (Sec. 4.2). Encoder embeddings are reshaped from (𝑁 , 𝑑) to (𝐾, 𝐿, 𝑑) and fed through 𝑛𝑙 transformer blocks with attention restricted to within-leaf token pairs (the Swinstyle local window of [Liu et al. 2021; Vaswani et al. 2017] adapted to a one-dimensional index range). Each block uses ℎ =8 heads, head dimension 𝑑/ℎ =16, a GELU-activated FFN with 4𝑑 hidden width, pre-norm LayerNorm, and residual connections. Attention logits carry a learned per-head bias produced by a two-layer MLP from the four-dimensional edge descriptor (Δx𝑖 𝑗 , 𝐴𝑖 𝑗 ). ACM Trans. Graph., Vol. 45, No. 6, Article 1. Publication date: December 2026.
1:2
•
Osborne et al.
Off-diagonal attention stack (Sec. 4.2). The same encoder embeddings are pooled along each tile’s row- and column-strip I𝑚 , J𝑚 to obtain per-leaf coarse summaries, then condensed along the in-strip axis to 𝐿𝑠 = 𝐿/𝑝 off = 32 tokens per tile via uniform mean-pooling. Tile-grouped tokens then arrive at the attention stack with shape (𝑀 H , 𝐿𝑠 , 𝑑), independent of the physical tile size 𝑆𝑚 𝐿, and are processed by 𝑛𝑙 transformer blocks identical in form to the diagonal stack (same head count, FFN width, normalization). Edge biases for each in-tile token pair are computed from the mean of (Δx𝑖 𝑗 , 𝐴𝑖 𝑗 ) over all node pairs spanning the corresponding sub-strips. Highway buffers (Sec. 4.3). After every attention sublayer in either stack, block-token embeddings are scatter-added into per-layer row, column, and global buffers rhw, chw ∈ R𝐵×𝑁 ×𝑑 and ghw ∈ R𝐵×𝑑 (offdiagonal tile tokens are first repeat-interleaved back to full 𝐿-leaf resolution so each tile contributes uniformly across its 𝑆 leaves). Before the FFN sublayer, each token’s row, column, and global highway slices are concatenated with its own 𝑑-dim embedding and the resulting 4𝑑-wide vector is fed through the FFN input projection. The four-channel mix (2D local + 1D row + 1D column + 0D broadcast) is repeated independently in every transformer block; the buffers themselves are not residual across blocks. Decoder heads (Sec. 4.4). Three lightweight heads project token embeddings to the factor tensors that the apply path consumes. The leaf head is a two-layer MLP per diagonal-stream token followed by a reshape into 𝐹𝑘 ∈ R𝐿×𝐿 . The two off-diagonal heads are single linear layers producing 𝑈𝑚 , 𝑉𝑚 ∈ R𝐿𝑠 ×𝐿𝑠 from the off-diagonal-stream tokens (𝐿𝑠 tokens per tile, one head emits 𝑈𝑚 , the other 𝑉𝑚 ). The two node heads are single linear layers producing the bridge matrices 𝑈˜𝑘 , 𝑉˜𝑘 ∈ R𝐿×𝐿𝑠 from the diagonal-stream per-node embeddings. A final scalar gate produces the Jacobi-residual weight 𝜆𝑖 . Apply-path tensor layout. The decoder outputs are concatenated into a single packed tensor of width 𝑃 = 𝐾𝐿 2 + 𝑀 H 𝐿𝑠2 + 2𝑁 𝐿𝑠 + 𝑁 with strides aligned to the H -matrix partition. With the weakadmissibility H -matrix partition at admissibility parameter 𝜂 =1, the number of unique off-diagonal tiles satisfies 𝑀 H =𝐾 −1 (geometric sum 𝐾/2+𝐾/4+· · ·+1 over the strict upper triangle); the symmetric c = lower triangle reuses the same 𝐵𝑚 via the transposed apply t𝑚 ⊤ sr from the main-paper Sec. 5 equations. Concrete factor-tensor 𝐵𝑚 𝑚 sizes at the scales of the main paper’s performance table (main paper Table 4) are listed in Sec. 7. Explicit off-diagonal apply equations (main paper Sec. 5). For leaf residual r𝑘 ∈ R𝐿 and tile leaf-ranges R𝑚 , C𝑚 , the off-diagonal contribution follows the same FMM-style chain used in the main text: û𝑘 = 𝑈˜𝑘⊤ r𝑘 , v̂𝑘 = 𝑉˜𝑘⊤ r𝑘 (restriction) (1) ∑︁ ∑︁ r c s𝑚 = û𝑘 , s𝑚 = v̂𝑘 (strip aggregation) (2) 𝑘 ∈ R𝑚
share the same launch shape), plus two partition-indexed scatteradds for strip aggregation and redistribution.
3
Training Procedure
Optimizer and schedule. We train with AdamW (PyTorch defaults for 𝛽 1, 𝛽 2, 𝜖, weight decay 10−4 ) at initial learning rate 2 × 10−4 . Learning rate is reduced on plateau with ReduceLROnPlateau (factor 0.5, patience 5 log steps, relative threshold 5×10−3 , minimum max(lr× 10−3, 10−6 )). Gradients are globally clipped to ℓ2 norm 1 before each AdamW step. Auto-stop. Training proceeds for at most 105 optimizer steps but is terminated early once the LR scheduler reaches min_lr and the cosine-Hutchinson loss has failed to improve by the same 5×10−3 relative threshold over 10 consecutive log steps (twice the LR-scheduler patience). At our default configuration this fires at roughly 24 300 steps (main paper, §7.3). √ Probe vectors. At each step we draw a batch of 𝐾𝑧 =max(64, ⌈ 𝑁 ⌉) probe vectors Z ∼ N (0, 𝐼 ) and apply two damped-Jacobi smoothing sweeps with 𝜔 =0.6 (main paper, Eq. 8) to redistribute probe energy toward lower spatial frequencies. Probes are detached after smoothing, so no gradient flows back through the smoother. The constant 64 floor √ on 𝐾𝑧 keeps gradient noise bounded at the smallest scales; the 𝑁 growth keeps the per-step gradient signal-to-noise ratio approximately constant as 𝑁 grows. Gradient accumulation. Each optimizer step aggregates gradients over four random cached training contexts (graph, 𝐴, masks, padded sizes, smoothed probes). Contexts are precomputed once and cached on disk so the per-step cost is dominated by the forward/backward through the network rather than dataset assembly. H -matrix partition. The weak-admissibility partition is constructed once per scale at training start with admissibility parameter 𝜂 =1 and the leaf count 𝐾 =𝑁 /𝐿 determined by the scale. The partition is keyed to 𝑁 and reused unchanged on every test frame at that scale; no part of the partition depends on 𝐴 or the right-hand side. Wall-clock budget. End-to-end training of the default configuration is ∼ 16 min on one H200 at 𝑁 =8 192, with the model overtaking GPU Jacobi at step ∼ 2 000 (∼ 1.3 min) and dropping under 200 PCG iterations by step ∼ 12 000 (∼ 8 min). The training profile is plotted in main paper Fig. 4.
4
Multiphase Poisson Benchmark Generation
This appendix gives the precise procedure for generating each frame of the multiphase Poisson benchmark used throughout the main paper.
𝑘 ∈ C𝑚
c ⊤ r r c t𝑚 = 𝐵𝑚 s𝑚 t𝑚 = 𝐵𝑚 s𝑚 , Í r + 𝑉˜ Í c Δy𝑘 = 𝑈˜𝑘 𝑚:𝑘 ∈ R𝑚 t𝑚 𝑘 𝑚:𝑘 ∈ C𝑚 t𝑚
(coarse coupling)
(3)
(prolongation).
(4)
The stage-wise shape progression for one tile of span 𝑆𝑚 𝐿 is 𝑆𝑚 𝐿 → 𝑆𝑚 𝐿𝑠 → 𝐿𝑠 → 𝐿𝑠 → 𝑆𝑚 𝐿𝑠 → 𝑆𝑚 𝐿. In implementation this is four batched GEMMs with shapes (𝐾, 𝐿, 𝐿𝑠 ), (𝑀 H , 𝐿𝑠 , 𝐿𝑠 ), (𝐾, 𝐿, 𝐿𝑠 ) (restriction, coupling, prolongation; two directional bridge passes ACM Trans. Graph., Vol. 45, No. 6, Article 1. Publication date: December 2026.
Grid and ordering. Each frame targets 𝑁 degrees of freedom on a 2D structured grid of dimensions 𝑊 ×𝐻 with 𝑊 ·𝐻 ≥ 𝑁 and 𝑊 , 𝐻 chosen as close to square as possible. All 𝑊 ·𝐻 cells are sorted by Morton (Z-order) code and truncated to the first 𝑁 cells, giving a contiguous Morton-ordered index set. Edges and density values follow the same ordering. We use a structured grid (rather than a particle cloud) so that the assembled 𝐴 has the canonical 5-point
Supplementary Material: Hierarchical Transformer Preconditioning for Interactive Physics Simulation
Laplacian sparsity pattern, isolating the effect of the heterogeneous coefficient field on conditioning.
•
1:3
distribution as ours; the row reports mean±std across the same per-scale test-frame evaluation as the other methods.
Density field. Per frame the heavy density 𝜌 heavy is drawn log6 Generalization Robustness uniformly from [5, 100] with 𝜌 light =1. We then sample 𝑛𝑏 ∼ Uniform{1, 2, 3} Table 6 in the main paper reports four cells of a wider train × eval barriers; each barrier is independently grid that probes within-family robustness at 𝑁 =4 096. Table 2 re(1) assigned an orientation (vertical or horizontal, equiprobable); ports the full grid for our method, adding both the in-distribution (2) given a center coordinate drawn uniformly from [0.2, 0.8] row (Full training, evaluated on closed-only and high-contrast subalong the cross axis, in normalized grid coordinates; sets) and a fourth, compositional training distribution that excludes (3) given a thickness drawn uniformly from [0.05, 0.20]; closed barriers and restricts contrast to 𝜌 heavy ∈ [5, 25]. (4) given a gap topology drawn uniformly from {top, bottom, middle_hole, closed}. A cell is heavy iff it lies in the heavy region of any barrier, so multiple barriers can overlap to form cross- or L-shaped inclusions. Each cell’s density is finally multiplied by 1+N (0, 0.052 ) for fine-grain symmetry breaking. Operator assembly. 𝐴 is the standard 5-point pressure-Poisson Laplacian with harmonic-mean face conductance 𝑤𝑖 𝑗 =2𝜌𝑖 𝜌 𝑗 /(𝜌𝑖 + 𝜌 𝑗 ), evaluated on cell-cell faces. Boundary cells use zero-flux (Neumann) boundary conditions, as is standard for incompressible pressureprojection in graphics-grade fluid simulators. We do not apply a zero-mean constraint to the resulting null-space mode; CG handles it implicitly when the right-hand side is consistent. Right-hand side. For each frame we draw b ∼ N (0, 𝐼 ) projected onto the orthogonal complement of the constant vector (1⊤ b =0 to ensure consistency with the Neumann null-space). This is a worstcase right-hand-side distribution from a preconditioner’s perspective: it weights all spatial frequencies equally, including the low ones that are hardest for local preconditioners. Per-scale dataset sizes. For each of five scales, spanning 𝑁 = 1 024 to 𝑁 = 16 384, we generate 100 training frames and 20 test frames with disjoint seeds. Training and test split the same generator distribution.
5
Per-frame Iteration-count Variability
The consolidated per-scale times and mean iteration counts have moved to the main paper, Table 4. This section reports the acrossframe variability (±1𝜎 over the same 20 test frames per scale) that the main table omits to fit the column budget. Table 1. Per-frame mean±std PCG iteration counts on the multiphase Poisson benchmark (rtol = 10 −8 , 20 frames per scale). Bold row marks the proposed method. Method
𝑁 =1 024 𝑁 =2 048 𝑁 =4 096
Unprec. CG (GPU) 497±144 Jacobi (GPU) 325±83 IC / DILU (GPU) 11±1 AMG / AMGX (GPU) 1±0 Neural SPAI (GPU) [Yang et al. 2025] 118±6 IC (CPU) 59±18 AMG (CPU) 5±6 Ours (GPU) 47±9
𝑁 =8 192
𝑁 =16 384
650±240 1153±434 1765±1744 2103±1113 429±112 839±476 968±261 1543±1137 15±1 22±4 30±5 40±4 1±0 1±1 2±0 3±0 167±8 246±41 338±24 496±107 96±23 113±31 170±47 254±61 9±6 5±7 7±8 10±8 66±16 80±24 168±26 394±68
The neural SPAI baseline is re-trained per scale (one network per 𝑁 ) using the SAI loss of [Yang et al. 2025] on the same training
Table 2. Within-family generalization grid for the proposed method at 𝑁 =4 096 (rtol =10 −8 , 20 frames per cell). Each cell reports mean±std PCG iterations; dashes mark cells that are in-distribution for the corresponding training set (so the entry would duplicate the leftmost column). The fourth training row is a compositional distribution that excludes closed barriers and restricts contrast to 𝜌 heavy ∈ [5, 25]; the rightmost column evaluates it on the joint shift (closed barriers + high contrast). Eval distribution Train distribution
Full (id)
Closed only
High contrast
Closed + high
Full Excl. closed barriers Low contrast Excl. closed + low 𝜌
82±28 78±19 78±26 95±29
— 68±16 — 89±20
— — 142±126 153±97
— — — 147±102
Mean total solve (ms) ΔOOD (iters)
8.2–9.5 1.00 (def.)
7.4–9.0 0.87–0.94
13.1–13.9 1.82, 1.61
13.5 1.55
ΔOOD is the ratio of iterations in the cell to the in-distribution cell of the same training row (lower is better). Row 1 (Full) is included so the in-distribution operating point is on the same page as the OOD numbers; we omit its trivial OOD columns. The grid shows that topology shift is inexpensive (ΔOOD =0.87– 0.94), while contrast shift is the dominant OOD cost (ΔOOD =1.55– 1.82). Even in the compositional shift cell (closed + high contrast), the model remains at 1.9× Jacobi-relative speedup.
7
Precision and Hardware Efficiency
This section expands the precision and apply-path arguments referenced from main paper §5 and §7.1. Mixed precision split and its motivation. The apply path consists of dense batched GEMMs of shapes (𝐾, 𝐿, 𝐿), (𝑀 H , 𝐿𝑠 , 𝐿𝑠 ), and (𝐾, 𝐿, 𝐿𝑠 ) (main paper §5). At 𝐿 =128, 𝐿𝑠 =32 each individual matmul has condition number bounded by the singular-value spread of its factor (typically ≤ 103 in our trained models), well within float32’s ∼ 107 representable range. PCG’s scalar accumulators, by contrast, run over the whole residual vector and accumulate 𝑂 (𝑁 ) terms per iteration; on stiff multiphase systems we observed residual norms drift by 1–2 orders of magnitude over a few hundred iterations when those accumulators were held at float32, which can cause spurious early termination at rtol =10−8 . Holding only the scalar accumulators (dot products, residual norms, step sizes) at float64 eliminates that drift without paying float64’s 2× memory/throughput penalty on the bulk GEMMs. This is the same precision split used in CG implementations in cuSPARSE and AMGX. ACM Trans. Graph., Vol. 45, No. 6, Article 1. Publication date: December 2026.
1:4
•
Osborne et al.
Factor-tensor memory footprint. The packed factor tensor of width 𝑃 =𝐾𝐿 2 + 𝑀 H 𝐿𝑠2 + 2𝑁 𝐿𝑠 + 𝑁 (main paper §4.4) determines the applypath memory budget. With 𝐿 =128, 𝐿𝑠 =32, and 𝑀 H =𝐾 −1 at 𝜂 =1 weak admissibility, Table 3 lists the absolute numbers at the scales in Table 1.
Table 3. Factor-tensor breakdown for the default configuration (𝐿 =128, 𝐿𝑠 = 32, 𝑝 off =4). 𝐾 =𝑁 /𝐿 is the leaf count; 𝑀H =𝐾 −1 is the unique off-diagonal tile count under weak admissibility (𝜂 =1); 𝑃 is the total packed width in float32 elements; “MB” is the resulting on-device footprint. 𝑁
𝐾
𝑀H
𝐾𝐿 2
𝑃
MB
1 024 2 048 4 096 8 192 16 384
8 16 32 64 128
7 15 31 63 127
0.131 M 0.262 M 0.524 M 1.049 M 2.097 M
0.20 M 0.41 M 0.82 M 1.65 M 3.29 M
0.82 1.64 3.29 6.58 13.2
(a) H-matrix partition.
(b) Highway channels.
Fig. 1. Left: weak-admissibility H-matrix partition. Dense diagonal leaves along the main diagonal; admissible off-diagonal tiles double in size with separation. Right: per-layer highway channels. One red overlay marks the row- and column-index sets of a representative off-diagonal tile; this tile contributes scatter-adds to those strips and to a single global token — four communication channels of dimensions 2D / 1D / 1D / 0D per layer.
The factor tensor scales linearly in 𝑁 at fixed 𝐿 and stays below 10 MB through 𝑁 =8 192; at 𝑁 =16 384 it reaches 13.2 MB and still fits comfortably in a single H200’s HBM. The dominant term at every scale is 𝐾𝐿 2 (the dense diagonal leaves), not 𝑀 H 𝐿𝑠2 : distant tiles are compressed so aggressively that the entire off-diagonal contribution costs less than the diagonal even though it covers 1 − 1/𝐾 of the matrix. Apply-path kernel breakdown. Main paper Table 3 reports the kernel-family split at 𝑁 =8 192: CUTLASS Tensor-Op GEMMs (31%) plus CUTLASS fused attention (14%) account for ∼ 45% of device time; the remaining ∼ 55% is elementwise operations, layout reshapes, and small bookkeeping kernels. What would have to change for tighter tolerance. The float32 apply path supports rtol ≤ 10−8 across all scales we evaluate. For graphics applications this is two to three orders of magnitude tighter than typical pressure-projection tolerances and is therefore sufficient. For scientific applications requiring rtol ≤ 10−10 or stricter, the apply path would either need a full float64 pass (doubling device-memory and halving Tensor-Core throughput) or an outer iterative refinement loop that uses the float32 apply as an inner preconditioner. We did not implement either since they fall outside the regime we target.
8 Additional Figures and Tables Referenced from the Main Paper This section provides optional figures and detailed-breakdown tables for readers who want additional clarity on the method and results. Main-paper cross-references resolve to the Figure/Table number shown next to each caption below. ACM Trans. Graph., Vol. 45, No. 6, Article 1. Publication date: December 2026.
Fig. 2. The sparse operator 𝐴, its dense inverse 𝐴 −1 , and the assembled learned 𝑀 ≈ 𝐴 −1 (top-left 256 × 256) on a multiphase pressure-Poisson frame, on a shared rank-normalized color scale. All three share the weakadmissibility block-and-tile pattern: full-rank diagonal blocks plus offdiagonal tiles whose magnitude decays with separation. This pattern is exactly the prior quantified by main paper Fig. 3.
Fig. 3. A representative frame from our multiphase pressure-Poisson benchmark at 𝑁 =16 384. Left: heterogeneous density 𝜌 with 27× contrast, separated by closed and partially-closed barriers (yellow). Middle: buoyancydriven right-hand side 𝑏 = −𝜕𝑦 𝜌, concentrated on horizontal density interfaces. Right: recovered pressure 𝑝 with velocity −∇𝑝/𝜌 overlaid.
Supplementary Material: Hierarchical Transformer Preconditioning for Interactive Physics Simulation
Table 4. Hyperparameter ablations on multiphase Poisson (rtol = 10 −8 , 20 frames per scale, neural preconditioner only). Train = wall-clock to convergence; Infer. = preconditioner forward pass; Iters = mean PCG iters; Total = inference + PCG solve; Δ solve time = fractional change vs. the “—” row. Total (ms)
Δ solve time
Width (avg. 𝑁 ∈ {2048, 4096, 8192} ): 𝑑=64, 𝑛𝑙 =3, hw 9.5 3.0 147 𝑑=128, 𝑛𝑙 =3, hw 10.7 3.1 105 𝑑=256, 𝑛𝑙 =3, hw 12.3 3.2 191
15.4 12.0 19.3
+28% — +61%
Depth (𝑁 =8192): 𝑑=128, 𝑛𝑙 =1, hw 𝑑=128, 𝑛𝑙 =2, hw 𝑑=128, 𝑛𝑙 =3, hw
6.1 9.6 16.2
1.3 2.3 3.4
421 218 168
36.9 20.8 17.9
+106% +16% —
Highways (𝑁 =2048): 𝑑=128, 𝑛𝑙 =3, hw 𝑑=128, 𝑛𝑙 =3, no-hw
6.0 4.8
3.3 2.2
66 149
8.8 13.8
— +58%
Configuration
Train Infer. (min) (ms)
Iters
•
1:5
Table 5. Neural-method runtime components by scale: preconditioner inference and PCG solve time (ms), plus total time and iterations aligned with the main-paper performance table. Neural SPAI uses the re-trained-per-scale CUDA apply path of [Yang et al. 2025]. Scale 𝑁
Total (ms)
Iters
Neural SPAI (GPU, CUDA apply) [Yang et al. 2025] 1 024 4.6 13.8 18.4 2 048 3.8 22.2 26.0 4 096 6.6 31.7 38.3 8 192 4.4 43.8 48.1 16 384 4.9 66.0 70.9
Infer. (ms)
118 167 246 338 496
Ours (GPU) 1 024 2 048 4 096 8 192 16 384
47 66 80 168 394
1.9 1.9 2.2 4.0 17.0
Solve (ms)
5.1 6.9 7.0 13.9 30.6
7.0 8.8 9.2 17.9 47.6
References Thomas N. Kipf and Max Welling. 2017. Semi-Supervised Classification with Graph Convolutional Networks. In International Conference on Learning Representations (ICLR). Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. 2021. Swin Transformer: Hierarchical Vision Transformer using Shifted Windows. arXiv preprint arXiv:2103.14030 (2021). https://arxiv.org/abs/2103.14030 Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Łukasz Kaiser, and Illia Polosukhin. 2017. Attention Is All You Need. In Advances in Neural Information Processing Systems 30 (NIPS 2017). 5998–6008. Zherui Yang, Zhehao Li, Kangbo Lyu, Yixuan Li, Tao Du, and Ligang Liu. 2025. Learning Sparse Approximate Inverse Preconditioners for Conjugate Gradient Solvers on GPUs. arXiv preprint arXiv:2510.27517 (2025). https://arxiv.org/abs/2510.27517
ACM Trans. Graph., Vol. 45, No. 6, Article 1. Publication date: December 2026.