Exact SO(3)-Equivariant Isotropic Kernels for Rotation-Robust Neural Dynamics

The gist: The Invariant-Conditioned Isotropic Kernel Neural Operator (IKNO) demonstrates that a compact, PDE-specialized model can remove coordinate dependence without sacrificing forecasting accuracy in three-dimensional Navier–Stokes dynamics.

Problem Setting

The study investigates the failure of neural surrogates for vector-valued partial differential equations when expressed in a rotated coordinate frame The relevant transformation rotates coordinates and vector components together The physical system under consideration is the incompressible Navier–Stokes equation on a periodic three-dimensional domain

IKNO Architecture

The Invariant-Conditioned Isotropic Kernel Neural Operator (IKNO) is introduced as a compact graph model that builds local interactions from scalar quantities unchanged by rotation and vector directions that rotate with the data Each edge coefficient is produced from scalar rotation invariants The construction has three properties relevant to the experiment: it is exactly SO(3)-equivariant, it annihilates constant fields in its difference channels, and it contains only 4,869 trainable parameters.

Key Components of IKNO

The method separates each message into invariant coefficients and equivariant vector generators The coefficient network receives six scalar invariants These six invariants are defined as the relative coordinate, unit direction, and velocity difference. The equivariant message is constructed using a subset of admissible SO(3)-equivariant generators Specifically, the message is defined as mij = aijδuij + bijrbij (rb T ijδuij) + cijui + dij (ui × rbij).

Experimental Results and Comparison

