MolHIT: Advancing Molecular-Graph Generation with Hierarchical Discrete Diffusion Models

Hojung Jung; Rodrigo Hormazabal; Jaehyeong Jo; Youngrok Park; Kyunggeun Roh; Se-Young Yun; Sehui Han; Dae-Woong Jeong

cs.AIFeb 19, 2026

MolHIT: Advancing Molecular-Graph Generation with Hierarchical Discrete Diffusion Models

Authors: Hojung Jung, Rodrigo Hormazabal, Jaehyeong Jo, Youngrok Park, Kyunggeun Roh, Se-Young Yun, Sehui Han, Dae-Woong Jeong

Organizations: KAIST AI · LG AI Research · Seoul National University

Abstract

Molecular generation with diffusion models has emerged as a promising direction for AI-driven drug discovery and materials science. While graph diffusion models have been widely adopted due to the discrete nature of 2D molecular graphs, existing models suffer from low chemical validity and struggle to meet the desired properties compared to 1D modeling. In this work, we introduce MolHIT, a powerful molecular graph generation framework that overcomes long-standing performance limitations in existing methods. MolHIT is based on the Hierarchical Discrete Diffusion Model, which generalizes discrete diffusion to additional categories that encode chemical priors, and decoupled atom encoding that splits the atom types according to their chemical roles. Overall, MolHIT achieves new state-of-the-art performance on the MOSES dataset with near-perfect validity for the first time in graph diffusion, surpassing strong 1D baselines across multiple metrics. We further demonstrate strong performance in downstream tasks, including multi-property guided generation and scaffold extension.

Figures & tables

Figure 1: MolHIT achieves SOTA result on MOSES dataset. (Top) Near-perfect validity, outperforming existing graph diffusion models. (Bottom) Pareto-optimal in quality-novelty trade-off.

Figure 2: Overview of MolHIT . (a) Markov chain of Hierarchical Discrete Diffusion Model (HDDM). Clean states ( $S_{0}$ ) are transited to the mid-level states ( $S_{1}$ ) and finally to the masked state ( $S_{2}$ ). (b) Generation process of MolHIT. From the masked prior, atoms are denoised into mid-level states and then to atomic tokens in a coarse-to-fine manner. (c) Phase diagram of HDDM showing the transition probability of the forward process. (d) Decoupled atom encoding scheme, separately encoding the aromatic and charged atom types.

Figure 3: Existing atom encoding for molecular graph is ill-posed. (Left) Reconstruction success rate on the Moses dataset with previous encoding and our decoupled atom encoding. (Right) Proportion of generated molecules containing pyrrolic nitrogen

[nH]

Category	Model	Quality $\uparrow$	Scaf-Novel $\uparrow$	Scaf-Ret. $\uparrow$	Valid $\uparrow$	Unique $\uparrow$	Novel $\uparrow$	Filters $\uparrow$	FCD $\downarrow$	SNN $\uparrow$	Scaf $\uparrow$
-	Training set	95.4	—	—	100.0	100.0	—	100.0	0.48	0.59	-
1D Sequence	VAE[ 25 ]	92.8	0.22	0.031	97.7	99.7	69.5	99.7	0.57	0.58	5.9
	CharRNN [ 42 ]	92.6	0.29	0.035	97.5	99.9	84.2	99.4	0.52	0.56	11.0
	SAFE-GPT [ 33 ]	92.8	0.12	0.015	99.8	98.9	43.7	97.7	0.72	0.57	6.3
	GenMol [ 27 ]	62.1	0.05	0.012	99.7	64.0	68.9	98.1	16.4	0.64	1.6
2D Graph	DiGress [ 47 ]	82.5	0.26	0.031	87.1	100.0	94.2	97.5	1.25	0.53	12.8

Table 1: Comprehensive MOSES benchmark results. Scaffold Novelty (Scaf-Novel) measures the ratio of novel scaffold molecules to the number of generated molecules, while Scaffold Retrieval (Scaf-Ret.) quantifies test scaffold retrievals. All of the results are the averaged value over 3 runs of 25,000 samples. Bold denotes the best in each category, and underline indicates SOTA performance within the 2D Graph models. Empty values are due to the absence of publicly available checkpoints or samples.

Model	Val.	V.U.	V.U.N.	KL Div.	FCD
Training set	100.0	100.0	—	99.9	92.8
DiGress [ 47 ] (org.)	85.2	85.2	85.1	92.9	68.0
DiGress [ 47 ] (full)	74.7	74.6	74.0	92.4	61.1
DiGress + DAE	65.2	65.2	64.9	87.0	49.2
MolHIT (Ours)	87.1	87.1	86.0	96.7	54.9

