🤖 machine learning

ConRetroBert: EMA Stabilized Dual Encoders for Template-Based Single-Step Retrosynthesis

ConRetroBert is a dual-encoder framework that enhances template-based single-step retrosynthesis by reframing it as a dense retrieval and ranking task, utilizing contrastive pretraining and EMA-stabilized adaptation to achieve state-of-the-art accuracy on USPTO benchmarks.

Original authors: Mohammad Jahid Ibna Basher, Ali Khodabandeh Yalabadi, Ivan Garibay, Ozlem Ozmen Garibay

Published 2026-05-14

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Mohammad Jahid Ibna Basher, Ali Khodabandeh Yalabadi, Ivan Garibay, Ozlem Ozmen Garibay

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 a master chef trying to recreate a complex, delicious dish (the Product) just by looking at the final plate. Your goal is to figure out exactly which ingredients (Reactants) you started with and what cooking steps (Templates) you used to get there. This is the job of retrosynthesis in chemistry: working backward from a finished molecule to its building blocks.

For a long time, computers trying to do this had two main approaches:

The "Free-Style" Chef: Guesses the ingredients from scratch without a recipe book. This is flexible but hard to explain why a guess was made.
The "Recipe Book" Chef: Looks up a specific, pre-written rule (a Template) that says, "If you have this shape, cut it here and add that." This is very clear and explainable, but the paper argues these chefs have been getting worse scores because their "recipe books" are too huge and messy to search through efficiently.

ConRetroBert is a new system designed to make the "Recipe Book" chef just as good as the free-style one, while keeping the clear, explainable rules.

Here is how it works, broken down into simple steps:

1. The Problem: The "Long Tail" Library

Imagine a library with millions of recipe cards. Most of the time, you need one of the top 100 popular recipes (like "Chocolate Cake"). But sometimes, you need a very rare, specific recipe (like "Fermented Moss Soup").
Old computer systems tried to pick the right recipe by looking at every single card in the library at once and asking, "Is this the one?"

The Issue: The computer gets overwhelmed by the popular cards and ignores the rare ones. It also wastes time comparing your dish to recipes that are chemically impossible.

2. The Solution: A Two-Stage Search Engine

ConRetroBert changes the game. Instead of asking the computer to pick from the whole library, it uses a two-step process, like a smart librarian.

Stage 1: The "Vibe Check" (Contrastive Pretraining)
First, the system learns to understand the "vibe" of a dish and the "vibe" of a recipe.

It takes a picture of a dish and a picture of a recipe card and learns to match them up.
It creates a shared mental space where similar dishes and similar recipes are standing close together.
The Result: When you give it a new dish, it can quickly find a small pile of "likely" recipe cards that match the vibe, ignoring the millions that don't.

Stage 2: The "Taste Test" (Listwise Ranking)
Now, the system has a small pile of candidates (say, 64 recipes). It doesn't just pick the first one; it ranks them carefully.

It looks at the pile and asks, "Which of these 64 is the best fit?"
Crucially, it learns to distinguish between recipes that look similar but are slightly wrong (called Hard Negatives). It's like a judge tasting two very similar soups and deciding which one is actually the right one.
This step is where the biggest improvement happens. It stops the computer from just guessing the most popular recipe and forces it to find the correct one for the specific dish.

3. The Secret Sauce: The "Slow-Moving Shadow" (EMA)

Here is the tricky part. As the computer learns, it wants to update its understanding of what a "recipe card" looks like. But if it changes its mind too fast, the "pile of candidates" it found in Stage 1 becomes outdated and useless. It's like a librarian who keeps moving the books around while you are still trying to find them.

ConRetroBert solves this with a Slow-Moving Shadow (EMA):

Imagine the computer has a Live Librarian who is very active, learning fast, and updating the rankings.
But the Book Shelves (the database of recipes) are managed by a Shadow Librarian who moves very slowly.
The Live Librarian does the hard work of ranking, but the Shadow Librarian only updates the shelves once a day (or once per "epoch").
Why this matters: This keeps the search stable. The computer can learn and adapt without breaking the search engine it's using to find the answers.

The Results

The paper tested this on a standard chemistry dataset (USPTO-50k).

The Old Way: Got about 50% of the answers right.
ConRetroBert Stage 1: Improved to about 50.5% (just the "vibe check").
ConRetroBert Stage 2: Jumped to 61.3% just by adding the "taste test" ranking.
With the Shadow Librarian: Reached 62.4%.

Even more impressively, this method is especially good at finding rare recipes (the "long tail"). While other methods struggle with uncommon chemical reactions, ConRetroBert's search-and-rank approach handles them much better.

Why This Matters

The paper argues that you don't have to choose between accuracy and explainability.

Free-style models are accurate but hard to trust because you don't know why they guessed those ingredients.
ConRetroBert is accurate and transparent. Because it picks a specific, explicit rule (a template), a human chemist can look at the output and say, "Ah, I see. It chose this rule because it matches this pattern."

In short, ConRetroBert proves that if you build a smart search engine and a careful ranking system, you can make the "Recipe Book" approach just as powerful as the "Free-Style" approach, giving scientists the best of both worlds: high accuracy and clear, inspectable logic.

Technical Summary of ConRetroBert

