TB-Transolver: Parallelized and Coherent Long-Horizon Rollouts via Temporal Bundling for Time-Dependent PDEs

Sidharth S. Menon; Ameya D. Jagtap

cs.LGNov 2, 2026

TB-Transolver: Parallelized and Coherent Long-Horizon Rollouts via Temporal Bundling for Time-Dependent PDEs

Authors: Sidharth S. Menon, Ameya D. Jagtap

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 ( ${\rm C(K)}$ ) 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 $29\times$ 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 $45.6\times$ at optimal bundle widths for systems exhibiting high spatiotemporal coherence compared to single-GPU implementations.

Figures & tables

Panel of Figure 1 — Figure 1: (a) Overview of the spatial physics-aware interactions and temporal parallelism in TB-Transolver. The Physics-Attention mechanism aggregates spatial information into physics-aware slices, which serve as macro-tokens for capturing physical correlations across the spatial domain. Attention is then applied to these tokens to model interactions among the physical slices. (b) Comparison of sequential and parallel discrete-time computations. Sequential computation involves time-conditioned or autoregressive rollouts , whereas temporal bundling enables multiple time steps to be predicted in parallel. (c) In TB-Transolver, autoregressive time stepping corresponds to the special case of temporal bundling with bundle width $\rm K=1$ , whereas $\rm K>1$ enables parallel prediction across multiple time steps. The illustration shows an example with 12 temporal snapshots.

Figure 2: Spatiotemporal coherence ( $C(\rm K)$ ) for dynamical systems: (a) low-viscosity Navier–Stokes equation, $\nu=1\mathrm{e}{-5}$ ; (b) high-viscosity Navier–Stokes equation, $\nu=1\mathrm{e}{-3}$ ; (c) reaction–diffusion equation; (d) plasticity; and (e) sea surface temperature (SST). The curves show the degree of correlation (or decorrelation) as a function of bundle width ( $\rm K$ ). When $C(\rm K)\rightarrow 1$ , the first and last slices of a bundle of width $\rm K$ are highly correlated. Conversely, when $C(\rm K)\rightarrow 0$ , these slices exhibit low correlation, indicating reduced temporal coherence and potentially making temporal bundling less favorable.

$\rm C(K)$ Trend	Higher accuracy	Higher efficiency	Source of Error
Flat (short/long rollouts)	when $\rm K\rightarrow T_{out}$	when $\rm K\rightarrow T_{out}$	$\varepsilon_{\rm accumulation}$
Decays (short rollouts)	when $\rm K\rightarrow 1$		$\varepsilon_{\rm accumulation},~\varepsilon_{\rm decoherence}~(\uparrow)$
Decays (long rollouts)	at optimal $\rm K^{}$ ( $\rm 1<K^{}<T_{out}$ )		$\varepsilon_{\rm accumulation}~(\uparrow),~\varepsilon_{\rm decoherence}$

Table 1: Performance trade-offs based on the structural coherence metric

\rm C(K)

. Note that

(\uparrow)

represents the dominating error for time rollouts.

Model	Plasticity	Navier–Stokes
Model	Plasticity	( $\nu=1\mathrm{e}{-5}$ )
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

Table 2: Comparative performance (Relative

\rm L_{2}

Error) of baseline operator learning models and Transolver across standard Plasticity and Navier-Stokes PDE benchmarks by Wu et al. [ 33 ] . Note that here, bundle width with

\rm K=1

corresponds to Transolver. Also, plasticity is trained by time-conditioning while Navier-Stokes is trained by auto-regressive rollouts.

Figure 3: Performance of TB-Transolver on the Plasticity Benchmark : As bundle width ${\rm K}$ increases toward ${\rm T_{out}}$ , TB-Transolver achieves a speedup of up to $29\times$ while relative $\rm L_{2}$ error increases by $125\%$ compared to time-conditioned rollouts at ${\rm K}=1$ .

Bundle Width ( $\rm K$ )	Mean Error	Sec/Epoch	Total Time (500 Ep.)	Peak VRAM	$\rm C(K)$	Speed-up
1	$0.00120^{[1]}$	122.00	17.0 h	1.98 GB	1.0	$\rm Baseline$
2	$0.00160^{[1]}$	59.31	8.2 h	1.99 GB	0.894	$2\times$
4	$0.00192^{[2]}$	30.17	4.2 h	1.98 GB	0.761	$4\times$
5	$0.00194^{[2]}$	22.50	3.1 h	1.98 GB	0.709	$5\times$
10	$0.00229^{[2]}$	6.18	52 min	1.99 GB	0.499	$20\times$
20	$0.00270^{[1]}$	4.15	35 min	1.98 GB	0.135	$29\times$

