Equivariance and Augmentation for Bayesian Neural Networks

2026-06-24

AI Summary

This work investigates how Bayesian neural networks under variational inference can achieve exact equivariance through data augmentation. For exponential-family variational distributions, the paper establishes, for the first time, theoretical conditions under which data augmentation induces equivariance. Building on this foundation, three symmetrization techniques are proposed, among which the “orbit expansion” method substantially enhances both model equivariance and overall performance. Theoretical analysis provides bounds on equivariance error, and extensive experiments demonstrate that the proposed approaches outperform existing baselines across multiple benchmarks, thereby validating the efficacy and practicality of the theoretical framework.

Abstract

Symmetries are important for many deep learning tasks, ranging from applications in the sciences to medical imaging. However, there is an ongoing debate about whether to impose symmetry constraints on the neural network architecture (yielding equivariant neural networks) or learn them from augmented training data. Although equivariant networks are well-studied theoretically, much less is known about data augmentation, since analyzing augmentation requires control over the training dynamics. Inspired by recent results that show that augmented infinite deep ensembles are exactly equivariant, we study data augmentation for Bayesian neural networks (BNNs) trained with variational inference. We focus on variational distributions in the exponential family and derive conditions under which exact equivariance is reached. We furthermore obtain bounds on the equivariance error and introduce three novel symmetrization techniques which boost the effect of data augmentation in this setting. We conduct extensive numerical experiments which show that one of our symmetrization methods (orbit expansion) outperforms the baseline in both equivariance and overall performance. Our code is available at github.com/dmw1998/augment-BNNs

Problem

• equivariance
• data augmentation
• Bayesian neural networks
• symmetry
• variational inference

Innovation

• equivariance
• data augmentation
• Bayesian neural networks
• variational inference
• symmetrization