Equivariant Neural Networks Geometric Deep Learning
# Equivariant Neural Networks & Geometric Deep Learning
## Introduction & Motivation
Many of the most important application domains for machine learning, including molecular property prediction, protein structure modeling, particle physics, robotics, and 3D computer vision, involve data with an inherent geometric structure: point clouds, molecular graphs embedded in 3D space, or physical systems whose governing laws do not depend on an arbitrary choice of coordinate frame. A standard neural network, trained on such data without any special architectural consideration, must learn to recognize the same underlying physical or geometric pattern separately for every possible rotation, translation, or reflection in which that pattern might appear, since nothing in its architecture encodes the fact that these transformed inputs represent fundamentally the same object.
Geometric deep learning is the study of neural network architectures that build the relevant geometric symmetries directly into the network's structure, rather than requiring the network to learn them from data through extensive augmentation. Equivariant neural networks are the primary architectural tool for this purpose: layers designed so that transforming the input by some symmetry (a rotation, translation, permutation, or reflection) produces a correspondingly and predictably transformed output, rather than an arbitrary or unpredictable one. This equivariance property is distinct from, but closely related to, invariance, in which the network's output is required to remain completely unchanged under the relevant transformations.
The practical significance of equivariant architectures has grown enormously with their central role in modern computational structural biology, most notably in AlphaFold2's use of an SE(3)-equivariant attention mechanism (the Invariant Point Attention module) to reason directly about 3D atomic coordinates, and in the broader adoption of equivariant graph neural networks for molecular dynamics simulation, drug discovery, and materials science, where respecting physical symmetries is not merely a convenient inductive bias but a requirement for physically consistent and sample-efficient learning from comparatively small and expensive-to-collect datasets.
## Core Concepts & Theory
A function f is said to be equivariant with respect to a group of transformations G if, for every transformation g in G and every input x, transforming the input by g and then applying f produces the same result as applying f first and then transforming the output by a corresponding transformation of g (which may act differently on the output space than it does on the input space, depending on the type of output). Invariance is the special case of equivariance in which the transformation acting on the output space is always the identity, meaning the output does not change at all regardless of how the input was transformed.
The groups most relevant to geometric deep learning include the permutation group S_n (relevant to graphs and point clouds, where the ordering of nodes or points is arbitrary and should not affect the learned representation of the underlying structure), the translation group (relevant whenever absolute position in space should not matter, only relative position), the rotation group SO(3) in three dimensions (relevant to molecules and rigid bodies, where physical properties do not depend on the arbitrary orientation of the coordinate system used to describe them), and the special Euclidean group SE(3), which combines rotations and translations and is the most commonly targeted symmetry group for 3D molecular and structural biology applications.
Graph neural networks provide a natural starting point for building in permutation equivariance, since the standard message-passing update, in which each node aggregates information from its neighbors using a permutation-invariant aggregation function such as sum or mean, automatically produces node representations that transform correctly (permute correspondingly) under any relabeling of the graph's nodes. Extending this permutation equivariance to also respect rotational and translational symmetry, when node features include 3D spatial coordinates, requires more specialized architectural machinery, since ordinary neural network layers (linear transformations followed by pointwise nonlinearities) do not automatically preserve equivariance under continuous rotation groups the way permutation-equivariant aggregation does under discrete relabeling.
## Mathematical Formulation
The formal definition of equivariance for a function f mapping an input space X to an output space Y, with respect to a group G acting on X via a representation rho_X and on Y via a representation rho_Y, requires that applying f after transforming the input equals transforming the output of f by the corresponding group element:
$$ f\big( ho_X(g) \cdot x\big) = ho_Y(g) \cdot f(x), \qquad \forall \, g \in G, \; \forall \, x \in X $$
Invariance is the special case where rho_Y(g) is the identity transformation for every g, meaning the output is completely unaffected by the input transformation:
$$ f\big( ho_X(g) \cdot x\big) = f(x), \qquad \forall \, g \in G, \; \forall \, x \in X $$
For SE(3)-equivariant graph neural networks operating on point clouds with associated scalar and vector features, a common and simple construction achieves equivariance by restricting all learned transformations of vector-valued features to operations that commute with rotation, such as scaling a vector feature by a scalar computed from rotation-invariant quantities (like inter-point distances), or taking linear combinations of relative position vectors, both of which are guaranteed to transform correctly under any global rotation R applied to the input coordinates:
$$ R \cdot \big(\alpha(\|\mathbf{r}_{ij}\|) \, \mathbf{r}_{ij}\big) = \alpha(\|\mathbf{r}_{ij}\|) \, \big(R \cdot \mathbf{r}_{ij}\big) $$
where r_ij is the relative position vector between two points, alpha is any scalar function computed from rotation-invariant inputs (such as the distance between the points), and R is an arbitrary 3D rotation matrix; this identity holds because scaling a vector by a rotation-invariant scalar and then rotating it gives the same result as rotating the vector first and then scaling it by the same (unchanged, since it depends only on the invariant distance) scalar.
More general equivariant architectures, such as Tensor Field Networks and SE(3)-Transformers, represent intermediate features as spherical harmonics, whose rotation behavior under SO(3) is captured exactly by Wigner D-matrices, and combine features of different rotational orders using Clebsch-Gordan coefficients, which specify precisely how to combine two rotationally-equivariant tensors of given orders into an output tensor of a target order while preserving equivariance throughout the combination.
## Advanced Theory & Extensions
Steerable CNNs generalize the translation-equivariance of standard convolutional networks to additionally respect rotational symmetry, by constraining convolutional filters to be built from a restricted basis of functions (steerable filters) whose response to a rotated input is exactly predictable, rather than learning arbitrary, unconstrained filters that would only respect translation symmetry and not rotation; this approach has been applied both to 2D image data (rotation-equivariant CNNs) and to 3D volumetric and point-cloud data.
E(n)-Equivariant Graph Neural Networks (EGNN) offer a comparatively lightweight alternative to full spherical-harmonic-based equivariant architectures, achieving E(n) equivariance (rotations, translations, and reflections in n dimensions) using only simple operations on scalar distances and relative position vectors, without requiring the more mathematically involved spherical harmonic and Clebsch-Gordan machinery, at some cost in the expressiveness of higher-order geometric interactions the network can represent, but with substantially simpler implementation and often comparable empirical performance on molecular property prediction tasks.
AlphaFold2's Invariant Point Attention (IPA) module extends the standard Transformer attention mechanism to operate directly on 3D atomic coordinates in a manner invariant to the choice of global reference frame, by having each residue's local coordinate frame perform attention computations using geometric quantities (distances between locally-transformed query and key points) that are invariant under any global rigid-body transformation applied consistently to the entire structure, allowing the network to reason about 3D spatial relationships between residues without being sensitive to how the overall structure happens to be oriented in space.
Equivariant diffusion models extend denoising diffusion probabilistic models to generate 3D structures (such as molecular conformations) by ensuring that both the forward noising process and the learned reverse denoising network respect the relevant symmetry group, typically SE(3) or E(3), so that the generative model does not implicitly favor any particular orientation of the generated structure and produces samples whose distribution is properly invariant to arbitrary global rotations and translations, a critical property for physically meaningful molecular generation.
## Computational Considerations
Higher-order equivariant architectures based on spherical harmonics and Clebsch-Gordan tensor products incur substantially greater computational cost than simpler scalar/vector architectures like EGNN, since the tensor product operations combining features of different rotational orders scale unfavorably with the maximum order retained in the network, creating a direct tradeoff between the network's ability to represent complex, higher-order angular relationships and its computational and memory efficiency, an important consideration when scaling equivariant architectures to large molecular systems or high-resolution 3D structures.
Achieving strict, exact equivariance (as opposed to approximate equivariance obtained through data augmentation, in which a standard non-equivariant network is trained on many randomly rotated copies of the same input in hopes it learns approximate rotational invariance) guarantees that the network's predictions are perfectly consistent regardless of the arbitrary orientation in which an input happens to be presented, which is particularly valuable in low-data regimes (common in molecular and materials science applications) where a non-equivariant network would need to see many rotated examples of essentially the same underlying structure to learn approximate rotational robustness, effectively wasting model capacity and training data on learning a symmetry that could instead be built in architecturally at no data cost.
Efficient implementation of equivariant tensor operations benefits substantially from specialized software libraries (such as e3nn) that precompute and cache the relevant Clebsch-Gordan coefficients and provide optimized implementations of spherical harmonic evaluation and equivariant tensor products, since naive implementations of these operations can be a significant performance bottleneck relative to the comparatively simple matrix multiplications that dominate standard (non-equivariant) neural network layers.
## Practical Implementation Strategies
Selecting the appropriate symmetry group and equivariance strength for a given application requires balancing the benefit of a stronger inductive bias against the risk of over-constraining the model: for applications where the relevant physical symmetry is exact (such as molecular energy prediction, where physical energy truly does not depend on the arbitrary orientation of the coordinate system), strict equivariance is almost always beneficial, whereas for applications where the symmetry is only approximate or where non-geometric contextual cues genuinely correlate with absolute orientation (such as certain computer vision tasks where "up" has real semantic meaning), imposing strict equivariance may unnecessarily discard useful signal.
Combining scalar (invariant) and vector or higher-order tensor (equivariant) feature channels within the same network, as EGNN and similar architectures do, allows the network to simultaneously process rotation-invariant quantities (such as bond lengths, angles, and chemical identity) alongside genuinely directional, rotation-equivariant quantities (such as relative position vectors and predicted force directions), which is important because many practically relevant target quantities include both invariant components (such as total energy) and equivariant components (such as per-atom forces, which are vectors that must rotate correctly with the input structure).
Validating that an implemented architecture is truly equivariant, rather than merely approximately so due to a subtle implementation bug, is best done empirically as a standard unit test: applying a random rotation (and translation, for SE(3) equivariance) to a given input, passing both the original and transformed input through the network, and verifying that the corresponding transformation applied to the original output exactly matches the network's output on the transformed input, up to numerical precision, is a simple and highly effective way to catch equivariance-breaking bugs that might otherwise be difficult to detect purely through downstream task performance.
## Benchmark Datasets & Evaluation
QM9, a dataset of roughly 134,000 small organic molecules with a range of computed quantum-mechanical properties (including total energy, dipole moment, and orbital energy gaps), is among the most widely used benchmarks for evaluating equivariant graph neural networks on molecular property prediction, with mean absolute error on each target property serving as the standard evaluation metric, and has served as a key testbed distinguishing the performance of increasingly sophisticated equivariant architectures over time.
The Open Catalyst Project (OCP) dataset, comprising millions of density functional theory calculations of molecular adsorption on catalytic surfaces, provides a substantially larger-scale benchmark relevant to materials science and catalysis discovery applications, evaluating both energy and per-atom force prediction accuracy, with force prediction accuracy being particularly sensitive to whether a model architecture correctly respects equivariance, since forces are inherently vector-valued, rotation-equivariant quantities.
Protein structure prediction benchmarks, most notably the biennial CASP (Critical Assessment of Structure Prediction) competition, provided the primary evaluation venue in which AlphaFold2's SE(3)-equivariant-attention-based architecture demonstrated a dramatic leap in structure prediction accuracy over prior methods, measured using metrics such as the Global Distance Test (GDT) score comparing predicted and experimentally determined 3D atomic coordinates, and has since motivated widespread adoption of geometric and equivariant architectural principles throughout structural biology and drug discovery pipelines.
## Key Challenges & Limitations
Strictly equivariant architectures can be more difficult to implement correctly than standard neural network layers, since even small implementation errors (such as an unintended non-equivariant nonlinearity applied directly to vector-valued features) can silently break the equivariance guarantee without necessarily causing an obvious error or crash, potentially only manifesting as subtly degraded generalization to inputs presented in orientations underrepresented during training, making the empirical equivariance-testing practice described above an important safeguard.
The computational overhead of higher-order equivariant architectures can become a limiting factor for scaling to very large molecular systems or high-resolution 3D structures, motivating an active tradeoff space between representational expressiveness (retaining higher-order spherical harmonic features capable of representing more complex angular relationships) and computational tractability (restricting to lower-order or purely scalar/vector features for efficiency), with different applications settling at different points along this tradeoff depending on their accuracy requirements and computational budget.
Not all real-world data respects the idealized symmetry group a practitioner might wish to impose: real molecular dynamics data, for instance, may include subtle external field effects or boundary conditions that break perfect rotational symmetry, and real-world 3D scanning or imaging data may include sensor-specific biases correlated with absolute orientation, meaning that imposing strict equivariance in such settings, while often still beneficial as a strong prior, may not perfectly match the true underlying data-generating process and can occasionally slightly underperform a sufficiently well-regularized non-equivariant model trained with extensive data augmentation.
## Hyperparameter Tuning
The maximum spherical harmonic order (often denoted l_max) retained by a spherical-harmonic-based equivariant architecture is a primary tunable hyperparameter controlling the tradeoff between representational capacity for angular relationships and computational cost, with common choices ranging from l_max = 1 (capturing only scalar and simple vector relationships) to l_max = 3 or higher for applications requiring finer angular resolution, such as modeling complex bonding geometries in molecular systems.
The number of equivariant message-passing layers and the width (number of feature channels at each rotational order) of an equivariant graph neural network interact with the receptive field and representational capacity of the model in a manner broadly analogous to depth and width in standard graph neural networks, but with the added consideration that increasing the number of distinct rotational orders retained at each layer compounds the computational cost of the tensor-product operations combining them, making joint tuning of depth, width, and maximum order an important and computationally expensive hyperparameter search in practice.
For applications combining invariant and equivariant loss terms (such as jointly predicting a scalar energy and vector-valued forces), the relative weighting between these loss components requires careful tuning, since force prediction accuracy and energy prediction accuracy can respond differently to a given architecture and training regime, and poorly balanced loss weighting can cause the network to disproportionately optimize one target at the expense of the other.
## Real-World Applications & Case Studies
AlphaFold2 and its successors represent the most prominent real-world application of SE(3)-equivariant architectural principles, using the Invariant Point Attention mechanism to reason about 3D protein backbone and side-chain geometry in a manner invariant to arbitrary global rigid-body transformations, contributing substantially to a step-change improvement in computational protein structure prediction accuracy that has had a significant downstream impact on structural biology and drug discovery research workflows.
Equivariant graph neural networks are widely used in molecular dynamics simulation as machine-learned interatomic potentials, trained to predict per-atom forces and total system energy directly from atomic positions and species, serving as a substantially faster (though approximate) replacement for expensive quantum-mechanical density functional theory calculations in molecular dynamics simulations, enabling simulation of larger systems over longer timescales than would be computationally feasible using first-principles methods alone, with equivariance being essential to ensuring the predicted forces behave physically consistently under arbitrary rotation of the simulated system.
In robotics and 3D computer vision, equivariant architectures have been applied to point-cloud-based object pose estimation and robotic grasp planning, where the ability to predict a rotation-equivariant output (such as a predicted grasp orientation that rotates correspondingly with a rotated input point cloud of an object) is a natural and directly useful property, reducing the amount of training data and augmentation needed to achieve robust performance across the full range of possible object orientations encountered in real-world robotic manipulation settings.
## Integration with Other Methods
Equivariant architectures are frequently combined with standard graph neural network message-passing frameworks, using equivariant operations specifically for the geometric (3D coordinate-dependent) components of the message-passing update while using standard, non-equivariant operations for purely categorical or non-geometric node and edge features, allowing a single unified architecture to handle both geometric and non-geometric aspects of a given prediction task.
Equivariant denoising diffusion models combine the generative modeling framework of diffusion models with equivariant architectural principles to generate 3D molecular structures or conformations, using an equivariant network to parameterize the reverse denoising process, ensuring that the resulting generative model produces a properly rotation-and-translation-invariant distribution over generated structures rather than implicitly favoring particular arbitrary orientations.
Equivariant architectures also compose with attention mechanisms, as in the SE(3)-Transformer and Invariant Point Attention, replacing or augmenting standard scaled dot-product attention with geometry-aware attention computations that respect the relevant symmetry group, combining the long-range, content-based reasoning strengths of attention mechanisms with the sample-efficiency and physical consistency benefits of built-in geometric equivariance.
## Future Research Directions
Reducing the computational overhead of higher-order equivariant architectures, through more efficient tensor-product implementations, sparsity-exploiting techniques, or hybrid architectures that selectively apply expensive higher-order equivariant operations only where most beneficial, remains an active area of systems and algorithms research, aiming to make the accuracy benefits of higher-order equivariance accessible at the scale required for very large molecular systems, materials simulations, or high-resolution 3D scenes.
Extending equivariant architectures to more general and higher-dimensional symmetry groups beyond the standard rotation and permutation groups, including gauge symmetries relevant to certain physics applications and more general Lie group symmetries relevant to other scientific domains, is an active theoretical research direction seeking to broaden the applicability of geometric deep learning principles beyond the currently dominant 3D molecular and structural biology use cases.
Combining equivariant architectural principles with large-scale foundation-model-style pretraining, analogous to how large language models are pretrained on broad text corpora before fine-tuning on specific tasks, is an emerging direction in molecular and materials science, aiming to build broadly capable, equivariant foundation models pretrained across diverse chemical and structural data that can then be efficiently fine-tuned or adapted to a wide range of downstream molecular property prediction and structure generation tasks.
## Summary & Key Takeaways
Equivariant neural networks build geometric symmetries, such as rotation, translation, and permutation invariance, directly into a network's architecture, guaranteeing that transforming an input by a symmetry operation produces a correspondingly transformed output, rather than requiring the network to learn this consistency approximately from data through augmentation. This property is especially valuable in molecular, structural biology, and physical simulation applications, where the relevant symmetries are exact physical facts about the data-generating process rather than merely convenient approximations, and where training data is often too limited to allow a non-equivariant architecture to learn robust approximate symmetry through augmentation alone.
The field spans a spectrum of architectural sophistication, from lightweight scalar/vector architectures like EGNN, through steerable CNNs respecting rotational symmetry in convolutional settings, to fully general spherical-harmonic-and-Clebsch-Gordan-based architectures like Tensor Field Networks and SE(3)-Transformers capable of representing arbitrarily high-order angular relationships at increasing computational cost. AlphaFold2's Invariant Point Attention stands as the most prominent real-world demonstration of these principles' practical impact, and ongoing research continues to seek more computationally efficient equivariant architectures, broader symmetry group coverage, and large-scale equivariant foundation model pretraining.
Keywords: equivariant neural networks, geometric deep learning, SE(3) equivariance, SO(3) rotation group, permutation equivariance, Tensor Field Networks, SE(3)-Transformer, E(n)-Equivariant Graph Neural Network, EGNN, steerable CNN, spherical harmonics, Clebsch-Gordan coefficients, Invariant Point Attention, AlphaFold2, equivariant diffusion model, molecular dynamics potential, QM9 benchmark, Open Catalyst Project, protein structure prediction, Wigner D-matrix
---
## Appendix: Practical Labs
### Lab 1: Verifying permutation equivariance of a graph neural network layer
import numpy as np
def gnn_message_passing_layer(node_features, adjacency, W):
"""A minimal permutation-equivariant GNN layer: each node's updated
feature is a linear transform of the SUM of its neighbors' features
(sum aggregation is permutation-invariant to the order neighbors are
visited, which is what makes the overall layer permutation-equivariant
with respect to node relabeling)."""
aggregated = adjacency @ node_features # sum over neighbors
return aggregated @ W
def test_gnn_layer_is_permutation_equivariant():
rng = np.random.RandomState(0)
n_nodes, feat_dim, out_dim = 6, 4, 3
node_features = rng.randn(n_nodes, feat_dim)
adjacency = (rng.uniform(0, 1, (n_nodes, n_nodes)) > 0.5).astype(float)
adjacency = np.triu(adjacency, 1)
adjacency = adjacency + adjacency.T # symmetric adjacency matrix
W = rng.randn(feat_dim, out_dim)
output_original = gnn_message_passing_layer(node_features, adjacency, W)
# Apply a random permutation to the nodes
perm = rng.permutation(n_nodes)
permuted_features = node_features[perm]
permuted_adjacency = adjacency[np.ix_(perm, perm)]
output_on_permuted_input = gnn_message_passing_layer(permuted_features, permuted_adjacency, W)
# Equivariance requires: f(permute(x)) == permute(f(x))
output_original_permuted = output_original[perm]
max_diff = np.max(np.abs(output_on_permuted_input - output_original_permuted))
print(f"Max difference between f(permute(x)) and permute(f(x)): {max_diff:.2e}")
assert max_diff < 1e-10, "GNN layer with sum-aggregation should be exactly permutation-equivariant"
print("Permutation equivariance test passed.")
if __name__ == "__main__":
test_gnn_layer_is_permutation_equivariant()### Lab 2: Verifying SO(3) rotation equivariance of a simple vector-feature update
import numpy as np
def random_rotation_matrix(rng):
"""Generates a uniformly random 3D rotation matrix via QR decomposition
of a random Gaussian matrix, with a sign correction to ensure a proper
rotation (determinant +1, not a reflection)."""
A = rng.randn(3, 3)
Q, R = np.linalg.qr(A)
Q = Q @ np.diag(np.sign(np.diag(R)))
if np.linalg.det(Q) < 0:
Q[:, 0] *= -1
return Q
def equivariant_vector_update(positions_i, positions_j):
"""A simple SO(3)-equivariant update (as used in EGNN-style architectures):
the updated vector feature is a scalar function of the (rotation-invariant)
distance, multiplied by the (rotation-equivariant) relative position vector."""
rel_vec = positions_i - positions_j
dist = np.linalg.norm(rel_vec, axis=-1, keepdims=True)
scalar_weight = 1.0 / (1.0 + dist) # any function of the invariant distance
return scalar_weight * rel_vec
def test_vector_update_is_so3_equivariant():
rng = np.random.RandomState(1)
n_points = 10
positions_i = rng.randn(n_points, 3)
positions_j = rng.randn(n_points, 3)
output_original = equivariant_vector_update(positions_i, positions_j)
R = random_rotation_matrix(rng)
rotated_positions_i = positions_i @ R.T
rotated_positions_j = positions_j @ R.T
output_on_rotated_input = equivariant_vector_update(rotated_positions_i, rotated_positions_j)
# Equivariance requires: f(R @ x) == R @ f(x)
output_original_rotated = output_original @ R.T
max_diff = np.max(np.abs(output_on_rotated_input - output_original_rotated))
print(f"Max difference between f(R*x) and R*f(x): {max_diff:.2e}")
assert max_diff < 1e-10, "Distance-weighted relative-vector update should be exactly SO(3)-equivariant"
# Sanity check: verify a NON-equivariant baseline (e.g. treating each raw
# coordinate independently through an asymmetric linear map) actually
# FAILS this test, confirming our test methodology can detect violations
W_bad = rng.randn(3, 3) # arbitrary, non-equivariant linear mixing of xyz
bad_output_original = positions_i @ W_bad
bad_output_on_rotated = rotated_positions_i @ W_bad
bad_output_original_rotated = bad_output_original @ R.T
bad_max_diff = np.max(np.abs(bad_output_on_rotated - bad_output_original_rotated))
assert bad_max_diff > 1e-6, "Arbitrary linear mixing should NOT be equivariant, confirming the test can detect violations"
print(f"(Sanity check) non-equivariant baseline max difference: {bad_max_diff:.4f}")
print("SO(3) rotation equivariance test passed.")
if __name__ == "__main__":
test_vector_update_is_so3_equivariant()### Lab 3: Invariant vs. equivariant feature construction for molecular-style inputs
import numpy as np
def random_rotation_matrix(rng):
A = rng.randn(3, 3)
Q, R = np.linalg.qr(A)
Q = Q @ np.diag(np.sign(np.diag(R)))
if np.linalg.det(Q) < 0:
Q[:, 0] *= -1
return Q
def compute_invariant_features(positions):
"""Computes rotation-and-translation-INVARIANT features from a point
cloud: pairwise distances only depend on relative geometry, not on the
absolute position or orientation of the point cloud."""
diffs = positions[:, None, :] - positions[None, :, :]
distances = np.linalg.norm(diffs, axis=-1)
return distances
def compute_equivariant_features(positions):
"""Computes rotation-EQUIVARIANT (but translation-invariant) features:
pairwise relative position vectors, which rotate along with the input
but do not depend on absolute translation."""
diffs = positions[:, None, :] - positions[None, :, :]
return diffs
def test_invariant_features_unchanged_equivariant_features_rotate():
rng = np.random.RandomState(2)
n_atoms = 8
positions = rng.randn(n_atoms, 3) * 3.0
invariant_original = compute_invariant_features(positions)
equivariant_original = compute_equivariant_features(positions)
R = random_rotation_matrix(rng)
translation = rng.randn(3) * 5.0
transformed_positions = positions @ R.T + translation
invariant_transformed = compute_invariant_features(transformed_positions)
equivariant_transformed = compute_equivariant_features(transformed_positions)
# Invariant features (pairwise distances) should be COMPLETELY unchanged
# by rotation AND translation
invariant_diff = np.max(np.abs(invariant_transformed - invariant_original))
print(f"Max change in invariant (distance) features under rotation+translation: {invariant_diff:.2e}")
assert invariant_diff < 1e-10, "Pairwise distances must be exactly invariant to rotation and translation"
# Equivariant features (relative position vectors) should rotate exactly
# with R, and be unaffected by translation (since they are differences)
equivariant_original_rotated = np.einsum('ijc,dc->ijd', equivariant_original, R)
equivariant_diff = np.max(np.abs(equivariant_transformed - equivariant_original_rotated))
print(f"Max difference between rotated equivariant features and true transformed features: {equivariant_diff:.2e}")
assert equivariant_diff < 1e-10, "Relative position vectors must rotate exactly with R and be translation-invariant"
print("Invariant vs. equivariant feature construction test passed.")
if __name__ == "__main__":
test_invariant_features_unchanged_equivariant_features_rotate()### Lab 4: Sample efficiency benefit of equivariance vs. augmentation-based approximate invariance
import numpy as np
def random_rotation_matrix_2d(theta):
"""A 2D rotation matrix by angle theta, used for a simplified illustrative
comparison (full 3D rotation groups require more machinery, but the core
sample-efficiency argument is identical and easier to see clearly in 2D)."""
c, s = np.cos(theta), np.sin(theta)
return np.array([[c, -s], [s, c]])
def true_invariant_target(positions):
"""A genuinely rotation-invariant target function: the sum of squared
distances from the origin, which does not depend on orientation."""
return np.sum(positions ** 2)
def fit_invariant_model(train_positions, train_targets):
"""An exactly rotation-invariant model: predicts based only on the
(invariant) norm of the input, fit via simple least-squares on the
single invariant feature norm^2."""
features = np.sum(train_positions ** 2, axis=1, keepdims=True)
# Closed-form least squares: target = a * feature (single coefficient)
a = np.sum(features[:, 0] * train_targets) / np.sum(features[:, 0] ** 2)
return a
def fit_noninvariant_model(train_positions, train_targets):
"""A generic (non-invariant) linear model over raw xy coordinates and
their squares, which must learn approximate rotational invariance purely
from the training data it happens to see."""
X = np.stack([train_positions[:, 0] ** 2, train_positions[:, 1] ** 2,
train_positions[:, 0] * train_positions[:, 1]], axis=1)
coeffs, _, _, _ = np.linalg.lstsq(X, train_targets, rcond=None)
return coeffs
def predict_noninvariant(positions, coeffs):
X = np.stack([positions[:, 0] ** 2, positions[:, 1] ** 2,
positions[:, 0] * positions[:, 1]], axis=1)
return X @ coeffs
def test_invariant_model_generalizes_better_with_limited_orientation_coverage():
rng = np.random.RandomState(3)
# Training data: a small number of NOISY observations sampled only within
# a very NARROW range of orientations (0 to 15 degrees), simulating a
# realistic scenario where training data is limited in both quantity and
# orientation coverage, and labels carry measurement noise. Under these
# realistic conditions, a generic model without the correct invariance
# built in cannot reliably identify the true underlying invariant function
# from such a narrow slice of orientations, while the invariant model,
# having only one free parameter constrained by the correct functional
# form, remains robust to the noise.
n_train = 12
radii = rng.uniform(1, 5, n_train)
angles = rng.uniform(0, np.pi / 12, n_train) # very narrow angular range
train_positions = np.stack([radii * np.cos(angles), radii * np.sin(angles)], axis=1)
noise_std = 0.5
train_targets = np.array([true_invariant_target(p) for p in train_positions]) \
+ rng.normal(0, noise_std, n_train)
invariant_coeff = fit_invariant_model(train_positions, train_targets)
noninvariant_coeffs = fit_noninvariant_model(train_positions, train_targets)
# Test data: points at the SAME radii but rotated to a completely different,
# UNSEEN range of orientations (e.g. 90 to 120 degrees)
n_test = 100
test_radii = rng.uniform(1, 5, n_test)
test_angles = rng.uniform(np.pi / 2, np.pi / 2 + np.pi / 6, n_test)
test_positions = np.stack([test_radii * np.cos(test_angles), test_radii * np.sin(test_angles)], axis=1)
test_targets = np.array([true_invariant_target(p) for p in test_positions])
invariant_preds = invariant_coeff * np.sum(test_positions ** 2, axis=1)
noninvariant_preds = predict_noninvariant(test_positions, noninvariant_coeffs)
invariant_mse = np.mean((invariant_preds - test_targets) ** 2)
noninvariant_mse = np.mean((noninvariant_preds - test_targets) ** 2)
print(f"Invariant model test MSE (unseen orientations): {invariant_mse:.6f}")
print(f"Non-invariant model test MSE (unseen orientations): {noninvariant_mse:.6f}")
assert invariant_mse < noninvariant_mse * 0.1, (
"A model with the correct invariance built in should generalize far better "
"to unseen orientations than a generic model that must infer invariance "
"purely from a narrow slice of training orientations"
)
print("Equivariance sample-efficiency test passed.")
if __name__ == "__main__":
test_invariant_model_generalizes_better_with_limited_orientation_coverage()