arXiv:2605.18606v1 [cs.LG] 18 May 2026
Physics-Aligned Canonical Equivariant Fourier Neural Operator under Symmetry-Induced Shifts
Jiaxiao Xu Shanghai University of Finance and Economics [email protected]
Changhong Mou Utah State University [email protected]
Yeyu Zhang Shanghai University of Finance and Economics [email protected]
Fengxiang He University of Edinburgh [email protected]
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 12× 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.
1
Introduction
Neural operators learn maps between function spaces and approximate PDE solution operators [1, 2, 3, 4, 5]. We study maps G † : A → U from input to output fields for parametric PDEs, fluid simulation, and weather modeling [6, 7, 8, 9]. Once trained, these models generalize across parameters and discretizations without rerunning a solver, but out-of-distribution (OOD) generalization remains a central challenge. The difficulty comes from a geometric mismatch. Many PDE solution operators are equivariant under continuous symmetry groups: translated, rotated, or co-moving inputs induce outputs transformed by the same group action [10, 11, 12]. If training data lie near a canonical slice of the orbit [a] = {Cg (a) | g ∈ G}, physically equivalent inputs can fall outside the sampled coordinate range even though the equation is unchanged. Spectral truncation can amplify this mismatch when transformed fields move energy outside retained modes [13]. Thus small in-distribution (ID) error on canonical samples need not imply small error along the group orbit. Preprint.
Solution Space M Out-of-distribution orbit
aaug
Orbit on M
C gã c
ãc
Equivalenceclass mapping
[x] = {g·x | g ∈ G}g ·x 2 g1 ·x
1 ug Canonical slice C−g a a
x
(ID)
canonicalization
orbit / class [x]
x̃
canonical slice
x̃
Cĝ
orbit-class view
FNO ãc
canonicalization
ũc
ûpred
[x̃] ∈ M/G
(a)
(b)
Figure 1: Geometric view of PACE-FNO. (a) The input is pulled from its observed orbit to a canonical slice, evolved there, and pushed forward to the target frame. (b) Points connected by the group action share an orbit [x] = {g · x | g ∈ G}; canonicalization chooses one representative x̃ from each equivalence class.
We address OOD generalization by exploiting continuous PDE symmetries, including translation, rotation, and Galilean boost. These transformations preserve the solution operator up to induced input and output actions, yet move observations along group orbits not covered by the canonical training measure. For this regime we propose the Physics-Aligned Canonical Equivariant Fourier Neural Operator (PACE-FNO), which decouples symmetry alignment from dynamics prediction. PACE-FNO trains a Lie-algebra coordinate estimator jointly with a standard Fourier Neural Operator (FNO), so the canonicalizer adapts to the target dynamics. Its default inference path is one-shot; per-sample group optimization is only an optional frozen-weight refinement. In practice, PACE-FNO maps each OOD input to a canonical frame, applies the FNO there, and pushes the prediction back to the target frame. We study equations on tori, i.e. periodic domains Td , where OOD samples come from bounded group actions that preserve the equation and boundary conditions. The method targets symmetry-induced shifts [11, 14, 15, 16, 17, 18], not arbitrary distribution shift. Contributions. • A Lie-algebra canonicalization layer for solution operators on fields over tori. It uses spectral group actions and terminal-frame correction, and wraps a standard FNO predictor without changing its architecture. • An error decomposition for symmetry-induced OOD shifts that separates canonical operator approximation from input and output alignment errors, showing when canonicalization helps (Theorem 3.1). • Empirical evaluations on 1-D and 2-D Burgers, shallow-water, and Navier-Stokes equations on tori. One-shot canonicalization reduces OOD error by up to 12× on shifts dominated by translation and by 3–9% on coupled rotation-translation shifts, while TTA gives an optional accuracy-latency trade-off. Related works. Existing approaches address symmetry-induced OOD generalization in three main ways, but each leaves part of the frame alignment to the predictor or to inference-time optimization: (1) Data augmentation. [19, 20, 21] exposes the model to transformed samples, but the learned operator still represents frame alignment and PDE evolution in one map, tying performance to the sampled range. PACE-FNO instead removes the estimated group action before prediction and restores the terminal frame afterward, so the FNO operates mainly on the canonical slice. (2) Equivariant architectures. [22, 12, 23, 24] enforce commutation by design and give strong inductive bias when the representation and discretization are fixed. They can also restrict architectural choices or require problem-specific group representations. PACE-FNO keeps the standard FNO architecture and imposes equivariance through input/output group actions, separating symmetry handling from the spectral 2
predictor. (3) Canonicalization methods. [15, 16, 17, 18] estimate a frame correction, predict in a standard frame, and restore the output frame. The closest prior work, LieLAC [14], adds an energy-based Lie-group canonicalizer to a pretrained model and solves a group optimization problem at inference, with optional fine-tuning. PACE-FNO learns the Lie-algebra coordinate estimator jointly with the operator, using PDE-specific terminal-frame correction.
2
Preliminaries
Operator learning. Let A and U be Banach spaces of input and output fields [4, 2, 25, 26, 27, 28]. The physical solution operator is G † : A → U, and Gθ is its learned finite-dimensional approximation. Empirical training minimizes ∥G † − Gθ ∥2L2µ (A;U ) = Ea∼µ ∥G † (a) − Gθ (a)∥2U , while approximation theory uses the compact-set metric supa∈K ∥G † (a) − Gθ (a)∥U . Out-of-distribution (OOD) construction. The training measure is supported near a canonical slice, not the full orbit of symmetry-equivalent states. Let µc be the canonical data distribution and PG a distribution over bounded group elements. The induced OOD measure is aOOD = Cg (ac ),
ac ∼ µc ,
g ∼ PG .
The target operator is unchanged up to the induced group action; only the input-output frame moves. Thus the OOD split tests frame generalization, not viscosity, forcing, solver, or boundary-condition variation. Lie groups and induced actions. Let G act smoothly on the spatial domain Ω [29], inducing [Cg a](x) = a(g −1 · x) on fields. Since G is generally nonlinear [30], we use Lie-algebra coordinates ξ ∈ g = Te G and recover g = exp(ξ). Pk For a basis {v1 , . . . , vk } of g, a generator expands as ξ = j=1 αj vj , so the estimator regresses the local coordinates αj before applying the exponential map. PACE-FNO uses this representation for translations, rotations, and problem-specific kinematic components. Symmetries of PDEs and neural operators. Many PDEs carry continuous symmetries: transformed solutions remain valid under G [10]. Equivariant operators enforce consistency under admissible coordinate changes [12, 31, 23]. The Fourier Neural Operator (FNO) [2] alternates local maps with spectral convolutions and is widely used fordiscretization-invariant operator approximation [2, 4, 32]: vt+1 (x) = σ W vt (x) + F −1 Rϕ · F (vt ) (x) . Here Rϕ is a learnable tensor on retained Fourier modes. Spectral truncation is efficient, but a transformed input can redistribute phase and high-frequency content relative to the canonical training frame, leaving the predictor outside its training regime. Equivariance in operator learning. For symmetry group G, predictions should satisfy Gθ (Cg a) = Cg (Gθ (a)). Hard-coding this constraint can reduce flexibility [12, 31]. PACE-FNO instead canonicalizes the input, applies the learned solution operator, and maps the output back [14, 15, 16, 17].
3
Methodology
In OOD regimes with frame changes, a standard neural operator must handle both frame alignment and physical evolution. PACE-FNO separates the two: a Lie-algebra estimator canonicalizes the input, a standard FNO solves the PDE on the canonical slice, and a terminal group action maps the output back. This design builds on recent work in Lie-equivariant learning, canonicalization, and test-time adaptation [12, 20, 14, 17]. Unlike post-hoc energy canonicalization such as LieLAC [14], PACE-FNO learns alignment jointly during training. One-shot inference uses one estimator pass, one inverse action, one FNO pass, and one output-frame action; iterative refinement is only residual correction. We omit a direct 3
(a) PACE-FNO Pipeline Input (OOD Field)
Lie-Alg. Coord. (Encoder) Predicts ξb0 C −1 b0 ) exp(ξ
aOOD (x)
1
FNO
Canonical State
P
(Lifting)
Fourier layer 1
Fourier layer 2
···
Fourier layer T
Canonical Output
Q
(Proj.)
ãc (x)
ũc (x)
Estimate & Canonicalize (Observation → Canonical)
2
FNO in Canonical Coordinates
Lie-Alg. Coord. (Decoder) Predicts ξbT Cexp(ξb ) T
Output (Prediction)
b pred (x) u
Decode & Reconstruct (Canonical → Observation)
3
(b) Fourier Layer R
F
(Spectral) Operator
(FFT)
Input
F −1 (IFFT)
σ
+
v(x)
(Activation)
(Sum)
W
(Linear) Branch
Data / Field
FNO / Operator
Lie-Alg. Coord. Spectral Domain
Neural Network
Activation / Sum Linear W Operation
Figure 2: PACE-FNO pipeline. (a) The Lie-algebra encoder predicts ξˆ0 , pulls the OOD input aOOD (x) to the canonical state ãc (x), applies the FNO to obtain ũc (x), and restores the final prediction ûpred (x). Dashed connectors indicate that panel (b) expands a representative FNO layer. (b) One Fourier layer with spectral branch, linear branch W, summation, and activation σ. LieLAC comparison because no public implementation is available and its setting does not match our PDE/OOD protocol (Appendix B.6). For each problem, G acts through a bounded subset GB ⊆ Td ⋊ SO(d) ⋊ Gal, where Gal (≡ Galilean) appears only for co-moving frames. We use T1 ⋊ Gal for 1-D Burgers, T2 ⋊ Gal for 2-D Burgers and SWE, and SE(2) = T2 ⋊ SO(2) for Navier-Stokes. When spatial and kinematic components are active, g = gS ⊕ gK . ID samples come from a canonical distribution µc with no macroscopic drift; OOD samples apply bounded group actions to canonical states. The solution operator stays fixed up to equivariance; only the orbit coordinate changes. Given OOD input aOOD , the estimator predicts the input-frame generator ξˆ0 = ξˆS,0 ⊕ ξˆK,0 . The neural estimator learns ξˆS,0 . When present, ξˆK,0 is computed from the observed field, e.g. by mean background velocity, rather than taken from the test generator. The output-frame generator follows ξˆT = ΓB (ξˆ0 ; T ), where ΓB is the identity for conserved rigid spatial actions and adds the analytically induced comoving translation for Galilean shifts. The overall forward map is −1 ûpred = Cexp(ξ̂T ) ◦ Gθ ◦ Cexp( (a ), ξ̂ ) OOD
(1)
0
where Gθ is the FNO predictor [2]. Eq. (1) is the one-shot rule. Ground-truth generators are used only to create perturbed training pairs and supervise alignment; at evaluation, aOOD alone determines ξˆ0 and ξˆT . Pullback and pushforward are spectral: translations use the shift theorem and rotations use three-shear decomposition [33, 34], avoiding interpolation dissipation. Appendix A gives details. Training uses bounded group perturbations, following restricted Lie-group augmentation [20]. The total loss combines three terms: Ltotal (θ, ϕ) = LPDE + λ1 Lequiv + λ2 Lanchor . 4
LPDE is a relative L2 field loss, optionally with a Fourier-domain H 1 term. Lequiv matches the estimated generator to the training perturbation, and Lanchor regresses Lie coordinates with a decayed weight to stabilize TTA. We balance terms with gradient-normalization heuristics [35, 36]; full definitions are in Appendix A. Algorithm 1 summarizes the training loop and the two inference modes. Algorithm 1 PACE-FNO training and inference Input: Training pairs (x, y) ∈ D; Gθ , Eϕ , group action C; weights λeq , λanc . Output: Trained (θ, ϕ) and one-shot/TTA prediction rules. 1: for each training step do 2: Sample (x, y) and a bounded training generator ∆0 ∈ g. 3: Derive the training target generator ∆T = ΓB (∆0 ; T ) and form x+ = Cexp(∆0 ) x, y + = Cexp(∆T ) y. 4: Estimate ξˆ0+ = Eϕ (x+ ) and derive ξˆT+ = ΓB (ξˆ0+ ; T ). 5: Canonicalize xc = C −1 + x+ and predict ûc = Gθ (xc ). exp(ξ̂0 )
6: Push forward ŷ = Cexp(ξ̂+ ) ûc . T 7: Update (θ, ϕ) with Lpde + λeq Leq + λanc Lanc . 8: end for −1 x). 9: One-shot: estimate ξˆ0 = Eϕ (x), set ξˆT = ΓB (ξˆ0 ; T ), and return Cexp(ξ̂T ) Gθ (Cexp( ξ̂0 ) ˆ 10: Optional TTA: freeze networks and refine only the spatial part of ξ0 by minimizing JTTA (ξˆ0 ) for K steps; recompute ξˆT = ΓB (ξˆ0 ; T ). 11: Return the same pushforward prediction with coordinates taken modulo the spatial period.
The next result separates canonical approximation from residual input and terminal alignment error under compactness, equivariance, and Lipschitz assumptions (Appendix E). Theorem 3.1 (Error Decomposition under Symmetry-Induced Shifts). Assume the hypotheses of Lemma E.1 (compact canonical support and continuous group action), Theorem E.2 (universal FNO approximation on the canonical compact set), and Lemmas E.3–E.4 (isometric/Lipschitz group action and Lipschitz learned operator). Then the prediction error of PACE-FNO over the bounded OOD orbit satisfies sup a∈AOOD
CĝT ◦ Gθ ◦ Cĝ−1 (a) − G † (a) 0
U
≤ ϵ + M0 ϵgeo,0 + ϵgeo,T ,
where ĝ0 = exp(ξˆ0 ), ĝT = exp(ξˆT ), and M0 = KC Lθ with KC the group-action Lipschitz constant and Lθ the FNO Lipschitz constant. Here ϵ is the canonical approximation error, while ϵgeo,0 and ϵgeo,T are residual input and terminal alignment errors. This is an analysis decomposition, not a distribution-free bound: for symmetry-induced OOD, PACEFNO reduces prediction to canonical-set approximation when residual alignment errors are small. The proof is in Appendix E. Orbit averaging projects onto an equivariant subspace without increasing covering numbers [15, 17, 18]; Appendix D gives the resulting bound. Theorem 3.2 (Symmetry-Induced Complexity Reduction for Canonicalized Operators). Let F be a class of neural operators from A to U, and let GB ⊂ G be a compact bounded symmetry group whose actions on A and U are bijective isometries. Let FG be the GB -equivariant projection of F under orbit averaging. Define d∞ (F1 , F2 ) := supa∈AOOD ∥F1 (a) − F2 (a)∥U . The projection is non-expansive; hence, for any r > 0, N (FG , r, d∞ ) ≤ N (F, r, d∞ ). Thus the covering-number bound for FG is no larger than that for the original class F. Approximate equivariance is measured by the worst-case commutation residual ηG (F ) :=
sup g∈GB a∈AOOD
∥(F ◦ Cg )(a) − (Cg ◦ F )(a)∥U ≤ η.
When the loss ℓ is Lℓ -Lipschitz and ηG (F ) ≤ η, the generalization gap of F exceeds that of its equivariant projection ΠG F by at most 2Lℓ η. 5
The one-shot estimate can have residual error near support boundaries or ambiguous canonical frames, consistent with spectral bias in neural networks [37, 38, 39]. Optional TTA refines only the lowdimensional input-frame generator during evaluation, with all network weights frozen [40, 41, 42]. For aOOD , it minimizes 2 −1 JTTA (ξˆ0 ) = PS ◦ Eϕ Cexp( (a ) g, (2) ξ̂ ) OOD 0
using AdamW over ξˆ0 , where PS selects spatial generator components and ∥ · ∥g is the Euclidean norm. If Eq. (2) removes the residual transformation, the pulled-back input should appear canonical to the estimator, up to symmetry-equivalent minima. Freezing network weights avoids catastrophic forgetting [43, 44]; the cost is slower inference, reported separately. Under local regularity, the continuous-time gradient flow of Eq. (2) decreases the energy and approaches the minima manifold, though finite-step AdamW may miss the best non-convex basin. Empirically, one-shot canonicalization accounts for most gains; TTA corrects residual mismatch. Proofs are in Appendix E.
4
Experiments
We evaluate PACE-FNO on Burgers, shallow-water (SWE), and Navier-Stokes (NS) equations on tori. Equations and cross-resolution protocols follow standard operator-learning evaluations [2, 4, 45, 46]; the symmetry-induced OOD splits are constructed here. Tables report ID and OOD relative errors. Each OOD input is aOOD = Cg (ac ) with held-out canonical field ac and g sampled beyond the training augmentation range. Baselines are FNO [2], FNO with augmentation (FNO+Aug) [19], and GFNO [31]; runtime and shift ranges are in Appendix B. 4.1
Burgers equations
The Burgers experiments isolate translation and Galilean drift. The 1-D equation is 1 ∂t u + ∂x (u2 ) = ν ∂xx u, 2 u(x, 0) = u0 (x), u(0, t) = u(1, t),
(x, t) ∈ T × (0, T ), x ∈ T, t ∈ (0, T ).
The corresponding OOD initial state is u0,OOD (x) = u0 (x − ∆x) + c. Here ∆x changes the spatial frame, while c induces the co-moving terminal correction. Under the group shift, FNO has high error (1.0352), augmentation reduces it to 0.3222, and one-shot PACE-FNO lowers it by about 12× to 0.0267. TTA gives only a marginal gain (0.0259) at much higher inference cost, and GFNO does not improve OOD error. Table 1: 1-D Burgers under cross-resolution group shift (resolution 1024 → 2048). Model FNO (no aug.) FNO+Aug PACE-FNO (one-shot) PACE-FNO + TTA GFNO
ID Err. ↓
OOD Err. ↓
Train (s/ep.)
OOD Inf. (s)
0.0662 ± 0.0028 0.0293 ± 0.0005 0.0623 ± 0.0015 0.0623 ± 0.0015 0.1406 ± 0.0064
1.0352 ± 0.0001 0.3222 ± 0.0077 0.0267 ± 0.0014 0.0259 ± 0.0007 1.0244 ± 0.0027
0.4212 ± 0.0095 0.7698 ± 0.0062 0.4560 ± 0.0030 0.4560 ± 0.0030 1.0365 ± 0.0411
0.3943 ± 0.0031 1.2826 ± 0.0340 0.5287 ± 0.0171 5.4611 ± 0.1030 1.0148 ± 0.0091
The 2-D variant uses 1 1 ∂t u + ∂x (u2 ) + ∂y (u2 ) = ν(∂xx u + ∂yy u), 2 2 u(x, y, 0) = u0 (x, y),
(x, y, t) ∈ T2 × (0, T ), (x, y) ∈ T2 .
with translated and boosted OOD inputs u0,OOD (x, y) = u0 (x − tx , y − ty ) + Ubg . Table 2 shows that one-shot PACE-FNO reduces OOD error by 2.7× relative to FNO+Aug (0.076 vs. 0.20). TTA further reduces error to 0.064 but adds 4.3 s over one-shot inference. 6
Table 2: 2-D Burgers under cross-resolution group shift (resolution 642 → 1282 ). ID Err. ↓
OOD Err. ↓
Train (s/ep.)
OOD Inf. (s)
0.0274 ± 0.0010 0.0261 ± 0.0009 0.0393 ± 0.0012 0.0393 ± 0.0012 0.0967 ± 0.0054
0.6137 ± 0.0362 0.2035 ± 0.0423 0.0755 ± 0.0139 0.0640 ± 0.0012 1.0131 ± 0.0119
0.6524 ± 0.0157 0.7559 ± 0.0177 1.5770 ± 0.0066 1.5770 ± 0.0066 4.8613 ± 0.3527
0.1244 ± 0.0118 0.1045 ± 0.0038 0.1399 ± 0.0023 4.4345 ± 0.0409 0.7462 ± 0.0080
Model FNO (no aug.) FNO+Aug PACE-FNO (one-shot) PACE-FNO + TTA GFNO
Figure 3: OOD prediction of PACE-FNO on 2-D Burgers under translation and Galilean drift (full comparison in Appendix Figure 11).
4.2
2-D shallow-water equations
The shallow-water system tests whether PACE-FNO separates wave evolution from transport in coupled conservative variables (h, m):
∂t m + ∇·
m⊗m h
∂t h + ∇·m = ν∇2 h, 1 2 + ∇ gh = ν∇2 m, 2 U(x, 0) = U0 (x),
(x, t) ∈ T2 × (0, T ), (x, t) ∈ T2 × (0, T ), x ∈ T2 .
Table 3: 2-D shallow water under cross-resolution group shift (resolution 642 → 1282 ). Model FNO (no aug.) FNO+Aug PACE-FNO (one-shot) PACE-FNO + TTA GFNO
ID Err. ↓
OOD Err. ↓
Train (s/ep.)
OOD Inf. (s)
0.0827 ± 0.0009 0.0769 ± 0.0010 0.0491 ± 0.0002 0.0491 ± 0.0002 0.0307 ± 0.0007
0.7260 ± 0.0003 0.5136 ± 0.0010 0.1069 ± 0.0168 0.1071 ± 0.0168 0.7593 ± 0.0071
0.7024 ± 0.0474 0.8502 ± 0.0082 2.6139 ± 0.0046 2.6139 ± 0.0046 6.7509 ± 0.0355
0.1501 ± 0.0068 0.1503 ± 0.0012 0.2952 ± 0.0349 8.1388 ± 0.0855 0.7161 ± 0.0401
OOD states apply sub-pixel translations and a Galilean background velocity. PACE-FNO estimates translation from the shifted height/momentum field and computes background velocity from conservative channels, reducing OOD error by 4.8× relative to FNO+Aug (0.107 vs. 0.514). TTA does not change the reported error at this precision. 7
Figure 4: OOD prediction of PACE-FNO on the 2-D shallow-water system under Galilean perturbation (full comparison in Appendix Figure 13). 4.3
(2+1)-D spatiotemporal Navier-Stokes equations
The Navier-Stokes experiment considers coupled translation and rotation shifts in spatiotemporal rollout: ∂t ω + u·∇ω = ν∇2 ω, (x, t) ∈ T2 × (0, Ttotal ), u = (∂y ψ, −∂x ψ),
∇2 ψ = −ω,
(x, t) ∈ T2 × (0, Ttotal ), x ∈ T2 .
ω(x, 0) = ω0 (x),
Table 4: Spatiotemporal Navier-Stokes under cross-resolution group shift (resolution 642 → 1282 ). Model FNO (no aug.) FNO+Aug PACE-FNO (one-shot) PACE-FNO + TTA GFNO
ID Err. ↓
OOD Err. ↓
Train (s/ep.)
OOD Inf. (s)
0.0377 ± 0.0043 0.0266 ± 0.0059 0.0129 ± 0.0006 0.0129 ± 0.0006 0.0459 ± 0.0016
0.6422 ± 0.0302 0.2840 ± 0.2001 0.2766 ± 0.0353 0.2593 ± 0.0522 0.4091 ± 0.0118
10.6448 ± 0.0392 10.7371 ± 0.0281 16.5571 ± 0.3631 16.5571 ± 0.3631 55.9359 ± 0.0496
1.8308 ± 0.6626 1.9761 ± 0.5997 3.3728 ± 0.1381 10.0376 ± 0.0418 9.2956 ± 0.0585
Figure 5: OOD prediction of PACE-FNO under SE(2) perturbations for (2 + 1)-D Navier-Stokes (full comparison in Appendix Figure 15). Each method observes Tin = 10 steps and predicts a T = 20 rollout. OOD samples apply SE(2) transformations to the whole trajectory, so the same rigid motion is removed and restored; here ξˆT = ξˆ0 . Table 4 shows that TTA reduces OOD error from 0.277 to 0.259 (6% relative) at 3× inference cost (10.0 s vs. 3.4 s). 4.4 Ablation and TTA sensitivity Appendix Table 10 isolates PACE-FNO components. Removing canonicalization eliminates most OOD gains; the full Lie generator matters most for 2-D Burgers, SWE, and NS. TTA is stable over tested step counts and learning rates (Figure 7 and Appendix Figure 16). Figure 6 shows a non-convex 2-D Burgers landscape with symmetry-induced low-energy basins, so the estimator need only produce an equivalent canonical field. 8
Figure 6: Empirical landscape of JTTA on 2-D Burgers. The landscape is non-convex with symmetryinduced low-energy basins, and descent from the estimator initialization converges to one valid basin.
Figure 7: TTA sensitivity for the (2 + 1)-D Navier-Stokes equation. The error reduction is stable across the tested learning rates and step counts.
5
Conclusions
PACE-FNO separates frame alignment from PDE evolution by estimating input-frame symmetry coordinates, applying a standard FNO on a canonical slice, and mapping the output back to the target frame. The error decomposition makes approximation and alignment terms explicit. Across four equations on tori, one-shot canonicalization reduces OOD error by up to 12× over FNO+Aug on translation-dominated shifts (Burgers, SWE) and by 3–9% on coupled rotation-translation shifts (NS), while matching ID accuracy. TTA gives a smaller correction at higher cost. Limitations. The symmetry group is specified a priori, so approximate, emergent, or unknown symmetries are outside the current scope. Spectral group actions assume fields on a torus; irregular domains would need graph or point-cloud transport instead of Fourier shifts and shear rotations [47, 48]. We also assume clean observations on regular grids; robustness to sensor noise, sparse measurements, and partial observability is not studied here. Future work. A natural next step is learning approximate symmetries from data rather than specifying them in advance. Extending canonicalization to irregular domains via graph or spherical neural operators would broaden the approach beyond periodic boundary conditions. Multi-task canonicalization across PDE families and discrete symmetries such as reflection equivariance are also open directions. Broader impact.
This work is foundational research. No negative social impact is identified. 9
References [1] George Em Karniadakis, Ioannis G. Kevrekidis, Lu Lu, Paris Perdikaris, Sifan Wang, and Liu Yang. Physics-informed machine learning. Nature Reviews Physics, 3(6):422–440, 2021. [2] Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Fourier neural operator for parametric partial differential equations. In International Conference on Learning Representations, 2021. [3] Lu Lu, Pengzhan Jin, and George Em Karniadakis. DeepONet: Learning nonlinear operators for identifying differential equations based on the universal approximation theorem of operators. Nature Machine Intelligence, 3:218–229, 2021. [4] Nikola Kovachki, Zongyi Li, Burigede Liu, Kamyar Azizzadenesheli, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Neural operator: Learning maps between function spaces with applications to PDEs. Journal of Machine Learning Research, 24(89):1–97, 2023. [5] Nicolas Boullé and Alex Townsend. A mathematical guide to operator learning. In Handbook of Numerical Analysis, volume 25, pages 83–125. Elsevier, 2024. doi: 10.1016/bs.hna.2024.05. 003. [6] Gege Wen, Zongyi Li, Kamyar Azizzadenesheli, Anima Anandkumar, and Sally M. Benson. U-FNO: An enhanced Fourier neural operator-based deep-learning model for multiphase flow. Advances in Water Resources, 163:104180, 2022. [7] Jaideep Pathak, Shashank Subramanian, Peter Harrington, Sanjeev Raja, Ashesh Chattopadhyay, Morteza Mardani, Thorsten Kurth, David Hall, Zongyi Li, Kamyar Azizzadenesheli, Pedram Hassanzadeh, Karthik Kashinath, and Anima Anandkumar. FourCastNet: A global datadriven high-resolution weather model using adaptive Fourier neural operators. arXiv preprint arXiv:2202.11214, 2022. [8] Johannes Brandstetter, Daniel Worrall, and Max Welling. Message passing neural PDE solvers. In International Conference on Learning Representations, 2022. [9] Phillip Lippe, Bastiaan Veeling, Paris Perdikaris, Richard Turner, and Johannes Brandstetter. PDE-refiner: Achieving accurate long rollouts with neural PDE solvers. In Advances in Neural Information Processing Systems, volume 36, pages 67398–67433, 2023. [10] Peter J. Olver. Applications of Lie Groups to Differential Equations, volume 107. Springer Science & Business Media, 1993. [11] Michael M. Bronstein, Joan Bruna, Taco Cohen, and Petar Veličković. Geometric deep learning: Grids, groups, graphs, geodesics, and gauges. arXiv preprint arXiv:2104.13478, 2021. [12] Marc Finzi, Samuel Stanton, Pavel Izmailov, and Andrew Gordon Wilson. Generalizing convolutional neural networks for equivariance to Lie groups on arbitrary continuous data. In International Conference on Machine Learning, pages 3165–3176. PMLR, 2020. [13] Jianwei Zheng, Wei Li, Ni Xu, Junwei Zhu, Xiaoxu Lin, and Xiaoqin Zhang. Alias-free mamba neural operator. In Advances in Neural Information Processing Systems, volume 37, pages 52962–52995, 2024. [14] Zakhar Shumaylov, Pavel Zaika, Jacob Rowbottom, Ferdia Sherry, Melanie Weber, and CarolaBibiane Schönlieb. Lie algebra canonicalization: Equivariant neural operators under arbitrary Lie groups. In International Conference on Learning Representations, 2025. [15] Omri Puny, Matan Atzmon, Heli Ben-Hamu, Ishan Misra, Aditya Grover, Edward J. Smith, and Yaron Lipman. Frame averaging for invariant and equivariant network design. In International Conference on Learning Representations, 2022. [16] Sékou-Oumar Kaba, Arnab Kumar Mondal, Yan Zhang, Yoshua Bengio, and Siamak Ravanbakhsh. Equivariance with learned canonicalization functions. In International Conference on Machine Learning, pages 15546–15566. PMLR, 2023. 10
[17] George Ma, Yifei Wang, Derek Lim, Stefanie Jegelka, and Yisen Wang. A canonicalization perspective on invariant and equivariant learning. In Advances in Neural Information Processing Systems, volume 37, 2024. [18] Nadav Dym, Hannah Lawrence, and Jonathan W. Siegel. Equivariant frames and the impossibility of continuous canonicalization. arXiv preprint arXiv:2402.16077, 2024. [19] Rui Wang, Robin Walters, and Rose Yu. Incorporating symmetry into deep dynamics models for improved generalization. arXiv preprint arXiv:2002.03061, 2020. [20] Johannes Brandstetter, Max Welling, and Daniel E. Worrall. Lie point symmetry data augmentation for neural PDE solvers. In International Conference on Machine Learning, pages 2241–2256. PMLR, 2022. [21] Tara Akhound-Sadegh, Laurence Perreault-Levasseur, Johannes Brandstetter, Max Welling, and Siamak Ravanbakhsh. Lie point symmetry and physics-informed networks. In Advances in Neural Information Processing Systems, volume 36, pages 42468–42481, 2023. [22] Taco Cohen and Max Welling. Group equivariant convolutional networks. In International Conference on Machine Learning, pages 2990–2999. PMLR, 2016. [23] Bart M. N. Smets, Jim Portegies, Erik J. Bekkers, and Remco Duits. PDE-based group equivariant convolutional neural networks. Journal of Mathematical Imaging and Vision, 65: 209–239, 2023. [24] Tobias Pfaff, Meire Fortunato, Alvaro Sanchez-Gonzalez, and Peter Battaglia. Learning meshbased simulation with graph networks. In International Conference on Learning Representations, 2021. [25] Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Andrew Stuart, Kaushik Bhattacharya, and Anima Anandkumar. Multipole graph neural operator for parametric partial differential equations. In Advances in Neural Information Processing Systems, volume 33, pages 6755–6766, 2020. [26] Shuhao Cao. Choose a transformer: Fourier or Galerkin. In Advances in Neural Information Processing Systems, volume 34, pages 15437–15450, 2021. [27] Tapas Tripura and Souvik Chakraborty. Wavelet neural operator for solving parametric partial differential equations in computational mechanics problems. Computer Methods in Applied Mechanics and Engineering, 404:115783, 2023. [28] Qianying Cao, Somdatta Goswami, and George Em Karniadakis. Laplace neural operator for solving differential equations. Nature Machine Intelligence, 6:631–640, 2024. [29] Alexander A. Kirillov. An Introduction to Lie Groups and Lie Algebras, volume 113. Cambridge University Press, 2008. [30] John M. Lee. Smooth manifolds. In Introduction to Smooth Manifolds, pages 1–29. Springer, 2003. [31] Jacob Helwig, Xuan Zhang, Cong Fu, Jerry Kurtin, Stephan Wojtowytsch, and Shuiwang Ji. Group equivariant Fourier neural operators for partial differential equations. In International Conference on Machine Learning, pages 13376–13411. PMLR, 2023. [32] Alasdair Tran, Alexander Mathews, Lexing Xie, and Cheng Soon Ong. Factorized Fourier neural operators. In International Conference on Learning Representations, 2023. [33] Alan W. Paeth. A fast algorithm for general raster rotation. In Graphics Gems, pages 179–195. Academic Press, 1990. [34] Michael Unser, Philippe Thevenaz, and Leonid Yaroslavsky. Convolution-based interpolation for fast, high-quality rotation of images. IEEE Transactions on Image Processing, 4(10): 1371–1381, 1995. 11
[35] Sifan Wang, Yujun Teng, and Paris Perdikaris. Understanding and mitigating gradient flow pathologies in physics-informed neural networks. SIAM Journal on Scientific Computing, 43 (5):A3055–A3081, 2021. [36] Zhao Chen, Vijay Badrinarayanan, Chen-Yu Lee, and Andrew Rabinovich. GradNorm: Gradient normalization for adaptive loss balancing in deep multitask networks. In International Conference on Machine Learning, pages 794–803. PMLR, 2018. [37] Nasim Rahaman, Aristide Baratin, Devansh Arpit, Felix Draxler, Min Lin, Fred A. Hamprecht, Yoshua Bengio, and Aaron Courville. On the spectral bias of neural networks. In International Conference on Machine Learning, pages 5301–5310. PMLR, 2019. [38] Zhi-Qin John Xu, Yaoyu Zhang, Tao Luo, Yanyang Xiao, and Zheng Ma. Frequency principle: Fourier analysis sheds light on deep neural networks. Communications in Computational Physics, 28(5):1746–1767, 2020. [39] Keyulu Xu, Mozhi Zhang, Jingling Li, Simon S. Du, Ken-ichi Kawarabayashi, and Stefanie Jegelka. How neural networks extrapolate: From feedforward to graph neural networks. In International Conference on Learning Representations, 2021. [40] Dequan Wang, Evan Shelhamer, Shaoteng Liu, Bruno Olshausen, and Trevor Darrell. Tent: Fully test-time adaptation by entropy minimization. In International Conference on Learning Representations, 2021. [41] Yu Sun, Xiaolong Wang, Zhuang Liu, John Miller, Alexei A. Efros, and Moritz Hardt. Test-time training with self-supervision for generalization under distribution shifts. In International Conference on Machine Learning, pages 9229–9248. PMLR, 2020. [42] Samarth Sinha, Peter Gehler, Francesco Locatello, and Bernt Schiele. TeST: Test-time selftraining under distribution shift. In IEEE/CVF Winter Conference on Applications of Computer Vision, 2022. [43] James Kirkpatrick, Razvan Pascanu, Neil Rabinowitz, Joel Veness, Guillaume Desjardins, Andrei A. Rusu, Kieran Milan, John Quan, Tiago Ramalho, Agnieszka Grabska-Barwinska, Demis Hassabis, Claudia Clopath, Dharshan Kumaran, and Raia Hadsell. Overcoming catastrophic forgetting in neural networks. Proceedings of the National Academy of Sciences, 114(13): 3521–3526, 2017. [44] Robert M. French. Catastrophic forgetting in connectionist networks. Trends in Cognitive Sciences, 3(4):128–135, 1999. [45] Makoto Takamoto, Timothy Praditia, Raphael Leiteritz, Dan MacKinlay, Francesco Alesiani, Dirk Pflüger, and Mathias Niepert. PDEBench: An extensive benchmark for scientific machine learning. In Advances in Neural Information Processing Systems, volume 35, 2022. [46] Zongyi Li, Miguel Liu-Schiaffini, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Learning chaotic dynamics in dissipative systems. In Advances in Neural Information Processing Systems, volume 35, pages 16768–16781, 2022. [47] Zongyi Li, Daniel Zhengyu Huang, Burigede Liu, and Anima Anandkumar. Fourier neural operator with learned deformations for PDEs on general geometries. Journal of Machine Learning Research, 24(388):1–26, 2023. [48] Boris Bonev, Thorsten Kurth, Christian Hundt, Jaideep Pathak, Maximilian Baust, Karthik Kashinath, and Anima Anandkumar. Spherical Fourier neural operators: Learning stable dynamics on the sphere. In International Conference on Machine Learning, pages 2806–2823. PMLR, 2023. [49] Stephen B. Pope. Turbulent Flows. Cambridge University Press, 2000. [50] Arnold Sommerfeld. Mechanics: Lectures on Theoretical Physics, volume 1. Academic Press, 1952. 12
[51] Luca Falorsi, Pim de Haan, Tim R. Davidson, Nicola De Cao, Maurice Weiler, Patrick Forré, and Taco S. Cohen. Explorations in homeomorphic variational auto-encoding. arXiv preprint arXiv:1807.04689, 2018. [52] Ronald N. Bracewell. The Fourier Transform and Its Applications. McGraw-Hill, New York, 1986. [53] Hamed Karimi, Julie Nutini, and Mark Schmidt. Linear convergence of gradient and proximalgradient methods under the Polyak-Łojasiewicz condition. In Joint European Conference on Machine Learning and Knowledge Discovery in Databases, pages 795–811. Springer, 2016.
13
A
Additional methodological details
This section gives the definitions and training details used in the main text. Let the problem-dependent continuous transformation group be denoted by GB ⊆ Td ⋊ SO(d) ⋊ Gal, where inactive factors are dropped for a given PDE. When both spatial and kinematic components are present, the corresponding Lie algebra admits the decomposition g = gS ⊕ gK . Kinematic reduction. Let v ∈ L1 (D; Rd ) denote the velocity field. The kinematic generator v bg ∈ gK and the centering operator Sv : A → A are defined by Z 1 v bg = Pphys (a) := v(x) dx, Sv (a) = a − v bg . |D| D This co-moving decomposition matches the usual interpretation of background transport and momentum removal [49, 50]. Neural parameterization. To parameterize the geometric subalgebra gS without singularities, we use a neural mapping Mϕ : A → R2k , followed by a smooth retraction Ret : R2k → gS , following the general idea of continuous latent parameterizations for rotational variables [51, 14]. The spatial generator is written as ξS = Ret(Mϕ (Sv (a))) . The full Lie-algebra coordinate estimate is then Eϕ (a) = ξS ⊕ v bg . Spectral group action. Given the full Lie algebra coordinate ξˆ ∈ g, we obtain the corresponding ˆ To avoid dissipation induced by interpolation, the group action group element through ĝ = exp(ξ). is implemented spectrally. For translations, the continuous Fourier shift theorem gives [52] F(Tτ trans ◦ a)(k) = exp − 2πi⟨k, τ trans ⟩ â(k). For rotations, we use a three-shear factorization in the spectral domain [33, 34]: F(Rθ ◦ a) = H̃x (α) ◦ H̃y (β) ◦ H̃x (α) F(a). Training-time perturbations. To expose the learned operator to controlled geometric shifts, we sample bounded perturbation generators ξ0 from a restricted training support Ωtrain ⊊ ΩOOD and generate aaug = Cexp(ξ0 ) (a), for canonical samples a ∼ µc , following the logic of restricted Lie-group augmentation [20]. The same sampled generator defines the perturbed training target through the kinematic map ξT = ΓB (ξ0 ; T ). These sampled values are used only to construct training pairs and alignment losses; PACE-FNO does not receive them during in-distribution (ID) or out-of-distribution (OOD) evaluation. Source of input and terminal generators. The estimator and the known kinematic relations play different roles. At inference, PACE-FNO receives only the observed field. The estimator infers the spatial frame, analytic reductions compute any kinematic background from the observed physical channels, and the terminal frame follows from the PDE symmetry relation. Table 5 summarizes this separation. 14
Table 5: Source of the generators used by PACE-FNO. Learned quantities come from Eϕ at inference; analytic quantities come from the observed field or the symmetry-induced transport rule. Sampled generators are used only for training perturbations and alignment losses. Test case
Learned at inference
Analytic at inference
Training-only sampled generator
1-D Burgers
Input translation phase δ̂x
2-D Burgers
Input translation t̂ = (t̂x , t̂y )
2-D SWE
ˆ Input translation ∆x
δx , vbg construct augmented pairs and alignment losses t, Ubg construct augmented pairs and alignment losses ∆x, Vbg construct boosted momentum and alignment losses
(2+1)-D NS
Input SE(2) frame (θ̂, t̂x , t̂y )
Background offset v̂bg from the field mean; terminal shift δ̂x,T = δ̂x + v̂bg T Scalar background Ûbg from the field mean; terminal translation t̂T = t̂ + (Ûbg T, Ûbg T ) Background velocity V̂bg from conservative channels, using the spatial mean of m/h; terminal ˆ T = ∆x ˆ + V̂bg T translation ∆x Rigid spatial symmetry is conserved, so ξˆT = ξˆ0
Detailed losses.
The total training objective is
(θ, tx , ty ) rotate/translate training trajectories and supervise alignment losses
Ltotal (θ, ϕ) = LPDE + λ1 Lequiv + λ2 Lanchor . The physical term matches the perturbed target in both the field and gradient domains: ∥ûpred − uaug ∥L2 ∥|K| ⊙ (F (ûpred ) − F (uaug ))∥L2 LPDE = E +γ . ∥uaug ∥L2 ∥|K| ⊙ F (uaug )∥L2 The equivariance regularizer is Lequiv = E ∥M(ξˆ0 ) − M(ξ0 )∥2F , and the anchor loss is Lanchor = E ∥ξˆ0 − ξ0 ∥2g . We balance λ1 and λ2 so that the geometric terms stay comparable to the physical gradient, following gradient-balancing heuristics [35, 36]. The anchor weight is gradually decayed during training.
B
Additional experimental details and results
We report the problem construction and additional results here. All problem settings follow a crossresolution operator-learning protocol: methods are trained on canonical low-resolution data and evaluated on geometrically shifted OOD data at inference. The equations are standard in operatorlearning evaluation [2, 4, 45]; the symmetry-induced OOD splits are new to this work. B.1
Implementation and reproducibility details
All reported numbers use three random seeds, {42, 1234, 3407}, with Python, NumPy, and PyTorch randomness reset before constructing each data loader. Tables report mean and standard deviation across seeds. Runtime entries are wall-clock training time per epoch and one full OOD evaluation pass, including iterative updates when TTA is enabled. All experiments ran on a single NVIDIA GeForce RTX 4080 GPU with 32GB memory, driver 580.105.08, and CUDA 13.0. Total training compute can be estimated from the listed epochs, per-epoch time, and three seeds. The FNO+Aug baselines use the same FNO architecture and optimizer settings as the corresponding FNO baselines and differ only in the training/evaluation data construction. The implementation and reproduction scripts are available at https://github.com/emoixiao/ PACE-FNO. 15
Code and asset provenance. The supplementary package contains author-written training, evaluation, and plotting code for the experiments in this paper, together with configuration files and figure assets generated from those runs. We do not redistribute third-party datasets, pretrained checkpoints, or external codebases as part of the public release. External software dependencies and benchmark sources remain under their respective original licenses and terms of use; these dependencies are cited in the paper and used only to the extent required for reproduction. Table 6: Core training hyperparameters used by the experimental scripts. “Batch” denotes train/OODevaluation batch size. For GFNO on NS we use time modes 10 in addition to spatial modes 10. For SWE, the FNO-model weight decay is 10−4 for FNO+Aug/PACE-FNO and 10−3 for FNO without augmentation and GFNO, matching the experiment configuration. Test case
Models
Train/Test
ID Grid
OOD Grid
Batch
Epochs
Optimizer, Scheduler, Architecture
1-D Burgers 2-D Burgers 2-D SWE (2+1)-D NS
FNO, FNO+Aug, GFNO, PACE-FNO FNO, FNO+Aug, GFNO, PACE-FNO FNO, FNO+Aug, GFNO, PACE-FNO FNO, FNO+Aug, GFNO, PACE-FNO
1000/200 1000/200 1000/200 1000/200
1024 642 642 642
2048 1282 1282 1282
20/1 20/10 20/10 10/10
500 500 500 100
Adam, lr 10−3 , wd 10−4 ; StepLR (50, 0.5); modes 20, width 64 Adam, lr 10−3 , wd 10−4 ; StepLR (100, 0.5); modes 12, width 32 Adam, lr 10−3 ; StepLR (100, 0.2); modes 16, width 32 Adam, lr 2.5×10−3 , wd 10−4 ; StepLR (100, 0.5); spatial modes 10, width 20
Table 7: PACE-FNO specific losses and TTA settings. One-shot PACE-FNO uses the same trained checkpoint and sets the number of TTA steps to zero. The Lie-algebra coordinate estimator parameter group uses the same learning rate as the FNO predictor and no additional weight decay unless explicitly listed. Test case
PACE-FNO Training Loss
TTA Variables and Optimizer
TTA Steps / LR
1-D Burgers 2-D Burgers 2-D SWE (2+1)-D NS
Split L2 /H 1 loss with EMA H 1 weight (α = 0.9) plus 10 Lequiv L2 + 5LH 1 + 1Lequiv + 5Lanchor Channel-weighted PDE loss plus scheduled H 1 , 10Lequiv , 5Lanchor , and 0.5 invariant-height loss L2 + 5LH 1 + 1Lequiv + λanc (e)Lanchor , λanc (e) = 20(0.85)e + 0.1
Translation phase, Adam (tx , ty ), Adam (tx , ty ), AdamW with wd 10−4 (θ, tx , ty ), Adam
20 / 0.05 50 / 0.01 80 / 0.05 5 / 0.05
B.2
1-D Burgers equation
This appendix only records the data-generation and adaptation details not shown in the main text. Canonical ID initial fields are sampled from a Gaussian random field, while OOD samples are generated by u0,OOD (x) = u0 (x − ∆x) + c, which combines spatial translation ∆x and a Galilean velocity offset c. This follows the restricted Lie-group augmentation framework [20]. Data Augmentation via Lie Group Actions. To test how much standard neural operators can learn from transformed examples alone, we sample measure-preserving perturbations from the Galilean group G = T1 ⋊ Gal. For a canonical initial field u0 (x), the augmented state at t = 0 is generated by the pushforward action of g0 = exp(ξ0 ): uaug (x, 0) = u0 (x − δx ) + vbg where the initial Lie algebra generator ξ0 ∈ g is defined by the superposition of the spatial translation and velocity boost differential operators: ∂ ξ0 = −δx + vbg Iu ∂x Here, δx ∼ U(−0.1, 0.1) parameterizes the spatial translation generator within the geometric subalgebra gS , Iu denotes the identity shift operator on the codomain, and vbg ∼ U(−0.2, 0.2) represents the Galilean kinematic velocity boost within gK . To maintain Galilean invariance during physical evolution, the target state at time T is shifted by the induced kinematic convection. The target group element evolves analytically as gT = ∂ exp(ξ0 + ∆x · ∂x ), where the spatial translation increment is ∆x = vbg T . Unsupervised Energy Functional via Canonical Self-Consistency. After canonicalization, the trained estimator Eϕ should produce a near-zero residual. We define the residual phase shift as −1 ϵres = PS ◦ Eϕ Cexp( (a ) , where PS projects onto the spatial translation subspace. The TTA OOD ξ̂S ) loss is the batch mean squared residual: B B 2 1 X (i) 2 1 X (i) −1 ˆ JTTA (ξS ) = ϵ = PS ◦ Eϕ C ) (i) (a exp(ξ̂S ) OOD B i=1 res B i=1 16
where B denotes the batch size.
Figure 8: Canonicalization on 1-D Burgers. The physical panels show the translated and boosted input being pulled back toward the canonical frame; the latent panels show that the estimator residual is small after alignment.
17
(a) PACE-FNO.
(b) FNO.
(c) FNO+Aug.
(d) GFNO.
Figure 9: OOD prediction on 1-D Burgers under translation and Galilean drift. PACE-FNO aligns the wave location and amplitude with the target, while the baselines show frame mismatch or phase error under the same shifted input.
18
B.3
2-D Burgers equation
For the two-dimensional Burgers setting, canonical states are centered and aligned, while OOD samples follow the same form as in the main text, u0,OOD (x, y) = u0 (x − tx , y − ty ) + Ubg , with translations (tx , ty ) and background flow Ubg . Data Augmentation via Two-Dimensional Lie Group Actions. For the 2-D setting, we sample measure-preserving perturbations from the Galilean group G = T2 ⋊ Gal. The corresponding initial Lie algebra generator ξ0 ∈ g is defined by the superposition of the two-dimensional spatial translation and velocity boost differential operators: ξ0 = −tx
∂ ∂ − ty + Ubg Iu ∂x ∂y
Here, the translation vector t = (tx , ty ) ∼ U(−0.1, 0.1)2 parameterizes the spatial generators within the geometric subalgebra gS , Iu denotes the identity shift operator on the codomain, and Ubg ∼ U(−0.5, 0.5) represents the Galilean velocity boost within gK . These shifts violate the mass and amplitude constraints of the in-distribution dataset. To preserve Galilean invariance during spatiotemporal evolution, the target state at time T undergoes kinematic convection. The target group element evolves analytically as gT = exp ξ0 +∆x·∇ , where the translation increment vector is driven by the background flow such that ∆x = (Ubg T, Ubg T ). Unsupervised Energy Functional via Two-Dimensional Canonical Self-Consistency. The estimator applied to the canonicalized state should return near-zero residual shifts when alignment succeeds. We define the residual two-dimensional phase shift vector as ϵres = (ϵres,x , ϵres,y ). The adaptation objective during evaluation is the empirical mean squared error of these two-dimensional phase residuals over a given batch: B
JTTA (ξˆS ) =
1 X (i) 2 2 (ϵres,x ) + (ϵ(i) ) res,y B i=1
where B denotes the batch size.
19
Figure 10: Canonicalization on 2-D Burgers. The learned pullback removes the dominant spatial shift and background-flow component before the FNO rollout, leaving a canonical field whose latent residual is close to the ID frame.
20
(a) PACE-FNO.
(b) FNO.
(c) FNO+Aug.
(d) GFNO.
Figure 11: OOD prediction on 2-D Burgers under translation and Galilean drift. PACE-FNO more accurately preserves the location and shape of the advected structures, while FNO, FNO+Aug, and GFNO retain visible frame errors under the same OOD shift.
21
B.4
2-D shallow-water equations
For the shallow-water setting, canonical conditions enforce zero macroscopic background velocity and centered mass distribution. OOD states are generated by U0,OOD (x) = BVbg U0 (x − ∆x), which combines sub-pixel translation and a Galilean background flow. For the conservative state U = [h, m]⊤ , the boost operator leaves height unchanged and shifts momentum by the transported background velocity, BV [h, m]⊤ = [h, m + hV]⊤ . Data Augmentation via Two-Dimensional Lie Group Actions. For shallow-water dynamics, we use the same Galilean group G = T2 ⋊ Gal to perturb the initial vector field. The pushforward action of g0 = exp(ξ0 ) is the boost-and-translation map above, with initial Lie algebra generator ξ0 = −tx
∂ ∂ − ty + Vbg Im ∂x ∂y
Here, the translation vector ∆x = (tx , ty ) ∼ U (−0.15, 0.15)2 parameterizes the spatial generators within the geometric subalgebra gS . To reduce interpolation-induced dissipation of high-frequency surface tension ripples, these continuous sub-pixel translations are executed analytically in the Fourier domain. Im denotes the boost direction acting on the momentum channels, and Vbg = (ubg , vbg ) ∼ U (−0.2, 0.2)2 represents the Galilean background flow within gK . These shifts violate the mass-centering and amplitude constraints of the in-distribution dataset. To maintain Galilean invariance during spatiotemporal evolution, the target state at time T undergoes kinematic convection. The target group element evolves analytically as gT = exp ξ0 + ∆xT · ∇ , where the translation increment vector is dynamically driven by the background flow such that ∆xT = Vbg T . Unsupervised Energy Functional via Two-Dimensional Canonical Self-Consistency. After canonicalization, the estimator should produce near-zero residual shifts. We define the residual two-dimensional phase shift vector as ϵres = (ϵres,x , ϵres,y ). The TTA objective is the empirical mean squared error over a batch: B 1 X (i) 2 2 JTTA (ξˆS ) = (ϵres,x ) + (ϵ(i) res,y ) B i=1
where B denotes the batch size.
22
Figure 12: Canonicalization behavior of PACE-FNO under Galilean perturbations for the 2-D shallowwater system. The aligned states remove the dominant translation and background transport before prediction, which explains the smaller OOD error in Table 3.
23
(a) PACE-FNO.
(b) FNO.
(c) FNO+Aug.
(d) GFNO.
Figure 13: OOD prediction on the 2-D shallow-water system under Galilean perturbation. PACE-FNO more accurately tracks the transported wave pattern, while the baselines show larger displacement and amplitude errors.
24
B.5
(2+1)-D Navier-Stokes equations
This setting evaluates spatiotemporal rollouts under coupled translation and rotation shifts. Models observe an input window of Tin = 10 steps and predict a future rollout of T = 20 steps. The OOD trajectory is generated by applying a single composite SE(2) action consistently across the trajectory; at t = 0 this gives ω0,OOD (x) = ω0 R−θ (x − ∆x) , which applies a translation and a spatial rotation to the canonical initial condition. Data Augmentation via SE(2) Lie Group Actions. For the Navier-Stokes task, we sample composite perturbations from the special Euclidean group G = SE(2) = T (2) ⋊ SO(2). The corresponding initial Lie algebra generator ξ0 ∈ se(2) is defined by the superposition of the twodimensional spatial translation and rigid rotation differential operators: ∂ ∂ ∂ ∂ ξ0 = −tx − ty −θ x −y ∂x ∂y ∂y ∂x Here, the translation vector ∆x = (tx , ty ) ∼ U (−0.05, 0.05)2 parameterizes the continuous T (2) spatial translations within the geometric subalgebra, and the angle θ ∼ U (−π/4, π/4) parameterizes the SO(2) rotations. Unlike kinematic background flows, these rigid geometric perturbations do not accumulate a macroscopic advection increment over time. On the discretized torus, arbitrary rotations should be read as the controlled rotated-field construction used in this experiment, rather than as a claim that every rotation angle is an exact automorphism of the square torus. Unsupervised Energy Functional via SE(2) Canonical Self-Consistency. After canonicalization, the estimator should return little residual translation or rotation. We write the residual translation as tres = (tx,res , ty,res ) and the residual rotation as θres . To avoid phase-wrapping discontinuities and singular gradients near the coordinate identifications of T2 × SO(2), the TTA loss uses a cosine distance rather than mean squared error: B
JTTA (ξˆS ) =
1 X (k) (k) 1 − cos(θres ) + 1 − cos(2πt(k) x,res ) + 1 − cos(2πty,res ) B k=1
This formulation maps the residual vectors onto a smooth hypersphere and avoids artificial jumps caused by angle wrapping, such as the discontinuity from −π to π.
25
Figure 14: Canonicalization on the (2 + 1)-dimensional Navier-Stokes setting. The learned frame correction reduces the imposed translation-rotation shift before rollout, so the latent representation is closer to the canonical training regime.
26
(a) PACE-FNO.
(b) FNO.
(c) FNO+Aug.
(d) GFNO.
Figure 15: OOD prediction under SE(2) perturbations for the (2 + 1)-dimensional Navier-Stokes equations. PACE-FNO more accurately preserves vortex placement and rollout geometry; the remaining gap reflects the coupled rotation-translation alignment error.
27
B.6
Comparison with other methods
Scope of the LieLAC comparison. LieLAC [14] is the closest prior work in spirit, and we discuss it in the main text as a post-hoc Lie-algebra canonicalization method. We do not report a head-to-head LieLAC number in the tables. At the time of writing, we could not find a public implementation or released PDE datasets for the published version. The paper also does not specify enough details of the data-generation pipeline for us to reproduce its protocol on our Burgers, SWE, and NavierStokes settings, including the train/test splits, solver choices, and symmetry ranges. Reporting our own reimplementation would therefore require several judgment calls and could be unfair to their method. We instead compare against baselines whose code paths and data protocols we can reproduce locally, and leave a direct empirical comparison to future work once an implementation or an exact experimental configuration is available.
Table 8: In-distribution autoregressive NMSE on the KdV equation under the LPSDA protocol [20], evaluated at 512 training samples. Lower NMSE indicates better accuracy. Number of samples Task
Solver
KdV (20s) FNO (AR)
Sym.
Architecture
−
g1 g2 g3 g4 windowed g1 KdV (20s) PACE-FNO (AR) + structured g2 –g4 temporal-window g1 + PACE g2 –g4 KdV (20s) PACE-FNO (AR)
KdV (20s) FNO (AR)
512
m=32, w=256 5 layers, 20 frames 0.0130 ± 0.0026 m=32, w=256 5 layers, 20 frames 0.0053 ± 0.0012 m=20, w=64 4 layers, 1 frame 0.0047 ± 0.0001 m=32, w=256 5 layers, 20 frames 0.0061 ± 0.0009
Table 8 compares PACE-FNO with locally reproduced LPSDA baselines on the KdV equation. PACE-FNO achieves the lowest NMSE at 512 training samples. The best PACE-FNO variant uses a single input frame with a compact FNO predictor (modes m = 20, width w = 64, four Fourier layers), while the strongest LPSDA-augmented FNO baseline uses a 20-frame history window and a larger FNO predictor (modes m = 32, width w = 256, five layers). The difference is consistent with structured canonicalization, rather than parameter count, being the source of the improvement.
Table 9: In-distribution autoregressive NMSE on the Kuramoto-Sivashinsky equation under the LPSDA protocol [20], evaluated at 512 training samples. Lower NMSE indicates better accuracy. Number of samples Task
Solver
KS (20s) FNO (AR)
Sym.
Architecture
−
g1 g2 g3 windowed g1 KS (20s) PACE-FNO (AR) + structured g2 –g3 windowed g1 KS (20s) PACE-FNO (AR) + PACE g2 –g3 KS (20s) FNO (AR)
512
m=32, w=256 5 layers, 20 frames 0.0010 ± 0.0002 m=32, w=256 5 layers, 20 frames 0.0009 ± 0.0001 m=20, w=64 4 layers, 1 frame 0.0008 ± 0.0002 m=20, w=64 4 layers, 20 frames 0.0003 ± 0.0001
Table 9 repeats the comparison on the KS equation with the same locally reproduced LPSDA baselines. All methods already operate at very low error in this in-distribution setting, leaving limited room for improvement. PACE-FNO still achieves the lowest NMSE with a smaller FNO predictor: a singleframe architecture with four Fourier layers at width 64 matches or surpasses the 20-frame, five-layer, width-256 LPSDA baselines. As in KdV, the comparison is consistent with canonicalization, rather than parameter count, explaining the improvement. 28
Table 10: Unified ablation study of PACE-FNO. The w/o Canonicalization variant disables learned pullback/pushforward during inference. For SWE, w/o Full Lie Gen. is numerically unstable (unclipped error reported). Dashes indicate entries with no independent training cost or evaluation. Dataset
Relative Error
Model
Runtime (s)
ID ↓
OOD ↓
Train (s/ep.)
OOD Inf. (s)
0.0621 ± 0.0028 0.0471 ± 0.0027 − −
0.0260 ± 0.0012 0.0193 ± 0.0007 0.0267 ± 0.0014 1.0384 ± 0.0003
0.4864 ± 0.0037 0.4902 ± 0.0034 − −
5.9590 ± 0.0963 6.0546 ± 0.2355 0.5287 ± 0.0171 0.4016 ± 0.0058
PACE-FNO + TTA 0.0393 ± 0.0012 w/o Structured Spatial Gen. 0.0392 ± 0.0008 0.0303 ± 0.0018 2-D Burgers w/o Full Lie Gen. PACE-FNO (one-shot) − w/o Canonicalization −
0.0640 ± 0.0012 0.0666 ± 0.0076 0.5478 ± 0.0303 0.0755 ± 0.0139 1.0085 ± 0.0035
1.5770 ± 0.0066 1.0761 ± 0.0268 1.1030 ± 0.0083 − −
4.4345 ± 0.0409 4.3969 ± 0.0336 4.3551 ± 0.0131 0.1399 ± 0.0023 0.1417 ± 0.0311
2.6139 ± 0.0046 2.5280 ± 0.0059 2.0817 ± 0.0222 − −
8.1388 ± 0.0855 7.8222 ± 0.0129 7.5391 ± 0.0296 0.2952 ± 0.0349 0.1426 ± 0.0260
1D Burgers
2-D SWE
(2+1)D NS
C
PACE-FNO + TTA w/o Full Lie Gen. PACE-FNO (one-shot) w/o Canonicalization
PACE-FNO + TTA 0.0491 ± 0.0002 0.1071 ± 0.0168 w/o Structured Spatial Gen. 0.0501 ± 0.0011 0.1008 ± 0.0142 w/o Full Lie Gen. 0.0586 ± 0.0058 2.11×104 ± 2.28×104 PACE-FNO (one-shot) − 0.1069 ± 0.0168 w/o Canonicalization − 0.7232 ± 0.0008 PACE-FNO + TTA w/o Full Lie Gen. PACE-FNO (one-shot)
0.0129 ± 0.0006 0.0133 ± 0.0006 −
0.2593 ± 0.0522 0.6280 ± 0.0076 0.2766 ± 0.0353
16.5571 ± 0.3631 10.0376 ± 0.0418 11.0809 ± 0.0649 2.7494 ± 0.0169 − 3.3728 ± 0.1381
Ablation and TTA sensitivity
The ablation above separates the one-shot canonicalizer from iterative test-time refinement. Figure 16 reports the sensitivity of the final OOD error to the tested TTA learning rates and adaptation-step budgets on the 2-D shallow-water system; the OOD relative error is nearly flat across the tested learning rates and adaptation steps.
Figure 16: TTA sensitivity on the 2-D shallow-water system. The OOD relative error is nearly flat across the tested learning rates and adaptation steps.
D
Generalization Error and Covering Numbers
We formalize the connection between covering numbers and generalization for operator learning. These definitions support the symmetry-induced complexity reduction in Theorem 3.2 and the subsequent remark on improved generalization. Let ρ denote the data distribution over A × U, and let ℓ : U × U → [0, B] be a bounded loss function. For a learned operator F ∈ F , the population risk and empirical risk are N X b N (F ) = 1 R ℓ(F (ai ), ui ). N i=1
R(F ) = E(a,u)∼ρ ℓ(F (a), u),
29
Definition D.1 (Generalization Error). The generalization error of an operator F ∈ F with respect to N i.i.d. samples is b N (F ). GenN (F ) := R(F ) − R The uniform generalization error over the class is supF ∈F |GenN (F )|. Definition D.2 (Covering Number). Let (F, d) be a metric space. For ϵ > 0, the ϵ-covering number N (F , ϵ, d) is the minimum cardinality of a set {F1 , . . . , Fm } ⊆ F such that for every F ∈ F there exists j ∈ [m] with d(F, Fj ) ≤ ϵ. For operator learning, the natural metric is the uniform deviation over the input domain of interest. In our setting this is the OOD-saturated set AOOD defined in Lemma E.1 below: d∞ (F1 , F2 ) :=
∥F1 (a) − F2 (a)∥U .
sup a∈AOOD
Lemma D.1 (Covering Number to Generalization Error). Let F be a class of operators from A to U, and let the loss ℓ be Lℓ -Lipschitz in its first argument and bounded in [0, B]. For any ϵ > 0 and δ ∈ (0, 1), with probability at least 1 − δ over the draw of N i.i.d. samples, s log 2 N (F, ϵ, d∞ )/δ b sup R(F ) − RN (F ) ≤ 2Lℓ ϵ + B . 2N F ∈F Proof. Take an ϵ-cover {F1 , . . . , Fm } of F under d∞ with m = N (F, ϵ, d∞ ). For any F ∈ F, let Fj be the cover element with d∞ (F, Fj ) ≤ ϵ. The Lℓ -Lipschitz property of ℓ gives |ℓ(F (a), u) − ℓ(Fj (a), u)| ≤ Lℓ ∥F (a) − Fj (a)∥U ≤ Lℓ ϵ, so both population and empirical risks differ by at most Lℓ ϵ from those of Fj . Hence |GenN (F )| ≤ 2Lℓ ϵ + |GenN (Fj )|. For a fixed Fj , Hoeffding’s inequality for losses bounded in [0, B] yields 2N t2 Pr |GenN (Fj )| ≥ t ≤ 2 exp − 2 . B q A union bound over the m cover elements and the choice t = B log(2m/δ) complete the proof. 2N Remark D.1 (Improved Generalization via Canonicalization). Theorem 3.2 establishes that the orbitaveraging projection ΠG is non-expansive under d∞ , which yields the covering-number inequality N (FG , ϵ, d∞ ) ≤ N (F, ϵ, d∞ ). Substituting this into Lemma D.1, the generalization bound for the equivariant class FG is at most the bound for the original (unconstrained) class F at the same sample size N . Moreover, when the learned operator F is approximately equivariant with commutation residual ηG (F ) ≤ η, the proof of Theorem 3.2 shows d∞ (F, ΠG F ) ≤ η, and the Lℓ -Lipschitz property of ℓ gives |R(F ) − R(ΠG F )| ≤ Lℓ η,
b N (F ) − R b N (ΠG F )| ≤ Lℓ η. |R
(3)
Combining (3) with Lemma D.1 applied to ΠG F ∈ FG and using the triangle inequality yields s log 2 N (FG , ϵ, d∞ )/δ |GenN (F )| ≤ 2Lℓ ϵ + B + 2Lℓ η. 2N | {z } generalization bound for FG
In PACE-FNO, canonicalization makes the overall pipeline approximately equivariant, so the residual η is bounded by the alignment error of the Lie-algebra coordinate estimator. As the estimator improves, η → 0 and the generalization gap approaches that of the fully equivariant class. Since N (FG , ϵ, d∞ ) ≤ N (F, ϵ, d∞ ), this bound is never worse than that of an unconstrained operator; the gain comes from the symmetry constraint reducing the covering number of the hypothesis class. 30
E
Proofs
The standard FNO approximation result requires a compact input set and a continuous target operator. We restate the needed assumptions in the notation used here, then prove the two bounds used in the main text. Induced terminal group element. If the physical solution operator is equivariant under the symmetry group GB , then for each initial group action g0 ∈ GB there exists a terminal group element gT ∈ GB satisfying the intertwining relation G † ◦ Cg0 = CgT ◦ G † . For the spatial transformations studied here (translation, rotation), equivariance implies gT = g0 because the geometric shift is conserved during temporal evolution on the torus. For kinematic transformations (Galilean boosts), the background velocity transports the field by vbg T over the time horizon T , inducing gT = exp(ξ0 + vbg T · ∇). This relation is applied in Theorem 3.1 to separate the prediction error into canonical approximation and geometric alignment components. Lemma E.1 (Compactness and Orbit Saturation). Let the physical space D ⊂ Rd be a bounded open set with a Lipschitz boundary. Let the input and output spaces be A = L2 (D; Rda ) and U = L2 (D; Rdu ). Assume that the canonical training measure µ is supported on a compact subset Ac ⊂ A. Let GB ⊂ G be the bounded set of group elements used to define the OOD shift family, and assume GB is compact. The OOD set is the orbit saturation of the canonical set under GB : AOOD := {Cg (a) : a ∈ Ac , g ∈ GB } ⊂ A. Assume that the action (g, a) 7→ Cg (a) is continuous, that the true physical evolution operator G † : AOOD → U is equivariant for the transformations considered, and that G † is locally Lipschitz on AOOD . Under these assumptions, the standard Fourier Neural Operator approximation result applies on the canonical compact set. Theorem E.2. Given the premises established in Lemma E.1, for any ϵ > 0, there exists a Fourier Neural Operator Gθ of sufficient depth or width such that sup ∥Gθ (a) − G † (a)∥U ≤ ϵ. a∈Ac
Remark E.1. Theorem E.2 summarizes the standard compact-set approximation setting for neural operators. Its relevance to OOD prediction is limited: when an input field is transformed by a continuous Lie group element g ∈ G, the transformed data aOOD = Cg (a) may fall outside the original canonical compact set Ac . The theorem then no longer applies directly to that transformed input. PACE-FNO addresses this specific geometric mismatch by estimating the generator and applying Cg−1 to map the observation back toward Ac . The resulting guarantee is conditional on successful alignment; estimator error remains an explicit term in Theorem 3.1. The input geometric alignment error is ϵgeo,0 :=
sup a∈AOOD
∥Cĝ−1 (a) − Cg−1 (a)∥A , 0 0
and the terminal alignment error is ϵgeo,T := sup ∥CĝT (G † (ac )) − CgT (G † (ac ))∥U . ac ∈Ac
Lemma E.3 (Isometry and Lipschitz Action of the Lie Group). Assume that the group action Cg on the output space U is an isometric representation. For the continuum translations and measurepreserving rotations modeled in our experiments, the L2 norm is preserved: ∥Cg (u)∥U = ∥u∥U for any g ∈ GB and u ∈ U. More generally, it is enough to assume that there exists a constant KC ≥ 1 such that, for any g ∈ GB and u1 , u2 ∈ U, ∥Cg (u1 ) − Cg (u2 )∥U ≤ KC ∥u1 − u2 ∥U . 31
Lemma E.4 (Lipschitz Continuity of the Learned Operator). Assume that the trained Fourier Neural Operator Gθ is Lipschitz on a neighborhood of Ac that contains the pulled-back estimates considered in the proof. That is, for some Lθ ≥ 1 and all a1 , a2 in this neighborhood, ∥Gθ (a1 ) − Gθ (a2 )∥U ≤ Lθ ∥a1 − a2 ∥A . Proof of Theorem 3.1. Let a ∈ AOOD be an arbitrary out-of-distribution input. By Lemma E.1, there exists a canonical state ac ∈ Ac and a true initial group element g0 ∈ GB such that a = Cg0 (ac ) =⇒ ac = Cg−1 (a) 0 Because G † is equivariant for the transformations considered, its value on the OOD state can be written using the canonical state and the induced terminal group element gT : G † (a) = G † (Cg0 (ac )) = CgT (G † (ac )) The triangle inequality separates the total prediction error into input-alignment, canonicalapproximation, and terminal-alignment terms: ∥CĝT (Gθ (Cĝ−1 (a))) − CgT (G † (ac ))∥U = CĝT (Gθ (Cĝ−1 (a))) − CĝT (Gθ (ac )) 0 0 + CĝT (Gθ (ac )) − CĝT (G † (ac )) + CĝT (G † (ac )) − CgT (G † (ac ))
U
≤ ∥CĝT (Gθ (Cĝ−1 (a))) − CĝT (Gθ (ac ))∥U 0 | {z } (I) Initial alignment error
+ ∥CĝT (Gθ (ac )) − CĝT (G † (ac ))∥U | {z } (II) Canonical-domain approximation error
+ ∥CĝT (G † (ac )) − CgT (G † (ac ))∥U . {z } | (III) Terminal group-action mismatch
For the first term, Lemmas E.3 and E.4 give (I) ≤ KC ∥Gθ (Cĝ−1 (a)) − Gθ (ac )∥U ≤ KC · Lθ ∥Cĝ−1 (a) − Cg−1 (a)∥A . 0 0 0 By the definition of ϵgeo,0 , this term is bounded by M0 ϵgeo,0 , where M0 := KC Lθ . For the canonicaldomain approximation term, the isometry case of Lemma E.3 and Theorem E.2 give (II) = ∥Gθ (ac ) − G † (ac )∥U ≤ ϵ. The third term is exactly bounded by the definition of the terminal alignment error: (III) ≤ ϵgeo,T . Combining the three bounds and taking the supremum over a ∈ AOOD gives sup a∈AOOD
∥CĝT ◦ Gθ ◦ Cĝ−1 (a) − G † (a)∥U ≤ M0 ϵgeo,0 + ϵ + ϵgeo,T . 0
Proof of Theorem 3.2. Let ν be the normalized Haar measure on the compact group GB . Define the orbit-averaging projection ΠG : F → FG by Z (ΠG F )(a) = Cg−1 F (Cg a) dν(g), GB
where the integral is understood in the Banach space U. For any h ∈ GB , the change of variables k = gh and the right invariance of ν give Z (ΠG F )(Ch a) = Cg−1 F (Cg Ch a) dν(g) GB Z −1 = Ckh −1 F (Ck a) dν(k) GB Z = Ch Ck−1 F (Ck a) dν(k) = Ch (ΠG F )(a). GB
32
Thus ΠG F is GB -equivariant. Next, because each Cg is an isometry and a bijection on AOOD , Z d∞ (ΠG F, ΠG H) = sup Cg−1 F (Cg a) − H(Cg a) dν(g) a∈AOOD G U Z B ≤ sup ∥F (Cg a) − H(Cg a)∥U dν(g) a∈AOOD
GB
≤ d∞ (F, H). m Hence ΠG is non-expansive. If {Fj }m j=1 is an r-cover of F under d∞ , then {ΠG Fj }j=1 is an r-cover of FG = ΠG (F). Therefore,
N (FG , r, d∞ ) ≤ N (F, r, d∞ ). b N be as defined in (D.1). It remains to connect this complexity reduction to the risk. Let R and R Applying Lemma D.1 to FG gives, with probability at least 1 − δ, s log 2 N (FG , ϵ, d∞ )/δ b sup |R(F ) − RN (F )| ≤ 2Lℓ ϵ + B . 2N F ∈FG Since N (FG , ϵ, d∞ ) ≤ N (F, ϵ, d∞ ), this bound is no larger than the analogous bound for the unconstrained class F. Finally, suppose F is approximately equivariant in the sense that the residual ηG (F ) :=
sup g∈GB a∈AOOD
∥(F ◦ Cg )(a) − (Cg ◦ F )(a)∥U
satisfies ηG (F ) ≤ η. By applying the same bound after Cg−1 and using isometry, ∥Cg−1 F (Cg a) − F (a)∥U ≤ η for every g ∈ GB and a ∈ AOOD . Averaging gives d∞ (F, ΠG F ) ≤ η. The Lℓ -Lipschitz property of the loss then implies |R(F ) − R(ΠG F )| ≤ Lℓ η,
b N (F ) − R b N (ΠG F )| ≤ Lℓ η. |R
Therefore the generalization gap of F is bounded by the equivariant-class gap of ΠG F plus 2Lℓ η. Lemma E.5 (Local Smoothness and Quadratic Growth Around the Minima Manifold). Assume that the zero set Ξ∗ is nonempty, and let JTTA : g → R≥0 be continuously differentiable on an open neighborhood N (Ξ∗ ) ⊂ g of Ξ∗ := {ξ ∈ g | JTTA (ξ) = 0}. Assume further that there exists L > 0 such that, for all ξˆ1 , ξˆ2 ∈ N (Ξ∗ ), ∥∇JTTA (ξˆ1 ) − ∇JTTA (ξˆ2 )∥g ≤ L∥ξˆ1 − ξˆ2 ∥g , and that there exists α > 0 such that, for all ξˆ ∈ N (Ξ∗ ), ˆ ≥ α dist(ξ, ˆ Ξ∗ )2 , JTTA (ξ) where
ˆ Ξ∗ ) := inf ∥ξˆ − ξ ′ ∥g . dist(ξ, ′ ∗ ξ ∈Ξ
Since Ξ∗ is the zero set of JTTA , we have min JTTA = 0. The quadratic-growth condition immediately yields the local error bound [53] ˆ Ξ∗ ) ≤ α−1/2 JTTA (ξ) ˆ 1/2 . dist(ξ, 33
Theorem E.6 (Monotone Energy Decay for Test-Time Refinement Near the Minima Manifold). Assume the hypotheses of Lemma E.5. Consider a global solution of the continuous-time gradient flow dξˆ ˆ = −∇JTTA (ξ(t)). dt Let the initial point ξˆ(0) satisfy o n S0 := ξ ∈ g JTTA (ξ) ≤ JTTA (ξˆ(0) ) ⊂ N (Ξ∗ ). Then, for all t ≥ 0, the energy is nonincreasing along the flow: d 2 ˆ ˆ JTTA (ξ(t)) = −∥∇JTTA (ξ(t))∥ g ≤ 0. dt Consequently, ˆ JTTA (ξ(t)) +
Z t
2 ˆ ˆ(0) ), ∥∇JTTA (ξ(s))∥ g ds = JTTA (ξ
0
and the trajectory remains in S0 ⊂ N (Ξ∗ ). There exists J∞ ∈ [0, JTTA (ξˆ(0) )] such that
ˆ lim JTTA (ξ(t)) = J∞ .
t→∞
In addition,
Z ∞
2 ˆ ˆ(0) ) − J∞ < ∞, ∥∇JTTA (ξ(t))∥ g dt = JTTA (ξ
0
and
ˆ lim ∥∇JTTA (ξ(t))∥ g = 0.
t→∞
The distance to the minima manifold is controlled by the energy: ˆ Ξ∗ ) ≤ α−1/2 JTTA (ξ(t)) ˆ 1/2 ≤ α−1/2 JTTA (ξˆ(0) )1/2 . dist(ξ(t), If, in addition, J∞ = 0, then ˆ Ξ∗ ) = 0. lim dist(ξ(t),
t→∞
If the sublevel set S0 is precompact, then the trajectory has accumulation points, and every accumulation point is a critical point of JTTA . Proof of Theorem E.6. Since dξˆ ˆ = −∇JTTA (ξ(t)), dt the chain rule gives d ˆ JTTA (ξ(t)) = dt
*
dξˆ ˆ ∇JTTA (ξ(t)), dt
+ 2 ˆ = −∥∇JTTA (ξ(t))∥ g ≤ 0. g
This proves monotonicity, and integrating over [0, t] yields Z t 2 ˆ ˆ ˆ(0) ). JTTA (ξ(t)) + ∥∇JTTA (ξ(s))∥ g ds = JTTA (ξ 0
Hence the trajectory remains in the initial sublevel set S0 . ˆ Because JTTA (ξ(t)) is nonincreasing and bounded below by 0, it converges to a limit J∞ ≥ 0. Passing to the limit t → ∞ in the identity above gives Z ∞ 2 ˆ ˆ(0) ) − J∞ < ∞. ∥∇JTTA (ξ(s))∥ g ds = JTTA (ξ 0
34
ˆ is absolutely continuous and ∇JTTA is L-Lipschitz on N (Ξ∗ ), the map Since ξ(·) ˆ t 7→ ∇JTTA (ξ(t)) is uniformly continuous along the trajectory. Together with the square-integrability established above, Barbalat’s lemma implies ˆ ∥∇JTTA (ξ(t))∥ g →0
as t → ∞.
The quadratic-growth assumption in Lemma E.5 gives ˆ ˆ Ξ∗ )2 , JTTA (ξ(t)) ≥ α dist(ξ(t), and therefore
ˆ Ξ∗ ) ≤ α−1/2 JTTA (ξ(t)) ˆ 1/2 . dist(ξ(t), Since the energy is nonincreasing, ˆ Ξ∗ ) ≤ α−1/2 JTTA (ξˆ(0) )1/2 . dist(ξ(t), ˆ Finally, if JTTA (ξ(t)) → 0, the same inequality implies ˆ Ξ∗ ) → 0. dist(ξ(t), If S0 is precompact, standard compactness arguments imply the existence of accumulation points. Since the gradient vanishes asymptotically, every accumulation point is a critical point of JTTA . Lemma E.7 (Lipschitz Continuity of Inverse Group Action). Let G be a finite-dimensional Lie group with Lie algebra g, and let C :G×A→A be a smooth group action on the physical field space A. Let U ⊂ g be an open neighborhood on which the exponential map is a diffeomorphism onto its image. Then, for any fixed out-of-distribution input aOOD ∈ AOOD , there exists a constant La > 0 such that for all ξˆ1 , ξˆ2 ∈ U , −1 −1 Cexp( (a ) − Cexp( (a ) ≤ La ∥ξˆ1 − ξˆ2 ∥g . ξ̂ ) OOD ξ̂ ) OOD A 1
2
In other words, the inverse group action parameterized through the exponential map is locally Lipschitz continuous with respect to the Lie algebra coordinates. Corollary E.8 (Prediction Error Bound Under Symmetry-Induced Equivalent Minima). Assume the hypotheses of Theorem E.6. Let aOOD ∈ AOOD be an arbitrary out-of-distribution input, and let ξˆ denote any finite-time estimator produced by the test-time adaptation dynamics within N (Ξ∗ ). Suppose that the inverse group action on the input space and the induced action on the output space are locally Lipschitz, so that the geometric perturbation argument in Theorem 3.1 yields ˆ Ξ∗ ), ∥ûpred − G † (aOOD )∥U ≤ ϵ + K dist(ξ, for some K > 0. Then ˆ 1/2 . ∥ûpred − G † (aOOD )∥U ≤ ϵ + K α−1/2 JTTA (ξ) ˆ such that In particular, along any adaptation trajectory ξ(t) ˆ JTTA (ξ(t)) → 0, we have
∥ûpred − G † (aOOD )∥U → ϵ.
Proof of Corollary E.8. Fix an arbitrary estimator ξˆ produced by the test-time adaptation dynamics. Since g is finite-dimensional and JTTA is continuous, the nonempty zero set Ξ∗ = {ξ ∈ g | JTTA (ξ) = 0} is closed. Hence the distance from ξˆ to Ξ∗ is attained by some ξ ′ ∈ Ξ∗ , so that ˆ Ξ∗ ). ∥ξˆ − ξ ′ ∥g = dist(ξ, 35
By Lemma E.7, the inverse group action on the input space is locally Lipschitz with respect to the Lie algebra coordinates. Hence, −1 −1 Cexp( (a ) − Cexp(ξ ′ ) (aOOD ) ξ̂) OOD
A
≤ La ∥ξˆ − ξ ′ ∥g .
Assuming analogously that the induced action on the output space U is locally Lipschitz with constant Lu > 0, we also have ϵgeo,T ≤ Lu ∥ξˆ − ξ ′ ∥g . Applying Theorem 3.1 with the symmetrically equivalent minimizer ξ ′ ∈ Ξ∗ yields −1 −1 ∥ûpred − G † (aOOD )∥U ≤ ϵ + M0 Cexp( (a ) − Cexp(ξ ′ ) (aOOD ) ξ̂) OOD
A
+ ϵgeo,T
≤ ϵ + (M0 La + Lu ) ∥ξˆ − ξ ′ ∥g . Defining K := M0 La + Lu , we obtain
ˆ Ξ∗ ). ∥ûpred − G † (aOOD )∥U ≤ ϵ + K dist(ξ,
Finally, applying the quadratic-growth error bound from Lemma E.5, we obtain ˆ Ξ∗ ) ≤ α−1/2 JTTA (ξ) ˆ 1/2 , dist(ξ, and therefore
ˆ 1/2 . ∥ûpred − G † (aOOD )∥U ≤ ϵ + K α−1/2 JTTA (ξ)
ˆ If JTTA (ξ(t)) → 0, then the right-hand side converges to ϵ. The above result should be interpreted as a local stability statement for the test-time refinement stage. It does not assert global convergence for arbitrary OOD inputs. Rather, it shows that when the estimator initialization lies in a neighborhood of a symmetry-equivalent minima manifold, reducing the residual energy controls the remaining geometric error. On the tori used in our experiments, symmetry-equivalent elements of Ξ∗ (e.g., shifts by a full spatial period) produce identical canonical fields under the group action; the bound is therefore independent of which element of the minima manifold attains the distance.
36