Buckingham $\pi$-Invariant Test‑Time Projection for Robust PDE Surrogate Modeling

ICLR 2026 Conference Submission18502 Authors

Keywords: Buckingham-pi, PDE, model-agnostic

TL;DR: We propose a π-invariant test-time projection that aligns test inputs with the training distribution by solving a log-space least squares prob- lem to preserve Buckingham π-invariants.

Abstract: PDE surrogate models such as FNO and PINN struggle to predict solutions across inputs with diverse physical units and scales, limiting their out-of-distribution (OOD) generalization. We propose a $\{pi}$-invariant test-time projection that aligns test inputs with the training distribution by solving a log-space least squares problem that preserves Buckingham π-invariants. For PDEs with multidimensional spatial fields, we use geometric representative $\{pi}$-values to compute distances and project inputs, overcoming degeneracy and singular points that limit prior $\{pi}$-methods. To accelerate projection, we cluster the training set into K clusters, reducing the complexity from O(MN) to O(KN) for the M training and N test samples. Across wide input scale ranges, tests on 2D thermal conduction and linear elasticity achieve an average MAE reduction up to $\approx\!91\%$ with minimal overhead. This training-free, model-agnostic method is expected to apply to more diverse PDE-based simulations.

Primary Area: applications to physical sciences (physics, chemistry, biology, etc.)

Submission Number: 18502