SymDrift: One-Shot Generative Modeling under Symmetries

2026-05-07

AI Summary

Existing single-step generative models struggle to accurately capture the distribution of physical systems exhibiting global symmetries—such as 3D rotational invariance—leading to degraded generation quality. This work proposes SymDrift, a novel framework that endows the drift field itself with symmetry awareness, thereby addressing symmetry mismatch in single-step generation without requiring explicit symmetrization of the empirical data distribution. SymDrift achieves this through two complementary strategies: symmetry-aware drift formulation in an optimally aligned coordinate space and group-invariant embeddings. By integrating equivariant generative modeling, optimal alignment, and single-step diffusion techniques, SymDrift substantially outperforms current single-step methods on molecular conformation and transition state generation tasks, attaining performance comparable to multi-step models while accelerating inference by up to 40×—making it well-suited for high-throughput applications such as drug screening.

Abstract

Generative modeling of physical systems, such as molecules, requires learning distributions that are invariant under global symmetries, such as rotations in three-dimensional space. Equivariant diffusion and flow matching models can incorporate such invariances effectively, even when trained on a non-invariant empirical distribution, but they typically rely on costly multi-step sampling. Recently, drifting models have emerged as an efficient alternative, enabling single-step generation and achieving state-of-the-art performance in generative modeling tasks. However, we show that drifting models face a symmetry-specific challenge, since an equivariant generator does not generally produce the same drifting field as the one obtained from the symmetrized target distribution. Addressing this issue would require expensive symmetrization of the empirical distribution. To avoid this cost, we propose SymDrift, a framework that makes the drifting field itself symmetry-aware. We introduce two complementary strategies: (i) a symmetrized drift in coordinate space based on optimal alignment, and (ii) a $G$-invariant embedding that removes symmetry ambiguity by construction. Empirically, SymDrift outperforms existing one-shot methods on standard benchmarks for conformer and transition state generation, while remaining competitive with significantly more expensive multi-step approaches. By enabling one-shot inference, SymDrift reduces computational overhead by up to 40$\times$ compared to existing baselines, making it promising for high-throughput applications such as virtual drug screening and large-scale reaction network exploration.

Problem

• generative modeling
• symmetry
• drifting models
• equivariance
• one-shot generation

Innovation

• SymDrift
• one-shot generative modeling
• symmetry-aware drift
• equivariance
• molecular conformation generation