3D Steerable CNNs: Learning Rotationally Equivariant Features in Volumetric Data

Maurice Weiler; Mario Geiger; Max Welling; Wouter Boomsma; Taco Cohen

Source-linked AI summary

3D Steerable CNNs: Learning Rotationally Equivariant Features in Volumetric Data

Maurice Weiler, Mario Geiger, Max Welling, Wouter Boomsma, Taco Cohen

arXiv:1807.02547v2Jul 6, 2018cs.LGstat.ML

TL;DR

The paper addresses how to build convolutional networks that respect rigid-body symmetries in 3D scientific data. It represents data with scalar, vector, and tensor fields and uses analytically derived steerable kernels for SE(3)-equivariant convolutions. The resulting networks improve amino acid propensity prediction and protein structure classification while providing a general equivariant mapping framework.

Problem
Machine-learning models for scientific data should respect exact or approximate symmetries, including the SE(3) symmetry present in protein-related tasks.
Method
The model uses fields over R3 and parameterizes equivariant convolution kernels as linear combinations of analytically derived steerable basis kernels.
Results
3D Steerable CNNs improve amino acid propensity prediction and consistently outperform a strong CNN baseline for protein structure classification across training-set sizes.
Takeaways & Limitations
Equivariant convolutions provide the most general equivariant linear maps between the paper’s fields, while requiring only a minor adaptation to conventional 3D-CNN code.
Takeaways & Limitations
The choice of feature types and multiplicities remains a hyperparameter, and its detailed investigation is left for future work.

Abstract

from arXiv · show

We present a convolutional network that is equivariant to rigid body motions. The model uses scalar-, vector-, and tensor fields over 3D Euclidean space to represent data, and equivariant convolutions to map between such representations. These SE(3)-equivariant convolutions utilize kernels which are parameterized as a linear combination of a complete steerable kernel basis, which is derived analytically in this paper. We prove that equivariant convolutions are the most general equivariant linear maps between fields over R^3. Our experimental results confirm the effectiveness of 3D Steerable CNNs for the problem of amino acid propensity prediction and protein structure classification, both of which have inherent SE(3) symmetry.

1 Introduction

The paper develops SE(3)-equivariant networks so models respect symmetries in scientific data, using steerable kernels within a conventional 3D CNN framework.

Motivation: SE(3)-equivariant networks address the need for machine-learning models to respect known symmetries in natural-science data.The paper connects symmetry-respecting models to physically meaningful predictions.
Approach: 3D Steerable CNNs parameterize each convolution kernel as a linear combination of pre-computed steerable basis kernels.This incorporates equivariance without substantially departing from standard 3D-CNN engineering practices.
Approach: The model represents inputs as fields over R3 and uses steerable filters to map between scalar, vector, and tensor-like representations.Its kernels are modeled as tensor fields satisfying an equivariance constraint.
Experiments: The network is evaluated on amino acid preference prediction and protein-structure classification, exploiting tasks with inherent SE(3) symmetry.The first task concerns atomic environments; the second concerns overall protein shape.

2 Related Work

Related equivariant-network research varies by symmetry group and feature geometry; this work targets full SE(3) equivariance on regular 3D grids.

Equivariant networks: Equivariant networks can be organized by the symmetry group they support and the geometrical features they transform.Steerable G-CNNs use general fields over homogeneous spaces, including R3.
Prior 3D models: Earlier voxel-based 3D G-CNNs achieved data-efficiency gains but were equivariant only to certain discrete rotations.This distinguishes them from models targeting full SE(3) equivariance.
Closest related work: Tensor Field Networks and N-Body Networks achieve full SE(3) equivariance on irregular point clouds, whereas 3D Steerable CNNs operate on regular 3D grids.Regular grids can be processed more efficiently on current hardware, while point clouds are more general.
Dimensional analogue: The two-dimensional analogue uses SO(2) irreducible features and circular harmonics, paralleling SO(3) features and spherical harmonics here.The correspondence is between the dimensional rotation group and the associated harmonic filters.
Earlier signal representations: Earlier spherical-tensor approaches used fixed rather than learned expansions of signals, unlike the learned representation in 3D Steerable CNNs.Related methods also addressed volumetric biomedical-image processing.

3 Convolutional feature spaces as fields

The paper models convolutional features as fields whose fibers contain geometrical quantities with specified rotation laws, then decomposes those laws into irreducible SO(3) components.

