Bayesian 3D Steerable CNNs: Enabling Equivariance and Uncertainty Quantification Simultaneously

Technical Summary: Bayesian 3D Steerable CNNs

Problem Statement
Steerable Convolutional Neural Networks (Steerable-CNNs) have established themselves as a powerful paradigm for 3D data processing, guaranteeing SE(3)-equivariance by parameterizing convolutional kernels as linear combinations of steerable basis functions. This architectural constraint ensures that transforming the input (e.g., via 3D rotation) produces a correspondingly transformed output, significantly reducing the hypothesis space and improving generalization in domains like shape recognition and molecular property prediction. However, existing Steerable-CNNs are entirely deterministic. They produce single point estimates without associated confidence measures, rendering them unsuitable for safety-critical applications (e.g., medical imaging, autonomous systems) where uncertainty quantification (UQ) is essential.

A primary challenge in developing Bayesian Steerable-CNNs is that naively introducing stochasticity into the network weights typically breaks the equivariance property. If convolution kernels are randomly sampled, they generally violate the "intertwiner" or steerability constraint, meaning the resulting stochastic realizations would not be equivariant, even if their mean is.

Methodology
The authors propose a Bayesian Steerable-CNN framework that resolves this conflict by placing posterior distributions over the scalar coefficients of the steerable basis functions, rather than the kernel matrices themselves.

1. Stochastic Equivariant Kernels: The convolutional kernel $K(x)$ is defined as a linear combination of deterministic, pre-computed steerable basis functions $\{\kappa_b(x)\}$ (derived from Wigner D-matrices for SO(3)):
   $$K(x; w) = \sum_{b=1}^{B} w_b \kappa_b(x)$$
   The basis functions $\kappa_b$ are fixed and satisfy the steerability constraint by construction. The coefficients $w = (w_1, \dots, w_B)$ are treated as random variables.
2. Preservation of Equivariance: The paper proves (Proposition 1) that if the basis functions satisfy the steerability constraint, any linear combination of them—regardless of whether the coefficients are deterministic or sampled from a probability distribution—remains equivariant. This allows for stochastic kernels without architectural modifications to enforce equivariance.
3. Variational Inference: The model employs the Bayes-by-Backpropagation approach. A mean-field Gaussian variational posterior $q(w|\theta)$ is assumed over the coefficients, parameterized by a mean $\mu$ and a log-variance $\rho$. The model is trained by maximizing the Evidence Lower Bound (ELBO), which consists of a cross-entropy loss term and a Kullback-Leibler (KL) divergence term regularizing the posterior against a Gaussian prior.
4. Uncertainty Decomposition: At inference, the model performs $N$ stochastic forward passes. The total predictive uncertainty is decomposed into:
   • Aleatoric Uncertainty: Inherent noise in the data (captured by the average entropy of individual predictions).
   • Epistemic Uncertainty: Model uncertainty due to limited data (captured by the variance of predictions across the $N$ samples).

Key Contributions

• Novel Bayesian Formulation: The first integration of stochastic inference with steerable convolutions, obtaining the posterior distribution of kernel basis coefficients while maintaining exact SE(3)-equivariance.
• Theoretical Proof: A formal proof demonstrating that stochasticity in the coefficient space does not violate the intertwiner constraint, ensuring that every sampled kernel remains equivariant.
• Decomposable Uncertainty: A framework that provides calibrated, decomposable uncertainty estimates (epistemic vs. aleatoric) for 3D equivariant models.

Experimental Results
The model was evaluated on the ModelNet10 dataset (3D CAD models) using an architecture built on the escnn library.

• Classification Performance: The Bayesian model achieved a test accuracy of 86.34%, which is competitive with deterministic counterparts.
• Calibration: The model demonstrated excellent calibration with an Expected Calibration Error (ECE) of 0.0263, indicating that predicted confidence scores closely match empirical accuracy.
• Robustness to Distribution Shift: Under additive Gaussian noise, the Bayesian model exhibited superior robustness compared to the deterministic baseline. At a high noise level ($\epsilon_{noise} = 0.9$), the Bayesian model retained 81.83% accuracy, outperforming the deterministic model (75.66%) by 6.17%.
• Uncertainty Utility:
  • Selective Prediction: By filtering out samples with high uncertainty, the model achieved a ~4% accuracy gain on 84% of the test dataset (samples with uncertainty below a specific threshold).
  • Semantic Meaning: A statistically significant negative Spearman rank correlation (-0.3989) was found between epistemic uncertainty and prediction error, confirming that higher uncertainty reliably predicts higher classification error.
  • Uncertainty Convergence: During training, total uncertainty decreased steadily, while epistemic uncertainty remained consistently low, suggesting the model became confident in its learned weight distribution.

Significance and Claims
The paper claims to unify Bayesian uncertainty quantification with the inductive bias of equivariant CNNs. The authors emphasize that this is the first work to successfully integrate stochastic inference and steerable convolutions in a single discriminative model. The significance lies in enabling risk-aware decision-making in 3D domains: the model not only respects geometric symmetries but also provides principled, decomposable confidence estimates. This allows practitioners to distinguish between ambiguity inherent in the data (aleatoric) and a lack of model knowledge (epistemic), facilitating targeted model refinement and safer deployment in safety-sensitive applications. The framework achieves these benefits without sacrificing the equivariance guarantees that make Steerable-CNNs effective for 3D data.