OctMole: Multiscale 3D Molecular Generation with Hierarchical Spatial Flows

Yu-Tong Jin; Yi-Yang Tian; Qiwei Ye; Peng-Shuai Wang

cs.LGOct 12, 2026

OctMole: Multiscale 3D Molecular Generation with Hierarchical Spatial Flows

Authors: Yu-Tong Jin, Yi-Yang Tian, Qiwei Ye, Peng-Shuai Wang

Organizations: Peking University · Beijing Academy of Artificial Intelligence

Abstract

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.

Figures & tables

Figure 1: Local spatial representation, illustrated in 2D. Left: octree subdivision for a three-atom molecule. Right: definition of the local offset within a leaf cell.

Figure 2: OctMole generates spatial support through hierarchical occupancy flows, then jointly generates atom types and continuous local coordinates. Intermediate occupancy states provide starting points for spatial conditioning.

Method	Atom Sta.	Mol. Sta.	Valid.	Unique	Connected
QM9
Data	99.900	99.300	99.500	–	–
EDM ( Hoogeboom et al., 2022 )	99.841	97.230	98.740	98.491	99.630
GeoLDM ( Xu et al., 2023 )	99.909	98.430	99.090	98.032	99.700
UniGEM ∗ ( Feng et al., 2025 )	99.815	98.550	99.500	96.010	99.350
Uni-3DAR ( Lu et al., 2025 )	99.976	99.620	99.800	95.842	99.680

Table 1: Chemical quality of unconditional 3D molecule generation under our Open Babel evaluation. All values are percentages and higher is better. UniGEM ∗ uses additional molecular property information. All models are evaluated on 10,000 final samples. Dark and light blue backgrounds mark the best and second-best results among generative models, respectively.

Figure 3: Open Babel molecule stability by atom-count bin on 10,000 GEOM-DRUG molecules per group. Marker positions are the mean atom count within each bin. OctMole is the most stable generator in every bin and degrades only mildly with size. Uni-3DAR has only two samples in the largest bin (red open marker), so its value there is not statistically informative.

Dataset	Method	Yield (%)	Energy	RMSD	Atoms
QM9	Data	–	0.0690	0.0574	18.0
	Uni-3DAR	99.47	0.2178	0.2349	17.2
	OctMole	99.53	0.1149	0.0628	18.4
GEOM-DRUG	Data	–	$\approx 0$	$\approx 0$	46.6
	Uni-3DAR	86.51	0.2988	0.6019	30.1
	OctMole	94.86	0.6335	0.6067	43.7

Table 2: GFN2-xTB comparison. Yield is relative to all 10,000 final samples; energy (kcal/mol/atom) and RMSD (Å) are medians over converged valid, connected molecules; atoms is the mean count in the corresponding xTB-evaluated sample set.

Dataset	Depth	IoU@ $d$	Chamfer (Å)	Diversity	Joint (%)
QM9	2	1.00	1.00	0.88	41.3
	4	0.97	0.08	0.01	4.2
	6	0.99	0.03	0.00	0.0
GEOM-DRUG	2	0.99	1.33	0.88	1.1
	4	0.98	0.27	0.28	53.3
	6	1.00	0.06	0.00	1.7

Table 3: Multiscale spatial conditioning. IoU@

d

is occupancy overlap at the clamped depth; Joint is the percentage of final samples satisfying chemical quality, shape, and novelty simultaneously. Full chemical statistics and confidence intervals are in Table 11 .

Figure 4: Multiscale spatial conditioning on GEOM-DRUG. Coarse conditions leave substantial chemical freedom, intermediate conditions retain shape while supporting diverse chemistry, and finest conditions approach reconstruction.

Table 8

Condition	IoU@4	Chamfer (Å)	Tanimoto	Valid	Stable
Depth 4 (coarse)	0.909	0.840	0.107	0.895	0.321
Depth 6 (fine)	0.968	0.131	0.663	0.954	0.831