Table 2: Full GuacaMol benchmark results using unfiltered dataset. Metric abbreviations: Val. (Validity), V.U. (Unique), V.U.N. (Novel). DiGress (org.) is original DiGress trained with filtered dataset and DiGress (full) values are from the re-implementation of DiGress on unfiltered, full GuacaMol dataset.

	MAE $\downarrow$					Pearson $r\uparrow$					Validity (%) $\uparrow$
Method	QED	SA	logP	MW	Avg.	QED	SA	logP	MW	Avg.
Marginal	0.117	0.115	0.067	0.272	0.143	0.489	0.570	0.802	0.396	0.564	75.03
Marginal + DAE	0.107	0.094	0.061	0.227	0.122	0.565	0.559	0.836	0.437	0.599	87.85
MolHIT (Ours)	0.061	0.040	0.049	0.081	0.058	0.804	0.790	0.950	0.685	0.807	96.31

Table 3: Multi-property guided generation on MOSES with four different conditions. We report mean absolute error (MAE), Pearson correlation (Pearson

r

), and validity. Avg. denotes the macro-average across four properties. Bold denotes best performances. All results are the averaged value over 3 runs of 10,000 samples.

Model	Validity (%) $\uparrow$	Diversity $\uparrow$	Hit@1 $\uparrow$	Hit@5 $\uparrow$
DiGress	50.8	44.8	2.07	6.41
Marginal + DAE	64.8	58.0	1.67	6.37
MolHIT (Ours)	83.9	57.4	3.92	9.79

Table 4: Scaffold extension results on the MOSES dataset. Results are averaged over 3 runs of 10,000 targets.

Method	Quality $\uparrow$	FCD $\downarrow$	Validity (%) $\uparrow$
DiGress [ 47 ]	82.5	1.25	87.1
+ DAE	87.6	0.89	96.2
+ PN Sampler	92.9	1.65	99.4
+ HDDM ( MolHIT )	94.2	1.03	99.1

Table 5: Incremental performance gains on the MOSES dataset by integrating DAE, the PN Sampler, and HDDM into the DiGress.

Figure 4: Effect of top-p sampling in MolHIT.

Appendix figures & tables10 assets

Supplementary material from the paper’s appendix.

Appendix

Method	Elemental Basis	Unique Tokens (Vocabulary)	Size
Standard Encoding	{C, N, S, O, F, Cl, Br}	C, N, S, O, F, Cl, Br	7
DAE (Ours)	{C, N, S, O, F, Cl, Br}	Aliphatic: C, N, S, O, F, Cl, Br	12
		Aromatic: c, n, nH , s, o

Appendix

Table 6: Comparison of Atom Vocabularies on the MOSES Dataset. DAE resolves structural ambiguities by decoupling elements into specific aromatic and hydrogen-locked states.

Category	Standard Encoding (Size: 12)	DAE Tokens (Size: 56)
Neutral Aliphatic	{C, N, O, F, B, Br, Cl, I, P, S, Se, Si}	C, N, O, F, B, Br, Cl, I, P, S, Se, Si
Aromatic States	(None / Implicit)	c, c+, c-, n, nH, n+, nH+, n-, s, s+, o, o+, se, se+, p
Charged &	(None / Implicit)	C+, C-, N+, NH+, NH2+, NH3+, N-, NH-, O+, O-, F+, F-,
Hypervalent		B-, Br+2, Br-, Cl+, Cl+2, Cl+3, Cl-, I+, I+2, I+3,
		P+, P-, S+, S-, Se+, Se-, Si-

Appendix

Table 7: Vocabulary expansion for the Guacamol dataset. DAE scales from 12 elemental types to 56 semantic tokens including aromatic and charged atoms.

Figure 5: The ratios of generated molecules having formal charge. MolHIT can reach to the training level proportion, while models with previous coarse encoding (left two) barely generate the charged atoms.

Figure 6: Reconstruction Fidelity and Identity Preservation. We measure the proportion of generated molecules that have at least one atom with formal charge.

Dataset	Group ID	Atom Elements ( $\mathcal{S}_{0}$ )
MOSES	Group 1	{C}
	Group 2	{N, O, S}
	Group 3	{F, Cl, Br}
	Group 4	{c, o, n, [nH], s}
GuacaMol	Group 1	{F, Cl, Br, I, F - , Cl - , Br - }
	Group 2	{C, N, O, P, S, Se}

Appendix

Table 8: Deterministic grouping kernels for node state space partitioning in MOSES and GuacaMol.

