Does Transolver really need a Transformer?
Organizations: Seminar for Applied Mathematics, ETH Zurich, Switzerland
Abstract
The widely used Transolver family of neural operators is based on physics-attention, which softly assigns the points of an unstructured mesh to a small number of slices, applies self-attention among the resulting tokens, and broadcasts the result back to the points. We provide a comprehensive empirical and theoretical analysis to elucidate the mechanisms which are responsible for model performance. To this end, we perform careful ablations on a challenging suite of nine 3D fluid dynamics benchmarks to find that replacing token attention with a constant linear map does not affect the accuracy. Thus, Transolver does not need a Transformer at all. However, removing the global mixing (slicing/deslicing) or doing it only once leads to performance collapse. We leverage the theory of averaging neural operators to explain and corroborate our findings by showing that just slicing/deslicing, in conjunction with pointwise MLPs, already suffices for universal approximation of continuous operators and attention is redundant in this context. Finally, we provide a novel FlashAttention-style efficient implementation of the key slicing/deslicing module of Transolver. This flashslice kernel streams slice and deslice over the points without materializing the heavy slice-weight tensor, while reproducing the best available implementation to floating point error. At the same time, it leads to very significant memory and compute savings, particularly at large slice counts.
Figures & tables
| Dataset | Baseline | NoTokenAttn | MlpOnly | Untied | FrozenSlice | SliceOnce |
|---|---|---|---|---|---|---|
| Taylor–Green | ||||||
| SHIFT-Wing | ||||||
| surface | ||||||
| volume | ||||||
| DrivAerNet++ | ||||||
| surface | ||||||
| time (ms) | peak memory (GB) | |||||||
| eager | ours | speedup | eager | ours | saved | |||
| full model, , single-tile kernels | ||||||||
| training | fp32 | k | ||||||
| training | bf16 | k | ||||||
| training | bf16 | M | ||||||
| inference | fp32 | M | ||||||
Appendix figures & tables34 assets
Supplementary material from the paper’s appendix.
Appendix
| variant | nonlocal budget | hard ceiling? | measured cost |
|---|---|---|---|
| Baseline / Untied | no (Prop. 1 ) | reference | |
| NoTokenAttn | no (Prop. 1 ) | to | |
| FrozenSlice | no (§ A.4 ) | to | |
| SliceOnce | yes, rank (Prop. 2 ) | – | |
| MlpOnly | yes, locality (Cor. 1 ) | – |
| Benchmark | steps | train/val/test | points/sample | subsample | in out ch. | inputs | outputs |
|---|---|---|---|---|---|---|---|
| Taylor–Green | 250k | 1.55M/19.3k/38.6k | state at (+ lead time) | state at | |||
| SHIFT-Wing surface | 250k | 8377/1000/1000 | 1.6M | normals, freestream (varying) | , | ||
| SHIFT-Wing volume | 140k ∗ | 8377/1000/1000 | 5.85M | SDF, freestream (varying) | , | ||
| DrivAerNet++ surface | 184k | 5825/1148/1154 | 459k | normals, freestream (const.) | , | ||
| DrivAerNet++ volume | 180k | 5825/1148/1154 | 11.0M | SDF, freestream (const.) | , | ||
| SHIFT-SUV surface | 50k | 798/100/100 | 2.57M | normals only | , |
| Benchmark | / | max rel. std | |
|---|---|---|---|
| DrivAerNet++ surface | |||
| DrivAerNet++ volume | |||
| DrivAerML surface | |||
| DrivAerML volume | |||
| SHIFT-Wing surface | |||
| SHIFT-Wing volume |
| Dataset | Variable | baseline | no-token-attn | MLP-only | untied | frozen | slice-once |
|---|---|---|---|---|---|---|---|
| Taylor–Green | |||||||
| DrivAerNet++ surf. | |||||||
| DrivAerNet++ vol. | |||||||
| Dataset | Variable | baseline | no-token-attn | MLP-only | untied | frozen | slice-once |
|---|---|---|---|---|---|---|---|
| Taylor–Green | |||||||
| DrivAerNet++ surf. | |||||||
| DrivAerNet++ vol. | |||||||
| step | 1k | 3k | 5k | 7k | 9k | 11k |
|---|---|---|---|---|---|---|
| ratio |
| level | logits dot | value dots | inputs | gate |
|---|---|---|---|---|
| ieee (default) | fp32 FMA | fp32 FMA | fp32 or bf16 | eager fp32 |
| tf32 | fp32 FMA | tf32 tensor cores | fp32 or bf16 | |
| bf16v | fp32 FMA | bf16 tensor cores | bf16 | eager bf16 autocast |
| bf16 | bf16 tensor cores | bf16 tensor cores | bf16 | eager bf16 autocast |
| tf32x3 | three tf32 products | three tf32 products | fp32 | ; measured – eager fp32 |
| inference: ms (GB) | training step: ms (GB) | |||
| eager | ours | eager | ours | |
| fp32 | ||||
| k | ( ) | ( ) | ( ) | ( ) |
| M | ( ) | ( ) | OOM | OOM |
| M | ( ) | ( ) | OOM | OOM |
| M | ( ) | ( ) | OOM | OOM |
| fp32, peak GB | bf16, peak GB | bf16 step | |||
| eager | ours | eager | ours | speedup | |
| slice count (depth ) | |||||
| fp32, k | bf16, k | bf16, M | |
|---|---|---|---|
| eager | |||
| eager compile | |||
| fused (exact fp32 dots) | |||
| fused compile fast dot | |||
| vs. strongest baseline |
| bf16, M | bf16, M | fp32, M | fp32, M | |
|---|---|---|---|---|
| ckpt | ||||
| ckpt compile | ||||
| fused ckpt compile fast dot | ||||
| vs. ckpt compile |
| inputs / dots | eager (ms) | ours (ms) | speedup | eager (GB) | ours (GB) | saved | ||
|---|---|---|---|---|---|---|---|---|
| bf16 / bf16 | k | |||||||
| bf16 / bf16 | M | |||||||
| bf16 / bf16 | k | |||||||
| bf16 / bf16 | M | OOM | — | — | — | |||
| bf16 / bf16 | k | |||||||
| bf16 / bf16, inference | M |
| step (ms) | Mpts/s | GB/GPU | |
|---|---|---|---|
| ckpt sharding | |||
| ckpt sharding compile | |||
| fused ckpt compile bf16 dot |
Neural Operators
Neural operators learn mappings between functions, enabling fast surrogates for PDEs and other continuous physical systems.
See all 304 papers