Dynamic Quadtree Tokenization and Transformer for Adaptive Mesh PDE Forecasting

Yilin Zhuang; Noah Zambrano; Karthik Duraisamy

cs.CVOct 2, 2026

Dynamic Quadtree Tokenization and Transformer for Adaptive Mesh PDE Forecasting

Authors: Yilin Zhuang, Noah Zambrano, Karthik Duraisamy

Organizations: Department of Aerospace Engineering University of Michigan Ann Arbor, MI 48105 · University of Michigan Ann Arbor, MI 48105

Abstract

The quadratic attention cost of Vision Transformers (ViTs) forces a trade-off between spatial resolution and rollout horizon, particularly for fine-scale PDEs where shocks, reaction fronts, and material interfaces occupy small, evolving regions of the domain. Conventional neural surrogates also lack mechanisms to adapt resolution dynamically. We propose WAMRViT, a ViT that tokenizes inputs as balanced quadtrees using a wavelet-inspired refinement criterion, jointly encodes position and refinement level with 3D rotary positional embeddings, and regrids in cell space during inference for stable long-horizon rollouts. A multi-scale variant retains each leaf at its native source resolution and lets the model learn across resolution levels. Unlike fixed-budget adaptive-tokenization methods, WAMRViT imposes no predetermined token count and supports fully adaptive topology throughout autoregressive rollout. To our knowledge, it is the first machine-learning surrogate to natively tokenize multi-level Adaptive Mesh Refinement (AMR) data. On uniform-grid benchmarks, uniform-patch WAMRViT improves finest-level region-of-interest VRMSE over a finest-patch uniform ViT while using substantially fewer tokens. The multi-scale variant achieves the lowest first-step full-field RMSE and VRMSE on both benchmarks and improves rollout-averaged full-field and refined-region accuracy at long horizons. With parallelized regridding, its end-to-end rollout cost lies between finest-patch and approximately token-matched coarser-patch ViTs. On a complex AMR combustion problem whose finest features cannot be represented natively by the evaluated uniform-grid baselines, WAMRViT operates directly on adaptive cells and substantially reduces finest-level error at matched transformer capacity. Code: https://github.com/tonyzyl/wamrvit

Figures & tables

Figure 1: Overview of WAMRViT. (a) Adaptive tokenization: the input field is partitioned into a 2:1-balanced quadtree by a wavelet-inspired local-range criterion, with one token per leaf. (b) Autoregressive inference: every $R$ steps a refine-protect-coarsen update revises the quadtree in cell space so the topology tracks moving features.

Figure 2: Parallel-layout transformer block. Attention and MLP branches are concatenated and projected into the residual stream.

Panel of Figure 3 — Figure 3: AMReX combustion prediction on the chronological holdout. (a) Full-domain WAMRViT prediction at rollout step $25$ ; the white box marks the spatial extent of the zoomed field examples. (b) Predictions (top row) and ground truth (bottom row): peak-pressure soot foils over $50$ rollout steps (left), and heat-release rate, density, and temperature at step $25$ (right). These longer qualitative diagnostics extend beyond the $20$ -step quantitative protocol. Ground-truth grids are shown for reference. Additional full-domain comparisons are shown in Appendix A . Animations are provided in the repository .

	WAMRViT (uniform)	WAMRViT (multi-scale)	ViT (finest)	ViT (mid)	SwinV2	IAP ∗ [ 14 ]
PLI, full-field	0.0661 $\pm$ 1.67E-3	0.0389 $\pm$ 2.52E-4	0.0539 $\pm$ 6.44E-3	0.0553 $\pm$ 6.68E-4	0.0920 $\pm$ 2.32E-4	/
PLI, finest ( $\ell\!=\!3$ )	2.6839 $\pm$ 0.166	2.5000 $\pm$ 0.046	3.4125 $\pm$ 0.305	3.3401 $\pm$ 0.121	6.3348 $\pm$ 0.054	/
PLI, tokens $\bar{N}$	3,327	3,339	8,960	3,500	8,960	/
TRL, full-field	0.1502 $\pm$ 2.81E-4	0.1476 $\pm$ 5.40E-4	0.1490 $\pm$ 6.61E-4	0.1629 $\pm$ 4.87E-4	0.1955 $\pm$ 6.46E-4	0.1577 / 0.1478
TRL, finest ( $\ell\!=\!2$ )	9.0411 $\pm$ 0.665	18.8458 $\pm$ 0.699	16.5135 $\pm$ 0.351	22.1965 $\pm$ 0.334	69.8422 $\pm$ 0.503	/
TRL, tokens $\bar{N}$	1,597	1,597	3,072	1,536	3,072	1,020 / 2,040

