SymDiff: Equivariant Diffusion via Stochastic Symmetrisation
Authors: Leo Zhang, Kianoosh Ashouritaklimi, Yee Whye Teh, Rob Cornish
ICLR 2025 | Venue PDF | Archive PDF | Plain Text | LLM Run Details
| Reproducibility Variable | Result | LLM Response |
|---|---|---|
| Research Type | Experimental | We show that this additional flexibility yields significant empirical benefit for E(3)-equivariant molecular generation. 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. |
| Researcher Affiliation | Academia | Leo Zhang Kianoosh Ashouritaklimi Yee Whye Teh Rob Cornish Department of Statistics, University of Oxford |
| Pseudocode | Yes | Algorithm 1 SYMDIFF training step 1: Sample z0 pdata(z0), t Unif({1, . . . , T}) and ϵ NU(0, I) 2: zt αtz0 + σtϵ 3: Sample R0 from the Haar measure on O(3) and η ν(dη) 4: R R0 fθ(RT 0 zt, η) 5: Take gradient descent step with θ 1 2w(t) ϵ R ϵθ(RT zt) 2 |
| Open Source Code | Yes | Our code is available at: https://github.com/leozhang ML/Sym Diff. |
| Open Datasets | Yes | QM9 (Ramakrishnan et al., 2014) is a common benchmark dataset used for evaluating molecular generation. GEOM-Drugs (Axelrod & Gomez-Bombarelli, 2022) is a larger and more complicated dataset than QM9, containing 430,000 molecules with up to 181 atoms. |
| Dataset Splits | Yes | We used the same train-val-test split of 100K-8K-13K as in Anderson et al. (2019). |
| Hardware Specification | Yes | In fact, when we tried to run the EDM model it resulted in out-of-memory errors on our NVIDIA H100 80GB GPU (Hoogeboom et al. (2022) avoid this by training EDM on 3 NVIDIA RTX A6000 48GB GPUs.) |
| Software Dependencies | No | Both components rely on Diffusion Transformers (Di Ts) (Peebles & Xie, 2023) using the official Py Torch implementation at https: //github.com/facebookresearch/Di T. We also state the hyperparameters that we kept fixed for both our QM9 and GEOM-Drugs experiments. |
| Experiment Setup | Yes | For the optimisation of SYMDIFF models, we followed Peebles & Xie (2023) and used Adam W (Loshchilov & Hutter) with a batch size of 256. We chose a learning rate of 2 10 4 and weight decay of 10 12 for our 31.2M parameter model by searching over a small grid of 3 values for each. To match the same number of steps as in Hoogeboom et al. (2022), we trained our model for 4350 epochs. |