TB-Transolver: Parallelized and Coherent Long-Horizon Rollouts via Temporal Bundling for Time-Dependent PDEs
Organizations: Aerospace Engineering Department, Worcester Polytechnic Institute, Worcester, MA 01609, USA.
Abstract
Transolver has recently gained widespread adoption for solving diverse physical problems on unstructured geometric meshes by leveraging its Physics-Attention mechanism, which seamlessly extends standard softmax attention to handle physical domains at scale. For time-dependent problems, existing literature relies on time-conditioned or autoregressive rollouts to preserve temporal causality while retaining Physics-Attention as the spatial backbone. However, a fundamental limitation of this formulation is its sequential execution model, wherein computational costs scale linearly with the temporal horizon, rendering standard Transolver inefficient for transient problems despite its exceptional spatial efficacy. This bottleneck has historically hindered its broader application in time-dependent physical modeling. To address this limitation, we introduce TB-Transolver (also referred to as Temporally Bundled Transolver), an efficient framework built on temporal bundling that scales Transolver for time-dependent PDEs. Furthermore, we propose spatiotemporal coherence () as a general metric to systematically quantify snapshot correlation within a temporal bundle, enabling the a priori selection of optimal bundle widths using training data alone, thereby avoiding expensive grid searches. We demonstrate the computational efficacy of TB-Transolver across four physically distinct domains, namely, plasticity, fluid dynamics, reaction-diffusion, and global weather forecasting; achieving training speedups of up to over conventional autoregressive and time-conditioned baselines. Finally, these computational gains are further amplified through distributed multi-node and multi-GPU parallelism, yielding aggregated speedups of up to at optimal bundle widths for systems exhibiting high spatiotemporal coherence compared to single-GPU implementations.
Figures & tables
| Trend | Higher accuracy | Higher efficiency | Source of Error |
|---|---|---|---|
| Flat (short/long rollouts) | when | when | |
| Decays (short rollouts) | when | ||
| Decays (long rollouts) | at optimal ( ) |
| Model | Plasticity | Navier–Stokes |
|---|---|---|
| ( ) | ||
| FNO [ 13 ] | – | 0.1556 |
| WMT [ 6 ] | 0.0076 | 0.1541 |
| U-FNO [ 30 ] | 0.0039 | 0.2231 |
| geo-FNO [ 15 ] | 0.0074 | 0.1556 |
| U-NO [ 25 ] | 0.0034 | 0.1713 |
| Bundle Width ( ) | Mean Error | Sec/Epoch | Total Time (500 Ep.) | Peak VRAM | Speed-up | |
|---|---|---|---|---|---|---|
| 1 | 122.00 | 17.0 h | 1.98 GB | 1.0 | ||
| 2 | 59.31 | 8.2 h | 1.99 GB | 0.894 | ||
| 4 | 30.17 | 4.2 h | 1.98 GB | 0.761 | ||
| 5 | 22.50 | 3.1 h | 1.98 GB | 0.709 | ||
| 10 | 6.18 | 52 min | 1.99 GB | 0.499 | ||
| 20 | 4.15 | 35 min | 1.98 GB | 0.135 |
| Bundle Width ( ) | Mean Error | Sec/Epoch | Total Time (500 Ep.) | Peak VRAM | Speed-up | |
|---|---|---|---|---|---|---|
| 1 | 79.7 | 12.0 h | 7.65 GB | 1.0 | ||
| 2 | 47.4 | 6.5 h | 3.94 GB | 0.935 | ||
| 5 | 31.6 | 4.2 h | 1.72 GB | 0.696 | ||
| 10 | 23.4 | 2.7 h | 0.98 GB | 0.430 |
| Bundle Width ( ) | Mean Error | Sec/Epoch | Total Time (500 Ep.) | Peak VRAM | Speed-up | |
|---|---|---|---|---|---|---|
| 1 | 216.00 | 30.0 h | 29.83 GB | 1.0 | ||
| 2 | 109.15 | 15.2 h | 15.05 GB | 0.999 | ||
| 5 | 44.20 | 6.1 h | 6.18 GB | 0.998 | ||
| 10 | 22.60 | 3.1 h | 3.22 GB | 0.994 |
| Bundle Width ( ) | Mean Error | Sec/Epoch | Total Time (500 Ep.) | Peak VRAM | Speed-up | |
| 1 | ———— Out Of Memory (OOM) ———— | 94.86 GB* | 1.0 | — | ||
| 2 | 4.22 | 35.2 min | 54.72 GB | 0.993 | ||
| 4 | 2.12 | 17.7 min | 27.53 GB | 0.954 | ||
| 8 | 1.07 | 8.9 min | 13.93 GB | 0.994 | ||
| *Maximum available GPU memory | ||||||
| Bundle Width ( ) | Mean Error | Sec/Epoch | Total Time (500 Ep.) | Peak VRAM | Speed-up | |
|---|---|---|---|---|---|---|
| 1 | 243.83 | 33.9 h | 22.49 GB | 1.0 | ||
| 2 | 123.60 | 17.2 h | 11.37 GB | 0.984 | ||
| 5 | 54.17 | 7.5 h | 4.69 GB | 0.819 | ||
| 10 | 38.43 | 5.3 h | 2.46 GB | 0.505 | ||
| 15 | 33.20 | 4.6 h | 1.72 GB | 0.665 | ||
| 30 | 20.16 | 2.8 h | 0.98 GB | 0.454 |
| Benchmark | Optimal Bundle | 1-GPU | 2-GPU DDP | Speed-up | Scaling Efficiency |
|---|---|---|---|---|---|
| Width ( ) | (Sec/Epoch) | (Sec/Epoch) | |||
| Plasticity | 2 | 62.89 | 32.58 | 1.93 | 96.5 |
| Low Viscosity | 2 | 47.40 | 25.09 | 1.89 | 94.5 |
| High Viscosity | 5 | 54.17 | 28.23 | 1.92 | 96.0 |
| Reaction-Diffusion | 10 | 22.60 | 14.23 | 1.59 | 79.5 |
| SST | 8 | 1.07 | 0.57 | 1.88 | 94.0 |
| Benchmark | Speed-up | ||
|---|---|---|---|
| Plasticity | 2 | 0.894 | |
| Navier-Stokes (Low Viscosity) | 2 | 0.935 | |
| Navier-Stokes (High Viscosity) | 5 | 0.819 | |
| Reaction-Diffusion | 10 | 0.994 | |
| Sea Surface Temperature (SST) | 8 | 0.994 |
| Benchmark | Source of error dominated by | ||
|---|---|---|---|
| Plasticity | Decays | - | ✓ |
| Navier-Stokes (Low Viscosity) | Decays | ✗ | ✓ |
| Navier-Stokes (High Viscosity) | Decays and Oscillates | ✓ | ✗ |
| Reaction-Diffusion | Flat | ✓ | - |
| Sea Surface Temperature (SST) | Flat | ✓ | - |
| Benchmark | Computational efficiency dominated by | ||
|---|---|---|---|
| Temporal bundling | GPU scaling | ||
| Plasticity | Decays | ✗ | ✓ |
| Navier-Stokes (Low Viscosity) | Decays | ✗ | ✓ |
| Navier-Stokes (High Viscosity) | Decays and Oscillates | ✓ | ✓ |
| Reaction-Diffusion | Flat | ✓ | ✗ |
| Sea Surface Temperature (SST) | Flat | ✓ | ✗ |
Appendix figures & tables2 assets
Supplementary material from the paper’s appendix.
Appendix
| Parameter | Plasticity | NS (Low ) | Reaction–Diffusion | SST | NS (High ) |
|---|---|---|---|---|---|
| Spatial domain | 2D structured | 2D structured | 2D structured | 2D irregular | 2D structured |
| Hidden dimension | 128 | 256 | 256 | 128 | 256 |
| Number of layers | 8 | 8 | 8 | 8 | 8 |
| Number of attention heads | 8 | 8 | 8 | 8 | 8 |
| Number of learned slices | 64 | 32 | 32 | 32 | 32 |
| Unified positional encoding | Disabled | Enabled | Enabled | Disabled | Enabled |
Neural PDE Solvers
Neural PDE solvers use trainable models to approximate solutions or solution operators for partial differential equations, often with physical constraints.
See all 286 papers