Out-of-Distribution Generalization of Neural Operators via Active Geometry Synthesis

Training neural operators for PDEs over complex geometries requires high-quality simulation datasets that are often designed by careful parameterization of a few specific shapes, yielding geometries that cover a range of physical dynamics close to those of the base shapes. Still, as seemingly small changes in the geometry can induce widely different physical dynamics, neural operators trained with such data generalize poorly to unseen geometric features that may arise in engineering and design optimization. In this work, we propose Active Geometry Synthesis (AGS), a novel geometric active learning method that improves the out-of-distribution (OOD) generalization of neural operators by synthesizing the geometries the operator is most uncertain about. Specifically, AGS equips a few base geometries with low-dimensional free-form deformations, and maximizes the predictive variance of the solution field, obtained by the Laplace approximation, with respect to the deformation control points using gradient ascent. As we demonstrate empirically, the Laplace predictive variance provides a surrogate for the operator errors over unseen shapes. As such, AGS samples diverse new geometries in regions of the geometry manifold where the physical dynamics are most challenging for the operator, expanding the coverage of the training dataset over the geometry manifold, and significantly improving the OOD generalization compared to i.i.d. uniform sampling. Meanwhile, by avoiding the manual collection and augmentation of geometries, AGS does not restrict the neural operator to a predetermined set of shapes.