Learning Invariances in Neural Networks

TL;DR

This paper introduces a method to learn invariances and equivariances in neural networks by optimizing over augmentation distributions, enabling models to adapt to the data's inherent symmetries across various tasks.

Contribution

It proposes a novel approach to automatically learn the appropriate invariances in neural networks through joint optimization of network and augmentation parameters.

Findings

Successfully recovers correct invariances in image tasks

Applies to regression, segmentation, and molecular prediction

Operates using only training data

Abstract

Invariances to translations have imbued convolutional neural networks with powerful generalization properties. However, we often do not know a priori what invariances are present in the data, or to what extent a model should be invariant to a given symmetry group. We show how to \emph{learn} invariances and equivariances by parameterizing a distribution over augmentations and optimizing the training loss simultaneously with respect to the network parameters and augmentation parameters. With this simple procedure we can recover the correct set and extent of invariances on image classification, regression, segmentation, and molecular property prediction from a large space of augmentations, on training data alone.