The main result shows that architectural equivariance and rotation augmentation produce measurably different operators Augmentation reduces the rotation sensitivity of two graph controls but does not eliminate it, whereas IKNO remains unchanged across 1,000 sampled coordinate frames. A Tensor Field Network matches IKNO in forecasting accuracy but uses 5.6 times more parameters. The normalized operator equivariance diagnostic, eeq = ER RMS(fθ(RU, RX) − Rfθ(U, X)) / (RMS(fθ(RU, RX)) + RMS(Rfθ(U, X)), indicates that a value near zero means that the learned derivative is coordinate-consistent.

Conclusion

IKNO combines invariant scalar edge functions with an equivariant vector basis to guarantee SO(3) equivariance On an untouched Navier–Stokes holdout, it matches a tensor-field network with 5.6 times fewer parameters, removes coordinate dependence left by augmentation, and beats a parameter-matched control. This supports compact, PDE-specialized exact operator consistency.

Discussion and Limitations

The scope of this conclusion is limited because the final test contains five trajectories from one low-resolution Navier–Stokes generator Furthermore, arbitrary rotations change the coordinate description of the periodic box, meaning the test is a joint coordinate-frame transformation of sampled states rather than a second solver run on an independently rotated grid. The comparison with Tensor Field Networks supports an efficiency claim on this specific task, not a general efficiency ordering. The two-frame extrapolator in Appendix B is much more accurate because the trajectories are temporally smooth, and the present result is therefore specific to one-frame operator identification.

References

Johannes Brandstetter, Rob Hesselink, Elise van der Pol, Erik J. Bekkers, and Max Welling. Geometric and physical quantities improve E(3) equivariant message passing Johannes Brandstetter, Max Welling, and Daniel E. Worrall. Lie point symmetry data augmentation for neural PDE solvers Johannes Brandstetter, Daniel E. Worrall, and Max Welling. Message passing neural PDE solvers Michael M. Bronstein, Joan Bruna, Taco Cohen, and Petar Velickovic. Geometric deep learning: Grids, groups, graphs, geodesics, and gauges Claudio Canuto, M. Yousuff Hussaini, Alfio Quarteroni, and Thomas A. Zang. Spectral Methods in Fluid Dynamics Taco Cohen and Max Welling. Group equivariant convolutional networks Congyue Deng, Or Litany, Yueqi Duan, Adrien Poulenard, Andrea Tagliasacchi, and Leonidas J. Guibas. Vector neurons: A general framework for SO(3)-equivariant networks Marc Finzi, Samuel Stanton, Pavel Izmailov, and Andrew Gordon Wilson. Generalizing convolutional neural networks for equivariance to lie groups on arbitrary continuous data Fabian Fuchs, Daniel Worrall, Volker Fischer, and Max Welling. SE(3)-transformers: 3d roto-translation equivariant attention networks Mario Geiger and Tess Smidt. e3nn: Euclidean neural networks Lingshen He, Yuxuan Chen, Zhengyang Shen, Yibo Yang, and Zhouchen Lin. Neural ePDOs: Spatially adaptive equivariant partial differential operator based networks Jacob Helwig, Xuan Zhang, Cong Fu, Jerry Kurtin, Stephan Wojtowytsch, and Shuiwang Ji. Group equivariant Fourier neural operators for partial differential equations Erik Jenner and Maurice Weiler. Steerable partial differential operators for equivariant neural networks G. F. Smith. On isotropic functions of symmetric tensors, skew-symmetric tensors and vectors Gilbert Strang. On the construction and comparison of difference schemes Nathaniel Thomas, Tess Smidt, Steven Kearnes, Lusann Yang, Li Li, Kai Kohlhoff, and Patrick Riley. Tensor field networks: Rotation- and translation-equivariant neural networks for 3d point clouds Soledad Villar, David W. Hogg, Kate Storey-Fisher, Weichi Yao, and Ben Blum-Smith. Scalars are universal: Equivariant machine learning, structured like classical physics V´ıctor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E(n) equivariant graph neural networks Kristof Sch¨utt, Oliver Unke, and Michael Gastegger. Equivariant message passing for the prediction of tensorial properties and molecular spectra G. F. Smith. On isotropic functions of symmetric tensors, skew-symmetric tensors and vectors Gilbert Strang. On the construction and comparison of difference schemes Nathaniel Thomas, Tess Smidt, Steven Kearnes, Lusann Yang, Li Li, Kai Kohlhoff, and Patrick Riley. Tensor field networks: Rotation- and translation-equivariant neural networks for 3d point clouds Soledad Villar, David W. Hogg, Kate Storey-Fisher, Weichi Yao, and Ben Blum-Smith. Scalars are universal: Equivariant machine learning, structured like classical physics V´ıctor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E(n) equivariant graph neural networks Kristof Sch¨utt, Oliver Unke, and Michael Gastegger. Equivariant message passing for the prediction of tensorial properties and molecular spectra G. F. Smith. On isotropic functions of symmetric tensors, skew-symmetric tensors and vectors Gilbert Strang.

1. No coordinate frame dependence in learned dynamics: The Invariant-Conditioned Isotropic Kernel Neural Operator (IKNO) builds local interactions from scalar quantities unchanged by rotation and vector directions that rotate with the data, ensuring that rotating the positions and velocities rotates the predicted velocity change in exactly the same way.

2. Improved generalization under unseen orientations: The system can perform better on unseen frames than models trained via random rotation augmentation, as IKNO remains unchanged across 1,000 sampled coordinate frames, whereas augmentation reduces but does not eliminate their coordinate dependence.

3. Parameter-efficient exact equivariance for PDE solving: The IKNO architecture achieves high accuracy while using significantly fewer resources; it matches the forecasting accuracy of a general rotation-aware Tensor Field Network with 5.6 times fewer parameters.

4. Distinguishing operator consistency from performance: The metric eeq is defined to measure if the learned derivative is coordinate-consistent, allowing researchers to determine if the learned derivative is coordinate-consistent, separating it from simple forecasting accuracy.

5. Robustness against numerical instability: The method demonstrates that exact architectures are stable, as the analysis shows that IKNO prediction error is unchanged to the displayed precision and the residual operator error is on the order of single-precision rounding, which contrasts with approximate methods where the sampled maximum of ∆1(R)/e1 over 1,000 rotations remains nonzero after augmentation.

Neural surrogates for vector-valued partial differential equations can fit training data yet change their predictions when the same physical state is expressed in a rotated coordinate frame. We study this failure on three-dimensional Navier--Stokes dynamics observed at irregularly placed points. We introduce the Invariant-Conditioned Isotropic Kernel Neural Operator (IKNO), a compact graph model that builds local interactions from scalar quantities unchanged by rotation and vector directions that rotate with the data. Consequently, rotating the positions and velocities rotates the predicted velocity change in exactly the same way. On a held-out test set fixed after model design, training unconstrained graph models on randomly rotated examples reduces but does not eliminate their coordinate dependence. In contrast, IKNO is consistent to numerical precision, matches the forecasting accuracy of a general rotation-aware Tensor Field Network with 5.6 times fewer parameters, and outperforms a parameter-matched graph simulator. These results show that a compact, PDE-specialized model can remove coordinate dependence without sacrificing forecasting accuracy.