Problem Formulation
Template-based single-step retrosynthesis aims to predict reactant sets for a target product by selecting and applying an explicit reaction template (encoded as a SMARTS transformation). While this approach offers inherent interpretability and traceability—crucial for synthesis planning—template-based methods have historically been viewed as less competitive than template-free models. The authors argue this limitation stems not from the use of templates themselves, but from the dominant learning formulation: global classification over a long-tailed library of reaction rules. This formulation creates three issues: frequent templates dominate the loss, the softmax denominator treats unobserved but chemically valid templates as competing negative labels (label incompleteness), and the output layer scales poorly with large template libraries.

Methodology: ConRetroBert
The paper presents ConRetroBert, a dual-encoder framework that reframes template-based retrosynthesis as a dense retrieval and ranking problem rather than global classification. The architecture employs two independent Transformer encoders (one for product SMILES, one for template SMARTS) that map inputs into a shared embedding space. The training process occurs in two stages:

Stage 1: Contrastive Pretraining
The model learns a joint product-template representation space using symmetric contrastive pretraining (InfoNCE loss). For a minibatch of paired examples, the product encoder aligns with its corresponding template, while other examples in the batch serve as negatives. This establishes a searchable embedding space where templates can be retrieved via nearest neighbor search.
Stage 2: Multi-Positive Listwise Ranking
This stage refines the retrieval space into a template proposal policy. Instead of classifying against the full vocabulary, the model ranks a candidate set $C(x)$ for each product. This set includes:
- Observed Positives: All templates recorded with the product in the training data.
- Hard Negatives: Chemically similar templates retrieved from a FAISS index.
- In-Batch Negatives: Templates associated with other products in the batch.
The model optimizes a multi-positive listwise negative log-likelihood objective, encouraging the separation of observed positives from chemically similar hard negatives. This formulation better matches the inference-time decision problem, where the model must choose among a small set of retrieved candidates rather than the entire library.

EMA-Stabilized Template Adaptation
A key technical challenge addressed is the instability of hard-negative mining when the template encoder is updated. If the template encoder changes too rapidly, the retrieval index (used to mine hard negatives) becomes stale. ConRetroBert resolves this using an Exponential Moving Average (EMA) mechanism:

A live template encoder is updated via gradients from the ranking loss.
A slow-moving EMA shadow encoder (updated with $\alpha = 0.999$ ) is used to rebuild the retrieval index and mine hard negatives at the epoch level.
This allows the template representations to adapt to the ranking objective without destabilizing the candidate construction process.

Key Contributions

Retrieval-Based Formulation: The authors reframe template-based retrosynthesis as dense product-template retrieval followed by candidate-set listwise ranking, replacing global classification with an objective that aligns with inference-time decisions.
EMA-Stabilized Training: They introduce a trainable template encoder stabilized by an EMA shadow, enabling controlled adaptation of template representations while maintaining a stable retrieval index for hard-negative mining.
Rigorous Evaluation Pipeline: The paper provides a comprehensive evaluation protocol including reactant-set ranking, template application, deduplication, and leakage-controlled transfer from USPTO-Full to USPTO-50k.
Empirical Validation: Through local and long-tail analyses, the authors demonstrate that candidate-set ranking substantially strengthens template-based retrosynthesis while retaining explicit rule traceability.

Results
Experiments were conducted on the USPTO-50k benchmark (local regime) and via transfer learning from a leakage-controlled USPTO-Full checkpoint.

Local USPTO-50k Performance:
- Moving from Stage 1 (contrastive retrieval) to Stage 2 (frozen template encoder ranking) improved top-1 reaction accuracy from 50.5% to 61.4%.
- Adding EMA-stabilized template adaptation (Stage 2 EMA Mid) further improved top-1 accuracy to 62.4%.
- The model achieved top-3, top-5, and top-10 accuracies of 81.6%, 85.3%, and 87.8%, respectively.
- The performance gain is primarily attributed to the Stage 2 candidate-set ranking formulation, which better separates positives from hard negatives than contrastive alignment alone.
Transfer Learning:
- Fine-tuning from a leakage-controlled USPTO-Full checkpoint reached 75.4% top-1 accuracy on USPTO-50k, demonstrating the scalability of the approach with larger pretraining data.
Long-Tail and Reliability:
- Retrieval-based prediction showed particular strength in the long tail of rare templates, outperforming classification-based scoring on rare classes.
- Analysis revealed that 42.8% of correct top-1 reactant predictions were "secondary successes," generated by a template different from the recorded ground-truth label but chemically equivalent after application and deduplication. This highlights that the model learns to retrieve chemically valid alternatives, not just the specific recorded label.

Significance and Claims
The paper argues that template-based retrosynthesis need not sacrifice predictive performance for interpretability. By shifting from global classification to retrieval-based ranking with hard-negative mining and EMA stabilization, ConRetroBert challenges the assumption that high accuracy requires abandoning explicit reaction templates. The authors claim that their method produces predictions that are not only competitive with state-of-the-art template-free and semi-template models but also remain chemically inspectable, traceable to specific transformation rules, and compatible with downstream synthesis planning constraints. The work suggests that the "weakness" of template-based methods is a formulation issue solvable through modern retrieval learning techniques, rather than an inherent limitation of using explicit rules.