Table 1: First-step VRMSE and rollout-mean token counts. Accuracy is mean

\pm

sample standard deviation over seeds

\{0,1,42\}

. Token counts

\bar{N}

are averaged over evaluation windows, rollout steps (

20

for PLI/AMReX;

1

for TRL), and seeds; SwinV2 reports first-stage tokens. Finest-level metrics use the regions refined to level

L

in the initial adaptive topology. IAP ∗ [ 14 ] reports published performance at two token budgets,

1{,}020

and

2{,}040

, respectively.

Panel of Figure 4 — Figure 4: Aggregated rollout error. PLI full-field and finest-level VRMSE for $K=1$ models (top), and AMReX finest-level RMSE and VRMSE for $K=4$ models (bottom). WAMRViT regrids every four rollout steps. At the PLI finest level, the single ViT-finest curve is evaluated on uniform WAMRViT’s refined region; multi-scale WAMRViT uses its own region. AMReX errors are dominated by heat-release rate (HRR); per-field results are reported in Appendix G.9 . Three-seed first-step statistics are given in Table 1 .

Dataset	Variant	Train (min)	$\bar{N}$	Forward (s)	Regrid (s)	Total (s)	Peak (GiB)
PLI	WAMRViT (uniform)	188	3601	1.454	0.126	$1.580\pm 0.328$	1.584
	WAMRViT (multi-scale)	333	3606	1.601	0.259	$1.860\pm 0.314$	1.802
	ViT-finest	236	8960	4.789	–	$4.789\pm 0.008$	1.979
	ViT-mid	128	3500	1.301	–	$1.301\pm 0.003$	1.523
	SwinV2	283	8960	0.469	–	$0.469\pm 0.001$	1.343
TRL	WAMRViT (uniform)	136	2111	0.774	0.159	$0.933\pm 0.095$	1.354

Table 2: Per-model training and inference cost. Inference uses ten matched starting windows, batch one, and

20

model calls on one H100, including four complete regrids for adaptive models. Forward and regrid are mean per-window sums. Peak is maximum PyTorch-allocated CUDA memory;

\bar{N}

averages tokens per call over the timing cohort. Train reports total minutes for the schedules in Appendix Table 4 .

	PLI		TRL		AMReX
$s_{\ell}$	full-field	finest ( $\ell\!=\!3$ )	full-field	finest ( $\ell\!=\!2$ )	finest ( $\ell\!=\!4$ )
$1$ (default)	0.0642	2.5107	0.1515	9.7444	1.631E10
$2$ (perimeter ratio)	0.0630	2.5250	0.1505	8.9478	1.805E10
$4$ (area ratio)	0.0632	2.5042	0.1506	9.8886	1.544E10
$8$	0.0666	2.5394	0.1502	9.1519	1.660E10

Table 3: Level-axis RoPE ablation: first-rollout-step channel-mean VRMSE (lower is better). Training depth follows each dataset’s default (

K=1

for PLI/TRL and

K=4

for AMReX). PLI/AMReX use a fixed seed; TRL is a single-run sweep. Full rollout curves and geometric interpretation are in Appendix G.5 .

Appendix figures & tables26 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 5: Qualitative rollout snapshot on PLI. Density at rollout step 20, trajectory index 950 (physical ID 00951), starting frame 20. In row-major order: ground truth, WAMRViT (uniform), WAMRViT (multi-scale), ViT-finest, ViT-mid, and SwinV2. All panels share a color range determined only from ground truth.

Figure 6: Qualitative rollout snapshot on TRL. Density at rollout step 20 on trajectory 0 (starting frame 20). In row-major order: ground truth, WAMRViT (uniform), WAMRViT (multi-scale), ViT-finest, ViT-mid, and SwinV2. All panels share a ground-truth-only color range. Both adaptive variants use initial and regrid tolerance $0.1$ per channel. This horizon is qualitative, beyond the one-step accuracy protocol.

Figure 7: Qualitative rollout snapshot on AMReX combustion, full domain. Temperature at rollout step 20 in the chronological test region, starting from frame 330. Top to bottom: ground truth, WAMRViT (uniform), and ViT-mid. Figure 3 shows a zoomed view of the leading shock region.

Figure 8: Per-level patchify cascade of the multi-scale variant.

	PLI	TRL	AMReX
