SymDiff: Equivariant Diffusion via Stochastic Symmetrisation

2024-10-08

arXiv.org

AI Summary

To address the parameterization complexity and high computational cost arising from intricate equivariant network architectures in E(3)-equivariant molecular generation, this paper proposes SymDiff—a lightweight equivariant diffusion model based on stochastic symmetrization. Its core innovation is the first integration of symmetrization into the diffusion generative framework: instead of relying on built-in equivariant layers or higher-order geometric features, SymDiff achieves strict E(3) equivariance solely through symmetry-aware data augmentation during sampling and equivariant adaptation of off-the-shelf backbone networks (e.g., Transformers or EGNNs). Crucially, it incurs zero additional parameters at inference time, enabling plug-and-play deployment and high computational efficiency. Experiments demonstrate that SymDiff significantly outperforms existing equivariant diffusion models across multiple molecular generation benchmarks, validating its strong generalization capability and scalability.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Machine Learning: Deep Generative Models & Autoencoders
• Search and Optimization: Sampling/Simulation-based Search

Application Category

• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Economics, Online Markets and Human Computation: Economic ramifications for generative AI infrastructure and applications

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.

Problem

• Constructs equivariant diffusion models using stochastic symmetrisation.
• Avoids complex parameterisations and higher-order geometric features.
• Enhances E(3)-equivariant molecular generation with scalable architectures.

Innovation

• Uses stochastic symmetrisation for equivariant diffusion models
• Lightweight, efficient, and easy to implement on existing models
• Leverages scalable architectures without requiring equivariant components