Table 6: Shape retention, chemical similarity, and molecular quality under coarse and fine PDBBind ligand occupancy conditions. IoU@4 measures generated–query occupancy overlap at depth 4.

Appendix figures & tables27 assets

Supplementary material from the paper’s appendix.

Appendix

Dataset	Flow	Blocks	Heads	Registers	Parameters
QM9	Occupancy	12	12	8	129.11M
QM9	Atom	12	12	8	129.74M
GEOM-DRUG	Occupancy	12	12	8	129.11M
GEOM-DRUG	Atom	16	12	8	172.26M

Appendix

Table 7: Backbone configurations used in the main experiments. Parameter counts include the task-specific input and output heads.

Figure 7: Sparse transformer backbone used in OctMole. Cell coordinates provide spatial positional encoding, while time and, where applicable, target-depth conditions modulate the transformer blocks. Task-specific input and output heads instantiate the hierarchical occupancy flow or joint atom flow. Coordinate rotary embeddings and the register tokens are omitted from the schematic for clarity.

Dataset	Flow	GPUs	Batch/GPU	Epoch	Cell corruption
QM9	Occupancy	4	256	4000	–
QM9	Atom	4	256	1000	No
GEOM-DRUG	Occupancy	8	128	100	–
GEOM-DRUG	Atom	8	128	10	Yes

Appendix

Table 8: Training configurations used for the main results. Epoch is the training epoch used at inference. All GPUs are NVIDIA A100-SXM4-40GB.

Method	Nodes $W_{1}$	Atom TV	Bond TV	Valency $W_{1}$	Length $W_{1}$	Angle $W_{1}$
QM9
EDM ( Hoogeboom et al., 2022 )	–	.013020	.000711	.011710	.000548	.441630
GeoLDM ( Xu et al., 2023 )	–	.009140	.000689	.004940	.000716	.398670
UniGEM ∗ ( Feng et al., 2025 )	–	.009920	.005794	.008080	.002367	1.472070
Uni-3DAR ( Lu et al., 2025 )	.859498	.027621	.006703	.001959	.001845	.579394
OctMole	.316405	.022172	.002656	.000929	.000805	.437003

Appendix

Table 9: Distributional fidelity under our Open Babel evaluation. Lower is better for every metric. TV denotes total variation distance and

W_{1}

denotes one-dimensional Wasserstein distance. Dark and light blue backgrounds mark the best and second-best results among generative models, respectively.

QM9
Method	$N_{\rm vc}$	Conv. (%)	Yield (%)	Energy	RMSD
Data	10,000	100.00	–	0.0690 (0.1431)	0.0574 (0.2589)
EDM	9,839	100.00	98.39	0.1613 (1.1166)	0.2117 (0.8272)
GeoLDM	9,880	100.00	98.80	0.1001 (0.5513)	0.1174 (0.6952)
UniGEM	9,888	100.00	98.88	0.1194 (0.8162)	0.1364 (0.7021)
Uni-3DAR	9,948	99.99	99.47	0.2178 (0.6282)	0.2349 (0.8173)

Appendix

Table 10: Full GFN2-xTB comparison from 10,000 final samples per model.

N_{\rm vc}

counts valid, connected inputs; convergence is relative to these inputs and yield to all final samples. Energy is in kcal/mol/atom and RMSD in Å; both report median (95th percentile).

Figure 8: GEOM-DRUG atom-count distributions, including hydrogens, for 10,000 raw samples from each generator and a uniform 10,000-molecule sample from the training split. Curves are Gaussian kernel-density estimates; Uni-3DAR is shifted toward smaller molecules, whereas OctMole tracks the data distribution.

Figure 9: Empirical cumulative distributions of total GFN2-xTB relaxation energy over the full 10,000-sample evaluations, restricted to valid, connected molecules with converged optimization. The horizontal axis is logarithmic; nonpositive reference values are displayed at

10^{-3}

kcal/mol.

Figure 10: GEOM-DRUG distributional distances stratified by molecular size. Across all five size bins, OctMole is closest to the reference in molecular size, atom type, bond type, valency, and bond angle; bond length is the sole exception. Lower values indicate closer agreement with the reference.