Category	Model	Quality $\uparrow$	Scaf-Novel $\uparrow$	Scaf-Ret. $\uparrow$	Valid $\uparrow$	Unique $\uparrow$	Novel $\uparrow$	Filters $\uparrow$	FCD $\downarrow$	SNN $\uparrow$	Scaf $\uparrow$
-	Training set	95.4	—	—	100.0	100.0	—	100.0	0.48	0.59	-
1D Sequence	VAE[ 25 ]	92.8 $\pm$ 0.2	0.22 $\pm$ 0.01	0.031 $\pm$ 0.003	97.7 $\pm$ 0.1	99.7 $\pm$ 0.0	69.5 $\pm$ 0.6	99.7 $\pm$ 0.0	0.57 $\pm$ 0.00	0.58 $\pm$ 0.01	5.9 $\pm$ 1.0
	CharRNN [ 42 ]	92.6 $\pm$ 2.5	0.29 $\pm$ 0.04	0.035 $\pm$ 0.003	97.5 $\pm$ 2.6	99.9 $\pm$ 0.0	84.2 $\pm$ 5.1	99.4 $\pm$ 0.3	0.52 $\pm$ 0.03	0.56 $\pm$ 0.01	11.0 $\pm$ 0.8
	SAFE-GPT [ 33 ]	92.8 $\pm$ 0.0	0.12 $\pm$ 0.00	0.015 $\pm$ 0.000	99.8 $\pm$ 0.0	98.9 $\pm$ 0.0	43.7 $\pm$ 0.3	97.7 $\pm$ 0.0	0.72 $\pm$ 0.02	0.57 $\pm$ 0.01	6.3 $\pm$ 0.7
	GenMol [ 27 ]	62.1 $\pm$ 0.0	0.05 $\pm$ 0.00	0.012 $\pm$ 0.001	99.7 $\pm$ 0.1	64.0 $\pm$ 0.5	68.9 $\pm$ 0.4	98.1 $\pm$ 0.1	16.36 $\pm$ 0.07	0.64 $\pm$ 0.01	1.6 $\pm$ 0.1
2D Graph	DiGress [ 47 ]	82.5 $\pm$ 0.7	0.26 $\pm$ 0.00	0.031 $\pm$ 0.000	87.1 $\pm$ 0.9	100.0 $\pm$ 0.0	94.2 $\pm$ 0.2	97.5 $\pm$ 0.0	1.25 $\pm$ 0.03	0.53 $\pm$ 0.00	12.8 $\pm$ 1.4

Appendix

Table 9: Unconditional generation on MOSES dataset with full statistics. We bring the reported value from Cometh and DeFoG from their work.

Model	Val. $\uparrow$	V.U. $\uparrow$	V.U.N. $\uparrow$	KL $\uparrow$	FCD $\downarrow$
Training set	100.0	100.0	—	99.9	92.8
DiGress (org.)	85.2	85.2	85.1	92.9	68.0
DiGress (full)	74.7 $\pm$ 0.4	74.6 $\pm$ 0.5	74.0 $\pm$ 0.4	92.4 $\pm$ 0.5	61.1 $\pm$ 0.2
DiGress+DAE	65.2 $\pm$ 0.4	65.2 $\pm$ 0.4	64.9 $\pm$ 0.4	87.0 $\pm$ 0.4	49.2 $\pm$ 0.6
MolHIT (Ours)	87.1 $\pm$ 0.5	87.1 $\pm$ 0.3	86.0 $\pm$ 0.5	96.7 $\pm$ 0.1	54.9 $\pm$ 0.2

Appendix

Table 10: Full statistics of GuacaMol benchmark results (unfiltered). Val.: Validity, V.U.: Unique, V.U.N.: Novel. All results are averaged over 3 runs.

	MAE $\downarrow$					Pearson $r\uparrow$					Validity (%) $\uparrow$
Method	QED	SA	LogP	MW	Avg.	QED	SA	LogP	MW	Avg.
Marginal	0.117 ( $\pm$ 0.002)	0.115 ( $\pm$ 0.003)	0.067 ( $\pm$ 0.001)	0.272 ( $\pm$ 0.009)	0.143 ( $\pm$ 0.004)	0.489 ( $\pm$ 0.003)	0.570 ( $\pm$ 0.012)	0.802 ( $\pm$ 0.003)	0.396 ( $\pm$ 0.001)	0.564 ( $\pm$ 0.005)	75.03 ( $\pm$ 0.74)
Marginal + DAE	0.107 ( $\pm$ 0.001)	0.094 ( $\pm$ 0.001)	0.061 ( $\pm$ 0.000)	0.227 ( $\pm$ 0.004)	0.122 ( $\pm$ 0.001)	0.565 ( $\pm$ 0.005)	0.559 ( $\pm$ 0.009)	0.836 ( $\pm$ 0.005)	0.437 ( $\pm$ 0.015)	0.599 ( $\pm$ 0.002)	87.85 ( $\pm$ 0.46)
MolHIT	0.061 ( $\pm$ 0.001)	0.040 ( $\pm$ 0.001)	0.049 ( $\pm$ 0.001)	0.081 ( $\pm$ 0.005)	0.058 ( $\pm$ 0.002)	0.804 ( $\pm$ 0.009)	0.790 ( $\pm$ 0.011)	0.950 ( $\pm$ 0.004)	0.685 ( $\pm$ 0.024)	0.807 ( $\pm$ 0.011)	96.31 ( $\pm$ 0.23)

