The Lie Derivative for Measuring Learned Equivariance

TL;DR

This work introduces the Lie derivative-based Local Equivariance Error (LEE) to quantify how much neural networks violate continuous symmetries, enabling precise, layerwise attribution of equivariance errors across CNNs, Vision Transformers, and Mixer architectures. By decomposing the Lie derivative through network layers, the authors reveal that aliasing arising from downsampling and pointwise nonlinearities is a major source of equivariance violations, and that larger, better-trained models tend to be more equivariant, sometimes making transformers more equivariant than CNNs after training. The study spans hundreds of pretrained models and multiple transformations (translation, rotation, shear), showing that training scale and data often matter more than architectural changes for learned equivariance, and that there remains an equivariance gap on out-of-distribution data. These findings suggest that while architectural priors for equivariance can help, data-driven learning and anti-aliasing strategies may be more impactful in practice, with ongoing relevance for when exact symmetry properties are critical (e.g., molecular rotation invariance).

Abstract

Equivariance guarantees that a model's predictions capture key symmetries in data. When an image is translated or rotated, an equivariant model's representation of that image will translate or rotate accordingly. The success of convolutional neural networks has historically been tied to translation equivariance directly encoded in their architecture. The rising success of vision transformers, which have no explicit architectural bias towards equivariance, challenges this narrative and suggests that augmentations and training data might also play a significant role in their performance. In order to better understand the role of equivariance in recent vision models, we introduce the Lie derivative, a method for measuring equivariance with strong mathematical foundations and minimal hyperparameters. Using the Lie derivative, we study the equivariance properties of hundreds of pretrained models, spanning CNNs, transformers, and Mixer architectures. The scale of our analysis allows us to separate the impact of architecture from other factors like model size or training method. Surprisingly, we find that many violations of equivariance can be linked to spatial aliasing in ubiquitous network layers, such as pointwise non-linearities, and that as models get larger and more accurate they tend to display more equivariance, regardless of architecture. For example, transformers can be more equivariant than convolutional neural networks after training.

Figures (9)

• Figure 1: (Left): The Lie derivative measures the equivariance of a function under continuous transformations, here rotation. (Center): Using the Lie derivative, we quantify how much each layer contributes to the equivariance error of a model. Our analysis highlights surprisingly large contributions from non-linearities, which affects both CNNs and ViT architectures. (Right): Translation equivariance as measured by the Lie derivative correlates with generalization in classification models, across convolutional and non-convolutional architectures. Although CNNs are often noted for their intrinsic translation equivariance, ViT and Mixer models are often more translation equivariant than CNN models after training.
• Figure 2: Non-linearities generate new high-frequency harmonics.
• Figure 3: Lie derivatives can be computed using automatic differentiation. We show how a Lie derivative for continuous rotations can be implemented in PyTorch NEURIPS2019_9015. The implementation in our experiments differs slightly, for computational efficiency and to pass second-order gradients through grid_sample.
• Figure 4: Contributions to equivariance shown cumulatively by layer, in the order the layers occur in the network. Left: Convolutional architectures. In all the CNNs, much of the equivariance error comes from downsampling and non-linearities. Middle-Left: Non-convolutional architectures. The initial patch embedding, a strided convolution, is the largest contributor for the ViTs and Mixers. The rest of the error is distributed uniformly across other nonlinear operations. Middle-Right: ResNet-50 across different transformations as a percentage. Despite being designed for translation equivariance, the fraction of equivariance error produced by each layer is almost identical for other affine transformations, suggesting that aliasing is the primary source of equivariance error. Right: Comparing LEE with alternative metrics for translation equivariance. Using integer translations misses key contributors to equivariance errors, such as activations, while using fractional translations can lead to radically different outcomes depending on choice of normalization ($N$ or $\sqrt{N}$). LEE captures aliasing effects and has minimal design decisions.
• Figure 5: Equivariance metrics evaluated on the ImageNet test set. Left: Non-LEE equivariance metrics display similar trends to \ref{['fig:title-fig']}, despite using larger, multi-pixel transformations. Right: Norm of rotation and shear Lie derivatives. Across all architectures, models with strong generalization become more equivariant to many common affine transformations. Marker size indicates model size. Error bars show one standard error over test set images used in the equivariance calculation.

Theorems & Definitions (1)

• Theorem 1