MeshTok: Efficient Multi-Scale Tokenization for Scalable PDE Transformers

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.

Many scientific and engineering systems, such as fluids, heat transfer, and chemical reactions, are modeled by partial differential equations that describe how quantities change over space and time. Neural networks can learn to predict these systems faster than traditional numerical solvers, but Transformer-based models often divide the whole simulation grid into equally sized pieces. This uniform division can waste computation in smooth regions while failing to focus enough detail on areas with sharp changes or complex local behavior. MeshTok addresses this problem by letting the model use different levels of detail in different parts of the same fixed grid. It represents smooth regions with larger, coarser tokens and complex regions with smaller, finer tokens, then processes all of these tokens together in one Transformer model. This design allows the model to keep broad global information while paying more attention to locally important details. Experiments across several PDE benchmarks show that MeshTok improves the balance between prediction accuracy and computational cost compared with uniform-grid Transformer baselines. This makes adaptive tokenization a practical way to build more efficient neural predictors for scientific simulations.