🤖 machine learning

How Molecular Generative Models Organize Molecular Identity

This paper reveals that molecular generative models organize chemical identities into fixed, piecewise-constant regions with recurring boundaries, demonstrating that their internal structure must be empirically characterized rather than assumed before the latent space can be reliably used for chemical navigation.

Original authors: Raul Ortega-Ochoa, Tejs Vegge, Jens S. Bakander, Luis Mantilla Calderon, Alan Aspuru-Guzik, Tonio Buonassisi

Published 2026-08-10

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Raul Ortega-Ochoa, Tejs Vegge, Jens S. Bakander, Luis Mantilla Calderon, Alan Aspuru-Guzik, Tonio Buonassisi

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine a world where computers don't just recognize pictures of cats or translate languages, but actually dream up new ones. In the field of artificial intelligence, there are special programs called "generative models" that act like digital alchemists. Instead of mixing chemicals in a flask, they mix numbers in a vast, invisible landscape to create new molecules—the tiny building blocks of medicines, plastics, and materials. To make these models work, scientists give them a map: a continuous coordinate system where every point represents a potential molecule. The hope is that if you walk a short distance on this map, you'll find a molecule that is chemically similar to the one you started with, like moving from a red car to a slightly different shade of red car. This idea of "navigability" is crucial; if the map is smooth and logical, scientists can use it to find new drugs by simply walking toward the right direction. But what if the map is actually a patchwork quilt of sharp, invisible cliffs? What if taking a tiny step forward doesn't give you a similar molecule, but suddenly drops you into a completely different chemical universe? This is the mystery scientists are trying to solve: how do these AI models actually organize the molecules they create inside their digital brains?

A team of researchers decided to peek behind the curtain of three different molecular AI models to see how they really organize their creations. They treated the AI's internal "map" not as a smooth, flowing river, but as a territory divided into distinct neighborhoods. By freezing the random choices the AI makes and tracing exactly which molecules come out of which coordinates, they discovered that these models don't create a smooth gradient of chemistry. Instead, they build a landscape of "piecewise-constant" regions. Think of it like a giant, multi-colored floor made of tiles. If you stand on one tile, you get one specific molecule. If you take a tiny step, you might stay on the same tile and get the exact same molecule. But if you cross the invisible line between tiles, you suddenly get a completely different molecule, even if your step was microscopic.

The researchers found that this "tile map" isn't random; it has a structure. In some models, like the autoregressive "MolMiner" and the hierarchical "HierVAE," the tiles are organized like a family tree. You might start in a broad region representing a general type of chemical skeleton (like a bicycle frame). As you move deeper into that region, the boundaries get finer, splitting the area into smaller tiles for specific variations (like a mountain bike vs. a road bike). This happens because the AI builds molecules step-by-step; early decisions create broad categories, and later decisions carve out the details. However, the study also revealed a surprising twist: just because the map looks organized doesn't mean the distance on the map matches the chemical distance. In one model, HierVAE, two points that are very close together on the computer's coordinate grid could actually produce molecules that are chemically very different. The "distance" the computer sees isn't always the "distance" a chemist would feel.

Furthermore, the team watched how these maps change while the AI is learning. They discovered that the AI learns to keep its neighborhoods chemically consistent (staying in the same "chemical family") very quickly. But the number of different specific molecules it can produce within those neighborhoods keeps changing for a long time. It's as if the AI first learns to keep its red cars in the red zone and blue cars in the blue zone, and only later figures out exactly how many shades of red it can make. The study suggests that before we trust these AI maps to help us design new life-saving drugs, we need to check the specific rules of the map we are using. We can't just assume that walking a short distance on the grid will always lead to a chemically similar molecule; sometimes, the map is a patchwork of sharp boundaries that only make sense if you look at them through the right lens.

Technical Summary: How Molecular Generative Models Organize Molecular Identity

Problem Statement
Generative models for matter are typically evaluated based on the quality of their outputs (validity, novelty, diversity) and their latent spaces are often treated as smooth, navigable chemical spaces. However, a critical structural question remains unaddressed: how are discrete molecular identities internally arranged within the continuous coordinate spaces of these models? Specifically, it is unclear where molecular identity is resolved, whether nearby coordinates decode to chemically related molecules, and how this organization emerges during decoding and training. Current evaluations often assume that a continuous latent space implies a chemically navigable space without verifying the internal partition of molecular identities.

