Hierarchical Transformer Preconditioning for Interactive Physics Simulation

TL;DR

This paper introduces a Hierarchical Transformer Preconditioner that improves real-time physics simulation by efficiently capturing long-range couplings with O(N) scaling, leading to faster convergence and better performance.

Contribution

It presents a novel neural preconditioner based on hierarchical transformers and a cosine-Hutchinson training objective, enabling efficient, scalable, and accurate approximate-inverse computations.

Findings

Achieves up to 2.7x speedup over neural SPAI on benchmarks.

Runs at 17.9 ms/frame for N=8,192, outperforming traditional methods.

Effectively captures long-range couplings in multiphase Poisson systems.

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.