Beyond Flat Walks: Compositional Abstraction for Autoregressive Molecular Generation

Introduction

Discovering new molecules — for drugs, materials, or chemical tools — is a slow and expensive process. Recent advances in generative AI offer a way to accelerate this: train a model on known molecules and let it propose entirely new ones. But molecules aren’t just sequences of characters. They have rich internal structure: rings, branches, and recurring building blocks that determine their chemical properties.

Most existing approaches flatten a molecular graph into a linear sequence of atoms and bonds, then generate one token at a time — much like a language model that can only produce text one letter at a time. It technically works, but the model has to rediscover higher-level patterns (like common ring structures) from scratch. MOSAIC is more like generating at the word level: it learns to compose molecules from meaningful building blocks rather than individual atoms.

Pipeline Overview

MOSAIC operates through a three-stage pipeline. First, molecular graphs are recursively coarsened into multi-level hierarchies where atoms are grouped into structurally meaningful clusters. Second, these hierarchies are serialized into structured token sequences that preserve parent-child relationships and inter-cluster connectivity. Third, a Transformer decoder is trained to autoregressively generate these structured sequences, enabling the model to compose molecules from coarse structure down to fine-grained atomic detail.

Coarsening Strategies

MOSAIC supports multiple graph coarsening strategies that recursively partition molecular graphs into hierarchical clusters. Each strategy offers different trade-offs between preserving chemical motifs and computational efficiency.

• Spectral Clustering: Uses the eigenvectors of the graph Laplacian to identify natural clusters in the molecular graph. Spectral methods capture global graph structure and produce balanced partitions.

• Hierarchical Agglomerative Clustering (HAC): A bottom-up approach that iteratively merges the most similar adjacent nodes based on bond-type-weighted distances. HAC tends to preserve local chemical structure.

• Motif-aware Coarsening (MC): First identifies known chemical motifs (rings, functional groups) using SMARTS pattern matching, then applies spectral or HAC clustering to the remaining atoms. This ensures that chemically meaningful substructures are preserved as intact units in the hierarchy.

• Motif-aware + Functional Group (MC+FG): Extends motif-aware coarsening with an expanded library of functional group patterns, providing finer-grained control over which substructures are preserved during coarsening.

Tokenization Schemes

The coarsened hierarchies are serialized into token sequences using one of four tokenization schemes, each capturing increasing levels of hierarchical information:

• SENT (Segmented Eulerian Neighborhood Trails): A flat baseline that serializes atom-bond sequences without hierarchy. Serves as the non-hierarchical reference point.

• H-SENT (Hierarchical SENT): Extends SENT by encoding multi-level partition structure. Tokens include partition boundaries and parent-child relationships, enabling the model to learn coarse-to-fine generation.

• HDT (Hierarchical Depth Tokenization): Encodes the full tree structure of the hierarchy using depth-first traversal. Each token carries depth information, making the hierarchical nesting explicit.

• HDTC (Hierarchical Depth Tokenization with Composition): The most expressive scheme, extending HDT with explicit compositional tokens at each internal node. This gives the model complete information about the branching structure of the hierarchy.

Results on COCONUT

We evaluate unconditional generation on the COCONUT dataset of complex natural products (~5K molecules, 30–100 heavy atoms, ≥4 rings). All models use a GPT-2 backbone (12 layers, 768 hidden, 12 heads) trained with identical hyperparameters. We generate 500 molecules per model and compare against the full reference dataset. Bold = best, underline = second best.

HDTC achieves the highest validity (91.8%) and fragment similarity (0.969), while HDT MC is a close second (89.2% validity). The flat-walk baseline SENT drops to 61.8% validity — a 30-point gap — showing that flat tokenization struggles with complex molecular graphs. Motif-community (MC) coarsening consistently outperforms both spectral (SC) and agglomerative (HAC) coarsening across all hierarchical tokenizers, confirming the importance of preserving known chemical motifs during coarsening.

HDT MC and HDTC dominate across all four MMD metrics. HDTC leads on motif rate (70.4% of valid molecules contain at least one recognized motif) and functional-group fidelity. SC and HAC variants lag significantly, reinforcing that the motif-community and compositional approaches excel at preserving complex substructure distributions in large molecules.

Generation Gallery

Generation Demo

The following animation shows MOSAIC’s autoregressive generation process, where the model builds a molecule token-by-token, progressively assembling the hierarchical structure from coarse partitions down to individual atoms and bonds.

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},
}