MeshTok: Adaptive PDE Tokenization

• MeshTok is an AMR-inspired tokenization framework that defines multiscale tokens on a fixed simulation grid to capture both global and local PDE features.
• It adaptively refines regions with sharp gradients using fine tokens while using coarse patches for smooth areas, balancing accuracy with computational cost.
• Experimental evaluations show that MeshTok achieves a superior efficiency-accuracy trade-off compared to uniform fine methods, reducing runtime dramatically.

Searching arXiv for MeshTok and closely related PDE transformer tokenization work.

MeshTok is an adaptive mesh refinement (AMR)-inspired tokenization and sequence modeling framework for PDE Transformers that replaces uniform patchification with a heterogeneous set of multiscale tokens defined on a fixed simulation grid. Instead of distributing computational effort evenly across the domain, it selectively refines spatial regions exhibiting sharp gradients, transient features, or multiscale structures, while retaining coarse patches in smooth regions. The resulting tokens are processed within a unified Transformer sequence, so that coarse-grained global context and fine-grained local details are handled without requiring specialized architectural components. Experimental evaluations across multiple PDE families and benchmark datasets report a consistent improvement in the efficiency-accuracy trade-off relative to uniform-grid baselines, positioning adaptive multiscale tokenization as a scalable and generalizable design principle for neural PDE modeling (Zhao et al., 3 Jun 2026).

1. Motivation and problem setting

Conventional patchified Transformers for PDEs partition a fixed high-resolution grid into uniform patches of size $p \times p$ and embed each patch as a token. MeshTok is motivated by the observation that PDE solutions are highly non-uniform: most of the domain is smooth and can be represented with large patches, whereas localized features such as shocks, boundary layers, and reaction fronts require much finer resolution (Zhao et al., 3 Jun 2026).

Within this framing, uniform “all-fine” tokenization yields a quadratic blow-up in attention cost, $\mathcal{O}(S^2)$ with $S$ tokens, while uniform “all-coarse” tokenization incurs large approximation errors in active regions. MeshTok transfers the central intuition of AMR from classical numerical PDE solvers to the tokenization layer: on a fixed simulation grid, the token stride is varied so that coarse patches are used where the solution is smooth and fine patches are used where it has sharp gradients or multiscale structure. Berger and Oliger (1984) and Berger and Colella (1989) are identified as the historical AMR lineage motivating this design (Zhao et al., 3 Jun 2026).

A common misconception is to treat the adaptive allocation of tokens as a formal optimality claim. The paper explicitly does not do so. It states that adaptive refinement moderately increases token count, but that the targeted allocation of computational resources to physically informative regions should be viewed as a practical inductive bias rather than a formal optimality guarantee (Zhao et al., 3 Jun 2026).

2. Multi-scale tokenization algorithm

At a time step $t$, MeshTok begins by partitioning the domain into

$N_c = (H/p)\cdot(W/p)$

non-overlapping coarse patches $\{R_i\}_{i=1\ldots N_c}$. For each coarse patch, an activity score is computed from finite-difference approximations of gradient and Laplacian energy: $$a(u,v) = \|\nabla x(u,v)\|^2 + \lambda \|\Delta x(u,v)\|^2,$$ followed by patchwise aggregation

$S_i = \frac{1}{|R_i|}\sum_{(u,v)\in R_i} a(u,v).$

The top $\lfloor k N_c \rfloor$ patches by descending $S_i$ are placed in the refinement index set $\mathcal{O}(S^2)$0 (Zhao et al., 3 Jun 2026).

If $\mathcal{O}(S^2)$1, the coarse patch $\mathcal{O}(S^2)$2 of size $\mathcal{O}(S^2)$3 is encoded by a shared coarse encoder $\mathcal{O}(S^2)$4 into one token. If $\mathcal{O}(S^2)$5, the patch is subdivided into $\mathcal{O}(S^2)$6 fine sub-patches of size $\mathcal{O}(S^2)$7, and each sub-patch is encoded by a shared fine encoder $\mathcal{O}(S^2)$8 into four tokens. The heterogeneous multiscale sequence is therefore

$\mathcal{O}(S^2)$9

This yields a sequence length

