Lie Group Decompositions for Equivariant Neural Networks (ICLR 2024)

Implementation of "Lie Group Decompositions for Equivariant Neural Networks"

Mircea Mironenco, Patrick Forré

Overview • Installation • Experiments • Citation

Overview

Installation

Assuming torch (>= 2.0) and torchvision are installed & configured, simply run:

git clone https://github.com/mirceamironenco/rgenn.git
cd rgenn
pip install -r requirements.txt

To reproduce the main experiments, the minimal requirements are scipy einops tqdm.

Experiments

To reproduce the affNIST results from Section 5, for the $\mathbb{R}^{2} \rtimes \mathrm{SL}(2, \mathbb{R})$ group:

python train.py --amp --liegroup=sl2 --gsamples=10

Alternatively, setting --liegroup=gl2 will reproduce our results for the $\mathbb{R}^{2} \rtimes \mathrm{GL}^{+}(2, \mathbb{R})$ (full affine) group.

If multiple GPUs are available, training with DDP has been configured for single-node use:

torchrun --nnodes=1 --nproc_per_node=gpu train.py --amp --liegroup=sl2 --gsamples=10

Citation

@inproceedings{
mironenco2024lie,
title={Lie Group Decompositions for Equivariant Neural Networks},
author={Mircea Mironenco and Patrick Forr{\'e}},
booktitle={The Twelfth International Conference on Learning Representations},
year={2024},
url={https://openreview.net/forum?id=p34fRKp8qA}
}