Equivariant Neural Tangent Kernels

TL;DR

This paper develops neural tangent kernels for equivariant neural networks, revealing their training dynamics and showing they can outperform non-equivariant models in image and quantum property prediction.

Contribution

It introduces a theoretical framework for neural tangent kernels of equivariant architectures, linking data augmentation and group convolutions, and demonstrates their practical advantages.

Findings

Equivariant NTKs share the same expected prediction as data-augmented models.

Equivariant NTKs outperform non-equivariant kernels in classification tasks.

The framework applies to roto-translations and 3D rotations, showing broad applicability.

Abstract

Little is known about the training dynamics of equivariant neural networks, in particular how it compares to data augmented training of their non-equivariant counterparts. Recently, neural tangent kernels (NTKs) have emerged as a powerful tool to analytically study the training dynamics of wide neural networks. In this work, we take an important step towards a theoretical understanding of training dynamics of equivariant models by deriving neural tangent kernels for a broad class of equivariant architectures based on group convolutions. As a demonstration of the capabilities of our framework, we show an interesting relationship between data augmentation and group convolutional networks. Specifically, we prove that they share the same expected prediction at all training times and even off-manifold. In this sense, they have the same training dynamics. We demonstrate in numerical experiments that this still holds approximately for finite-width ensembles. By implementing equivariant NTKs for roto-translations in the plane ($G=C_{n}\ltimes\mathbb{R}^{2}$) and 3d rotations ($G=\mathrm{SO}(3)$), we show that equivariant NTKs outperform their non-equivariant counterparts as kernel predictors for histological image classification and quantum mechanical property prediction.