Fields: A layer’s feature maps form compactly supported continuous functions from R3 to a fiber R^Kn, with each position assigned a feature vector.The feature space Fn is the vector space of these functions.
Fields: The fibers are structured as stacked scalars, vectors, and tensors, with each geometrical quantity forming a field across space.Channels belonging to different quantities transform independently.
Disentangling: Weyl’s principle distinguishes geometrical quantities because rigid motions mix channels within one quantity but not across different quantities.This formalizes a form of disentangling for field representations.
Transformations: Under rotation, a vector field both moves each vector to a new position and rotates the vector by a 3 × 3 matrix.The transformation is [π(r)f](x) := ρ(r)f(r−1x).
Transformations: A field type corresponds one-to-one with a representation ρ describing how each fiber transforms under rotations.Rigid motions combine rotations and translations through the induced SE(3) representation.
Irreducible features: Any SO(3) representation decomposes into Wigner-D irreducibles of dimensions 2l + 1, and the network uses only these irreducible features.For rank-2 tensors, the decomposition yields irreducible components of dimensions 1, 3, and 5.

4 SE(3)-Equivariant Networks

The paper constructs SE(3)-equivariant linear maps by characterizing steerable kernels and parameterizing them with a complete analytically derived basis. These kernels can be combined in standard 3D convolutions, while equivariant nonlinearities and low-pass filtering support equivariant networks in practice.

4.1 The Subspace of Equivariant Kernels: Equivariant linear maps between fields are exactly cross-correlations with rotation-steerable kernels.The kernel must satisfy a rotational transformation constraint, while cross-correlation already provides translation equivariance.
4.1 The Subspace of Equivariant Kernels: The equivariant kernel space is a vector space, so kernels can be parameterized as linear combinations of basis kernels.This basis expansion turns the equivariance constraint into fixed structure with learnable expansion coefficients.
4.2 Solving for the Equivariant Kernel Basis: For irreducible input and output features of orders j and l, the kernel decomposes into blocks indexed by angular frequencies |j − l| ≤ J ≤ j + l.The tensor-product representation D^j ⊗ D^l decomposes into irreducible representations, separating the kernel constraints into invariant subspaces.
4.2 Solving for the Equivariant Kernel Basis: For j = l = 1, the three J values produce 3 × 3 angular matrices corresponding to identity, curl, and gradient-of-divergence patterns.Each matrix contributes one learnable parameter per radial basis function ϕ_m.
4.2 Solving for the Equivariant Kernel Basis: Each basis kernel combines a spherical-harmonic angular component with an arbitrary radial function, instantiated using radial basis functions ϕ_m.The resulting basis is indexed by radial index m and frequency index J, then transformed back to the original feature basis.
4.3 Equivariant Nonlinearities: Equivariance must also hold across nonlinearities and discretization, requiring specialized nonlinearities and low-pass filtering before downsampling.The authors report significant improvements in rotational equivariance from smoothing feature maps before subsampling, compared with standard strided downsampling.

5 Experiments

Experiments test equivariance, 3D shape recognition, amino acid preference prediction, and protein structure classification. The model generalizes across rotations and remains accurate with substantially reduced training data and parameter count.

Equivariance tests: 99±2% accuracy on randomly rotated Tetris blocks demonstrates generalization beyond the fixed training orientation.The 4-layer 3D Steerable CNN was trained on fixed-orientation voxel grids and tested on randomly rotated blocks.
3D shape classification: SHREC17 evaluates rotated 3D shape classification across 51,300 models, 55 classes, and 64x64x64 voxel grids.Its perturbed category contains arbitrarily rotated images, making it a test of rotational handling on complex objects.
Amino acid propensity prediction: The model replaces conventional convolutional layers with equivalent 3D steerable layers for amino acid preference prediction.This experiment keeps the original network architecture while substituting the convolutional operators.
CATH protein structure classification: A 3D Steerable CNN outperforms a strong conventional CNN baseline on protein structure classification despite having 100 times fewer parameters.The CATH experiment compares a 143,560-parameter steerable network with a 15,878,764-parameter conventional baseline.
CATH protein structure classification: Accuracy improvements remain robust as the training splits are reduced, with the proposed model outperforming the baseline on fractions of the training data.The networks retain the same architecture while regularization parameters are re-optimized for each training-set reduction.

6 Conclusion

The paper presents 3D Steerable CNNs as SE(3)-equivariant networks with a general equivariant convolution theory. Experiments show equivariance, accuracy, and data efficiency in amino acid propensity prediction and protein structure classification.

6 Conclusion: 3D Steerable CNNs represent data using multiple field types over R3 and provide a class of SE(3)-equivariant networks.The architecture uses scalar-, vector-, and tensor-like fields to encode features.
6 Conclusion: SO(3)-steerable convolutions are proved to provide the most general equivariant mapping between fields, establishing a universal architecture class.The conclusion frames the result as a comprehensive theory of 3D Steerable CNNs.
6 Conclusion: The networks require only minor code adaptations, can be converted to conventional 3D CNNs after training, and show strong accuracy and data efficiency.These properties are reported across amino acid propensity prediction and protein structure classification.

