TreeRef-BFN: Equivariance-Free De Novo Molecule Generation based on 2D Topology and Internal 3D Geometry
Authors: Ruiqing Sun, Sen Yang, Dawei Feng, Bo Ding, Yijie Wang, Huaimin Wang
Organizations: College of Computer Science and Technology National University of Defense Technology Changsha 410073, China · Bioinformatics Center of AMMS Beijing 100080, China
De novo 3D molecular generation jointly models molecular size, topology, and geometry. Most methods pre-sample molecular size and generate Cartesian coordinates, limiting variable-size conditional tasks such as fragment completion and scaffold decoration while often relying on equivariant architectures. Internal-coordinate methods avoid rigid-body redundancy but typically require a known molecular graph or autoregressive construction, which may accumulate errors. We propose TreeRef, a tree-based molecular representation that assigns molecular topology and topology-dependent local 3D geometry to a naturally variable-size tree. RingRef nodes encode ring closures while preserving the tree structure, while Null nodes allow molecular size to emerge directly from node occupancy. Based on TreeRef, we develop TreeRef-BFN, a Bayesian Flow Network with a standard Transformer backbone that globally couples these locally defined variables and jointly generates discrete molecular variables and continuous local geometry. A single pretrained TreeRef-BFN supports unconditional generation and variable-size structure-conditioned 3D generation through masking alone, without retraining. Empirical studies demonstrate strong chemical validity, molecular stability, and diversity, accurate local geometric distributions, fast sampling, and competitive property-conditioned generation, establishing TreeRef-BFN as an efficient and flexible framework for 3D molecular generation.
Figures & tables
Figure 1: Illustration of the TreeRef representation. The left panel shows the Real atoms of the original 3D molecule. The middle column lists the Real atoms in BFS order, each serving as the parent of the k=3 child slots in the three columns on the right, which are the actual BFN-modeled tokens represented by zi . Unused child slots are represented as Null nodes, and all slots after the last occupied region are Null. During BFS traversal, edges excluded from the spanning tree (red edge in the molecule) are represented by RingRef nodes, which encode ring-closing relations by referring to the corresponding existing Real atoms (dashed arrow).
Validity ( ↑ )
Unique ( ↑ )
Stability ( ↑ )
W1 ( ↓ )
Model
Steps
Time
RDKit
SMILES
Atom
Mol.
Bond
Angle
Dih.
EDM ( Hoogeboom et al., 2022 )
1000
104.2
84.08
98.44
98.34
81.53
0.0041
0.762
5.624
GCDM ( Morehead and Cheng, 2024 )
1000
136
87.03
98.38
98.70
85.79
0.0042
0.843
4.247
GeoLDM ( Xu et al., 2023 )
1000
101.5
91.36
97.90
98.91
89.90
0.0059
0.745
4.202
ADiT ( Joshi et al., 2025 )
500
59.6
91.37
96.76
76.43
7.56
0.0438
1.599
5.535
Zatom-1 ( Morehead et al., 2026 )
100
11.4
91.56
96.69
98.16
84.94
0.0066
0.822
3.590
Table 1: Unconditional generation results on QM9 with explicit hydrogens. ∗ Results are taken from the corresponding papers because pretrained checkpoints are not publicly available.
Validity ( ↑ )
Unique ( ↑ )
Stability ( ↑ )
W1 ( ↓ )
Model
Steps
RDKit
SMILES
Atom
Mol.
Bond
Angle
Dih.
EDM ( Hoogeboom et al., 2022 )
1000
41.15
100.00
81.09
0.22
0.0238
5.613
22.418
GCDM ( Morehead and Cheng, 2024 )
1000
55.51
100.00
87.96
3.79
0.0227
1.479
21.504
GeoLDM ( Xu et al., 2023 )
1000
45.50
100.00
84.29
0.60
0.0187
3.297
21.222
ADiT ( Joshi et al., 2025 )
500
58.65
100.00
75.47
0.12
0.0512
2.048
19.374
Zatom-1 ( Morehead et al., 2026 )
100
79.89
99.90
88.36
5.98
0.0256
1.271
20.227
Table 2: Unconditional generation on GEOM-Drugs.
Method
Conditioning
Steps
Valid ↑
RMSD ↓
Bond W1↓
Angle W1↓
Dihedral W1↓
JODO ( Huang et al., 2023 )
RePaint-style
1000
88.50
1.223
0.026
0.953
7.823
MiDi ( Vignac et al., 2023 )
RePaint-style
500
4.40
2.512
1.392
38.801
23.480
SemlaFlow ( Irwin et al., 2025 )
Frozen
100
63.90
1.849
0.419
11.918
11.806
TreeRef-BFN-S
Frozen
100
93.00
0.902
0.014
0.719
2.966
TreeRef-BFN-M
Frozen
100
78.40
0.804
0.015
1.066
3.553
TreeRef-BFN-M-inv2
Frozen
100
94.55
0.664
0.007
0.383
2.378
Table 3: Structure-conditioned 3D generation on QM9. Lengths are in Å and angles in degrees.
Table 5
Appendix figures & tables19 assets
Supplementary material from the paper’s appendix.
Appendix
Model
Blocks
Hidden dimension
Attention heads
FFN dimension
Parameters (M)
QM9-S
8
256
8
1024
10.50
QM9-M
12
512
8
2048
60.74
QM9-L
16
768
12
3072
179.04
QM9-L (property-conditioned)
16
768
12
3072
179.64
GEOM-Drugs-XL
16
1024
16
4096
319.83
Appendix
Table 6: Transformer configurations and total trainable parameter counts. M denotes one million parameters. The property-conditioned QM9 configuration includes the scalar-property embedding network and learned null-condition vector. GEOM-Drugs-XL uses k=3 , explicit hydrogens, and a formal-charge channel.
Figure 2: Unconditional QM9 generation across model sizes and sampling budgets. Panels show (a) RDKit validity, (b) molecule stability, (c) uniqueness, and Wasserstein-1 distances for (d) bond lengths, (e) bond angles, and (f) dihedrals. Higher values are better in (a–c), and lower values are better in (d–f). S, M, and L markers summarize 10,000 samples. The horizontal axis is logarithmic, and lines connect evaluated configurations. No variability across independent runs is reported.
Figure 3: Unconditional QM9 molecules at 100 steps, drawn at random from valid samples. Red edges indicate RingRef closures.
Model
Steps
Valid ↑
MiDi stability ↑
EDM stability ↑
Unique ↑
Novel
W1↓
Atom
Mol.
Atom
Mol.
Bond
Angle
Dih.
S
10
75.65
96.63
70.94
93.25
58.50
93.15
67.82
0.0234
3.797
6.998
S
20
81.12
97.71
78.25
94.62
64.33
96.18
58.79
0.0183
2.181
3.853
S
50
84.48
98.27
82.63
95.30
67.69
97.02
54.80
0.0154
1.215
2.286
S
100
85.98
98.48
84.60
95.59
69.25
96.78
53.19
0.0135
0.888
1.682
S
200
85.91
98.47
84.43
95.52
68.69
96.78
50.64
0.0138
0.846
1.764
Appendix
Table 7: Unconditional QM9 results for TreeRef-BFN-S, M, L, and M-inv2. Valid denotes RDKit validity. MiDi stability uses the generated bond graph and MiDi allowed-valence rules; EDM stability uses bonds inferred from Cartesian distances and EDM valence rules. Atom and molecule stability are reported separately, with decoding failures counted as unstable molecules. Validity, stability, uniqueness (Unique), and novelty (Novel) are percentages. Bond-length W1 is in Å; bond-angle and dihedral W1 are in degrees.
Figure 4: Unconditional GEOM-Drugs samples from TreeRef-BFN at 100 steps. Twelve molecules, four per row. The label above each molecule is its heavy-atom formula. Red edges are RingRef closures.
Property
EGNN
TreeRef
α (Bohr 3 )
0.132
0.318
Δϵ (meV)
99.8
137
ϵHOMO (meV)
69.1
95.6
ϵLUMO (meV)
53.0
93.5
μ (D)
0.057
0.434
Cv
0.055
0.101
Appendix
Table 8: Validation mean absolute error of the two property regressors. Both are trained on the scorer half of the QM9 training set and evaluated on the QM9 validation split.
Property
EGNN
TreeRef
NExT-Mol
α (Bohr 3 )
1.18
1.14
1.16
Δϵ (meV)
225
174
297
ϵHOMO (meV)
142
136
205
ϵLUMO (meV)
169
136
235
μ (D)
0.531
0.495
0.507
Cv
0.513
0.445
0.512
Appendix
Table 9: Property mean absolute error at 100 steps. The generator is the same in both columns. NExT-Mol is the number in Table 5 .
Figure 5: Property error across sampling budgets. Solid lines are the frozen EGNN and TreeRef regressors on the same generated molecules. The dashed line is NExT-Mol from Table 5 . Orbital errors are converted from Hartree to meV. Each point is 10,000 samples. The horizontal axis is logarithmic.
Figure 6: Sampling errors for unconditional QM9. (a) Decode-failure rate and (b) exact-valence instability among decoded molecules. S, M, and L use 10,000 fast-decoded samples. (c–f) Composition of the failures at 100 steps: unmatched RingRefs, under-valent carbon, over-valent hydrogen, and the remaining valence errors. The number in the ring is their combined share of all samples. The horizontal axes in (a) and (b) are logarithmic.
Figure 7: QM9-L repairs. Left is the fast result, right is rules. The first two rows are RingRef samples 3777 (200 steps) and 5472 (100 steps). The diamond is the unmatched RingRef; the dotted segment joins it to its host. In 3777 rules changes the oxygen under the diamond to nitrogen. In 5472 the promoted nitrogen is still under-valent, so the second step writes it as fluorine and doubles the carbon–oxygen bond. The last row is a 100-step molecule whose nitrogens are over-valent. Rules replaces two of them by carbon and oxygen. Coordinates are not moved.
Figure 8: Frozen versus RePaint-style 2D-to-3D generation on QM9. Colors indicate QM9-S, QM9-M, QM9-L, and M-inv2; solid and dashed lines denote Frozen and RePaint-style conditioning. Panel (a) is xyz2mol validity: bonds are inferred from the generated coordinates, then RDKit sanitization and a single component are required. Panels (b–d) are mean absolute errors of bond lengths, bond angles, and dihedral angles against the reference connectivity. Panels (e–h) are mean aligned RMSD and Wasserstein-1 distances of the three local geometric distributions. The sampling-step axis is logarithmic. Each point uses 1,000 conditioning molecules, with two conformers for Frozen and one for RePaint-style. Errors are means over eligible conformers, not a best-of-two score.
Figure 9: Channel deletion. (a) One-step distance error at t=0.1 after replacing the named input channel by its prior; other channels stay at the training forward state. Sixty-four validation molecules, one noise seed. (b) Unconditional sampling with that channel’s input reset to the prior before every denoiser call (32 molecules, 20 steps). Curves are QM9-S, QM9-M, QM9-L, and M-inv2. Validity in (b) uses the generated bond table.
Figure 10: QM9-L 2D-to-3D generation at 100 steps. Each pair places the input graph on the left and one generated conformer on the right. The left drawing is the RDKit 2D depiction with hydrogens removed. The label on the 3D panel is the heavy-atom formula and the RMSD to the reference conformer, in angstroms.
Figure 11: Fragment editing on QM9 across sampling budgets. Top row, success. Bottom row, RDKit validity. Colors are S, M, L, and M-inv2. Solid lines are Frozen conditioning and dashed lines are RePaint-style. Columns are decoration, multi-fragment completion, scaffold replacement, and linker generation. The step axis is logarithmic. S, M, and L Frozen runs use 10,000 conditions. RePaint-style runs use 1,000. Linker generation has 296 eligible conditions.
Figure 12: Fragment-edit success at 100 steps for S, M, L, and M-inv2. (a) Frozen conditioning. (b) RePaint-style conditioning. S, M, and L Frozen use 10,000 conditions. RePaint-style runs use 1,000.
Figure 13: QM9-L scaffold decoration at 100 steps. The left panel is a 3D embedding of the conditioning scaffold. The right panel is the generated molecule.
Figure 14: QM9-L multi-fragment completion at 100 steps. Each case shows the two conditioning fragments and the generated molecule. Fragments are embedded from SMILES without adding hydrogens at the cut.
Figure 15: QM9-L scaffold replacement at 100 steps. Only cases with at least two side chains are shown. The left panels are those side chains; the right panel is the generated molecule.
Figure 16: QM9-L linker generation at 100 steps. The left panels are the two anchors that are kept. The right panel is the generated molecule that joins them.
Molecular size is intrinsically linked to composition, structure, and properties, but most 3D molecular generators treat size as a predefined input rather than a quantity to be generated. We introduce Equivariant-Free Transformer-Autoencoded Latent Flow Matching, a two-stage framework that samples a fixed-dimensional latent vector using flow matching and employs an autoregressive Transformer to determine molecular size, atom types, coordinates, and chemical attributes. Canonical atom ordering and rigid-pose alignment enable Transformers without equivariant layers, while decoded attributes guide bond reconstruction. On PCQM4Mv2, unconditional generation yields 87.9% unique and novel molecules passing sanitization and PoseBusters checks, exceeding baselines with lower end-to-end training and sampling time and higher end-to-end throughput. Across ten target HOMO-LUMO gaps, internal ranking retains 30% of the screened candidates and increases the density functional theory-verified hit rate within 0.1 eV of the target from 25.0% to 52.4%, while largely preserving novelty and diversity. These results demonstrate fixed-dimensional latent generation with autoregressive decoding as a practical approach to molecular design without prespecifying size.
Weichi Yao, Cameron Gruich, Bryan R. Goldsmith +1
Michigan Institute for Data & AI in Society · Department of Chemical Engineering, University of Michigan, Ann Arbor, USA · Department of Statistics, University of Michigan, Ann Arbor, USA
Three-dimensional (3D) molecule generation has been dominated by diffusion models, which achieve strong generation quality but typically require the molecular size to be specified a priori. Recent autoregressive approaches have substantially narrowed the performance gap while naturally supporting variable-length generation and conditioning on partial molecular context. However, balancing unconditional and context-conditioned generation remains challenging. We introduce KRONOS, a latent autoregressive diffusion framework that generates molecules in the latent space of a pre-trained autoencoder, jointly modeling molecular graph topology and geometry, while retaining the flexibility of autoregressive generation. We further introduce a mixed training strategy inspired by Fill-in-the Middle (FIM) paradigm, enabling both unconditional and fragment-conditioned molecular generation within a single left-to-right autoregressive model. Experiments on QM9 and GEOM-Drugs demonstrate that KRONOS achieves leading unconditional generation performance among autoregressive methods, while remaining competitive with diffusion models. Moreover, fragment-conditioned generation is achieved with negligible impact on unconditional generation performance, demonstrating that both generation paradigms can be supported within a single architecture.
Federico Ottomano, Gaopeng Ren, Yingzhen Li +2
Department of Chemistry Imperial College London · Department of Computing Imperial College London
Generating accurate 3D molecules requires jointly determining a variable number of atoms, their identities, and geometry across spatial scales: global shape emerges coarsely, whereas chemical bonds demand fine local precision. Existing atom-based diffusion and flow models operate largely at a single scale, while dense spatial fields become costly at high resolution. We introduce OctMole, a multiscale octree flow framework for 3D molecular generation. A hierarchical occupancy flow progressively resolves sparse molecular support from coarse to fine, and a joint atom flow assigns atom types and continuous in-cell offsets to occupied leaves. This decomposition combines efficient global structure with subcell geometric precision, while cell corruption improves robustness to generated support. On QM9 and GEOM-DRUG, OctMole achieves strong chemical, distributional, and geometric quality, with particularly clear gains over Uni-3DAR on GEOM-DRUG and substantial efficiency gains over dense 3D generation. As an additional capability, intermediate occupancy states and selected regions provide spatial conditioning without task-specific retraining, supporting multiscale shape conditioning, regional editing, linker generation, and shape interpolation. Code and pretrained checkpoints are available at https://github.com/octree-nn/octmole and https://huggingface.co/Disembo/octmole, respectively.
Yu-Tong Jin, Yi-Yang Tian, Qiwei Ye +1
Peking University · Beijing Academy of Artificial Intelligence