Compute-Adaptive Surrogate Modeling of Partial Differential Equations

Modeling dynamical systems governed by partial differential equations (PDEs) presents significant challenges for machine learning-based surrogate models. While vision transformers have shown potential in capturing complex spatial dynamics, their reliance on fixed-size patches limits flexibility and scalability. In this work, we introduce two convolutional architectural blocks—Convolutional Kernel Modulator (CKM) and Convolutional Stride Modulator (CSM)—designed for patch processing in autoregressive prediction tasks. These blocks unlock dynamic patching and striding strategies to balance accuracy and computational efficiency during inference. Furthermore, we propose a rollout strategy that adaptively adjusts patching and striding configurations throughout temporally sequential predictions, mitigating patch artifacts and long-term error accumulation while improving the capture of finer-scale structures of physics-based PDEs. We show that the use of these blocks improves predictive accuracy against fixed-patch baselines, while also enabling inference time scaling.