Controllable Patching for Compute-Adaptive Surrogate Modeling of Partial Differential Equations
This addresses the problem of compute-adaptive deployment for production systems using PDE surrogates, offering a plug-and-play solution that is broadly applicable across architectures.
The paper tackles the limitation of fixed patch sizes in patch-based transformer surrogates for modeling spatiotemporal dynamics by introducing lightweight modules (CKM and CSM) that enable dynamic patch size control at inference without retraining or accuracy loss, improving rollout fidelity and runtime efficiency on 2D and 3D PDE benchmarks.
Patch-based transformer surrogates have become increasingly effective for modeling spatiotemporal dynamics, but the fixed patch size is a major limitation for budget-conscience deployment in production. We introduce two lightweight, architecture-agnostic modules-the Convolutional Kernel Modulator (CKM) and Convolutional Stride Modulator (CSM)-that enable dynamic patch size control at inference in patch based models, without retraining or accuracy loss. Combined with a cyclic patch-size rollout, our method mitigates patch artifacts and improves long-term stability for video-like prediction tasks. Applied to a range of challenging 2D and 3D PDE benchmarks, our approach improves rollout fidelity and runtime efficiency. To our knowledge, this is the first framework to enable inference-time patch-size tunability in patch-based PDE surrogates. Its plug-and-play design makes it broadly applicable across architectures-establishing a general foundation for compute-adaptive modeling in PDE surrogate tasks.