Backbone (shared)
Layers / heads / per-head $d_{h}$	$6$ / $12$ / $192$
Transformer parameters	$286.7$ M
Complete parameters (uniform)	$290.2$ M	$287.5$ M	$407.1$ M
Complete parameters (multi-scale)	$339.0$ M	$292.0$ M	–
Tokenizer

Appendix

Table 4: Training specifications for the WAMRViT runs reported in Section 3.4 ; topology resampling is disabled. AMReX constructs its quadtree directly from the simulator hierarchy (Appendix C.6 ) and has no fixed root tile.

	PLI	TRL	AMReX
Backbone and decoder
Stage depths	$[8,\,12,\,12,\,6]$
Heads per stage	$[6,\,12,\,24,\,48]$
Transformer-stage parameters	$286.5$ M
Embedding dim / decoder channels	$192$ / $256$
Decoder refinement blocks	$2$

Appendix

Table 5: SwinV2 specifications for the regular-grid architecture-family control. The four conventional shifted-window stages approximately match WAMRViT’s transformer capacity. The refined FPN decoder uses two ConvNeXt blocks per non-bottleneck skip and a learned

3\times 3

correction after bilinear upsampling; it is inspired by Poseidon [ 36 ] . On AMReX, the hierarchy is projected onto the uniform grid of AMR level

2

(

1024\times 6144

), four times coarser per direction than the finest level.

Dataset	Variant	Adaptation channels	Tolerance $\varepsilon$	Radius $\delta$	Execution
PLI	WAMRViT (uniform)	$3$ of $10$	$[0.1,\,0.1,\,0.3]$	$2$	CPU array
	WAMRViT (multi-scale)	$3$ of $10$	$[0.1,\,0.1,\,0.3]$	$2$	CPU array
TRL	WAMRViT (uniform)	$4$ of $4$	$[0.1,\,0.1,\,0.1,\,0.1]$	$2$	CPU array
	WAMRViT (multi-scale)	$4$ of $4$	$[0.1,\,0.1,\,0.1,\,0.1]$	$2$	CPU array
AMReX	WAMRViT (uniform)	HRR	$0.01$	$2$	CPU array + GPU replay

Appendix

Table 6: Inference regrid settings for Table 2 . PLI and TRL tolerances are per selected channel; AMReX uses a scalar tolerance. “CPU array” denotes parallel packed-array topology and CPU payload operations; AMReX keeps CPU topology authoritative and replays payload operations on the GPU.

	WAMRViT (uniform)	WAMRViT (multi-scale)	ViT (finest)	ViT (mid)	SwinV2
PLI, full-field	0.0110 $\pm$ 6.232E-5	4.963E-3 $\pm$ 2.198E-5	6.092E-3 $\pm$ 2.338E-4	6.932E-3 $\pm$ 4.526E-5	0.0111 $\pm$ 6.474E-5
PLI, finest ( $\ell\!=\!3$ )	2.475E-3 $\pm$ 8.896E-5	1.909E-3 $\pm$ 1.167E-5	2.473E-3 $\pm$ 1.338E-4	2.858E-3 $\pm$ 3.869E-5	5.702E-3 $\pm$ 4.940E-5
TRL, full-field	0.5185 $\pm$ 4.802E-3	0.4979 $\pm$ 8.661E-4	0.5069 $\pm$ 1.483E-4	0.5972 $\pm$ 3.193E-3	0.7047 $\pm$ 3.396E-3
TRL, finest ( $\ell\!=\!2$ )	0.3283 $\pm$ 3.064E-3	0.3333 $\pm$ 1.103E-3	0.3382 $\pm$ 2.818E-4	0.4007 $\pm$ 2.046E-3	0.5193 $\pm$ 7.857E-4
AMReX, finest ( $\ell\!=\!4$ )	1.202E10 $\pm$ 8.293E8	/	/	3.064E10 $\pm$ 1.818E8	3.022E10 $\pm$ 1.498E8

Appendix

Table 7: First-rollout-step RMSE (channel mean), counterpart of Table 1 . Entries are mean

\pm

sample standard deviation over three training seeds; training depth is

K=1

for PLI/TRL and

K=4

for AMReX. Finest-level metrics are evaluated on the regions refined to level

L

in the initial adaptive topology. AMReX omits ViT-finest because the native finest grid cannot be materialized; ViT-mid and SwinV2 use

p=32

Panel of Figure 9 — Figure 9: Aggregated full-field error on PLI , $K\!=\!1$ .