Methodology
The authors propose a framework to make molecular identity explicit as an equivalence relation on the output representation space and "pull back" these identity classes through the generative process.

Formalism: The study defines molecular identity via an equivalence relation ( $\sim$ ) on the output space $X$ (e.g., strings, graphs). This creates a quotient space $X/\sim$ representing molecular objects. To handle stochastic generators, the authors introduce a "random tape" $\eta$ , treating the generation process as a deterministic map from $(z, \eta)$ to an output. This induces a partition of the joint space $Z \times E$ by molecular identity.
Equivalence Conventions: Six distinct conventions for molecular equivalence are tested, ranging from fine-grained to coarse:
- Canonical isomeric SMILES (exact structure).
- InChIKey-14 (connectivity, collapsing stereoisomers/tautomers).
- Molecular Formula.
- Heavy-atom Elements.
- Bemis–Murcko Scaffold.
- Murcko-generic (skeleton only).
Models Studied: The framework is applied to three distinct architectures trained on the ZINC dataset:
- MolMiner: A property-conditioned autoregressive transformer.
- HierVAE: A hierarchical graph variational autoencoder with an autoregressive decoder.
- GDSS: A score-based graph diffusion model.
Experimental Probes:
- Fixed-Randomness Sections: Probing 1D and 2D cross-sections of the latent space $Z$ while holding the random tape $\eta$ constant to visualize identity partitions.
- Stochastic Persistence: Varying $\eta$ at fixed $z$ to test if a specific coordinate consistently decodes to the same molecule.
- Neighborhood Analysis: Sampling balls in $Z$ to measure chemical cohesiveness (similarity of molecules within a neighborhood vs. across neighborhoods) using Tanimoto similarity on ECFP and MACCS fingerprints.
- Metric Correlation: Testing whether Euclidean distance or cosine similarity in $Z$ correlates with chemical similarity.
- Training Evolution: Tracking the development of these properties across training checkpoints.

Key Contributions and Results

Piecewise-Constant Identity Regions: Across all three architectures, fixing the random tape reveals that the latent space is partitioned into piecewise-constant regions separated by sharp boundaries. These regions correspond to specific molecular identities.
Dependence on Convention and Architecture:
- MolMiner and HierVAE: Exhibit strong chemical cohesiveness. Nearby coordinates decode to chemically related molecules, and distinct neighborhoods occupy distinct chemical territories. This organization holds under fine-grained conventions (SMILES, InChIKey-14).
- GDSS: Shows little persistence of exact molecular identity under stochastic decoding (99.95% unique SMILES for fixed $z$ ). Clear organization only emerges under coarse conventions (e.g., Murcko scaffold or elemental composition), indicating that its exact-molecule variation is largely driven by decoder randomness rather than coordinate structure.
Metric Limitations: For HierVAE, while the latent space is chemically organized (nearby points decode to similar molecules), Euclidean distance and cosine similarity between coordinates are poor proxies for chemical distance ( $R^2 \approx 0.05$ ). This suggests the latent space is curved and that standard geometric metrics do not reliably predict chemical proximity.
Coarse-to-Fine Boundaries: The partitions exhibit a recurring hierarchical structure where broad molecular territories are subdivided by finer boundaries. In MolMiner, this is traced to the sequential decoding process: early decoding steps define broad lineages (shared partial structures), and later decisions subdivide these regions.
Decoupling of Cohesiveness and Granularity: During training, chemical cohesiveness (the degree to which a neighborhood is chemically localized) stabilizes early. In contrast, the number of distinct molecular identities represented within a neighborhood (granularity) continues to evolve and refine long after cohesiveness has plateaued.

Significance and Claims
The paper argues that a continuous generative coordinate space is insufficient to establish a navigable chemical space. Before treating a latent space as chemically navigable for interpolation or optimization, its internal organization must be characterized relative to:

The specific representation being generated.
The equivalence convention used to define molecular identity.
The stochastic decoding process.
The metric used to compare coordinates.

The authors conclude that internal organization is not a single, inherent property of a generative space but a complex interaction of these factors. The framework provides a direct method to determine whether a generative space is organized and whether it can be meaningfully navigated for specific chemical questions. The authors note that while the study focuses on molecules, the framework applies to any generative model mapping to an output space with a meaningful equivalence relation, such as proteins or crystals.

Technical Summary: How Molecular Generative Models Organize Molecular Identity

More like this