# Dynamic Quadtree Tokenization and Transformer for Adaptive Mesh PDE Forecasting

arXiv: https://arxiv.org/abs/2610.04044 · HTML: https://mlsanity.com/paper/2610.04044/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 23 of its 36 references that are in the map.
- First posted on arXiv 2 Oct 2026.

## Builds on (most cited first, 23 of 23)

- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Scalable Diffusion Models with Transformers](https://mlsanity.com/paper/2212.09748/) (2022)
- [RoFormer: Enhanced Transformer with Rotary Position Embedding](https://mlsanity.com/paper/2104.09864/) (2021)
- [Fourier Neural Operator for Parametric Partial Differential Equations](https://mlsanity.com/paper/2010.08895/) (2020)
- [DeepONet: Learning nonlinear operators for identifying differential equations based on the universal approximation theorem of operators](https://mlsanity.com/paper/1910.03193/) (2019)
- [ViViT: A Video Vision Transformer](https://mlsanity.com/paper/2103.15691/) (2021)
- [Swin Transformer V2: Scaling Up Capacity and Resolution](https://mlsanity.com/paper/2111.09883/) (2021)
- [DGM: A deep learning algorithm for solving partial differential equations](https://mlsanity.com/paper/1708.07469/) (2017)
- [DeepXDE: A deep learning library for solving differential equations](https://mlsanity.com/paper/1907.04502/) (2019)
- [Ray: A Distributed Framework for Emerging AI Applications](https://mlsanity.com/paper/1712.05889/) (2017)
- [Message Passing Neural PDE Solvers](https://mlsanity.com/paper/2202.03376/) (2022)
- [Adaptive Fourier Neural Operators: Efficient Token Mixers for Transformers](https://mlsanity.com/paper/2111.13587/) (2021)
- [Rotary Position Embedding for Vision Transformer](https://mlsanity.com/paper/2403.13298/) (2024)
- [Poseidon: Efficient Foundation Models for PDEs](https://mlsanity.com/paper/2405.19101/) (2024)
- [CViT: Continuous Vision Transformer for Operator Learning](https://mlsanity.com/paper/2405.13998/) (2024)
- [Interfacing Finite Elements with Deep Neural Operators for Fast Multiscale Modeling of Mechanics Problems](https://mlsanity.com/paper/2203.00003/) (2022)
- [An Image is Worth More Than 16x16 Patches: Exploring Transformers on Individual Pixels](https://mlsanity.com/paper/2406.09415/) (2024)
- [Accelerating Vision Transformers with Adaptive Patch Sizes](https://mlsanity.com/paper/2510.18091/) (2025)
- [AMR-Transformer: Enabling Efficient Long-range Interaction for Complex Neural Fluid Simulation](https://mlsanity.com/paper/2503.10257/) (2025)
- [Adaptive physics-informed neural operator for coarse-grained non-equilibrium flows](https://mlsanity.com/paper/2210.15799/) (2022)
- [Adaptive Patching for High-resolution Image Segmentation with Transformers](https://mlsanity.com/paper/2404.09707/) (2024)
- [MATEY: multiscale adaptive foundation models for spatiotemporal physical systems](https://mlsanity.com/paper/2412.20601/) (2024)
- [DART: Differentiable Dynamic Adaptive Region Tokenizer for Vision Foundation Models](https://mlsanity.com/paper/2506.10390/) (2025)

## Abstract (arXiv)

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