Table 3: Plasticity Benchmark (

{\rm T_{out}=20}

) : Accuracy and computational efficiency for varying bundle widths.

{\rm C(K)}

follows the trend illustrated in Figure 2 . The model is trained using time-conditioned rollouts. The model comprises

2.8

million parameters.

Figure 4: Performance of TB-Transolver on the Navier-Stokes Benchmark : For low viscosity (top), as bundle width ${\rm K}$ increases toward ${\rm T_{out}}$ , TB-Transolver achieves a speedup of up to $3.4\times$ while relative $\rm L_{2}$ error increases by $115\%$ compared to autoregressive rollouts at ${\rm K}=1$ . On the other hand, for high viscosity (bottom), as bundle width ${\rm K}$ increases toward ${\rm T_{out}}$ , TB-Transolver achieves a speedup of up to $12\times$ while relative $\rm L_{2}$ error decreases by $99\%$ compared to autoregressive rollouts at ${\rm K}=1$ .

Bundle Width ( $\rm K$ )	Mean Error	Sec/Epoch	Total Time (500 Ep.)	Peak VRAM	$\rm C(K)$	Speed-up
1	$0.09618^{[1]}$	79.7	12.0 h	7.65 GB	1.0	$\rm Baseline$
2	$0.10638^{[1]}$	47.4	6.5 h	3.94 GB	0.935	$1.7\times$
5	$0.16571^{[1]}$	31.6	4.2 h	1.72 GB	0.696	$2.5\times$
10	$0.20670^{[1]}$	23.4	2.7 h	0.98 GB	0.430	$3.4\times$

Table 4: Navier-Stokes, Low Viscosity Benchmark (

\rm T_{out}=10

) : Accuracy and computational efficiency for varying bundle widths.

{\rm C(K)}

follows the trend illustrated in Figure 2 . The model is trained using autoregressive time rollouts. The model comprises

11.2

million parameters.

Figure 5: Performance of TB-Transolver on the Reaction–Diffusion Benchmark : As bundle width ${\rm K}$ increases toward ${\rm T_{out}}$ , TB-Transolver achieves a speedup of up to $10\times$ alongside substantial reduction in relative $\rm L_{2}$ error by $76.9\%$ compared to autoregressive rollouts at ${\rm K}=1$ .

Bundle Width ( $\rm K$ )	Mean Error	Sec/Epoch	Total Time (500 Ep.)	Peak VRAM	$\rm C(K)$	Speed-up
1	$0.00489^{[1]}$	216.00	30.0 h	29.83 GB	1.0	$\rm Baseline$
2	$0.00438^{[1]}$	109.15	15.2 h	15.05 GB	0.999	$2\times$
5	$0.00394^{[2]}$	44.20	6.1 h	6.18 GB	0.998	$5\times$
10	$0.00113^{[2]}$	22.60	3.1 h	3.22 GB	0.994	$10\times$

Table 5: Reaction-Diffusion Benchmark (

\rm T_{\text{out}}=10

) : Accuracy and computational efficiency for varying bundle widths.

{\rm C(K)}

follows the trend illustrated in Figure 2 . The model is trained using autoregressive time rollouts. The model comprises

11.2

million parameters.

Figure 6: Performance of TB-Transolver on the SST Benchmark : As bundle width ${\rm K}$ increases toward ${\rm T_{out}}$ , TB-Transolver achieves a speedup of up to $7.5\times$ alongside reduction in relative $\rm L_{2}$ error by $7.9\%$ compared to time-bundled rollouts at ${\rm K}=2$ . Note that ${\rm K}=1$ exhausts available GPU memory for a batch size of $8$ at float32 precision. The test dataset spans years 2021–2025, with additional evaluation results provided in Appendix B .

Bundle Width ( $\rm K$ )	Mean Error	Sec/Epoch	Total Time (500 Ep.)	Peak VRAM	$\rm C(K)$	Speed-up
1	———— Out Of Memory (OOM) ————			94.86 GB*	1.0	—
2	$0.08000^{[3]}$	4.22	35.2 min	54.72 GB	0.993	$\rm Baseline$
4	$0.07550^{[3]}$	2.12	17.7 min	27.53 GB	0.954	$1.9\times$
8	$0.07370^{[3]}$	1.07	8.9 min	13.93 GB	0.994	$3.9\times$
*Maximum available GPU memory

