Deep invariant networks with differentiable augmentation layers

TL;DR

This paper introduces a novel method for learning data invariances directly within neural networks using differentiable augmentation layers, offering a versatile, efficient alternative to bilevel optimization-based data augmentation techniques.

Contribution

It proposes learnable augmentation layers integrated into networks to learn invariances from training data, bypassing complex bilevel optimization methods.

Findings

Faster and easier to train than bilevel optimization methods.

Achieves comparable performance to state-of-the-art data augmentation techniques.

Produces invariances that are insensitive to model expressivity.

Abstract

Designing learning systems which are invariant to certain data transformations is critical in machine learning. Practitioners can typically enforce a desired invariance on the trained model through the choice of a network architecture, e.g. using convolutions for translations, or using data augmentation. Yet, enforcing true invariance in the network can be difficult, and data invariances are not always known a piori. State-of-the-art methods for learning data augmentation policies require held-out data and are based on bilevel optimization problems, which are complex to solve and often computationally demanding. In this work we investigate new ways of learning invariances only from the training data. Using learnable augmentation layers built directly in the network, we demonstrate that our method is very versatile. It can incorporate any type of differentiable augmentation and be applied to a broad class of learning problems beyond computer vision. We provide empirical evidence showing that our approach is easier and faster to train than modern automatic data augmentation techniques based on bilevel optimization, while achieving comparable results. Experiments show that while the invariances transferred to a model through automatic data augmentation are limited by the model expressivity, the invariance yielded by our approach is insensitive to it by design.