Hierarchical Transformer Preconditioning for Interactive Physics Simulation

Abstract: Neural preconditioners for real-time physics simulation offer promising data-driven 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 O(N) 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 MA on convergence-critical spectral subspaces, optimizing angular alignment of MAz 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, N = 1,024-16,384), the solver runs from ~143 to ~21 fps. At N = 8,192, it reaches 17.9 ms/frame, with 2.2x speedup over GPU Jacobi, ~28x over GPU IC/DILU (AMGX multicolor_dilu), and 2.7x over neural SPAI retrained per scale on the same benchmark.

• The paper introduces a hierarchical transformer preconditioner that integrates H-matrix partitioning and cosine-Hutchinson loss, achieving a 68× reduction in condition number.
• The method attains consistent Θ(N) scaling in both training and inference, enabling real-time interactive physics simulations with significantly fewer PCG iterations.
• The architecture employs highway connections and fixed-shape GEMMs to maximize GPU efficiency, outperforming traditional preconditioners such as Jacobi and AMG.

Hierarchical Transformer Preconditioning for Interactive Physics Simulation

Introduction

This paper introduces a Hierarchical Transformer Preconditioner (HTP) tailored for high-performance, interactive physics simulation tasks. The central objective is to overcome limitations in existing neural preconditioners, especially regarding their efficiency and capability to capture long-range couplings in large, frame-varying, sparse SPD systems commonly encountered in fluid dynamics and multiphysics simulations. The proposed solution integrates a Transformer-based architecture with a weak-admissibility $\mathcal{H}$-matrix partitioning scheme, coupled with a cosine-Hutchinson probe loss, enabling consistent $\Theta(N)$ scaling in both training and inference. This results in a preconditioning operator that can be applied effectively and efficiently in real-time settings, without incurring the high setup costs or memory overhead seen in classical approaches such as Incomplete Cholesky (IC), Algebraic Multigrid (AMG), and earlier neural alternatives.

Hierarchical Transformer Architecture and $\mathcal{H}$-Matrix Partitioning

HTP leverages a multiscale $\mathcal{H}$-matrix structure, where the degrees of freedom are partitioned into contiguous leaves, producing dense diagonal blocks and hierarchically coarsening off-diagonal tiles. The diagonal blocks are processed by a high-capacity attention stack emitting full-rank PSD factors per leaf. For off-diagonal blocks, the model emits a fixed number of coarse tokens per tile, allowing rank to naturally decay with spatial separation—thereby directly exploiting the compressibility of long-range interactions. This layout enables the preconditioner to be represented with $O(N)$ complexity at inference by ensuring block sizes and computation per tile are fixed, regardless of $N$.

Figure 1: The fixed architecture-provided rank fraction for off-diagonal blocks remains above the numerically required threshold for effective approximate inversion across all distance classes.

To facilitate information flow across the partition without sacrificing scaling properties, the architecture introduces highway connections—axial row/column buffers and a global summary token—that integrate context at multiple hierarchical levels. These design choices enable the transformer to propagate information across non-local regions while preserving linear complexity.

Cosine-Hutchinson Loss: Spectral Clustering via Angular Alignment

A pivotal contribution of the paper is the adoption of a cosine-Hutchinson probe loss for training the preconditioner. Rather than minimizing pointwise distances between $MA\mathbf{z}$ and $\mathbf{z}$ (as in SAI/Frobenius-objective losses), the cosine loss targets the angular alignment between these vectors, matching the scale invariance required for optimal PCG convergence. This strategy allows eigenvalue clusters in $MA$ to form wherever angular alignment is optimal, rather than forcing them to a prescribed location (e.g., near $\|A\|$), thus avoiding unnecessary spectral placement constraints that do not impact convergence.

Figure 2: Spectra of $\Theta(N)$0 for the multiphase Poisson system; the cosine-Hutchinson trained model achieves substantially tighter spectral clustering, unconstrained by target eigenvalue location, leading to a $\Theta(N)$1 reduction in condition number.

