Skip to content

Lorentz group equivariant autoencoders based on Lorentz Group Network

License

Notifications You must be signed in to change notification settings

zichunhao/lgn-autoencoder

Repository files navigation

Lorentz Group Equivariant Autoencoders

DOI

Overview

This autoencoder model exploits the symmetry of the Lorentz Group, a Lie group that represents a fundamental symmetry of spacetime and describes the dynamics of relativistic objects such as elementary particles in a particle physics experiment. The model is built using the architecture of the LGN introduced by Bogatskiy et al. in arXiv:2006.04780 (see the README file in each directory for more details). A preprint is available on arXiv:2212.07347. Published on The European Physical Journal C.

To achieve Lorentz equivariance, the model works on the irreducible representations of the Lorentz group $\mathrm{SO}(1,3)$. For instance, Lorentz scalars are $(0,0)$ representations, and 4-vectors, such as the particle 4-momenta, are $(1/2,1/2)$ representations. Each representation has its transformation rules. That the model is equivariant implies that each parameter in the model will transform according to its corresponding transformation rule if the input undergoes a Lorentz transformation. In this way, the model can always generate data that satisfy the special relativity, and the latent space, since all internal parameters are Lorentz tensors, can possibly be more physically interpretable. The diagram for a message passing is shown below.

Download Dataset

To download data:

  1. Install JetNet:
    pip3 install jetnet; 
    
  2. Run preprocess.py
    python utils/data/preprocess.py \
    --jet-types g q t w z \
    --save-dir "./data"
    

Train the model

An example training looks like this.

mkdir -p results;

python -u main.py \
--data-paths "./data/g_jets_30p_p4.pt" "./data/q_jets_30p_p4.pt" \
--test-data-paths "./data/g_jets_30p_p4.pt" "./data/q_jets_30p_p4.pt" \
-j QCD \
-e 50 \
-bs 512 \
--train-fraction 0.75 \
--lr 0.0005 \
--loss-choice chamfer \
--get-real-method sum \
--tau-latent-vectors 8 \
--tau-latent-scalars 1 \
--maxdim 2 \
--l1-lambda 1e-8 \
--l2-lambda 0 \
--map-to-latent "min&max" \
--mlp-width 6 \
--mlp-depth 6 \
--encoder-num-channels 3 3 4 4 \
--decoder-num-channels 4 4 3 3 \
--patience 1000 \
--plot-freq 100 \
--save-freq 200 \
--plot-start-epoch 50 \
--equivariance-test \
--num-test-batch 1024 \
--save-dir "results" \
| tee "results/training-log.txt"

Some important parameters for main.py

  • -bs (batch-size): batch size.
  • -tbs (--test-batch-size): test batch size.
  • -j (--jet-type): the jet type used for training (mainly used for naming files).
  • -e (--num-epochs): number of epochs to train.
  • --loss-choice: loss function to use.
  • --train-fraction: fraction of the data used for training.
  • --data-paths: paths to the training data.
  • --test-data-paths: paths to the test data.
  • --map-to-latent: aggregation method to the latent space.
    • For min&max, min and max will be concatenated.
    • For min+max, min and max will be added.
    • For mix, a Lorentz group equivariant linear layer will mix the irreps to the latent space.
  • --maxdim: maximum weight of representation to keep in training (recommended: 2 or 3).
  • --tau-latent-scalars: number of (complexified) Lorentz scalars to keep in the latent space.
  • --tau-latent-vectors: number of (complexified) 4-vectors to keep in the latent space.
  • --encoder-num-channels: the number of channels for each representation in the encoder.
  • --decoder-num-channels: the number of channels for each representation in the decoder.
  • --plot-freq: frequency of plotting (plotting / epoch).
  • --plot-start-epoch: epoch number at which to start plotting (for better GPU usage in the beginning).
  • --save-dir: directory to save the trained model and plots.
  • --equivariance-test: whether to test the model for equivariance.

Test the model

The test includes reconstruction and can include Lorentz group equivariance test and anomaly detection. An example test looks like this.

python test.py \
-tbs 512 \
--num-test-batch 200 \
-j QCD \
--maxdim 2 \
--tau-latent-vectors 8 \
--tau-latent-scalars 1 \
--loss-choice chamfer \
--get-real-method sum \
--map-to-latent "min&max" \
--mlp-width 6 \
--mlp-depth 6 \
--encoder-num-channels 3 3 4 4 \
--decoder-num-channels 4 4 3 3 \
--model-path "results/LGNAutoencoder_QCDJet_min&max_tauLS1_tauLV1_encoder3344_decoder4433" \
--data-paths "data/g_jets_30p_p4.pt" "data/q_jets_30p_p4.pt" \
--test-data-paths "data/g_jets_30p_p4.pt" "data/q_jets_30p_p4.pt" \
--anomaly-detection \
--signal-paths "data/t_jets_30p_p4.pt" "data/w_jets_30p_p4.pt" "data/z_jets_30p_p4.pt" \
--signal-types t w z \
--equivariance-test \
| tee -a "results/test-log.txt"

Results

Equivariance Tests

Boost and rotational equivariance tests were done on the model. The rotation angles range from $0$ to $2\pi$, and the Lorentz factors range from $0$ to $11013.2$. The model is equivariant with respect to rotation up to floating point errors and is equivariant with respect to boost in the physically relevant region (the errors increase as the Lorentz factor increases because of the floating point sensitivity of boost).

References

Relevant Group Equivariant Models

Background Knowledge

Group theory and group representations

The Lorentz group and its representations

  • K. Rao, "Linear Algebra and Group Theory for Physicists". Text and Readings in Physical Sciences, Springer, 2nd edition, 2006, doi: 10.1007/978-93-86279-32-3. ISBN 9789386279323.
  • R. Sexl and H. Urbantke, "Relativity, Groups, Particles: Special Relativity and Relativistic Symmetry in Field and Particle Physics". Springer, 2001, doi: 10.1007/978-3-7091-6234-7. ISBN 9783709162347.
  • Tung Wu-Ki, "Group Theory in Physics". World Scientific, 1985, doi: doi:10.1142/0097.

The connection between group theory and particle physics

Main Libraries