Dataset	Depth	IoU@ $d$	IoU@6	Chamfer (Å)	Sim.	Div.	Sta.	Joint
QM9	2	1.00	0.00	1.00	0.12	0.88	99.7	41.3 [33.6, 49.1]
	4	0.97	0.51	0.08	0.97	0.01	100.0	4.2 [0.4, 10.4]
	6	0.99	0.99	0.03	1.00	0.00	100.0	0.0 [0.0, 0.0]
GEOM-DRUG	2	0.99	0.01	1.33	0.13	0.88	85.2	1.1 [0.5, 1.9]
	4	0.98	0.34	0.27	0.74	0.28	95.4	53.3 [45.5, 61.5]
	6	1.00	1.00	0.06	1.00	0.00	99.1	1.7 [0.5, 3.5]

Appendix

Table 11: Complete multiscale conditioning results. Joint is the percentage of finals that are valid and stable, within 1 Å Chamfer of the reference, and chemically distinct from it, with a 95% cluster-bootstrap interval over references.

Figure 11: Joint success by conditioning depth, with 95% cluster-bootstrap intervals. The useful depth shifts from 2 on QM9 to 4 on GEOM-DRUG.

Dataset	Depth	Identity	Rotation	Shift	Both
QM9	2	99.65	99.80	99.70	99.70
	4	100.00	99.50	100.00	98.65
	6	100.00	100.00	100.00	99.95
GEOM-DRUG	2	85.20	86.90	85.40	85.00
	4	95.40	82.95	93.05	80.05
	6	99.10	97.05	98.15	94.65

Appendix

Table 12: Molecule stability (%) under coordinate transformations with cell corruption enabled. Shift changes the grid origin; Both combines shift and rotation.

Dataset	Representation	Params (M)	GPUs	Batch	Epochs	Time (h)	Speedup
QM9	Dense	131.31	4	32	4,000	118.84	$1.00\times$
	Octree	129.10	4	256	4,000	17.35	$6.85\times$
GEOM-DRUG	Dense	236.07	8	32	100	143.86	$1.00\times$
	Octree	129.10	8	128	100	28.02	$5.13\times$

Appendix

Table 13: Training efficiency of dense and octree spatial generators. Batch is per GPU and time is observed wall-clock training time.

Dataset	Repr.	Time	Atom Sta.	Mol. Sta.	Valid	Unique	Conn.	Nodes $W_{1}$
QM9	Dense	2.090	96.831	69.500	94.300	99.576	80.400	3.5609
	Octree	0.379	99.989	99.900	99.900	98.298	100.000	0.5004
GEOM-DRUG	Dense	1.928	49.197	4.911	99.896	22.176	34.169	42.4219
	Octree	0.443	99.543	84.400	98.200	100.000	96.800	2.6977

Appendix

Table 14: Single-GPU sampling efficiency and quality for approximately 1,000 reranked final samples. Time is seconds per retained molecule. Chemical metrics are percentages; lower Nodes

W_{1}

is better.

Figure 12: Tokens per molecule across equivalent resolutions for 1,000 validation molecules from each dataset. Octree counts sum occupied nodes over all levels; lines show means and shaded bands show the 50th–90th percentiles. Dense-grid token counts grow cubically, whereas sparse octree token counts grow approximately linearly.

Figure 13: Quantization strain for 300 ground-truth GEOM-DRUG validation conformers at

32^{3}

64^{3}

, and

128^{3}

resolution without continuous offsets. Collision-bearing conformers are counted in the collision-free rate and excluded from geometry and xTB summaries at that resolution. The working

64^{3}

grid has no collisions in this cohort but still incurs substantial bond error and GFN2-xTB relaxation energy.

Dataset	$D$	$k$	Atom Sta.	Mol. Sta.	Valid	Connected	Nodes $W_{1}$	Length $W_{1}$
QM9	6	1	99.935	99.080	99.560	99.700	0.9984	0.0015
	6	2	99.985	99.730	99.770	99.960	0.3183	0.0008
