TreeRef-BFN: Equivariance-Free De Novo Molecule Generation based on 2D Topology and Internal 3D Geometry

Ruiqing Sun; Sen Yang; Dawei Feng; Bo Ding; Yijie Wang; Huaimin Wang

cs.AISep 26, 2026

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

Abstract

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 $z_{i}$ . 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 ( $\uparrow$ )	Unique ( $\uparrow$ )	Stability ( $\uparrow$ )		$W_{1}$ ( $\downarrow$ )
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 ( $\uparrow$ )	Unique ( $\uparrow$ )	Stability ( $\uparrow$ )		$W_{1}$ ( $\downarrow$ )
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 $\uparrow$	RMSD $\downarrow$	Bond $W_{1}$ $\downarrow$	Angle $W_{1}$ $\downarrow$	Dihedral $W_{1}$ $\downarrow$
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 $\uparrow$	MiDi stability $\uparrow$		EDM stability $\uparrow$		Unique $\uparrow$	Novel	$W_{1}\downarrow$
			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

W_{1}

is in Å; bond-angle and dihedral

W_{1}

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
$\alpha$ (Bohr 3 )	0.132	0.318
$\Delta\epsilon$ (meV)	99.8	137
$\epsilon_{\mathrm{HOMO}}$ (meV)	69.1	95.6
$\epsilon_{\mathrm{LUMO}}$ (meV)	53.0	93.5
$\mu$ (D)	0.057	0.434
$C_{v}$	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
$\alpha$ (Bohr 3 )	1.18	1.14	1.16
$\Delta\epsilon$ (meV)	225	174	297
$\epsilon_{\mathrm{HOMO}}$ (meV)	142	136	205
$\epsilon_{\mathrm{LUMO}}$ (meV)	169	136	235
$\mu$ (D)	0.531	0.495	0.507
$C_{v}$	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.

3D Molecular Generation

3D molecular generation creates molecules by modeling their chemical structures and three-dimensional geometries.

See all 55 papers

Fixed-Dimensional Latent Flow for Generating Variable-Size 3D Molecules
Sep 8, 2026Weichi Yao, Cameron Gruich, Bryan R. Goldsmith +13D Molecular Generation Molecular Generation
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
Autoregressive latent diffusion for 3D molecule generation
Jul 10, 2026Federico Ottomano, Gaopeng Ren, Yingzhen Li +2Latent Diffusion Models Molecular Graph Generation
Department of Chemistry Imperial College London · Department of Computing Imperial College London
OctMole: Multiscale 3D Molecular Generation with Hierarchical Spatial Flows
Oct 12, 2026Yu-Tong Jin, Yi-Yang Tian, Qiwei Ye +13D Molecular Generation Generative Modeling
Peking University · Beijing Academy of Artificial Intelligence

Sep 8, 2026J/K move · Enter open · S save

Fixed-Dimensional Latent Flow for Generating Variable-Size 3D Molecules

Weichi Yao, Cameron Gruich, Bryan R. Goldsmith, Yixin Wang

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

3D Molecular Generation Molecular Generation Flow-Based Generative Modeling

Open paper AskPDF arXiv

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.

Abstract

3D Molecular Generation

Explore similar work

Fixed-Dimensional Latent Flow for Generating Variable-Size 3D Molecules

Autoregressive latent diffusion for 3D molecule generation

OctMole: Multiscale 3D Molecular Generation with Hierarchical Spatial Flows