$S$0

The construction is explicitly two-level: unrefined coarse patches remain single tokens, while refined patches expand into four fine tokens. The paper’s ablations report that the activity-based refinement policy outperforms random and a posteriori indicator variants for rollout, and that errors shrink as the refinement ratio $S$1 increases, with $S$2 capturing most of the gain at moderate cost (Zhao et al., 3 Jun 2026).

3. Token content and geometry-aware embedding

Each token $S$3 carries three kinds of information: encoded features from either $S$4 or $S$5, continuous center coordinates $S$6, and a scale level $S$7, where $S$8 denotes coarse and $S$9 denotes fine (Zhao et al., 3 Jun 2026).

MeshTok uses a geometry-aware positional encoding based on FiLM. A small MLP $t$0 maps $t$1 to FiLM parameters $t$2: $t$3 The token embedding is then modulated as

$t$4

In this formulation, spatial location and scale adaptively re-weight the features rather than merely adding a positional bias (Zhao et al., 3 Jun 2026).

An alternative multi-scale sinusoidal basis is also described: $t$5 and similarly for $t$6, followed by concatenation. However, the reported ablations state that FiLM(pos+depth) outperforms fixed sinusoidal and naïve learnable tables. This result localizes a significant part of MeshTok’s inductive bias in the combination of multiscale tokenization and geometry-aware positional modulation rather than in modified attention mechanics (Zhao et al., 3 Jun 2026).

4. Integration into the Transformer and reconstruction

Tokens from a history of $t$7 time steps are flattened into one sequence and processed with a block-causal spatio-temporal attention mask. The mask is defined by

$t$8

with full spatial attention within a frame. This preserves the standard Transformer attention pattern while imposing temporal causality across frames (Zhao et al., 3 Jun 2026).

Within each Transformer layer, multi-head self-attention is computed in the usual form

$t$9

with $N_c = (H/p)\cdot(W/p)$0, $N_c = (H/p)\cdot(W/p)$1, and $N_c = (H/p)\cdot(W/p)$2 projected from the input token matrix. The design implication emphasized by the paper is that MeshTok does not depend on specialized architectural components; the tokens remain part of a unified Transformer sequence (Zhao et al., 3 Jun 2026).

At the output corresponding to the final time step, decoding reverses the heterogeneous merge. A coarse decoder $N_c = (H/p)\cdot(W/p)$3 maps each coarse token back to a $N_c = (H/p)\cdot(W/p)$4 patch, and a fine decoder $N_c = (H/p)\cdot(W/p)$5 maps each fine token to a $N_c = (H/p)\cdot(W/p)$6 patch. A small CNN $N_c = (H/p)\cdot(W/p)$7 fuses the full coarse field and the sparse fine correction: $N_c = (H/p)\cdot(W/p)$8 This suggests a separation between global field reconstruction from coarse tokens and local correction from fine tokens, but the paper frames the method as a standard encoder or encoder-decoder integration rather than a new attention architecture (Zhao et al., 3 Jun 2026).

5. Computational scaling and empirical trade-offs

The theoretical cost argument is stated relative to a uniform all-fine discretization. Uniform “all-fine” tokenization uses $N_c = (H/p)\cdot(W/p)$9 tokens per frame and attention cost proportional to $\{R_i\}_{i=1\ldots N_c}$0. MeshTok with refinement ratio $\{R_i\}_{i=1\ldots N_c}$1 uses $\{R_i\}_{i=1\ldots N_c}$2 tokens, giving an attention cost

$\{R_i\}_{i=1\ldots N_c}$3

At $\{R_i\}_{i=1\ldots N_c}$4, the model retains approximately $\{R_i\}_{i=1\ldots N_c}$5 of full-fine attention cost, corresponding to a $\{R_i\}_{i=1\ldots N_c}$6 reduction (Zhao et al., 3 Jun 2026).