Appendix

Table 11: Full statistics of multi-property guided generation on MOSES with 4 different conditions. We report mean absolute error (MAE;

\downarrow

), Pearson correlation (

r

;

\uparrow

), and validity. Avg. is the macro-average across properties. Bold denotes best values.

Model	Valid (%) $\uparrow$	Diversity $\uparrow$	Hit@1 $\uparrow$	Hit@5 $\uparrow$
DiGress	50.8 $\pm$ 0.5	44.8 $\pm$ 1.8	2.07 $\pm$ 0.09	6.41 $\pm$ 0.21
Marginal + DAE	64.8 $\pm$ 0.2	58.0 $\pm$ 0.1	1.67 $\pm$ 0.10	6.37 $\pm$ 0.24
MolHIT (Ours)	83.9 $\pm$ 0.4	57.4 $\pm$ 0.6	3.92 $\pm$ 0.23	9.79 $\pm$ 0.09

Appendix

Table 12: Full statistics of scaffold extension results on MOSES. Results are averaged over 3 runs of 10,000 targets. Hit@k denotes the recovery of ground-truth within

k

samples.

Split	Statistic	QED	SA Score	logP	MW (Da)
Training	Min	0.1912	1.2694	-5.3940	250.017
Training	Max	0.9484	7.4831	5.5533	349.999
Test	Min	0.2265	1.3339	-4.2894	250.042
Test	Max	0.9484	6.6916	5.7255	349.990

Appendix

Table 13: Min and max values for molecular property conditioning in MOSES training / test split.

Explore similar work

Controllable Molecular Generative Foundation Models
May 14, 2026Yihan Zhu, Yuhan Liu, Weijiang Li +2Molecular Optimization Molecular Graph Generation
University of Notre Dame
FragFlex: Fragment-Based Flexible Molecular Generation for Exploring Vast Chemical Space
Oct 11, 2026Sihyun Park, David Ryan KoesMolecular Graph Generation Molecular Generation
Carnegie Mellon University · University of Pittsburgh
MolDA: Molecular Understanding and Generation via Large Language Diffusion Model
Apr 6, 2026Seohyeon Shin, HanJun Choi, Jun-Hyung Park +2Multimodal Large Language Models Molecular Graph Generation
Gwangju Institute of Science and Technology, Gwangju, South Korea · Hankuk University of Foreign Studies, Seoul, South Korea

May 14, 2026J/K move · Enter open · S save

Controllable Molecular Generative Foundation Models

Yihan Zhu, Yuhan Liu, Weijiang Li, Tengfei Luo+1

University of Notre Dame

Molecular Optimization Molecular Graph Generation Graph Foundation Models RL for Diffusion Models Molecular Generation Conditional Generative Modeling

Open paper AskPDF arXiv

Despite the success of foundation models in language and vision, molecular graph generation still lacks a unified framework for heterogeneous design tasks with reliable controllability. While reinforcement learning (RL) offers a natural post-training mechanism for task-specific optimization, applying it to graph generative models is hindered by the vast atom-wise action spaces and chemically invalid intermediate states. We propose Controllable Molecular Generative Foundation Models (CoMole), built with a unified motif-aware graph diffusion pipeline. By learning a motif-aware graph space, CoMole transfers pretrained structural priors into controllable generation, where RL optimizes conditional reverse policies over chemically meaningful decisions. We theoretically characterize the bottleneck of atom-level RL and justify motif-aware policy optimization. Across three heterogeneous benchmarks spanning materials and drug discovery, CoMole ranks first in controllability on all nine targets, reduces MAE by up to 48.2% relative to the strongest baselines, and maintains validity above 0.94 without rule-based correction or post-hoc filtering. We further show that CoMole transfers controllability to unseen properties by optimizing only task embeddings with the generator frozen, achieving performance competitive with strong task-specific baselines.

Abstract

Explore similar work

Controllable Molecular Generative Foundation Models

FragFlex: Fragment-Based Flexible Molecular Generation for Exploring Vast Chemical Space

MolDA: Molecular Understanding and Generation via Large Language Diffusion Model