Adaptive Tokenization for Vision Transformer PDE Simulation

Hanwen Wang, Paris Perdikaris

Keywords: PDE surrogates, transformers, adaptive tokenization, scientific machine learning

TL;DR: An adaptive scheme that partitions a regular grid into variable-size patches under a fixed token budget.

Abstract: Vision Transformers have been increasingly adopted as neural-operator surrogates for simulating physical systems governed by partial differential equations (PDEs). However, the quadratic computational complexity of self-attention with respect to the number of tokens limits their scalability to fine discretizations in space and time. To address this challenge, we present Iterative Adaptive Patchification (IAP), a novel fixed-budget adaptive patching scheme tailored for ViT-based autoregressive PDE simulation. IAP dynamically allocates tokens to regions of high complexity, while enforcing a strict global token budget. This design ensures that training and inference costs can be flexibly adjusted during training or inference. We validate IAP on a suite of 2D and 3D PDE benchmarks, demonstrating that it achieves performance comparable to vanilla ViTs with fixed patch sizes across a range of token budgets, and outperforms the vanilla ViT baseline significantly on one 3D benchmark. The analysis of the compute-accuracy trade-off also shows that the model is robust against post-training budget adjustments.

Journal Opt In: No, I do not wish to participate

Submission Number: 59