Leo Zhang

A meta-point of this paper is that category theory has utility for reasoning about current problems of interest in mainstream machine learning. The theory is predictive, not just descriptive. 🧵(1/6)

GitHub - leozhangML/SymDiff: Code for the paper: SymDiff: Equivariant Diffusion via Stochastic Symmetrisation Code for the paper: SymDiff: Equivariant Diffusion via Stochastic Symmetrisation - leozhangML/SymDiff

Also, you can find our code at github.com/leozhangML/S...

For future work, we're excited about exploring the design space of SymDiff for improving molecular/protein diffusion models

Me and my co-authors will also be at ICLR presenting this work, so feel free to come chat with us to discuss more about SymDiff!

In addition, as a consequence of the flexibility allowed by SymDiff, we report much lower computational costs compared to message-passing based EGNNs
(8/9)

We note that the generality of our methodology should also allow us to apply SymDiff as a drop-in replacement to these other methods as well

We show substantial improvements over EDM on QM9, GEOM-Drugs, and competitive performance with much more sophisticated recent baselines (which all use intrinsically equivariant architectures)

To validate our methodology, we apply SymDiff to molecular generation with E(3)-equivariance

We take the basic molecule generation framework of EDM (Hoogeboom et al. 2022) and use SymDiff with Diffusion Transformers as a drop-in replacement for the EGNN they use in their reverse process

✅ We show how to bypass previous requirements of a intrinsically equivariant sub-network with the Haar measure, for even more flexibility
✅ We also sketch how to extend SymDiff to score and flow-based models as well

✅ We derive a novel objective to train our new model
✅ We overcome previous issues with symmetrisation concerning pathologies of cannonicalisation, and the computational cost and errors involved in evaluating frame averaging/probabilistic symmetrisation

This provides us with the mathematical foundations to build SymDiff

Our key idea is to "symmetrise" the reverse kernels of a diffusion model to build equivariance with unconstrained architectures

This is where recent work from Cornish (2024) (arxiv.org/abs/2406.11814) comes in. This generalises all previous work and extends it to the stochastic case using category-theoretic arguments, under the name of "stochastic symmetrisation"

A key limitation of this line of work however is that they only consider deterministic functions - i.e. they cannot convert the stochastic kernels involved in generative modelling to be equivariant

This decouples prediction from equivariance constraints, benefiting from the flexibility of the arbitrary base model, which doesn't suffer from the constraints of intrinsically equivariant architectures.

As an alternative, people have proposed cannonicalisation/frame averaging/probabilistic symmetrisation which convert arbitrary networks to be equivariant through a learned group averaging

One common criticism of (intrinsically) equivariant architectures is that due to the architecture constraints required to ensure equivariance, they suffer from worse expressivity, greater computational cost and implementation complexity

There's been a lot of recent interest in whether equivariant models are really needed for data containing symmetries (think AlphaFold 3)

SymDiff: Equivariant Diffusion via Stochastic Symmetrisation We propose SymDiff, a method for constructing equivariant diffusion models using the framework of stochastic symmetrisation. SymDiff resembles a learned data augmentation that is deployed at sampling ...

Joint work with Kianoosh Ashouritaklimi @yeewhye.bsky.social @rob-cornish.bsky.social

📄Paper link: arxiv.org/abs/2410.06262

In our new paper (accepted at ICLR), we propose the first framework for constructing equivariant diffusion models via symmetrisation

This allows us to ensure E(3)-equivariance with just highly scalable standard architectures such as Diffusion Transformers, instead of EGNNs, for molecular generation