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 time, and is lightweight, computationally efficient, and easy to implement on top of arbitrary off-the-shelf models. In contrast to previous work, SYMDIFF typically does not require any neural network components that are intrinsically equivariant, avoiding the need for complex parameterisations or the use of higher-order geometric features. Instead, our method can leverage highly scalable modern architectures as drop-in replacements for these more constrained alternatives. We show that this additional flexibility yields significant empirical benefit for E(3)-equivariant molecular generation. To the best of our knowledge, this is the first application of symmetrisation to generative modelling, suggesting its potential in this domain more generally. Published as a conference paper at ICLR 2025 tractable optimisation objective. Our model is stochastically E(3)-equivariant overall without needing any intrinsically E(3)-equivariant neural networks as subcomponents. We also sketch how to extend SYMDIFF to score and flow-based generative models (Song et al., 2020; Lipman et al., 2022) . To validate our framework, we implemented SYMDIFF for de novo molecular generation, and evaluated it as a drop-in replacement for the E(3)-equivariant diffusion of Hoogeboom et al. ( 2022 ), which relies on intrinsically equivariant neural networks. In contrast, our model is able to leverage highly scalable off-the-shelf architectures such as Diffusion Transformers (Peebles & Xie, 2023) for all of its subcomponents. We demonstrate this leads to significantly improved empirical performance for both the QM9 and GEOM-Drugs datasets. BACKGROUND We provide here an overview of the underlying theory behind equivariant diffusion modelling. This theory is most conveniently developed in terms of Markov kernels, whose definition we recall first. EQUIVARIANT MARKOV KERNELS Markov kernels At a high level, a Markov kernel k : X → Y may be thought of as a conditional distribution or stochastic map that, when given an input x ∈ X , produces a random output in Y with distribution k(dy|x). For example, given a function f : X × E → Y and a random element η of E, there is a Markov kernel k : X → Y for which k(dy|x) is the distribution of f (x, η) 1 . As a special case, every deterministic function f : X → Y may be thought of as a Markov kernel X → Y also. When k(dy|x) has a density (or likelihood), we will denote this by k(y|x), although we note that we can still reason about Markov kernels even when they do not admit a likelihood in this sense.