Skip to content

Official PyTorch and JAX Implementation of "Harmonics of Learning: Universal Fourier Features Emerge in Invariant Networks"

Notifications You must be signed in to change notification settings

sophiaas/spectral-universality

Repository files navigation

Harmonics of Learning

Rotational harmonics

Description

This repository accompanies the paper:

Harmonics of Learning: Universal Fourier Features Emerge in Invariant Networks (2024)

Here, we provide an implementation of a complex-valued Power-Spectral Network trained via contrastive learning on an invariance objective for a finite group. As shown in the companion paper, at convergence the network learns all the irreducible unitary representations of the group. In particular, the multiplication table can be extracted from its weights.

We provide implementations of the model and its training in both PyTorch and JAX.

Setup

python 3.8 
pip install -r requirements.txt

Groups

The file groups.py provides implementations of various finite groups, including cyclic, dihedral and symmetric.

Training

In order to train the models in PyTorch and in JAX, run the files train_torch.py and train_JAX.py respectively. The training parameters are set at the beginning of these files.

About

Official PyTorch and JAX Implementation of "Harmonics of Learning: Universal Fourier Features Emerge in Invariant Networks"

Topics

Resources

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published

Languages