GEOM-DRUG	6	1	99.614	86.790	98.180	96.960	3.6906	0.0023
	6	2	99.624	86.820	98.200	96.900	2.9844	0.0020
	8	2	98.229	52.700	92.440	82.160	4.9498	0.0039

Appendix

Table 15: Hierarchy ablation on 10,000 reranked final samples. Chemical metrics are percentages; lower

W_{1}

is better.

Figure 14: Matched-strength random perturbations over 128 references and 256 generations per arm. Deletion and pruning remove the same depth-2-derived physical region at both depths, and matched shifts move cells by 7.68 Å. The hatched depth-4 one-step shift is the unmatched 1.92 Å reference. Bars report changes from same-depth controls.

Dataset	Setting	Time	Mol. Sta.	Valid	Connected	Unique
QM9	Ours 50	327.61	100.00	100.00	100.00	98.40
	Ours 25	184.85	99.80	99.80	99.80	99.40
	Ours 10	98.50	98.20	99.40	98.80	99.80
	Uni-3DAR	218.90	100.00	100.00	100.00	99.60
GEOM-DRUG	Ours 50	342.82	87.80	98.00	96.20	100.00
	Ours 25	181.17	86.00	97.60	96.40	100.00

Appendix

Table 16: Sampling-step pilot with 500 final samples per setting. Time is pure generation time per retained final molecule in milliseconds.

Figure 15: Molecule stability by reranking-score decile for 10,000 final samples per dataset. Higher-confidence bins generally improve stability on GEOM-DRUG, whereas QM9 is close to saturation throughout.

Figure 16: Ligand rediscovery from shape alone on 354 PDBBind ligands, with 20 candidates per condition. Bars show the fraction of all query ligands whose best valid candidate reaches each Tanimoto threshold; unavailable query fingerprints count as failures.

Figure 17: PDBBind queries (left) and generated depth-4 shape mimics (right) with a shared camera per row. Shapes are similar while elemental composition and inferred chemistry differ.

Figure 18: Regional fragment editing examples. Each row shows a different source molecule and its edited output. Dark gray marks atoms occupying the retained finest-level cells, while the colored region is regrown from its depth-4 support; atom identities and within-cell coordinates are generated throughout.

	Connected	Valid	Atom Sta.	Mol. Sta.
Generated	37.0	82.0	95.5	27.0
All candidates	25.6	79.1	93.1	13.5
Ground truth	100.0	94.0	–	92.0

Appendix

Table 17: Linker generation quality over 100 targets using best-of-eight candidates. Ground-truth molecules are re-evaluated under the same protocol.

Figure 19: Three successful linker-generation cases. End-fragment occupancy is fixed in dark gray, and the regenerated linker is shown in element colors.

Figure 20: Linker-generation diagnostics on 100 targets. The left panel shows attachment-span errors for the 92 best-of-eight outputs with matched attachment atoms; the right panel classifies all selected outputs. Disconnection is the dominant outcome.

Figure 21: Shape interpolation for one representative GEOM-DRUG pair. Intermediate frames are best-of-eight generations from interpolated depth-4 occupancy.

Figure 22: Depth-4 occupancy IoU to the two endpoints across 20 interpolation pairs. Thin curves show individual pairs and bold curves show their means; the mean curves cross near

t=0.5

, forming a geometric intermediate.

3D Molecular Generation

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

See all 53 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
VecMol: Vector-Field Representations for 3D Molecule Generation
Mar 13, 2026Yuchen Hua, Xingang Peng, Jianzhu Ma +13D Molecular Generation Molecular Representation Learning
Institute for Artificial Intelligence, Peking University · School of Intelligence Science and Technology, Peking University · Department of Electronic Engineering, Tsinghua University +2
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

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

VecMol: Vector-Field Representations for 3D Molecule Generation

Autoregressive latent diffusion for 3D molecule generation