Masked Autoencoders are PDE Learners

TL;DR

This paper introduces masked autoencoders for physics problems, enabling neural PDE solvers to learn generalizable, rich representations that improve performance across diverse equations and conditions.

Contribution

It adapts masked pretraining to PDE modeling, allowing self-supervised learning of heterogeneous physics representations that enhance generalization and solver performance.

Findings

Learned representations generalize to unseen equations and parameters.

Conditioned neural solvers improve time-stepping and super-resolution.

Masked autoencoders effectively consolidate diverse physics data.

Abstract

Neural solvers for partial differential equations (PDEs) have great potential to generate fast and accurate physics solutions, yet their practicality is currently limited by their generalizability. PDEs evolve over broad scales and exhibit diverse behaviors; predicting these phenomena will require learning representations across a wide variety of inputs which may encompass different coefficients, boundary conditions, resolutions, or even equations. As a step towards generalizable PDE modeling, we adapt masked pretraining for physics problems. Through self-supervised learning across PDEs, masked autoencoders can consolidate heterogeneous physics to learn rich latent representations. We show that learned representations can generalize to a limited set of unseen equations or parameters and are meaningful enough to regress PDE coefficients or the classify PDE features. Furthermore, conditioning neural solvers on learned latent representations can improve time-stepping and super-resolution performance across a variety of coefficients, discretizations, or boundary conditions, as well as on certain unseen PDEs. We hope that masked pretraining can emerge as a unifying method across large, unlabeled, and heterogeneous datasets to learn latent physics at scale.