The reported experimental setup spans PDEBench, The Well, and PDENNEval, with datasets including CNS, SWE, Reaction-Diffusion, Gray–Scott, Shear Flow, Allen–Cahn, Burgers, and BSB. The main resolution is $\{R_i\}_{i=1\ldots N_c}$7, with ablations at $\{R_i\}_{i=1\ldots N_c}$8 and $\{R_i\}_{i=1\ldots N_c}$9. Model scales are SMALL (4-layer, $$a(u,v) = \|\nabla x(u,v)\|^2 + \lambda \|\Delta x(u,v)\|^2,$$0), BIG (8-layer, $$a(u,v) = \|\nabla x(u,v)\|^2 + \lambda \|\Delta x(u,v)\|^2,$$1), and LARGE (12-layer, $$a(u,v) = \|\nabla x(u,v)\|^2 + \lambda \|\Delta x(u,v)\|^2,$$2). Baselines include DeepONet, FNO, ViT, AFNO/DPOT, MPP, MoE-POT, and BCAT. Evaluation uses relative $$a(u,v) = \|\nabla x(u,v)\|^2 + \lambda \|\Delta x(u,v)\|^2,$$3 error for 1-step and 10-step autoregressive rollout, wall-clock runtime, and MACs (Zhao et al., 3 Jun 2026).

For the BIG-model pretraining setting, representative benchmark results are reported as follows.

Setting	CNS*(1.0,0.01)	SWE
Uniform Coarse	1.138	0.647
MeshTok ($$a(u,v) = \|\nabla x(u,v)\|^2 + \lambda \|\Delta x(u,v)\|^2,$$4)	1.056	0.443
Uniform Fine	0.911	0.395

The corresponding averages are 1.086 for Uniform Coarse, 0.962 for MeshTok ($$a(u,v) = \|\nabla x(u,v)\|^2 + \lambda \|\Delta x(u,v)\|^2,$$5), and 0.840 for Uniform Fine (Zhao et al., 3 Jun 2026).

Wall-clock inference time for the BIG model is reported separately.

Setting	Time
coarse	18.0 ms
MeshTok ($$a(u,v) = \|\nabla x(u,v)\|^2 + \lambda \|\Delta x(u,v)\|^2,$$6)	39.7 ms
full-fine	106.9 ms

These numbers underwrite two distinct comparisons. Relative to uniform coarse tokenization, MeshTok lowers error by refining only selected regions. Relative to uniform fine tokenization, it approaches fine-grid accuracy while operating at substantially lower runtime and attention cost (Zhao et al., 3 Jun 2026).

6. Findings, scope, and limitations

The paper’s principal finding is that adaptive multi-scale tokenization yields consistently lower relative $$a(u,v) = \|\nabla x(u,v)\|^2 + \lambda \|\Delta x(u,v)\|^2,$$7 errors than uniform coarse grids under the same token budget. It also reports that, compared to uniform fine grids, MeshTok achieves comparable accuracy at a fraction of the computational cost, often with $$a(u,v) = \|\nabla x(u,v)\|^2 + \lambda \|\Delta x(u,v)\|^2,$$8 wall-clock speedups. The activity-based indicator combining gradient and Laplacian energy is presented as a practical surrogate for more expensive a posteriori or learned policies (Zhao et al., 3 Jun 2026).

A second major finding concerns architectural generality. MeshTok is described as integrating seamlessly into any ViT-style encoder or encoder-decoder; the core inductive bias resides in the tokenization and positional FiLM rather than in custom attention layers. The paper therefore treats adaptive multiscale tokenization as a design principle for scalable PDE foundation models rather than as a domain-specific architectural exception (Zhao et al., 3 Jun 2026).

The limitations are explicit. The current design is tied to structured Cartesian grids and a two-level refinement hierarchy. Extension to unstructured meshes or deeper multilevel refinements would require new mesh graph encodings and connectivity-aware positional maps. This suggests that the present formulation is best understood as an AMR-inspired tokenization scheme for structured-grid PDE learning, not as a complete solution to mesh generality (Zhao et al., 3 Jun 2026).

Input noise is also treated as a practical consideration rather than a theoretical component of the method: a small relative noise, $$a(u,v) = \|\nabla x(u,v)\|^2 + \lambda \|\Delta x(u,v)\|^2,$$9, is reported to stabilize rollout and improve long-horizon accuracy. Code is available through the SCAILab-USTC repository linked from the paper (Zhao et al., 3 Jun 2026).