{"id":13477065,"url":"https://github.com/mfinzi/equivariant-MLP","last_synced_at":"2025-03-27T04:32:14.041Z","repository":{"id":47245384,"uuid":"294766103","full_name":"mfinzi/equivariant-MLP","owner":"mfinzi","description":"A library for programmatically generating equivariant layers through constraint solving","archived":false,"fork":false,"pushed_at":"2023-05-08T00:30:03.000Z","size":20780,"stargazers_count":254,"open_issues_count":5,"forks_count":21,"subscribers_count":9,"default_branch":"master","last_synced_at":"2024-10-15T03:55:14.325Z","etag":null,"topics":["deep-learning","equivariance"],"latest_commit_sha":null,"homepage":"","language":"Jupyter Notebook","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"mit","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/mfinzi.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null}},"created_at":"2020-09-11T17:40:49.000Z","updated_at":"2024-10-13T11:40:07.000Z","dependencies_parsed_at":"2024-01-07T11:15:19.039Z","dependency_job_id":null,"html_url":"https://github.com/mfinzi/equivariant-MLP","commit_stats":{"total_commits":313,"total_committers":4,"mean_commits":78.25,"dds":0.009584664536741228,"last_synced_commit":"da6986f7766697423e9c9e974a6dfd9d8ae6e5e6"},"previous_names":[],"tags_count":3,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/mfinzi%2Fequivariant-MLP","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/mfinzi%2Fequivariant-MLP/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/mfinzi%2Fequivariant-MLP/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/mfinzi%2Fequivariant-MLP/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/mfinzi","download_url":"https://codeload.github.com/mfinzi/equivariant-MLP/tar.gz/refs/heads/master","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":245785494,"owners_count":20671622,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["deep-learning","equivariance"],"created_at":"2024-07-31T16:01:37.849Z","updated_at":"2025-03-27T04:32:14.034Z","avatar_url":"https://github.com/mfinzi.png","language":"Jupyter Notebook","funding_links":[],"categories":["Jupyter Notebook","Libraries"],"sub_categories":["New Libraries"],"readme":"\u003cdiv align=\"center\"\u003e\n\u003cimg src=\"docs/_static/emlp_logo4x.png\" width=\"350\" alt=\"logo\"/\u003e\n\u003c/div\u003e\n\n# A Practical Method for Constructing Equivariant Multilayer Perceptrons for Arbitrary Matrix Groups\n[![Documentation](https://readthedocs.org/projects/emlp/badge/)](https://emlp.readthedocs.io/en/latest/) | [![Paper](https://img.shields.io/badge/arXiv-2104.09459-red)](https://arxiv.org/abs/2104.09459) | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/mfinzi/equivariant-MLP/blob/master/docs/notebooks/colabs/all.ipynb) | \n[![codecov.io](https://codecov.io/github/mfinzi/equivariant-MLP/coverage.svg)](https://codecov.io/github/mfinzi/equivariant-MLP)\n| [![PyPI version](https://img.shields.io/pypi/v/emlp)](https://pypi.org/project/emlp/) \n\u003c!-- [![tests](https://github.com/mfinzi/equivariant-MLP/actions/workflows/python-package.yml/badge.svg?branch=dev)](https://github.com/mfinzi/equivariant-MLP/actions/workflows/python-package.yml) |  --\u003e\n\u003c!-- | [![codecov](https://codecov.io/gh/mfinzi/equivariant-MLP/branch/master/graph/badge.svg?token=DYAFHK68JX)](https://codecov.io/gh/mfinzi/equivariant-MLP) --\u003e\n\n\n*EMLP* is a jax library for the automated construction of equivariant layers in deep learning based on the ICML2021 paper [A Practical Method for Constructing Equivariant Multilayer Perceptrons for Arbitrary Matrix Groups](https://arxiv.org/abs/2104.09459). You can read the documentation [here](https://emlp.readthedocs.io/en/latest/).\n\u003c!-- #and paper [here](https://arxiv.org/abs/2104.09459).  --\u003e\n\n## What EMLP is great at doing\n\n- Computing equivariant linear layers between finite dimensional\nrepresentations. You specify the symmetry group (discrete, continuous,\nnon compact, complex) and the representations (tensors, irreducibles, induced representations, etc), and we will compute the basis of equivariant\nmaps mapping from one to the other.\n\n- Automatic construction of full equivariant models for small data. E.g.\nif your inputs and outputs (and intended features) are a small collection of elements like scalars, vectors, tensors, irreps with a total dimension less than 1000, then you will likely be able to use EMLP as a turnkey solution for making the model or atleast function as a strong baseline.\n\n- As a tool for building larger models, but where EMLP is just one component in a larger system. For example, using EMLP as the convolution kernel in an equivariant PointConv network.\n\n## What EMLP is not great at doing\n\n- An efficient implementation of CNNs, Deep Sets, typical translation + rotation equivariant GCNNs, graph neural networks.\n\n- Handling large data like images, voxel grids, medium-large graphs, point clouds.\n\nGiven the current approach, EMLP can only ever be as fast as an MLP. So if flattening the inputs into a single vector would be too large to train with an MLP, then it will also be too large to train with EMLP.\n\n--------------------------------------------------------------------------------\n\n# Showcasing some examples of computing equivariant bases\n\nWe provide a type system for representations. With the operators ρᵤ⊗ρᵥ, ρᵤ⊕ρᵥ, ρ* implemented as `*`,`+` and `.T` build up different representations. The basic building blocks for representations are the base vector representation `V` and tensor representations `T(p,q) = V**p*V.T**q`. \n\nFor any given matrix group and representation formed in our type system, you can get the equivariant basis with [`rep.equivariant_basis()`](https://emlp.readthedocs.io/en/latest/package/emlp.reps.html#emlp.reps.equivariant_basis) or a matrix which projects to that subspace with [`rep.equivariant_projector()`](https://emlp.readthedocs.io/en/latest/package/emlp.reps.html#emlp.reps.equivariant_projector). \n\nFor example to find all O(1,3) (Lorentz) equivariant linear maps from from a 4-Vector Xᶜ to a rank (2,1) tensor Mᵇᵈₐ, you can run\n\n```python\nfrom emlp.reps import V,T\nfrom emlp.groups import *\n\nG = O13()\nQ = (T(1,0)\u003e\u003eT(2,1))(G).equivariant_basis()\n```\n\nor how about equivariant maps from one Rubik's cube to another?\n```python\nG = RubiksCube()\n\nQ = (V(G)\u003e\u003eV(G)).equivariant_basis()\n```\n\nUsing `+` and `*` you can put together composite representations (where multiple representations are concatenated together). For example lets find all equivariant linear maps from 5 node features and 2 edge features to 3 global invariants and 1 edge feature of a graph of size n=5:\n```python\nG=S(5)\n\nrepin = 10*T(1)+5*T(2)\nrepout = 3*T(0)+T(2)\nQ = (repin(G)\u003e\u003erepout(G)).equivariant_basis()\n```\n\nFrom the examples above, there are many different ways of writing a representation like `10*T(1)+5*T(2)` which are all equivalent.\n`10*T(1)+5*T(2)` = `10*V+5*V**2` = `5*V*(2+V)` \n\u003c!-- Feel free to go wild:\n```python\nW=V(O13())\nrepin = (W+2*W**2)*(W.T+1*W).T + W.T\nrepout = 3*W**0 + W + W*W.T\nQ = (repin\u003e\u003erepout).equivariant_basis()\n``` --\u003e\n\nYou can even mix and match representations from different groups. For example with the cyclic group ℤ₃, the permutation group 𝕊₄, and the orthogonal group O(3)\n\n```python\nrep = 2*V(Z(3))*V(S(4))+V(O(3))**2\nQ = (rep\u003e\u003erep).equivariant_basis()\n```\n\nOutside of these tensor representations, our type system works with any finite dimensional linear representation and you can even build your own bespoke representations following the instructions [here](https://emlp.readthedocs.io/en/latest/notebooks/4new_representations.html).\n\nYou can visualize these equivariant bases with [`vis(repin,repout)`](https://emlp.readthedocs.io/en/latest/package/emlp.reps.html#emlp.reps.vis), such as with the three examples above\n\n\u003cimg src=\"https://user-images.githubusercontent.com/12687085/115313228-e19be000-a140-11eb-994f-d4eae4057eba.png\" width=\"200\"/\u003e \u003cimg src=\"https://user-images.githubusercontent.com/12687085/115312972-6afee280-a140-11eb-82f0-603748694645.png\" width=\"360\"/\u003e \u003cimg src=\"https://user-images.githubusercontent.com/12687085/111226510-a0e7fe80-85b7-11eb-913b-09776cdaa92e.png\" width=\"200\"/\u003e \n\u003c!-- ![basis B](https://user-images.githubusercontent.com/12687085/111226517-a2192b80-85b7-11eb-8dba-c01399fb7105.png \"title2\")\n![basis A](https://user-images.githubusercontent.com/12687085/111226510-a0e7fe80-85b7-11eb-913b-09776cdaa92e.png \"title1\") --\u003e\n\n\nCheckout our [documentation](https://emlp.readthedocs.io/en/latest/) to see how to use our system and some worked examples.\n\n# Simple example of using EMLP as a full equivariant model\n\nSuppose we want to construct a Lorentz equivariant model for particle physics data that takes in the input and output 4-momentum of two particles\nin a collision, as well as a some metadata about these particles like their charge, and we want to classify the output\nas belonging to 3 distinct classes of collisions. Since the outputs are simple logits, they should be unchanged by\nLorentz transformation, and similarly with the charges.\n\n```python\nimport emlp\nfrom emlp.reps import T\nfrom emlp.groups import Lorentz\nimport numpy as np\n\nrepin = 4*T(1)+2*T(0) # 4 four vectors and 2 scalars for the charges\nrepout = 3*T(0) # 3 output logits for the 3 classes of collisions\ngroup = Lorentz()\nmodel = emlp.nn.EMLP(repin,repout,group=group,num_layers=3,ch=384)\n\nx = np.random.randn(32,repin(group).size()) # Create a minibatch of data\ny = model(x) # Outputs the 3 class logits\n```\n\nHere we have used the default Objax EMLP, but you can also use our [PyTorch](https://emlp.readthedocs.io/en/latest/notebooks/pytorch_support.html), [Haiku](https://emlp.readthedocs.io/en/latest/notebooks/haiku_support.html), or [Flax](https://emlp.readthedocs.io/en/latest/notebooks/flax_support.html) versions of the models. To see more examples, or how to use your own representations or symmetry groups, check out the documentation.\n\n# Installation instructions\n\nTo install as a package, run \n```bash\npip install emlp\n```\n\nTo run the scripts you will instead need to clone the repo and install it locally which you can do with\n\n```bash\ngit clone https://github.com/mfinzi/equivariant-MLP.git\ncd equivariant-MLP\npip install -e .[EXPTS]\n```\n\n# Experimental Results from Paper\n\nAssuming you have installed the repo locally, you can run the experiments we described in the paper. \n\nTo train the regression models on one of the `Inertia`, `O5Synthetic`, or `ParticleInteraction` datasets found in [`emlp.datasets.py`](https://github.com/mfinzi/equivariant-MLP/blob/master/emlp/datasets.py) you can run the script [`experiments/train_regression.py`](https://github.com/mfinzi/equivariant-MLP/blob/master/experiments/train_regression.py) with command line arguments specifying the dataset, network, and symmetry group. For example to train [`EMLP`](https://emlp.readthedocs.io/en/latest/package/emlp.nn.html#emlp.nn.EMLP) with [`SO(3)`](https://emlp.readthedocs.io/en/latest/package/emlp.groups.html#emlp.groups.SO) equivariance on the `Inertia` dataset, you can run\n\n```\npython experiments/train_regression.py --dataset Inertia --network EMLP --group \"SO(3)\"\n```\n\nor to train the MLP baseline you can run\n\n```\npython experiments/train_regression.py --dataset Inertia --network MLP\n```\nOther command line arguments such as `--aug=True` for data augmentation or `--ch=512` for number of hidden units and others are available, and you can browse the options and their defaults with `python experiments/train_regression.py -h`. If no group is specified, EMLP will automatically choose the one matched to the dataset, but you can also go crazy with any of the other groups implemented in [`groups.py`](https://github.com/mfinzi/equivariant-MLP/blob/master/emlp/groups.py) provided the dimensions match the data (e.g. for the 3D inertia dataset you could do `--group=` [`\"Z(3)\"`](https://emlp.readthedocs.io/en/latest/package/emlp.groups.html#emlp.groups.Z) or [`\"DkeR3(3)\"`](https://emlp.readthedocs.io/en/latest/package/emlp.groups.html#emlp.groups.DkeR3) but not [`\"Sp(2)\"`](https://emlp.readthedocs.io/en/latest/package/emlp.groups.html#emlp.groups.Sp) or [`\"SU(5)\"`](https://emlp.readthedocs.io/en/latest/package/emlp.groups.html#emlp.groups.SU)).\n\nFor the dynamical systems modeling experiments you can use the scripts\n[`experiments/neuralode.py`](https://github.com/mfinzi/equivariant-MLP/blob/master/experiments/neuralode.py) to train (equivariant) Neural ODEs and [`experiments/hnn.py`](https://github.com/mfinzi/equivariant-MLP/blob/master/experiments/hnn.py) to train (equivariant) Hamiltonian Neural Networks.\n\n\nFor the dynamical system task, the Neural ODE and HNN models have special names. [`EMLPode`](https://emlp.readthedocs.io/en/latest/package/emlp.nn.html#emlp.nn.EMLPode) and [`MLPode`](https://emlp.readthedocs.io/en/latest/package/emlp.nn.html#emlp.nn.MLPode) for the Neural ODEs in `neuralode.py` and [`EMLPH`](https://emlp.readthedocs.io/en/latest/package/emlp.nn.html#emlp.nn.EMLPH) and [`MLPH`](https://emlp.readthedocs.io/en/latest/package/emlp.nn.html#emlp.nn.MLPH) for the HNNs in `hnn.py`. For example,\n\n```\npython experiments/neuralode.py --network EMLPode --group=\"O2eR3()\"\n```\nor \n\n```\npython experiments/hnn.py --network EMLPH --group=\"DkeR3(6)\"\n```\n\nThese models are trained to fit a double spring dynamical system. 30s rollouts of the dataset, along with rollout error on these trajectories, and conservation of angular momentum are shown below.\n\n\u003cimg src=\"https://user-images.githubusercontent.com/12687085/114937183-759d3d00-9e0b-11eb-9310-bbfc606e6bda.gif\" width=\"230\"/\u003e \u003cimg src=\"https://user-images.githubusercontent.com/12687085/114937167-703ff280-9e0b-11eb-8421-d8408b31908a.PNG\" width=\"280\"/\u003e \u003cimg src=\"https://user-images.githubusercontent.com/12687085/114937171-71711f80-9e0b-11eb-885e-a541ae1d28cc.PNG\" width=\"240\"/\u003e \n\n\u003c!-- # \n\u003cp align=\"center\"\u003e\n  \u003cimg src=\"https://user-images.githubusercontent.com/12687085/94081992-e75d5d00-fdcd-11ea-9df0-576af6909944.PNG\" width=1000\u003e\n\u003c/p\u003e --\u003e\n\nIf you find our work helpful, please cite it with\n```bibtex\n@article{finzi2021emlp,\n  title={A Practical Method for Constructing Equivariant Multilayer Perceptrons for Arbitrary Matrix Groups},\n  author={Finzi, Marc and Welling, Max and Wilson, Andrew Gordon},\n  journal={Arxiv},\n  year={2021}\n}\n```\n\u003c!-- \nTop quark tagging dataset: https://zenodo.org/record/2603256#.YAoEPehKiUl --\u003e\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fmfinzi%2Fequivariant-MLP","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fmfinzi%2Fequivariant-MLP","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fmfinzi%2Fequivariant-MLP/lists"}