Physics-Aligned Canonical Equivariant Fourier Neural Operator under Symmetry-Induced Shifts

TL;DR

The paper introduces PACE-FNO, a neural operator that respects physical symmetries to improve generalization in solving PDEs, especially out-of-distribution, by aligning inputs and outputs with known symmetries.

Contribution

It proposes a novel symmetry-aware neural operator architecture that separates coordinate alignment from physical evolution, enhancing OOD performance without changing the core FNO architecture.

Findings

PACE-FNO matches in-distribution accuracy of standard neural operators.

Reduces out-of-distribution error by up to 12x compared to FNO with symmetry augmentation.

Input alignment and output restoration account for most OOD gains.

Abstract

Neural operators approximate PDE solution maps, but they need not respect the symmetries of the governing equation. In out-of-distribution (OOD) regimes, a standard neural operator must often learn coordinate alignment and physical evolution within a single map, which can hurt generalization. We use known continuous symmetries of evolution equations on periodic domains to separate these two roles. We propose the Physics-Aligned Canonical Equivariant Fourier Neural Operator (PACE-FNO), which estimates the input frame with a Lie-algebra coordinate estimator, maps the field to a reference frame, applies a standard Fourier Neural Operator (FNO), and restores the prediction to the target frame. We train alignment and operator prediction jointly using bounded symmetry perturbations, with an optional low-dimensional refinement step that updates the estimated frame at inference. Equivariance is enforced by the input and output transformations, while the FNO architecture remains unchanged. Across 1-D and 2-D Burgers, shallow-water, and Navier-Stokes equations on periodic domains, PACE-FNO matches the in-distribution (ID) accuracy of standard neural operators and reduces out-of-distribution (OOD) relative error by up to 12x over FNO with symmetry augmentation (FNO+Aug) under translations and Galilean shifts, with smaller gains for coupled rotation-translation shifts. Ablations show that aligning the input and restoring the output frame account for most OOD gains; inference-time refinement provides a smaller correction.