💻 computer science

BiMol-Diff: A Unified Diffusion Framework for Molecular Generation and Captioning

BiMol-Diff is a unified diffusion framework that employs a token-aware noise schedule to overcome the limitations of standard diffusion and autoregressive models, achieving state-of-the-art performance in both text-conditioned molecular generation and molecule captioning by preserving structurally informative substructures.

Original authors: Aditya Hemant Shahane, Anuj Kumar Sirohi, Devansh Arora, Nitin Kumar, Prathosh A P, Sandeep Kumar

Published 2026-04-28

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Aditya Hemant Shahane, Anuj Kumar Sirohi, Devansh Arora, Nitin Kumar, Prathosh A P, Sandeep Kumar

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 you are trying to teach a robot two very different skills at once:

Reading a chemical blueprint and describing it in plain English (like a caption for a photo).
Reading a text description (like "a molecule that stops cancer cells") and drawing the exact chemical blueprint.

For a long time, scientists tried to teach these skills using "Autoregressive" models. Think of these like a person writing a sentence one word at a time, left to right. If they make a mistake in the first word, the whole sentence can go off the rails, and they can't easily go back to fix it. Also, these models treat every part of the chemical structure the same way, even though some parts (like the "skeleton" of the molecule) are much harder to get right than others.

The paper introduces BiMol-Diff, a new way to teach the robot using a technique called Diffusion.

The Core Idea: The "Dusty Sketch" Analogy

Imagine you have a perfect drawing of a molecule.

Standard Diffusion is like taking that drawing and throwing a bucket of dust over it. It throws the dust evenly everywhere. The delicate, important lines get covered just as much as the empty white space. When the robot tries to clean the dust off to see the drawing again, it struggles because the important lines are buried under too much dust.
BiMol-Diff is smarter. It knows which parts of the drawing are the "critical lines" (the chemical structure) and which parts are just "background noise."
- It throws less dust on the critical lines so they stay visible.
- It throws more dust on the easy parts.
- When the robot cleans the drawing, it can easily see the important structure because it was protected, allowing it to reconstruct the molecule perfectly.

How It Works (The Two-Way Street)

The authors built a single "brain" that can do both directions of the task:

From Molecule to Text (Captioning): The robot looks at a chemical structure, sees the "protected" parts, and writes a clear description.
From Text to Molecule (Generation): The robot reads a description, figures out the "hard" chemical parts it needs to protect, and draws the molecule.

To do this, they didn't just use the standard chemical text format (SMILES). They broke the molecule down into a specific format: Atom-Bond-Atom triplets. Imagine taking a molecule apart and listing every connection like a recipe: "Carbon connected to Oxygen, Oxygen connected to Hydrogen." This helps the robot understand the shape better than just a string of letters.

The Secret Sauce: "Token-Aware Noise"

The biggest innovation is how they handle the "dust" (noise).

Old Way: "I will mess up every part of the molecule equally."
BiMol-Diff Way: "I will look at each part of the molecule. If a part is hard to guess (like a complex ring structure), I will keep it very clean. If a part is easy to guess, I can mess it up more."

This is like a teacher grading a test. If a student struggles with a specific math concept, the teacher gives them extra practice on that specific part, rather than giving them the same amount of practice on everything.

The Results: Did It Work?

The team tested this on two big datasets (collections of molecules and their descriptions).

For Describing Molecules: BiMol-Diff wrote better captions than any previous model. It was more accurate and used better vocabulary.
For Drawing Molecules: This is where it really shined. When asked to draw a molecule from a description, it got the structure right 15.4% more often than the best previous methods. It also produced molecules that were chemically valid (they actually exist in the real world) and looked very similar to the target.

The Trade-off (The Catch)

The paper is honest about one downside: Speed.
Because the robot has to "clean" the drawing step-by-step (iteratively), it takes longer than the "write one word at a time" models.

If you need a result instantly (low latency), the old models are still faster.
If you need high quality and accuracy (like in drug discovery where getting the shape wrong is bad), BiMol-Diff is the winner.

Summary

BiMol-Diff is a unified system that treats chemical structures and language as two sides of the same coin. By being "smart" about how it corrupts and cleans data—protecting the hard, important chemical parts while letting the easy parts get messy—it creates a much more reliable bridge between text and molecular design.

1. Problem Statement

The paper addresses the challenge of bridging molecular structures (graphs of atoms and bonds) and natural language descriptions. This bidirectional capability is crucial for controllable molecular design:

Molecule Generation (Text $\to$ Molecule): Converting natural language descriptions (e.g., "a kinase inhibitor with good solubility") into valid molecular structures.
Molecule Captioning (Molecule $\to$ Text): Generating natural language descriptions that explain a molecule's structure and properties.

Limitations of Existing Approaches:

Autoregressive (AR) Models: While popular (e.g., GPT, T5), they generate tokens sequentially (left-to-right). This makes them prone to error propagation, struggle with long-range dependencies (e.g., ring closures in SMILES), and fail to enforce global structural validity easily.
Standard Diffusion Models: While non-autoregressive and capable of holistic updates, standard diffusion applies uniform noise across all tokens. This is suboptimal for chemistry because it corrupts chemically critical tokens (e.g., ring indices, branching markers) just as aggressively as redundant tokens, hindering the recovery of valid structures.