Table 6: Sea Surface Temperature (SST) Benchmark (

\rm T_{\text{out}}=8

months) : Accuracy and computational efficiency for varying bundle widths.

{\rm C(K)}

follows the trend illustrated in Figure 2 . The model is trained using autoregressive time rollouts. The model comprises 710k parameters.

Bundle Width ( $\rm K$ )	Mean Error	Sec/Epoch	Total Time (500 Ep.)	Peak VRAM	$\rm C(K)$	Speed-up
1	$0.54380$	243.83	33.9 h	22.49 GB	1.0	$\rm Baseline$
2	$0.08000$	123.60	17.2 h	11.37 GB	0.984	$2\times$
5	$0.00530$	54.17	7.5 h	4.69 GB	0.819	$4.5\times$
10	$0.00687$	38.43	5.3 h	2.46 GB	0.505	$6.4\times$
15	$0.00603$	33.20	4.6 h	1.72 GB	0.665	$7.3\times$
30	$0.00714$	20.16	2.8 h	0.98 GB	0.454	$12\times$

Table 7: Navier-Stokes, High-Viscosity Benchmark (

\rm T_{\text{out}}=30

) : Accuracy and computational efficiency for varying bundle widths.

{\rm C(K)}

follows the trend illustrated in Figure 2 . The model is trained using autoregressive time rollouts. The model comprises 11.2 million parameters.

Benchmark	Optimal Bundle	1-GPU	2-GPU DDP	Speed-up	$\%$ Scaling Efficiency
Benchmark	Width ( $\rm K^{*}$ )	(Sec/Epoch)	(Sec/Epoch)	Speed-up	$\%$ Scaling Efficiency
Plasticity	2	62.89	32.58	1.93 $\times$	96.5
Low Viscosity	2	47.40	25.09	1.89 $\times$	94.5
High Viscosity	5	54.17	28.23	1.92 $\times$	96.0
Reaction-Diffusion	10	22.60	14.23	1.59 $\times$	79.5
SST	8	1.07	0.57	1.88 $\times$	94.0

Table 8: Distributed Data Parallel (DDP) Scaling : Scaling efficiency up to 2 GPUs evaluated across diverse physical regimes using optimal temporal bundle widths (

\rm K^{*}

). Near-linear scaling is achieved on compute-heavy domains where Physics-Attention overhead dominates communication latency.

Panel of Figure 7 — Figure 7: Multi-GPU DDP Scaling with Temporal Bundling : (a) Illustration comparing regular time stepping on a single GPU and multi-GPU training with temporal bundling for a 4-GPU configuration. Gradients are synchronized across GPUs using the AllReduce operation during backpropagation, before the optimizer update. (b) Multi-GPU scaling up to 8 GPUs across diverse physical domains using optimal bundle width $\rm K^{*}$ with compute load dominating communication latency.

Benchmark	$\rm K^{*}$	$\rm C(K^{*})$	Speed-up
Plasticity	2	0.894	$2\times$
Navier-Stokes (Low Viscosity)	2	0.935	$1.7\times$
Navier-Stokes (High Viscosity)	5	0.819	$4.5\times$
Reaction-Diffusion	10	0.994	$10\times$
Sea Surface Temperature (SST)	8	0.994	$3.9\times$

Table 9: Optimal bundle widths and their accuracy-efficiency trade-offs : Following the selection criterion in Equation ( 4 ):

\rm K^{*}=\max\left\{K\;:\;C(K)\gtrsim 0.8\right\}

, the selected optimal bundle widths (

\rm K^{*}

) provide computational speed-ups while maintaining favorable predictive accuracy across the benchmarks.

Benchmark	$\rm C(K)$	Source of error dominated by
Benchmark	$\rm C(K)$	$\rm\varepsilon_{accumulation}$	$\rm\varepsilon_{decoherence}$
Plasticity	Decays	-	✓
Navier-Stokes (Low Viscosity)	Decays	✗	✓
Navier-Stokes (High Viscosity)	Decays and Oscillates	✓	✗
Reaction-Diffusion	Flat	✓	-
Sea Surface Temperature (SST)	Flat	✓	-

Table 10: Sources of error in temporal bundling : Error is governed by either accumulation or decoherence for

1<{\rm K}<{\rm T_{out}}

\rm\varepsilon_{decoherence}

can be neglected in regimes where

{\rm C(K)}

