Atoms as 3D tokens
Every atom is a token, 44.4 on average in GEOM-Drugs; ring geometry is learned atom by atom.
NeurIPS 2026
3D Molecule Generation from Rigid Motifs via $\boldsymbol{\mathrm{SE}(3)}$ Flows
1Technical University of Munich2Munich Data Science Institute3Munich Center for Machine Learning4Imperial College London
TL;DRGenerate 3D molecules from rigid motifs, not atoms: one joint flow over motif types and their $\mathrm{SE}(3)$ poses, with each motif's inner geometry fixed by its canonical pose.
t = 0.00
noise · t = 0t = 0.5t = 1 · data
Molecule. We start from the usual all-atom representation: $N$ atoms with coordinates $\mathbf{y}_j$ and types $h_j$ from the atom vocabulary $\mathcal{V}_a$.
\(\{(\mathbf{y}_j,\,h_j)\}_{j=1}^{N}\)
\(\mathbf{y}_j\in\mathbb{R}^{3},\;\;h_j\in\mathcal{V}_a\)
Rigid motifs. Keep planar rings, double and triple bonds and bonds to hydrogens; cut the rest, and split motifs rarer than a fraction $\alpha$ of the data. Isolated atoms get dummy atoms towards their neighbours, which fix their orientation. The molecule $\boldsymbol{\mathcal{M}}$ becomes $K$ motifs $m_k$, each with a frame $\mathbf{T}_k$: a rotation $\mathbf{R}_k$ and a translation $\mathbf{x}_k$.
\(\boldsymbol{\mathcal{M}}=\{(\mathbf{T}_k,\,m_k)\}_{k=1}^{K}\)
\(\mathbf{T}_k=(\mathbf{R}_k,\mathbf{x}_k)\in\mathrm{SE}(3)\)
Motif tokens. Each token pairs a frame $\mathbf{T}_k$ with a type $m_k$ from the motif vocabulary $\mathcal{V}_m$. A type stores its canonical pose $\mathbf{P}_k$, its atom types $\mathbf{h}_k$ and its symmetry rotations $\mathcal{S}_k$.
\(m_k=(\mathbf{P}_k,\mathbf{h}_k,\mathcal{S}_k)\in\mathcal{V}_m\)
\(\mathbf{P}_k\in\mathbb{R}^{N_k\times 3},\;\;\mathcal{S}_k\subset\mathrm{SO}(3)\)
Noise. Every type $m_0$ starts masked, rotations $\mathbf{R}_0$ are drawn uniformly from $\mathrm{SO}(3)$ and translations $\mathbf{x}_0$ from a Gaussian.
\(m_0=[\mathrm{MASK}]\)
\(\mathbf{R}_0\sim\mathcal{U}(\mathrm{SO}(3)),\;\;\mathbf{x}_0\sim\mathcal{N}(\mathbf{0},\mathbf{I})\)
Joint flow. Each type $m_t$ unmasks by a discrete flow [2]; each rotation $\mathbf{R}_t$ follows a geodesic and each translation $\mathbf{x}_t$ a straight line [1]. One IPA network predicts all of them.
\(p_t(m_t\mid m_1)=(1-t)\,\delta_{[\mathrm{MASK}]}+t\,\delta_{m_1}\)
\(\mathbf{R}_t=\exp_{\mathbf{R}_0}\!\big(t\,\log_{\mathbf{R}_0}\mathbf{R}_1\big),\;\;\mathbf{x}_t=(1-t)\,\mathbf{x}_0+t\,\mathbf{x}_1\)
Deterministic lift. The atoms $\mathbf{Y}_k$ of each motif are its canonical pose $\mathbf{P}_k$ placed by its frame $(\mathbf{R}_k,\mathbf{x}_k)$, with no learned decoder or assembly.
\(\mathbf{Y}_k=\mathbf{P}_k\mathbf{R}_k+\mathbf{1}\,\mathbf{x}_k^{\top}\)
\(\mathbf{Y}_k\in\mathbb{R}^{N_k\times 3},\;\;k=1,\dots,K\)
1Why rigid motifs?
Atom-level models learn even the most rigid parts of a molecule, such as aromatic rings, atom by atom. Earlier protein structure generators already treated residues as rigid frames, and graph-based generators build molecules from fragments. MotiFlow brings the two ideas together for general drug-like molecules in 3D.
Protein residues as rigid 3D frames
Every backbone residue is an $\mathrm{SE}(3)$ frame with a fixed vocabulary of 20 side-chains.
Motifs as 2D graphs
Chemically meaningful fragments, but generated as 2D graphs without any 3D geometry.
Motifs as rigid 3D frames
52 atoms → 11 motif tokens with $\mathrm{SE}(3)$ frames; each motif keeps its canonical geometry.
2From molecules to motif tokens
A fragmentation has to give rigid pieces, at least three non-collinear points per piece to fix a frame, and pieces frequent enough to learn. MotiFlow keeps double and triple bonds, bonds to hydrogens and all bonds within approximately planar rings and fused ring systems, and cuts the rest. Motifs rarer than a fraction $\alpha$ of the dataset are split further; $\alpha$ is the only tunable parameter. Isolated atoms and collinear motifs get dummy atoms that lock their orientation.
Fragmentation ablation · 100 sampling steps
Treating planar rings as motifs raises atom stability on both datasets, at a small cost in validity × connectivity. A smaller $\alpha$ gives a larger vocabulary with rarer motifs.
Representation size on QMugs
Rigid motifs compress the all-atom and heavy-atom representations by 3.4× and 1.9×, with no learned latent space and no expensive reconstruction.
3Canonical poses and symmetries
Each motif type keeps one arbitrarily chosen and fixed canonical pose $\mathbf{P}_k$. Every instance in the data gets its rotation by Kabsch alignment to that pose, using a graph automorphism for the atom correspondence, which avoids searching over all $N!$ orderings. Symmetric motifs, such as a para-substituted benzene, have several rotations $\mathcal{S}_k$ that map them onto themselves, so their frame is only defined up to $\mathcal{S}_k$. During training the target is the equivalent rotation closest to the current noisy frame.
\(\tilde{\mathbf{R}}_1=\mathbf{S}^*\mathbf{R}_1,\quad \mathbf{S}^*=\operatorname*{arg\,max}_{\mathbf{S}\in\mathcal{S}_k}\operatorname{Tr}\big(\mathbf{R}_t^{\top}\mathbf{S}\mathbf{R}_1\big)\)
The dashed frame is a noisy rotation $\mathbf{R}_t$; solid frames are the $|\mathcal{S}_k|$ equivalent targets, with their geodesic distance to $\mathbf{R}_t$. Regressing to the closest one avoids conflicting gradients from equivalent but distant targets at small $t$.
4Two flows, one network
Following the multimodal flows of Campbell et al. [2], the conditional path factorises over motifs and modalities: a discrete flow decides what each motif is, an $\mathrm{SE}(3)$ flow decides where it sits. The factorisation only defines the path; one network sees the whole noisy molecule and predicts all of it.
\(p_t(\boldsymbol{\mathcal{M}}_t\mid\boldsymbol{\mathcal{M}}_1)=\prod\limits_{k=1}^{K}\underbrace{p_t(m_t^k\mid m_1^k)}_{\text{discrete flow}}\;\underbrace{p_t(\mathbf{T}_t^k\mid\mathbf{T}_1^k)}_{\mathrm{SE}(3)\text{ flow}}\)
What · discrete flow
Every type starts as $[\mathrm{MASK}]$. A masked motif unmasks to its true type at rate $(1-t)^{-1}$, so on the conditional path each one reveals at a uniform time and then stays fixed.
Where · $\mathrm{SE}(3)$ flow
Translations move on straight lines, rotations along geodesics of $\mathrm{SO}(3)$. Shown for the highlighted motif, with its frame at evenly spaced times.
t = 0.62
Network · 12.4M parameters
Inputsnoisy frames $\mathbf{T}_t$tokens $m_t$, time $t$self-conditioning $\hat m_1$
Edgesall pairs of motifs64 RBFs of centre distances
4 × blocktriangle updatesIPA, 8 heads × 64transition, 2 transformer layers
Headsrotation updatetranslation MLPlogits over $\mathcal{V}_m$
Built on FoldFlow-Base [1], developed for protein backbones. Its residue-index edge features are replaced by purely geometric ones, since motifs have no order. Trained with $\mathcal{L}=\mathbb{E}_t\big[\mathcal{L}_{\text{trans}}+\mathcal{L}_{\text{rot}}+\tfrac{1}{1-t}\,\mathcal{L}_{\text{DFM}}\big]$.
5Results
6Scope
- Setup
- Following EDM [3], bonds are inferred from interatomic distances.
- Works best
- Medium and large drug-like molecules.
- Limitations
- Gains shrink for small or highly flexible molecules, which offer little rigid structure.
- Vocabulary
- A larger vocabulary holds larger motifs but makes each one rarer, so uncommon motifs are under-sampled as $\alpha$ decreases. Rare substructures stay representable by splitting them into frequent motifs.
Citation
@inproceedings{poletukhin2026d,
title = {3D Molecule Generation from Rigid Motifs via $\mathrm{SE}(3)$ Flows},
author = {Roman Poletukhin and Marcel Kollovieh and Eike Eberhard and Stephan G{\"u}nnemann},
booktitle = {The Fortieth Annual Conference on Neural Information Processing Systems},
year = {2026},
url = {https://openreview.net/forum?id=10JJ9uxd6J}
}References
- Bose et al. SE(3)-stochastic flow matching for protein backbone generation. ICLR 2024.
- Campbell et al. Generative flows on discrete state-spaces. ICML 2024.
- Hoogeboom et al. Equivariant diffusion for molecule generation in 3D. ICML 2022.