Evaluating Equivariance for Reconstruction

Savannah Thais, Columbia University

Daniel Murnane, Lawrence Berkeley National Lab

Why do some ML models outperform others?

We often highly fine-tune models and attribute their performance to a certain design choice: size, input features, hyperparemeters, network design, etc

Physical Inductive Biases

Data Structure

Model Constraints

Task Formulation

Relational structure, ordering, feature selection, pre-processing, etc

Restricting model weights, learned function, propagated information, etc

Physics informed neural networks, incorporating conservation laws or equations through loss function design, etc

Symmetries are (potentially) powerful physical inductive biases

ɸ

ɸ

Image from arxiv:2102.09844

Equivariant networks are popular, partially because they are (in some formulations) easy to implement

Diagram from Daniel Murnane

Potential Benefits of Equivariance

Model Efficiency

• Most published ML results achieve SotA accuracy
• Often the main theoretical advantage of equivariant models (inspired by CNNs and GNNs)

• Should learn complete symmetry orbit from one example
• Learns reduced representation for easier generalization

• Models may have an ‘easier’ time learning optimal functions
• Equivariance allows models to use parameters more efficiently

• Do not need to rely on data augmentation to learn symmetries
• Should reduce training data size (and computational resource usage)

Accuracy

Generalizability

Data Efficiency

Physics as an ML Sandbox

Particle Tracking

• Using TrackML Dataset
• Build an SO(2) rotation equivariant model (in x-y plane)

• Using Top Quark Tagging Reference Dataset
• Build a Lorentz equivariant model (rotations and boosts in spacetime)

Jet Tagging

Image from Nelson

Image from arxiv:2109.12636

Baseline Tagging Models

ParticleNet

Message passing dynamic graph GNN on particle graph

​

arXiv:1902.08570

Deep 2D CNN on jet images

​

arXiv:1611.05431

​

​

ResNexT

Deep set network on particle features

​

arXiv:1810.05165

​

​

Particle Flow

Linear discriminant on EFP complete linear basis

​

arXiv:1712.07124

​

​

​

​

Energy Flow Polynomials

Equivariant Tagging Models

LorentzNet

Message passing GNN with Lorentz equivariant message, on particle graph

​

arXiv:2201.08187

​

​

NN with CG-layers that take tensor products and decompose into irreps using Clebsch-Gordan map, on particle features

​

arXiv:2006.04780

​

​

​

​

Lorentz Group Network

Deep set-esq network using all totally symmetric Lorentz invariants and full set of 15 rank 2 to rank 2 maps as aggregators

​

arXiv:2211.00454

​

​

PELICAN

Message passing GNN with Lorentz equivariant message and (optionally) unconstrained message, on particle graph

​

arXiv:2202.06941

​

​

​

​

​

​

VecNet

Tracking Models

EuclidNet

Interaction Network

Message passing GNN with node and edge updates, on hit graph (with physics-based edge construction)

​

arXiv:2103.16701

​

​

​

​

​

​

Message passing GNN with SO(2)-equivariant message construction, on hit graph (with physics-based edge construction)

​

arXiv:2304.05293

​

​

​

​

Evaluating Equivariance

• Jet tagging: highest accuracy model is equivariant, but not all equivariant models perform well
• Tracking: for small models equivariant models have highest accuracy, but performance plateaus as models grow
• Overall, relationship between equivariance and accuracy is unclear (confounding factors remain)

Accuracy

• Jet tagging: regression model with physics inputs is most efficient. Semi-equivariant model is also efficient.
• Tracking: relationship changes with model size
• Overall, equivariance does not seem to contribute directly to model efficiency

Model Efficiency

Ant factor = 10^5/[(1-AUC)*N_p]

Tracking

Tagging

Evaluating Equivariance

• Jet tagging: equivariant models generalize, but not all to the same extent
• Tracking: both equivariant and sufficiently large non-equivariant models generalize
• Overall, equivariance provides a good amount of generalization, but other models can too (tradeoffs)

Generalizability

• Jet tagging: clear benefit from equivariance in very small data regimes: achieves 99% of full accuracy with just 0.5% of training dataset
• Compared to 97% for non-equivariant model
• Overall, seems to be the most replicable benefit of equivariance. This is demonstrated in other papers, such as NequIP

Data Efficiency

Tagging

Tagging

Tracking

So, What Now?

What kinds of inductive biases are useful? How are they useful?

Over-constraint?

Is full equivariance the right approach for HEP tasks?

• Unconstrained models can learn to generalize under symmetry transformations
• VecNet studies show optimal accuracy and model efficiency are achieved with mixed equivariant and non-equivariant information
• While the underlying physics is obeys symmetries, observed data is likely NOT fully symmetric

​

Expressivity?

Equivariance is a bit of an ill-defined model characteristic

• Different methods of enforcing equivariance yield substantially different results
• Not proven that it is easier to find the global optimum in a constrained optimization space
• Highly performing HEP models have other design choices besides equivariance

​

Fundamental work on network expressivity indicates that feature choice and message construction largely determine expressivity

Transformers?

Transformers provide excellent performance with minimal or no inductive bias…

​

In a many parameter, big data regime does physics matter at all?

​

arXiv:2202.03772

Three Takeaways

Consider the Physics Goals

Explore Other Inductive Biases

There is no one-size-fits-all solution to building the optimal model for a physics task. Trade-offs between model size, compute resources, data availability, robustness, etc are key

Consider inductive biases that modify data structures or task design, rather than constraining optimization space

Conduct apples-to-apples model comparisons that are better able to isolate the impact of design choices and carefully designed ablation studies

More Systematic Studies

What’s Next?

Let’s discuss!

​

​

st3565@columbia.edu

@basicsciencesav

Our paper: Equivariance Is Not All You Need

CREDITS: This presentation template was created by Slidesgo, and includes icons by Flaticon, and infographics & images by Freepik