Panel of Figure 10 — Figure 10: Aggregated coarsest-level error on PLI , $K\!=\!1$ . Each coarsest leaf spans $8^{2}$ times the finest-leaf area. Uniform-patch WAMRViT stores $50$ pixels per leaf, whereas multi-scale WAMRViT retains an $80\times 40$ raster. The displayed regular-grid controls are block-mean projected onto uniform WAMRViT’s leaf rasters and can benefit from error-cancelling averages at this scope. The coarse-level curves therefore differ in retained resolution as well as adaptive topology.

Panel of Figure 11 — Figure 11: Aggregated finest-level error on PLI , $K\!=\!1$ .

	full-field		finest ( $\ell\!=\!3$ )
$K$	WAMRViT	ViT-finest	WAMRViT	ViT-finest
$1$	0.0642	0.0494	2.5107	3.1023
$2$	0.0638	0.0501	2.5252	3.1863
$4$	0.0680	0.0593	2.4921	3.3610

Appendix

Table 8: First-step VRMSE on PLI under the push-forward sweep, channel-mean over the ten PLI fields.

Panel of Figure 12 — Figure 12: Push-forward ablation on PLI. WAMRViT (solid, blue family) and the ViT-finest baseline (dashed, green family) at $K\in\{1,2,4\}$ , with light-to-dark shading for $K=1\to 4$ within each family. Top row: full-field aggregate. Bottom row: finest refinement level ( $\ell=3$ ).

Variant	full-field	finest ( $\ell\!=\!3$ )
WAMRViT (no resample, default)	0.0642	2.5107
WAMRViT (with resample)	0.0683	2.7734

Appendix

Table 9: First-step VRMSE on PLI under the random-resample augmentation, channel-mean over the ten PLI fields.

Panel of Figure 13 — Figure 13: Random-resample ablation on PLI. WAMRViT (uniform tokenization) trained without the augmentation (solid, dark blue, the main-text default) vs. with the augmentation of Appendix C.8 (solid, light blue), at $K\!=\!1$ , against the ViT-finest baseline (dashed, green). Top row: full-field aggregate. Bottom row: finest refinement level ( $\ell\!=\!3$ ).

Variant	coarsest ( $\ell\!=\!0$ )	finest ( $\ell\!=\!4$ )
WAMRViT (no resample, default)	1.525E9	1.631E10
WAMRViT (with resample)	1.648E9	1.668E10

Appendix

Table 10: First-step VRMSE on AMReX combustion under the random-resample augmentation, channel-mean over the seventeen combustion fields. AMReX has no full-field row (Section 3.1 ). Resampling increases coarsest- and finest-level VRMSE by

8.1\%

and

2.3\%

, respectively.

Panel of Figure 14 — Figure 14: Random-resample ablation on AMReX combustion. WAMRViT (uniform tokenization) trained without the augmentation (solid, dark blue, the no-resampling control) vs. with the augmentation (solid, light blue), at $K\!=\!4$ . Top row: coarsest level ( $\ell\!=\!0$ ). Bottom row: finest refinement level ( $\ell\!=\!4$ ).

Panel of Figure 15 — Figure 15: RoPE level-axis scale ablation on PLI. Rollouts for WAMRViT (uniform tokenization, no resample) at $s_{\ell}\in\{1,2,4,8\}$ . Top row: full-field aggregate. Bottom row: finest refinement level ( $\ell=3$ ). The ordering varies with scope, metric, and horizon.

Panel of Figure 16 — Figure 16: RoPE level-axis scale ablation on AMReX combustion. WAMRViT rollouts at $s_{\ell}\in\{1,2,4,8\}$ . Top row: coarsest level ( $\ell=0$ ). Bottom row: finest level ( $\ell=4$ ). Full-field materialization is omitted.

	WAMRViT (uniform)			WAMRViT (multi-scale)
Mesh policy	$\bar{N}$	RMSE	VRMSE	$\bar{N}$	RMSE	VRMSE
Default	3,324	0.0403	0.1749	3,336	0.0306	0.1405
Radius $\delta=1$	2,717	0.0486	0.2238	2,739	0.0414	0.1950
Radius $\delta=3$	3,981	0.0427	0.1836	3,998	0.0329	0.1476
Tolerances $\times 0.5$	3,745	0.0408	0.1780	3,717	0.0316	0.1435
Tolerances $\times 1.5$	3,269	0.0404	0.1753	3,299	0.0308	0.1428

Appendix