2. Methodology: BiMol-Diff

BiMol-Diff is a unified diffusion framework that handles both generation and captioning tasks under a single conditional denoising formulation.

A. Representation Pipeline

Instead of treating molecules purely as character-level strings, the framework uses a structured approach:

Canonical SMILES: Serves as the standard interface.
Atoms-in-SMILES (AIS): Converts SMILES into a chemistry-aware token sequence.
Serialized Graph Triplets ( $\tilde{G}$ ): For graph-conditioned modeling, the molecule is represented as a sequence of atom-bond-atom triplets using special tokens ([HEAD], [REL], [TAIL], [SEP]).
- Example: A bond between C1 and C2 is serialized as [HEAD] C1 [REL] SINGLE [TAIL] C2.
- This allows the model to view the molecule as a serialized Knowledge Graph (KG), reducing dependency on specific SMILES syntax rules.

B. Core Innovation: Token-Aware Noise Schedule

The central contribution is a Token-Aware Noise Schedule that adapts the corruption process based on the difficulty of recovering specific tokens.

Rationale: Chemically salient tokens (high information density) are harder to recover than syntactic or redundant tokens. Uniform noise destroys these critical tokens too early.
Mechanism:
1. Difficulty Estimation: During training, the model estimates the per-token reconstruction loss ( $\ell^i_t$ ) at each timestep. High loss indicates a "hard" token.
2. Adaptive Scheduling: A mapping function (specifically a piecewise-linear map) translates the difficulty profile into a custom noise schedule ( $\bar{\alpha}^i_{t,new}$ ) for each token $i$ .
3. Result: Tokens that are harder to denoise receive a more conservative corruption schedule (less noise added early on), preserving their integrity during the forward process and improving recoverability during the reverse process.

C. Unified Training Objective

The framework uses a Variational Lower Bound (VLB) derived objective with a $z_0$ -prediction reparameterization. The model is trained to directly predict the clean data ( $z_0$ ) rather than just the noise. The loss function combines:

Denoising: Predicting clean data across timesteps.
Consistency: Ensuring the first denoising step aligns with the true data embedding.
Rounding: A specific term to handle the conversion from continuous latent space back to discrete tokens (SMILES or text).

3. Key Contributions

Token-Aware Noising Strategy: A novel mechanism that assigns position-dependent noise levels based on token recovery difficulty, preserving chemically critical substructures.
Unified Bidirectional Framework: A single model architecture capable of performing both Text-to-Molecule (S2G) and Molecule-to-Text (G2S) tasks using the same token-aware mechanism.
Serialized Graph Representation: The use of atom-bond-atom triplets allows for better generalization and reduced reliance on rigid SMILES syntax compared to pure character-level models.

4. Experimental Results

The model was evaluated on two datasets: ChEBI-20 (Molecule Generation) and M3-20M (Molecule Captioning).

Molecule Captioning (G2S) on M3-20M

Performance: BiMol-Diff achieved state-of-the-art results, outperforming the best autoregressive baseline (Text+Chem T5) by 4.6% in BLEU and 4.7% in ChrF++.
Metrics: It also achieved the best BERTScore (0.843) and MAUVE (0.925) among all baselines.
Efficiency: Despite being a diffusion model, it was 3.6x faster than the DiffuSeq baseline while maintaining higher quality.

Molecule Generation (S2G) on ChEBI-20

Reconstruction: BiMol-Diff achieved an Exact Match accuracy of 0.262, a 15.4% relative improvement over the best autoregressive and diffusion baselines (both at 0.227).
Structural Fidelity: Significant gains in fingerprint similarity (MACCS: 0.894, Morgan: 0.762), indicating better preservation of chemical substructures.
Validity: Maintained high chemical validity (0.901), ensuring generated molecules are chemically plausible.

Ablation Studies

Noise Schedule: Replacing the uniform "sqrt" schedule with the token-aware schedule yielded substantial gains (e.g., BLEU increased from 0.495 to 0.567).
Tokenization: The Atoms-in-SMILES (AIS) tokenizer outperformed regex-based and atom-level tokenizers, confirming the benefit of chemistry-aware segmentation.

5. Significance and Conclusion

Scientific Impact: BiMol-Diff demonstrates that token-aware denoising is critical for high-fidelity molecular language modeling. By explicitly protecting chemically important tokens during the diffusion process, the model overcomes the structural fragility often seen in generative AI for chemistry.
Unified Approach: It proves that a single diffusion framework can effectively handle both generation and captioning, promoting consistency in iterative design loops.
Trade-offs: While diffusion models generally incur higher inference costs than autoregressive models, BiMol-Diff shows that with optimized sampling (e.g., 1000–2000 steps), it can outperform AR models in quality while remaining competitive in speed. However, for ultra-low latency applications, AR models remain superior.

In summary, BiMol-Diff represents a significant step forward in controllable molecular design, offering a robust, unified, and structurally aware alternative to traditional autoregressive methods.