PICL: Physics Informed Contrastive Learning for Partial Differential Equations

TL;DR

This work tackles generalization of neural PDE operators across multiple governing equations by introducing Physics Informed Contrastive Learning (PICL). PICL combines a Generalized Contrastive Loss with a physics-informed similarity and a system-aware distance between latent representations and model updates, guiding a shared neural operator to capture diverse PDE behaviors. A two-stage training pipeline pretrains the Fourier Neural Operator on a mixture of 1D/2D Heat, Burgers', and Advection data, then fine-tunes per equation, yielding strong gains in both fixed-future and autoregressive tasks—most notably in 1D. The approach produces meaningful latent-space clustering by system behavior and demonstrates that physics-informed pretraining can enhance cross-system generalization, albeit with added pretraining compute and reliance on known governing equations.

Abstract

Neural operators have recently grown in popularity as Partial Differential Equation (PDE) surrogate models. Learning solution functionals, rather than functions, has proven to be a powerful approach to calculate fast, accurate solutions to complex PDEs. While much work has been done evaluating neural operator performance on a wide variety of surrogate modeling tasks, these works normally evaluate performance on a single equation at a time. In this work, we develop a novel contrastive pretraining framework utilizing Generalized Contrastive Loss that improves neural operator generalization across multiple governing equations simultaneously. Governing equation coefficients are used to measure ground-truth similarity between systems. A combination of physics-informed system evolution and latent-space model output are anchored to input data and used in our distance function. We find that physics-informed contrastive pretraining improves accuracy for the Fourier Neural Operator in fixed-future and autoregressive rollout tasks for the 1D and 2D Heat, Burgers', and linear advection equations.

Figures (10)

• Figure 1: Our pretraining distance function measures distance between the difference of successive frames of input data, and the difference of model output and model output with a physics-informed update.
• Figure 2: Our two-step training procedure first pretrains on all three equations simultaneously (a), then fine-tunes on each equation individually (b). Darker green points represent higher similarity to our Heat sample. Bolder arrows represent stronger attraction between samples in embedding space.
• Figure 3: 1D comparison of fixed-future performance between FNO and FNO pretrained using PICL.
• Figure 4: Comparison of autoregressive rollout performance between FNO and FNO pretrained using PICL.
• Figure 5: t-SNE of latent embeddings after PICL pretraining. We see clear clustering of similar systems, denoted by color and transparency. Advection systems are clustered separately from Heat and Burgers systems, strongly diffusive systems are clustered, strongly advective Burgers systems are clustered, and weakly to moderately diffusive and advective systems are clustered.