remains flat. Plasticity benchmark has no contribution from

\rm\varepsilon_{accumulation}

since it is trained by time-conditioned rollouts. Note that this is a qualitative interpretation of the dominant error mechanisms.

Benchmark	$\rm C(K)$	Computational efficiency dominated by
Benchmark	$\rm C(K)$	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	✓	✗

Table 11: Computational efficiency across domains for temporal bundling : Computational efficiency across distinct physical domains, governed by the combined effects of temporal bundling and multi-GPU scaling.

Appendix figures & tables2 assets

Supplementary material from the paper’s appendix.

Appendix

Parameter	Plasticity	NS (Low $\nu$ )	Reaction–Diffusion	SST	NS (High $\nu$ )
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

Appendix

Table 12: Transolver architecture and training configurations across the plasticity, low- and high-viscosity Navier–Stokes, reaction–diffusion, and sea surface temperature (SST) benchmarks.

Figure 8: Performance of TB-Transolver on the SST Benchmark : As bundle width ${\rm K}$ increases toward ${\rm T_{out}}$ , TB-Transolver achieves a speedup of up to $7.5\times$ alongside reduction in relative $\rm L_{2}$ error by $7.9\%$ compared to time-bundled rollouts at ${\rm K}=2$ . Note that ${\rm K}=1$ exhausts available GPU memory for a batch size of $8$ at float32 precision.

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

Transolver- $σ$ : Joint Spectral-Physical Subspace Modeling for Neural PDE Solving
Sep 29, 2026Haonan Shangguan, Hang Zhou, Haixu Wu +3Fourier Neural Operator Neural PDE Solvers
School of Software, BNRist, Tsinghua University · MIT CSAIL
Transolver-3: Scaling Up Transformer Solvers to Industrial-Scale Geometries
Feb 4, 2026Hang Zhou, Haixu Wu, Haonan Shangguan +4Efficient Neural Network Inference Neural PDE Solvers
School of Software, BNRist, Tsinghua University, China.
When Attention Beats Fourier: Multi-Scale Transformers for PDE Solving on Irregular Domains
May 8, 2026Brandon Yee, Pairie Koh, Jack Rodriguez +1Neural PDE Solvers PDE Solving
Physics Lab, Yee Collins Research Group

Sep 29, 2026J/K move · Enter open · S save

Transolver- $σ$ : Joint Spectral-Physical Subspace Modeling for Neural PDE Solving

Haonan Shangguan, Hang Zhou, Haixu Wu, Yuezhou Ma+2

School of Software, BNRist, Tsinghua University · MIT CSAIL

Fourier Neural Operator Neural PDE Solvers PDE Surrogate Modeling PDE Operator Learning

Open paper AskPDF arXiv

Neural solvers offer efficient surrogates for numerical simulation of partial differential equations (PDEs). For time-dependent problems, strong one-step accuracy does not necessarily translate into reliable autoregressive rollout. We observe that a solver based only on physical-state modeling can achieve lower one-step error, whereas its spectral-only counterpart can become more accurate at later rollout steps. Motivated by this observation, we present Transolver-

σ

, a neural PDE solver based on joint spectral--physical subspace modeling. Within each block, adaptive physical-state interactions and spectral transformations are modeled in dedicated latent subspaces, whose responses are recomposed to enable information exchange between the two representations. Within the physical subspace, we introduce Slice-Residual Physics-Attention (SRPA), which preserves an explicit slice-space identity path while retaining learnable cross-slice interaction. In parallel, an axis-factorized Fourier operator captures global spectral structure. Across five well-established PDE benchmarks spanning steady-state prediction and time-dependent dynamics, Transolver-

σ

achieves state-of-the-art with a benchmark-averaged relative error reduction of 33.4% over the strongest baseline for each metric, while consistently improving autoregressive rollout over single-operator counterparts. Transolver-

σ

further delivers strong gains on coupled multiphysics systems and real-world fluid and combustion measurements from RealPDEBench, demonstrating its effectiveness beyond standard simulation benchmarks.

Abstract

Neural PDE Solvers

Explore similar work

Transolver-σσσ: Joint Spectral-Physical Subspace Modeling for Neural PDE Solving

Transolver-3: Scaling Up Transformer Solvers to Industrial-Scale Geometries

When Attention Beats Fourier: Multi-Scale Transformers for PDE Solving on Irregular Domains

Transolver- $σ$ : Joint Spectral-Physical Subspace Modeling for Neural PDE Solving