PaperlayerView exact revision on arXiv

Source-linked AI summary

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

arXiv:2602.17602v1cs.AI

TL;DR

Molecular graph diffusion seeks to combine graph models’ structural exploration with the chemical validity of sequence models. MolHIT addresses this gap using hierarchical discrete diffusion and decoupled atom encoding, and reports state-of-the-art performance across molecular benchmarks and downstream tasks. The paper also notes that larger models, architectural improvements, and additional GuacaMol training remain untested or incomplete.

  • Problem

    Graph diffusion models can generate invalid or chemically unstable molecules, while graph atom encodings ignore chemical states such as aromaticity and formal charge.

  • Method

    MolHIT combines Hierarchical Discrete Diffusion Model, which adds chemical-group states, with Decoupled Atom Encoding that explicitly represents aromaticity and formal charge.

  • Results

    MolHIT achieves state-of-the-art performance across molecular benchmarks and tasks, surpassing existing graph diffusion models and 1D baselines while attaining near-perfect reconstruction with DAE.

  • Takeaways & Limitations

    MolHIT supports end-to-end atom-level generation with formal charges and explicit nH, while improving performance on multi-property guided generation and scaffold extension.

  • Takeaways & Limitations

    The study does not test increased model size or architectural improvements, and GuacaMol training was not continued to performance saturation.

Abstract

from arXiv · show

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.

1. Introduction

MolHIT targets the validity–novelty gap in molecular graph generation with hierarchical diffusion and chemically informed atom encoding. It reports state-of-the-art performance across benchmarks and downstream tasks, including MOSES results that surpass graph and 1D baselines.

  • Motivation: Molecular generation must explore chemical space while preserving chemical validity, because minor atom-level errors can produce impossible or synthetically inaccessible structures.
  • Existing approaches: 1D SMILES models often generate valid molecules but can memorize training patterns, limiting exploration and creating a validity–novelty trade-off.
  • Existing approaches: Graph diffusion models support structural exploration but remain prone to invalid or chemically unstable samples compared with optimized 1D models.
  • Limitations: Existing discrete diffusion models neglect chemical relationships among atom categories and use naive encodings that make fine-grained atomic states difficult to reconstruct.
  • MolHIT: MolHIT combines Hierarchical Discrete Diffusion Model, which adds chemical-group states, with Decoupled Atom Encoding based on roles such as charge and aromaticity.
  • Results: Across MOSES, GuacaMol, scaffold extension, and multi-property guided generation, MolHIT improves over graph diffusion models and surpasses 1D models, achieving state-of-the-art results.

2. Preliminaries

Discrete diffusion defines noising and denoising over categorical states, with transition choices including uniform, marginal, and absorbing processes. For molecular graphs, atom and bond matrices are independently noised and a neural network learns to reconstruct the clean graph.

  • Discrete diffusion: Discrete diffusion models define noising and denoising as a Markov process over a discrete state space with K categories.
  • Discrete diffusion: The marginal distribution at timestep t can be written as q(x_t|x_0) = Cat(x_t; x_0 Q̄_t), where Q̄_t is the product of transition matrices.
  • Transition mechanisms: Uniform transition interpolates clean data with a uniform prior through a monotonically decreasing diffusion scheduler.
  • Transition mechanisms: Marginal transition uses a data prior approximating the empirical training distribution, whereas absorbing transition adds a masked state that receives all categories.
  • Molecular graph diffusion: For molecular graphs G = (X, E), the forward process independently transforms atom and bond matrices using their cumulative transition matrices.
  • Molecular graph diffusion: A neural network estimates the clean atoms and bonds from a noisy graph using cross-entropy loss, with λ balancing node and edge contributions.

3. MolHIT Framework

