Beyond Flat Walks: Compositional Abstraction for Autoregressive Molecular Generation

Introduction

Molecular graphs are a natural representation for drug discovery and property prediction: nodes are atoms, edges are bonds, and validity and function depend heavily on recurring substructures — rings, functional groups, and their connectivity. While discriminative graph neural networks excel at predicting properties from such graphs, generating novel, valid molecules remains challenging. Autoregressive transformers, which have scaled successfully in language and code, offer an attractive backbone for graph generation when graphs are first tokenized into sequences. Existing tokenization schemes typically flatten the graph into a single linear order — a random walk or depth-first traversal — so that a causal transformer can predict the next token. These flat walks capture local adjacency and some long-range edges via back-references, but they do not expose the graph’s compositional hierarchy: which atoms belong to which ring or functional group, and how those units connect. The model must therefore infer motif structure implicitly from the token stream, which can lead to broken rings, inconsistent motif statistics, and limited control over high-level chemistry.

MOSAIC addresses this gap by introducing an explicit hierarchical graph (H-graph) abstraction and a three-stage pipeline for autoregressive molecular generation. First, we coarsen the molecular graph into a set of communities and inter-community edges, optionally using domain knowledge (ring and functional-group detection) to guarantee motif cohesion. Second, we tokenize the H-graph into a flat sequence using a coarse-before-fine ordering: the model sees community-level structure before atom-level detail, mirroring the compositional nature of chemistry. Third, we flatten the generated token sequence back into a molecular graph. All tokenization schemes support a lossless roundtrip from graph to tokens and back.

Pipeline Overview

Coarsening Strategies

The coarsening step determines how atoms are grouped into communities. We explore two families: flexible methods that optimize a graph-theoretic objective without domain knowledge, and constraint methods that leverage chemical structure to define communities. Spectral clustering uses the eigenvectors of the graph Laplacian to identify natural clusters, selecting the number of partitions that maximizes Newman–Girvan modularity. Hierarchical agglomerative clustering (HAC) takes a bottom-up approach, iteratively merging the most similar adjacent nodes based on bond-type-weighted distances, which tends to preserve local chemical neighborhoods. Motif-aware coarsening (MC) first identifies known chemical motifs (rings, functional groups) via SMARTS pattern matching, then clusters the remaining atoms with spectral or HAC methods. This ensures that chemically meaningful substructures are preserved as intact units. MC with functional groups (MC+FG) extends this with an expanded library of functional-group patterns for finer-grained control over which substructures remain intact during coarsening.

Tokenization Schemes

Given a coarsened hierarchy, we serialize it into a flat token sequence for autoregressive modeling. A shared design principle across all tokenizers is coarse-before-fine ordering: community-level structure always precedes atom-level detail, so the model first commits to the high-level decomposition before filling in local connectivity. H-SENT (Hierarchical SENT) encodes multi-level partition structure with explicit partition boundaries and bipartite cross-community edge blocks; each community’s atoms are serialized via a SENT walk with back-edge brackets. HDT (Hierarchical DFS Tokenization) encodes the full tree structure using depth-first traversal with ENTER/EXIT nesting tokens, capturing both intra- and inter-community edges as back-edges during the DFS, with no separate bipartite reconstruction needed. HDTC (HDT with Composition) is the most expressive scheme: it types each community node (ring, functional group, singleton), encodes a super-graph of inter-community bonds, and serializes atom-level detail within each typed community block.

Results

We evaluate unconditional generation on MOSES (~1M drug-like molecules, 10–26 heavy atoms) and COCONUT (~5K complex natural products, 30–100 heavy atoms, ≥4 rings). All models use a GPT-2 backbone (12 layers, 768 hidden, 12 heads, ~85M parameters) trained with identical hyperparameters. We generate 500 molecules per model and evaluate against 5,000 randomly sampled references from the combined train+test set. We report four representative tokenizers: SENT (flat baseline), H-SENT MC (hierarchical flat-walk), HDT MC (motif-community hierarchy), and HDTC (typed compositional). Bold = best among the four per dataset.

On MOSES (simple drug-like molecules), hierarchy provides modest but consistent improvements: SENT achieves 86.8% validity while HDT MC reaches 89.1% with the best FCD (2.30) and PGD (0.000). H-SENT MC achieves the best scaffold similarity (0.896). HDTC leads on motif rate (0.836). All four achieve near-identical motif-level fidelity (FG MMD ≤ 0.002). On COCONUT (complex natural products), the gap is dramatic: SENT drops to 58.1% validity while all hierarchical tokenizers achieve 87–90%, a 31-point improvement. No single model dominates distributional metrics: H-SENT MC leads SNN (0.783) and scaffold similarity (0.143), HDT MC achieves perfect PGD (0.000) and the best FCD (3.25), and HDTC leads motif rate (0.708) and fragment similarity (0.972). All hierarchical models achieve comparable distributional fidelity (FCD 3.1–4.6, all MMDs ≤ 0.004), suggesting that the choice of tokenizer and coarsening is less about overall quality than about which specific aspects of the distribution to optimize.

Full Results

Generation Gallery

Discussion

Hierarchy is necessary for complex molecules. On MOSES (~20 atoms), flat walks achieve 86.8% validity and hierarchy adds only 2.3 points. On COCONUT (30–100 atoms), flat walks drop to 58% validity while all hierarchical tokenizers reach 87–90%, a 31-point gap. The benefit scales with molecular complexity: larger molecules have real compositional structure that hierarchy can exploit.

Enumeration and traversal are complementary. H-SENT serializes molecules as a catalog of parts plus a wiring diagram (explicit index blocks and bipartite edge lists), excelling at distributional precision (best FCD 3.12, best Subst TV/KL). HDT serializes molecules as one continuous DFS walk with nesting, excelling at valid and diverse generation (best validity 89.9%, best PGD 0.000). The coarsening strategy amplifies this: spectral coarsening pairs with H-SENT’s enumeration, HAC pairs with HDT’s traversal.

Coarsening robustness scales with molecular size. On small MOSES molecules, only chemistry-aware MC coarsening works (generic methods drop to ≤22% validity). On complex COCONUT molecules, all coarsening methods achieve 87–90% validity, because larger molecules contain genuine hierarchical motifs that even naive algorithms can discover.

Typed decomposition gives motif precision, not distributional dominance. HDTC’s R/F/S type labels yield the highest motif rate (0.708) and fragment similarity (0.972), but untyped models match or beat it on FCD, SNN, and PGD with greater diversity. The diversity cost of hierarchy is a one-time flat-vs-hierarchical penalty, not a gradient that worsens with more constraint.

Substructure metrics generalize better than molecule-level similarity. SNN and scaffold similarity drop 2–5× between full-reference and test-only evaluation, while fragment similarity and MMD metrics remain stable. This gap is uniform across all models, driven by the small training set (5K molecules) rather than tokenizer-specific memorization.

BibTeX Citation

@article{bian2025mosaic,
  title     = {Beyond Flat Walks: Compositional Abstraction for
               Autoregressive Molecular Generation},
  author    = {Bian, Kaiwen and Yang, Andrew H. and Parviz, Ali
               and Mishne, Gal and Wang, Yusu},
  year      = {2025},
  url       = {https://github.com/KevinBian107/MOSAIC},
}