Empirical analysis demonstrates that, under the same architecture and data distribution, the cosine loss produces markedly better spectral compression than the SAI loss, with a $\Theta(N)$2 improvement in condition number reduction attributable solely to the objective function. This is directly reflected in dramatically reduced PCG iterations and solve times.

Computational Implementation: Dense GEMMs and CUDA Graphs

By structuring the entire PCG inner loop as a sequence of batched GEMMs with fixed shapes—enabled by the static $\Theta(N)$3-matrix partition—HTP allows the entire solve loop to be captured within a single CUDA Graph. This eliminates kernel launch overhead, enables efficient use of GPU compute resources (favoring compute-bound over memory-bound workloads), and further distinguishes the method from alternatives like neural SPAI and IC, whose sparse or triangular solves induce scattered memory access and limited parallelism.

Numerical and Performance Results

Extensive benchmarks on 2D multiphase Poisson systems (with random per-frame variation in topology and coefficient contrast) confirm the efficacy of HTP in comparison to both classical and learned preconditioners. Across a range of problem sizes ($\Theta(N)$4–$\Theta(N)$5), HTP consistently achieves interactive rates ($\Theta(N)$6–$\Theta(N)$7 fps), with iteration counts improving over GPU Jacobi by factors of $\Theta(N)$8–$\Theta(N)$9 and outpacing retrained neural SPAI baselines by up to $\mathcal{H}$0 in wall-clock time, surpassing even those classical methods that require heavy setup amortization.

Figure 3: End-to-end solve times demonstrate consistent interactive rates and significant speedup over Jacobi and neural SPAI baselines.

Figure 4: Visualizations of residual propagation highlight the rapid, multiscale error correction enabled by HTP, comparable to AMG but at substantially lower per-iteration cost.

Training dynamics confirm that the cosine-Hutchinson loss offers stable gradient flow and closely tracks PCG convergence, while SAI-based models plateau prematurely. Ablations make clear that both the architectural adaptations (highway connections, block sizes, model width) and the loss formulation are critical to performance.

Figure 5: Training profile reveals that the cosine-Hutchinson loss maintains improvement in PCG iteration count after SAI loss stagnates.

Generalization and Robustness

Generalization tests indicate that HTP provides robust performance across significant OOD shifts, including unseen barrier topologies and increased coefficient contrast. Performance remains competitive, as the $\mathcal{H}$1-matrix partition is agnostic to barrier geometry and the loss’s scale invariance absorbs amplitude variations.

Theoretical and Practical Implications

Theoretical implications center on demonstrating that preconditioners operating in the quotient space of linear operators modulo positive scalar can optimize subspace clustering rather than pointwise action, yielding superior spectral properties for conjugate gradient solvers. Practically, this enables deployment in time-critical simulation pipelines where matrix structure evolves every frame and setup amortization is infeasible.

HTP’s approach can be extended beyond 2D pressure-Poisson systems to 3D, soft-body, contact dynamics, and other SPD systems provided there exists a spatial ordering admitting low-rank off-diagonal compressibility. However, performance gains diminish when applied to operators lacking this geometric structure or when $\mathcal{H}$2 is reused sufficiently for classical setup amortization to take effect.

Future Directions

Open avenues include:

• Incorporating dynamic partitioning to allow for variable problem sizes without retraining (moving from fixed leaf size to fixed leaf count).
• Extending the structural priors and loss strategies to indefinite or non-SPD systems.
• Unifying this framework with physics-informed or differentiable simulation environments for end-to-end-optimized solvers.

Conclusion

HTP provides a Transformer-based, $\mathcal{H}$3-matrix-anchored, strictly $\mathcal{H}$4, and hardware-tailored preconditioner for SPD systems in interactive simulation. Its architectural innovations, coupled with the scale-invariant cosine-Hutchinson loss, enable superior spectral clustering, dramatically reduced solver iteration counts, and consistent real-time performance. The method’s structural and theoretical adaptability positions it as a robust alternative for next-generation neural solvers in scientific computing.