MolHIT combines hierarchical discrete diffusion, decoupled atom encoding, and specialized sampling to model chemically meaningful molecular graphs. Its design introduces chemical groups during diffusion and explicitly represents atom properties such as aromaticity and formal charge.

  • Hierarchical Discrete Diffusion Models: HDDM adds mid-level states between clean atom categories and masking, enabling a multi-stage corruption process rather than standard single-space or absorbing diffusion.The augmented state space partitions clean states, mid-level states, and a masked state; the forward process uses transition kernels satisfying Chapman–Kolmogorov consistency.
  • Hierarchical Discrete Diffusion Models: The transition operator Φ maps clean atom categories into chemically meaningful mid-level groups before later masking and denoising.Φ is a row-stochastic K×G matrix, while S1 states absorb clean states under the first transition.
  • Decoupled Atom Encoding: DAE explicitly encodes aromaticity and formal charge, producing near-perfect reconstruction on MOSES and GuacaMol and recovering motifs such as pyrrolic nitrogen.Coarse atomic-number encodings create one-to-many mappings across physical states and can fail to reconstruct important motifs.
  • Forward and Reverse Process: MolHIT uses HDDM for atom transitions and a uniform kernel for edges because the authors observe different transition preferences for atoms and bonds.The forward process decouples atom and bond dynamics; uniform transitions are reported as essential for edge generation, while HDDM performs better for atom types.
  • Forward and Reverse Process: The PN sampler projects denoising predictions onto the clean one-hot manifold before adding noise, addressing restricted structural exploration during standard posterior updates.Temperature sampling is applied to atom predictions, and the authors report that this configuration gives the best performance.

4. Experiments

Experiments evaluate MolHIT on unconditional and conditional molecular-generation benchmarks, including MOSES, GuacaMol, multi-property control, and scaffold extension. MolHIT reports strong graph-generation results while retaining validity, quality, and structural novelty across these settings.

  • Experimental Setup: MolHIT is evaluated on MOSES and GuacaMol, using decoupled vocabularies of 12 and 56 tokens respectively, with a DiGress graph transformer of unchanged model size.Reported results average three independent runs.
  • Unconditional Generation on MOSES: On MOSES, MolHIT achieves state-of-the-art Quality of 94.2%, Scaffold Novelty of 0.39, and near-perfect Validity of 99.1%.It outperforms graph baselines across nearly all key metrics, while 1D models show higher validity but lower novelty-related scores.
  • Unconditional Generation on GuacaMol: On GuacaMol, MolHIT ranks highest on every reported metric except FCD.Original DiGress has strong FCD, while DAE increases molecules with charged or special atoms; MolHIT results use 40 training epochs versus 1,000 for original DiGress.
  • Multi-property Guided Generation: For multi-property generation, MolHIT achieves macro-averaged MAE of 0.058, a 52.4% reduction versus Marginal+DAE, Pearson r of 0.807, and validity above 95%.The evaluated properties are QED, SA, MW, and logP.
  • Scaffold Extension: For scaffold extension, MolHIT significantly outperforms DiGress across all reported metrics, while applying DAE improves validity and diversity but lowers Hit@1.DAE produces higher diversity and matches the original DiGress on Hit@5.
  • Component Analysis: An ablation study attributes incremental gains to integrating DAE, the PN sampler, and HDDM, which together yield the highest Quality, FCD, and Validity among graph diffusion models.The study is conducted on the MOSES dataset.

5. Related Works

Related work spans discrete diffusion mechanisms for categorical state spaces and molecular-generation models built with continuous, discrete, and sequence-based representations.

  • Discrete Diffusion Models: Discrete diffusion research includes uniform transitions, absorbing transitions, and related transition mechanisms for discrete state spaces.The paper positions HDDM as a hierarchical extension with intermediate states.
  • Diffusion Models for Molecular Generation: Molecular-generation studies apply continuous score-based diffusion, discrete graph diffusion, CTMC-based methods, and 1D sequence modeling.The cited graph methods include GDSS, DiGress, and later CTMC-based approaches.

6. Conclusion

MolHIT is presented as a hierarchical discrete diffusion model for molecular generation, achieving state-of-the-art performance on large molecular datasets and enabling direct atom-level generation with formal charges or explicit nH.

  • MolHIT uses a hierarchical discrete diffusion framework for molecular generation.
  • The algorithm achieves state-of-the-art performance on large molecular datasets.
  • MolHIT directly generates atoms with formal charges or explicit nH, supporting more realistic molecule generation.

Impact Statement

The paper aims to advance molecule generation to accelerate discovery of useful drugs and materials, while acknowledging potential misuse to generate harmful substances.

  • The work aims to accelerate the discovery of useful drugs and materials and improve human lives.
  • The model could be maliciously used to generate substances harmful to humans and environments.

A. Limitation and future directions