1 Design choices

The design combines equivariant field representations with specialized nonlinearities and normalization. Feature types and multiplicities control representational capacity, while gated nonlinearities were preferred in practice.

Feature types and multiplicities: Feature types and multiplicities are hyperparameters analogous to channel choices in conventional CNNs.The network doubles multiplicities when downsampling, while input and output types are prescribed by the task.
Feature types and multiplicities: Using only scalar fields restricts kernels to isotropic forms, whereas higher-order representations permit more complex kernels.The choice of representation therefore controls the geometric structure of allowable filters.
Nonlinearities: Scalar fields use conventional ReLU nonlinearities, while nonscalar fields require equivariant alternatives.Scalar features do not transform under rotations, so elementwise nonlinearities preserve equivariance for them.
Nonlinearities: Norm nonlinearities preserve equivariance because orthogonal representations preserve feature norms under rotations.Applying a nonlinearity to the norm therefore commutes with the group transformation.
Nonlinearities: Norm nonlinearities converged more slowly and were omitted from final experiments because suitable bias initialization lacked a proper scale.The bias acts as a threshold on feature norms while preserving the orientation of larger vectors.
Nonlinearities: Tensor-product nonlinearities remain equivariant and produce outputs transforming under the tensor-product representation of the inputs.A change of basis can express the tensor-product representation in the framework’s irreducible feature basis.
Nonlinearities: Gated nonlinearities scale feature vectors using data-dependent scalar gates and performed better in practice than the other tested options.Gate fields use sigmoid, while scalar output fields can use ReLU.

2 Reduced parameter cost of 3D Steerable CNNs

The paper examines whether the steerable network’s performance advantage persists despite its much smaller parameter count. A reduced-capacity conventional CNN comparison indicates that the relative gain increases when the baseline is restricted.

Reduced parameter cost: The authors trained conventional CNNs with fewer filters to test whether the steerable model’s advantage depended on conventional-network overfitting.The main experiments had already shown better performance for the steerable CNN with many fewer parameters.
Reduced parameter cost: The relative performance gain of the 3D Steerable CNN increases dramatically when the conventional CNN is restricted to fewer parameters.This comparison directly addresses whether parameter count alone explains the observed advantage.

3 The Tetris experiment

The Tetris experiment tests whether the 3D Steerable CNN generalizes from fixed-orientation training to randomly rotated 3D inputs. Its architecture uses constrained equivariant kernels and far fewer parameters than the conventional CNN.

The Tetris network uses four hidden layers with kernel size 5, padding 4, and no batch normalization.
The SE3 network has 41k parameters, whereas the conventional CNN has 6M because its kernels are unconstrained.
The architecture applies stride-2 downsampling between layers 1–2 and 2–3, followed by global average pooling before output.
Table 2 reports test accuracy for classifying rotated Tetris pieces, averaged with standard deviations over 17 runs.

4 3D Model classification

The 3D model classification experiment evaluates architectures for voxelized 3D shapes using a network-selection process and compares results with competition and published baselines. The experiment uses 64 × 64 × 64 voxel inputs.

The model-selection process varied depth, multiplicities, dropout, low-pass filtering or stride, and initialization method across 10 models.
The architecture uses kernel size 5, padding 4, batch normalization, and no low-pass filters, with 142k parameters.
The experiment converts 3D models into voxel grids of size 64 × 64 × 64.
Table 4 compares the model's results with the original competition and two other articles.
Table 3 documents the architecture used for the 3D Model experiment, including stride-2 reductions and global average pooling.
Table 4 presents the results of the SHREC17 experiment.

5 The CATH experiment

The CATH experiment represents protein structures as bounded voxelized Cα-atom density fields and compares a minimal equivariant network with a conventional baseline. The steerable model uses the same 3D channel counts while containing substantially fewer parameters.

Protein structures are simplified to Cα atoms and represented by Gaussian densities on a 503-voxel grid with 0.2 nm voxel spacing.
The input is limited to a 5 nm sphere around each protein's center of mass, excluding structures that extend beyond that boundary.
The dataset uses 10 splits with no cross-split overlap at the CATH superfamily level, while covering all 10 architectures.
The conventional CATH baseline is a ResNet34-like 3D architecture with half the original channels and global pooling for translational invariance.
The steerable network follows the same ResNet template, replacing each layer with an equivariant equivalent and matching the baseline's number of 3D channels.
The minimal steerable model contains 143,560 parameters, more than a factor hundred fewer than the baseline.