Table 11: PLI mesh-policy sensitivity. Full-field rollout-mean errors and token counts for the fixed checkpoints. Token counts

\bar{N}

are averaged over the same evaluation windows and rollout steps.

Figure 17: PLI full-field errors ,

K=1

. Rows show the four field groups described above; columns show VRMSE (left) and RMSE (right).

Figure 18: PLI coarsest-level errors (

\ell=0

K=1

. Rows show the same four field groups; columns show VRMSE (left) and RMSE (right).

Figure 19: PLI finest-level errors (

\ell=3

K=1

. Rows show the same four field groups; columns show VRMSE (left) and RMSE (right).

Figure 20: Aggregated AMReX combustion errors ,

K=4

. Rows show the coarsest (

\ell=0

) and finest (

\ell=4

) levels; columns show VRMSE (left) and RMSE (right), averaged over the

17

channels. As on PLI, the regular-grid predictions are block-mean-downsampled onto the adaptive coarsest-level leaves and can benefit from error-cancelling averages at this scope. The finest row corresponds to the AMReX panels of Figure 4 .

Figure 21: AMReX coarsest-level errors (

\ell=0

K=4

. Each row corresponds to one of the six representative fields; columns show VRMSE (left) and RMSE (right).

Figure 22: AMReX finest-level errors (

\ell=4

K=4

. Rows show the same six fields; columns show VRMSE (left) and RMSE (right).

Explore similar work

MeshTok: Efficient Multi-Scale Tokenization for Scalable PDE Transformers
Jun 3, 2026Yanshun Zhao, Xiaoyu Peng, Jiamin Jiang +2Neural PDE Solvers PDE Surrogate Modeling
School of Mathematical Sciences, University of Science and Technology of China, Hefei 230026, China · Suzhou Institute for Advanced Research, University of Science and Technology of China, Suzhou 215123, China.
Less Tokens, Better Forecasts: Sparse Residual Routing for Efficient Weather Prediction
Jul 2, 2026Janet Wang, Yunbei Zhang, Lin Zhao +3Efficient ViTs Dense Prediction
Tulane University · Northeastern University · University of Alabama at Birmingham +1
AE-ViT: Stable Long-Horizon Parametric Partial Differential Equations Modeling
Apr 7, 2026Iva Mikuš, Boris Muha, Domagoj VlahPDE Surrogate Modeling Reduced-Order Modeling
University of Zagreb Faculty of Electrical Engineering and Computer Science · University of Zagreb Faculty of Science Department of Mathematics

Jun 3, 2026J/K move · Enter open · S save

MeshTok: Efficient Multi-Scale Tokenization for Scalable PDE Transformers

Yanshun Zhao, Xiaoyu Peng, Jiamin Jiang, Congcong Zhu+1

School of Mathematical Sciences, University of Science and Technology of China, Hefei 230026, China · Suzhou Institute for Advanced Research, University of Science and Technology of China, Suzhou 215123, China.

Neural PDE Solvers PDE Surrogate Modeling Adaptive Mesh Refinement

Open paper AskPDF arXiv

Conventional patchified Transformers operate on uniform spatial partitions, distributing computational effort evenly across the domain irrespective of local features. This inflexible tokenization scheme is inherently limited in its ability to efficiently represent and process solutions to complex PDEs. To address this, we propose MeshTok, an adaptive mesh refinement (AMR)-inspired tokenization and sequence modeling framework. This method selectively refines spatial regions exhibiting sharp gradients, transient features, or multiscale structures, generating a heterogeneous set of multiscale tokens defined on a fixed simulation grid. These tokens are processed within a unified Transformer sequence, enabling the model to simultaneously capture coarse-grained global context and fine-grained local details without requiring specialized architectural components. Although adaptive refinement moderately increases token count, it promotes a more targeted allocation of computational resources to physically informative regions, which we view as a practical inductive bias rather than a formal optimality guarantee. Experimental evaluations across multiple PDE families and benchmark datasets demonstrate that MeshTok consistently improves the efficiency-accuracy trade-off compared to uniform-grid baselines. This suggests adaptive multiscale tokenization as a scalable and generalizable design principle for neural PDE modeling. Code is available at https://github.com/SCAILab-USTC/MeshTok.

Abstract

Explore similar work

MeshTok: Efficient Multi-Scale Tokenization for Scalable PDE Transformers

Less Tokens, Better Forecasts: Sparse Residual Routing for Efficient Weather Prediction

AE-ViT: Stable Long-Horizon Parametric Partial Differential Equations Modeling