The paper identifies implementation and evaluation limits while outlining future directions for hierarchical diffusion, richer tokenization, and related discrete-diffusion analysis.

  • Limitations: The models were not tested with increased model size or architectural improvements, leaving stated room for better performance.
  • Limitations: Performance on GuacaMol was not fully trained to saturation, so further training may improve results.
  • Future directions: Future work includes applying Hierarchical Discrete Diffusion Models to language and image domains and combining them with different sampling schemes.
  • Future directions: The framework can be further improved through tokenization that incorporates motifs or functional groups.
  • Generalized HDDM forward process: The generalized forward process partitions the state space into clean atomic states and multiple intermediate hierarchical levels.
  • Generalized HDDM forward process: Each hierarchical stage uses a row-stochastic kernel mapping cumulative lower-level states to specific higher-level states.
  • Generalized HDDM forward process: The induced transition matrices have block structure, with hierarchical transitions in active blocks and identity mappings elsewhere.
  • Generalized HDDM forward process: The transition kernel satisfies Chapman–Kolmogorov consistency under the stated conditions.

C.2. Proof of Theorem 3.2

The proof develops posterior and loss expressions for hierarchical discrete diffusion by combining forward transition kernels, masking, Bayes’ rule, and neural parameterization of clean-state probabilities.

  • The conditional forward transition is categorical, and Bayes’ rule yields the posterior q(zs|zt, x).
  • The posterior is analyzed across three cases determined by whether zt belongs to the clean, intermediate, or masked state spaces.
  • When zt is masked, the posterior combines the clean state, its stochastic projection, and the mask prior.
  • A neural network is parameterized to estimate only probabilities in the clean final state S0, producing closed-form posterior expressions.
  • The derivation proceeds to the continuous-time NELBO of HDDM and a general NELBO theorem.
  • The theorem applies to any stochastic row matrix Φ, allowing arbitrary stochastic mappings from S0 to S1.
  • Deterministic grouping, where rows of Q(1) are one-hot, enables parameterization around categories sharing a mid-level state and simplifies the NELBO.

D.1. Decoupled Atom Encoding (DAE)

Decoupled Atom Encoding (DAE) expands molecular atom representations beyond atomic numbers by separating chemically meaningful descriptors, resolving structural ambiguities in graph generation.

  • DAE decouples aromaticity, hydrogen saturation, and formal charge magnitude into separate chemical descriptors.
  • MolHIT reaches the training-level proportion of charged molecules, whereas prior coarse encodings barely generate charged atoms.
  • DAE resolves pyrrolic-hydrogen ambiguities such as [nH] in heteroaromatic rings by explicitly separating aromaticity and hydrogen counts.
  • DAE expands the GuacaMol vocabulary from 12 elemental types to 56 semantic tokens, including aromatic and charged atoms.
  • On GuacaMol, coarse encoding achieves 1.88% success on [nH] reconstruction and 0.09% identity preservation, illustrating the reconstruction bottleneck DAE targets.

D.3. Full experimental results with standard deviations

The experiments evaluate unconditional and conditional molecular generation using standardized benchmarks, repeated runs, and consistent model-training settings.

  • Three experiments are run for each setting, with standard deviations reported for MOSES, GuacaMol, guided generation, and scaffold extension.
  • Baselines use released checkpoints or official implementations, with GenMol retrained on MOSES for fair comparison.
  • The graph-transformer backbone uses 12 transformer blocks and approximately 16.2M trainable parameters across experimental settings.
  • MOSES evaluation covers Validity, Uniqueness, Novelty, Filters, FCD, SNN, and Scaffold Similarity.
  • Validity measures the fraction of generated molecules passing RDKit sanitization and basic chemical valency checks.
  • Uniqueness measures valid molecules that are not duplicates, while Novelty measures valid unique molecules absent from the training set.

D.6. Structure novelty metric

The paper evaluates molecular novelty alongside scaffold retrieval and conditional extension, distinguishing exploration beyond training data from recovery of known structures.

  • Scaffold Novelty measures generated Bemis–Murcko scaffolds absent from the training distribution.
  • Scaffold Retrieval measures unique test-set scaffolds successfully generated from the held-out distribution.
  • Scaffold retrieval tests distributional accuracy by checking whether generated structures capture the valid, drug-like test manifold rather than noise.
  • Scaffold extension fixes the input Bemis–Murcko scaffold and generates a completed molecule with the same atom count as the ground truth.
  • Validity is computed from chemically valid generated extensions, while diversity is computed on the unique valid set using Morgan-fingerprint Tanimoto similarity.
  • Hit@K records whether the canonical ground-truth molecule appears among K generated candidates.
Loading 2602.17602v1…