mirceamironenco / rgenn

Official implementation of Lie Group Decompositions for Equivariant Neural Networks (ICLR 2024)

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}
}