Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

SymDiff: Equivariant Diffusion via Stochastic Symmetrisation

Leo Zhang, Kianoosh Ashouritaklimi, Yee Whye Teh, Rob Cornish

TL;DR

SymDiff introduces stochastic symmetrisation to diffusion models to achieve $E(3)$-equivariance without relying on intrinsically equivariant neural networks. By upgrading an $\mathcal{H}$-equivariant kernel with a learnable $\gamma_\theta$ via a symmetrisation operator, the method yields a $(\mathcal{H}\times\mathcal{G})$-equivariant reverse process with a tractable training objective. The authors instantiate this for $N$-body systems on the CoM-free space, using scalable architectures like Diffusion Transformers and Set Transformers, and demonstrate strong gains on QM9 and GEOM-Drugs while reducing computational demands. They also show data augmentation is a special case of SymDiff, and provide extensions to score and flow models. Overall, SymDiff broadens the toolkit for equivariant generative modelling by combining stochastic symmetrisation with off-the-shelf components, enabling flexible, efficient, and accurate $N$-body molecule generation with practical impact in computational chemistry and related domains.

Abstract

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 $\mathrm{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.

SymDiff: Equivariant Diffusion via Stochastic Symmetrisation

TL;DR

SymDiff introduces stochastic symmetrisation to diffusion models to achieve -equivariance without relying on intrinsically equivariant neural networks. By upgrading an -equivariant kernel with a learnable via a symmetrisation operator, the method yields a -equivariant reverse process with a tractable training objective. The authors instantiate this for -body systems on the CoM-free space, using scalable architectures like Diffusion Transformers and Set Transformers, and demonstrate strong gains on QM9 and GEOM-Drugs while reducing computational demands. They also show data augmentation is a special case of SymDiff, and provide extensions to score and flow models. Overall, SymDiff broadens the toolkit for equivariant generative modelling by combining stochastic symmetrisation with off-the-shelf components, enabling flexible, efficient, and accurate -body molecule generation with practical impact in computational chemistry and related domains.

Abstract

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 -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.
Paper Structure (65 sections, 3 theorems, 53 equations, 8 tables, 4 algorithms)

This paper contains 65 sections, 3 theorems, 53 equations, 8 tables, 4 algorithms.

Table of Contents

  1. Introduction
  2. Background
  3. Equivariant Markov Kernels
  4. Markov kernels
  5. Stochastic equivariance
  6. Stochastic invariance
  7. Equivariant diffusion models
  8. Denoising diffusion models
  9. Invariant and equivariant diffusion
  10. N-body systems and $E(3)$-equivariant diffusion
  11. $N$-body data
  12. Center of mass free space
  13. CoM-free diffusions
  14. Invariance and equivariance
  15. Equivariant Diffusion via Stochastic Symmetrisation
  16. ...and 50 more sections

Key Result

Theorem 1

Suppose $\mathcal{H} \times \mathcal{G}$ acts on $\mathcal{X}$ and $\mathcal{Y}$, and $\gamma: \mathcal{X} \to \mathcal{G}$ is an equivariant base case. Then every Markov kernel $k : \mathcal{X} \to \mathcal{Y}$ that is equivariant with respect to the induced action of $\mathcal{H}$ gives rise to a

Theorems & Definitions (6)

  • Theorem 1
  • Proposition 1
  • Proposition 2
  • proof
  • proof
  • proof