Learning Symmetries via Weight-Sharing with Doubly Stochastic Tensors

TL;DR

This paper introduces a method to learn symmetry-based weight sharing in neural networks using doubly stochastic matrices, enabling dynamic discovery of symmetries and improving model flexibility.

Contribution

It proposes a novel approach to learn weight-sharing schemes via learnable doubly stochastic matrices, generalizing group equivariance to discover symmetries adaptively.

Findings

Permutation matrices converge to regular group representations with strong dataset symmetries

Method effectively captures partial symmetries in data

Learned symmetries improve model generalization and robustness

Abstract

Group equivariance has emerged as a valuable inductive bias in deep learning, enhancing generalization, data efficiency, and robustness. Classically, group equivariant methods require the groups of interest to be known beforehand, which may not be realistic for real-world data. Additionally, baking in fixed group equivariance may impose overly restrictive constraints on model architecture. This highlights the need for methods that can dynamically discover and apply symmetries as soft constraints. For neural network architectures, equivariance is commonly achieved through group transformations of a canonical weight tensor, resulting in weight sharing over a given group $G$. In this work, we propose to learn such a weight-sharing scheme by defining a collection of learnable doubly stochastic matrices that act as soft permutation matrices on canonical weight tensors, which can take regular group representations as a special case. This yields learnable kernel transformations that are jointly optimized with downstream tasks. We show that when the dataset exhibits strong symmetries, the permutation matrices will converge to regular group representations and our weight-sharing networks effectively become regular group convolutions. Additionally, the flexibility of the method enables it to effectively pick up on partial symmetries.