[Paper Review] Equivariant Learning of Stochastic Fields: Gaussian Processes and Steerable Conditional Neural Processes

Peter Holderrieth, Michael Hutchinson+1 more|arXiv (Cornell University)|Nov 25, 2020

Gaussian Processes and Bayesian Inference46 references4 citations

TL;DR

This paper introduces Steerable Conditional Neural Processes (SteerCNPs), a novel equivariant model for learning stochastic fields such as vector fields in physics and engineering. By enforcing equivariance under rotations and reflections using group theory and steerable kernels, SteerCNPs achieve superior generalization and robustness on synthetic and real-world data, outperforming baseline models in both likelihood and transfer learning tasks, especially under data augmentation and distribution shift.

ABSTRACT

Motivated by objects such as electric fields or fluid streams, we study the problem of learning stochastic fields, i.e. stochastic processes whose samples are fields like those occurring in physics and engineering. Considering general transformations such as rotations and reflections, we show that spatial invariance of stochastic fields requires an inference model to be equivariant. Leveraging recent advances from the equivariance literature, we study equivariance in two classes of models. Firstly, we fully characterise equivariant Gaussian processes. Secondly, we introduce Steerable Conditional Neural Processes (SteerCNPs), a new, fully equivariant member of the Neural Process family. In experiments with Gaussian process vector fields, images, and real-world weather data, we observe that SteerCNPs significantly improve the performance of previous models and equivariance leads to improvements in transfer learning tasks.

Motivation & Objective

• To develop a machine learning model that respects physical symmetries such as rotation and reflection in stochastic fields.
• To formalize the conditions under which Gaussian processes are equivariant under spatial transformations.
• To design a new Neural Process variant—SteerCNPs—that enforces equivariance through steerable kernels and group representations.
• To demonstrate improved performance and generalization in transfer learning and out-of-distribution settings.
• To bridge equivariant deep learning with probabilistic modeling for scientific machine learning.

Proposed method

• The paper derives necessary and sufficient conditions for vector-valued Gaussian processes to be equivariant under orthogonal transformations using group representation theory.
• It introduces SteerCNPs, a conditional neural process that uses steerable convolutional layers with fiber representations to ensure equivariance under finite subgroups of O(n), such as C₄, C₈, D₄, and D₈.
• The model employs a group-equivariant message-passing mechanism where context points are transformed consistently with the prediction field under group actions.
• Equivariance is preserved by ensuring the model's output transforms in the same way as the input under group operations, using irreducible representations of the group H.
• The model is trained end-to-end with likelihood maximization, using a softplus activation for covariance and sigmoid for mean to maintain invariance and bounded outputs.
• Data augmentation with blank images and rotated digits is used to improve robustness to distribution shift and out-of-context generalization.

Experimental results

Research questions

• RQ1 Under what conditions is a stochastic process model equivariant under spatial transformations such as rotations and reflections?
• RQ2 How can we construct Gaussian processes that are equivariant to general orthogonal transformations in Rⁿ?
• RQ3 Can we design a Neural Process model that is fully equivariant to non-translation symmetries, such as rotations and reflections?
• RQ4 Does enforcing equivariance lead to improved generalization and robustness in few-shot and out-of-distribution settings?
• RQ5 How does equivariance affect performance on real-world scientific data such as weather fields and vector fields?

Key findings

• SteerCNPs achieve significantly higher log-likelihood scores than baseline models, with a mean test log-likelihood of 1.18 ± 0.02 on rotated MNIST, compared to 1.06 ± 0.04 for ConvCNP.
• The model trained with D₈ symmetry achieves the highest performance, indicating that richer symmetry groups improve generalization.
• Qualitative results show that SteerCNPs maintain consistent predictions under rotation, while non-equivariant models (e.g., ConvCNP) produce erratic outputs when context sets are rotated.
• On extrapolation tasks, SteerCNPs generalize better to unseen regions and larger canvases, with reduced noisy predictions far from context points.
• The inclusion of blank images during training improved generalization to test sets with large empty regions, reducing distribution shift issues.
• Equivariance leads to improved transfer learning performance, particularly under rotation and distribution shift, validating the physical inductive bias.