MolHIT: Hierarchical Molecular Graph Diffusion
- MolHIT is a molecular graph generation framework that integrates hierarchical discrete diffusion with decoupled atom encoding to improve chemical validity.
- It employs a coarse-to-fine corruption process and a Project-and-Noise sampler to preserve molecular structure while ensuring novelty.
- The framework supports unconditional, multi-property guided generation and scaffold extension, outperforming previous graph diffusion models.
Searching arXiv for papers relevant to “MolHIT” and closely related molecular generation / hit-identification work. MolHIT is a molecular graph generation framework introduced by Hojung Jung, Rodrigo Hormazabal, Jaehyeong Jo, Youngrok Park, Kyunggeun Roh, Se-Young Yun, Sehui Han, and Dae-Woong Jeong in “MolHIT: Advancing Molecular-Graph Generation with Hierarchical Discrete Diffusion Models” (Jung et al., 19 Feb 2026). It is designed for unconditional molecular generation, multi-property guided generation, and scaffold extension in 2D molecular graphs. The method combines a Hierarchical Discrete Diffusion Model (HDDM) with Decoupled Atom Encoding (DAE) in order to improve chemical validity, distributional fidelity, and structural novelty in graph diffusion, while addressing the persistent gap between graph-based and strong 1D sequence-based molecular generators (Jung et al., 19 Feb 2026).
1. Problem setting and design rationale
MolHIT is motivated by a specific limitation of prior molecular graph diffusion models: although discrete diffusion is naturally matched to categorical atom and bond labels, earlier graph diffusion systems were reported to suffer from low chemical validity and weaker practical performance than strong 1D baselines such as SMILES-based generators (Jung et al., 19 Feb 2026). The paper identifies invalid molecules as a central failure mode, since incorrect atom assignments, aromatic states, or valence patterns can render a generated graph unusable.
The method is situated within the distinction between 1D molecular generation and 2D graph generation. In 1D approaches, molecules are generated as strings, typically SMILES, which often yields high validity because the modeling problem is close to standard sequence generation. In 2D graph generation, atoms and bonds are modeled directly as a graph, which is more faithful to molecular topology but harder to denoise while preserving chemistry (Jung et al., 19 Feb 2026). MolHIT takes the position that graph diffusion remains attractive because molecules are naturally graphs and graph models should better capture topological and chemical structure, but that the state space itself must encode stronger chemistry-aware priors.
The paper studies three settings. The main setting is unconditional molecular graph generation on MOSES and GuacaMol. It also evaluates multi-property guided generation, where molecules are generated to satisfy target values of several molecular descriptors simultaneously, and scaffold extension, where a scaffold is fixed and the remaining region is generated (Jung et al., 19 Feb 2026). This makes MolHIT a general graph diffusion framework rather than a narrowly unconditional generator.
2. Hierarchical discrete diffusion model
The core probabilistic contribution is HDDM, which augments standard discrete diffusion with additional intermediate atom categories that encode coarse chemical group information (Jung et al., 19 Feb 2026). Let the clean atom state space be with categories. MolHIT introduces intermediate categories and one masked state , so the total state space has size .
A row-stochastic matrix maps clean atom categories to intermediate categories. This induces the projection-like transition
while the absorbing transition to the mask state is
Using schedules and 0 with 1, 2, and 3, MolHIT defines
4
with 5 and 6. The paper proves Chapman–Kolmogorov consistency,
7
yielding the cumulative kernel
8
This produces a coarse-to-fine corruption process. A clean atom can remain unchanged with mass 9, be coarsened to an intermediate chemical group with mass 0, or be fully masked with mass 1 (Jung et al., 19 Feb 2026). In the paper’s formulation, this is the main mechanism by which diffusion preserves broad chemical identity before exact atom identity is forgotten.
MolHIT applies HDDM only to atoms. Bonds use a simpler uniform transition, which the authors report works better empirically for edge generation (Jung et al., 19 Feb 2026). The schedules used in experiments are
2
The paper also derives an ELBO/NELBO for HDDM and states that, although the hierarchical process has a principled variational interpretation, the model is trained in practice with the simpler original cross-entropy loss from standard graph diffusion because the implied regularization did not improve performance (Jung et al., 19 Feb 2026).
3. Decoupled atom encoding and chemistry-aware state design
The second major component is DAE. MolHIT argues that previous graph diffusion models often encode nodes only by atomic number, which collapses chemically distinct states into the same token. The paper identifies aromatic versus aliphatic states, hydrogen saturation, and formal charge magnitude as the critical ambiguities (Jung et al., 19 Feb 2026). Under such coarse encoding, reconstruction becomes ill-posed because a single node token can correspond to multiple chemically distinct states.
DAE resolves this by expanding the atom vocabulary so that chemically specific atom roles become explicit categories. On MOSES, the original 7 heavy-atom vocabulary 3 is expanded to 12 tokens:
- aliphatic: 4
- aromatic: 5
On GuacaMol, DAE expands 12 coarse element types to 56 semantic atom tokens, including aromatic forms and charged states such as 6, and others (Jung et al., 19 Feb 2026). The paper emphasizes that this avoids introducing explicit hydrogens as separate graph nodes, which would increase graph size and cost.
The intermediate categories used by HDDM are defined over these DAE tokens. On MOSES, the 12 tokens are grouped into four categories:
- 7
- 8
- 9
- 0
On GuacaMol, six groups are used, spanning neutral halogens, neutral heavy atoms, aromatic states, cationic states, anionic or other charged states, and special elements such as boron and silicon (Jung et al., 19 Feb 2026).
This state redesign is central to the paper’s chemical-validity argument. The reconstruction analysis reports that, on GuacaMol, coarse encoding had only 1.88% success rate on the 1 group and negligible identity preservation, whereas DAE achieved 100% success and over 99.98% identity preservation (Jung et al., 19 Feb 2026). In the paper’s interpretation, DAE does not merely improve tokenization; it changes the denoising problem into one where chemically specific states are locally represented instead of being inferred indirectly from the full edge tensor.
4. Denoiser architecture, sampling, and conditional generation
MolHIT uses a graph transformer denoiser inherited from DiGress. At a high level, a molecular graph is written as 2, with 3 the atom feature matrix and 4 the bond tensor, and the denoiser predicts clean atoms and bonds from noisy graph states (Jung et al., 19 Feb 2026). The training objective is
5
The model keeps different forward processes for atoms and edges. Atoms use HDDM,
6
whereas edges use the uniform transition
7
The paper reports that the denoiser backbone remains essentially the same size as DiGress: 12 transformer blocks and about 16.2M parameters, with DAE and HDDM adding negligible overhead, reported as less than 0.01% of parameters (Jung et al., 19 Feb 2026).
Sampling uses the Project-and-Noise sampler rather than only the standard discrete-diffusion posterior update. At each step, the denoiser predicts clean atom and bond distributions, these are projected back to a discrete one-hot graph by categorical sampling, and the projected graph is then re-noised to the previous timestep with the forward kernel (Jung et al., 19 Feb 2026). The paper characterizes this as bypassing posterior constraints from the current noisy state in order to improve exploration. It also studies temperature and top-8 nucleus sampling, concluding that applying them to atom predictions only works best.
For multi-property guided generation, MolHIT modifies the graph transformer with Adaptive LayerNorm:
9
where 0 and 1 are affine maps of the conditioning vector 2 (Jung et al., 19 Feb 2026). The paper further uses property dropout during training to enable classifier-free guidance at inference.
For scaffold extension, scaffold atoms are held fixed throughout the reverse process,
3
while the remaining atoms and bonds are generated (Jung et al., 19 Feb 2026). This is the paper’s constrained-generation mechanism for medicinal-chemistry-style scaffold completion.
5. Empirical performance
The principal benchmark is MOSES. MolHIT reports the following metrics: Quality 94.2, Scaffold Novelty 0.39, Scaffold Retrieval 0.033, Validity 99.1%, Uniqueness 99.8%, Novelty 91.4%, Filters 98.0%, FCD 1.03, SNN 0.55, and Scaffold similarity 14.4 (Jung et al., 19 Feb 2026). The paper characterizes the 99.1% validity result as near-perfect validity for the first time in graph diffusion. In the direct comparison reported there, graph diffusion baselines had substantially lower validity: DiGress 87.1, DisCo 88.3, Cometh 87.2, and DeFoG 92.8 (Jung et al., 19 Feb 2026).
On GuacaMol, which the paper describes as harder because it includes charged chemistry and is used in its full unfiltered form, MolHIT reports Validity 87.1, Valid-unique 87.1, Valid-unique-novel 86.0, KL 96.7, and FCD 54.9 (Jung et al., 19 Feb 2026). The authors state that MolHIT outperforms DiGress trained on the full unfiltered dataset in all but FCD.
The ablation study isolates the contributions of DAE, the PN sampler, and HDDM. On MOSES, the reported sequence is:
- DiGress: Quality 82.5, FCD 1.25, Validity 87.1
- DAE: Quality 87.6, FCD 0.89, Validity 96.2
- PN sampler: Quality 92.9, FCD 1.65, Validity 99.4
- HDDM (full MolHIT): Quality 94.2, FCD 1.03, Validity 99.1
This pattern is interpreted as follows. DAE provides the largest direct gain in validity and fidelity; PN sampling pushes validity higher still but worsens FCD; HDDM then improves the quality–fidelity tradeoff while preserving near-perfect validity (Jung et al., 19 Feb 2026).
MolHIT also reports strong downstream results. In multi-property guided generation on MOSES, with target properties QED, SA, logP, and molecular weight, the reported MAEs are 0.061, 0.040, 0.049, and 0.081 respectively, with average 0.058; the corresponding Pearson correlations are 0.804, 0.790, 0.950, and 0.685, with average 0.807; and validity is 96.31% (Jung et al., 19 Feb 2026). In scaffold extension, MolHIT reports Validity 83.9%, Diversity 57.4, Hit@1 3.92, and Hit@5 9.79, exceeding both DiGress and the Marginal+DAE baseline reported in the paper (Jung et al., 19 Feb 2026).
6. Position within molecular machine learning
MolHIT belongs to the lineage of discrete molecular graph diffusion, but its novelty lies primarily in state-space design rather than in replacing the denoiser backbone (Jung et al., 19 Feb 2026). In contrast to earlier graph diffusion systems such as DiGress, DisCo, Cometh, and DeFoG, it injects chemistry into the diffusion process through hierarchical atom-state corruption and chemically explicit node vocabularies (Jung et al., 19 Feb 2026). This makes it a representation-level and transition-level intervention rather than a purely architectural one.
The method is distinct from several nearby lines of work. “Data-Efficient Molecular Generation with Hierarchical Textual Inversion” introduces HI-Mol, a low-shot SMILES-based generator built on hierarchical textual inversion over frozen text-to-molecule models, rather than graph diffusion (Kim et al., 2024). “TextOmics-Guided Diffusion for Hit-like Molecular Generation” introduces ToDi, a SELFIES diffusion model conditioned jointly on omics expressions and textual molecular descriptions for hit-like generation (Yuan et al., 14 Jul 2025). “Lo-Hi: Practical ML Drug Discovery Benchmark” is not a generator but a benchmark and splitter for novelty-constrained hit identification and lead optimization (Steshin, 2023). “MolProphecy” addresses molecular property prediction through gated cross-attention fusion of graph features and chemist-knowledge text, not graph generation (Zhao et al., 26 Jun 2025). A common source of confusion is therefore terminological rather than methodological: MolHIT, in the sense of (Jung et al., 19 Feb 2026), is a chemistry-aware molecular graph generation framework, not a hit-identification benchmark, not a human-in-the-loop predictor, and not the low-shot textual inversion method HI-Mol.
The paper’s limitations are also explicit. It does not claim to solve synthesizability beyond standard SA-like proxies, 3D conformational consistency, or experimental validation of generated molecules (Jung et al., 19 Feb 2026). The authors state that they have not fully explored scaling model size or architectural improvements beyond the DiGress backbone, and that GuacaMol training was not run to full saturation. The appendix identifies future directions including richer tokenization with motifs or functional groups, extension to 3D molecular generation and proteins, and combination with more advanced samplers (Jung et al., 19 Feb 2026). This suggests that MolHIT is best understood not as a complete endpoint for graph-based molecular generation, but as a demonstration that graph diffusion can approach the reliability previously associated with strong 1D generators once chemical priors are embedded directly into the discrete state space.