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

TL;DR

ConRetroBert introduces a dual encoder framework with EMA stabilization for improved template-based retrosynthesis, significantly enhancing top-1 accuracy and handling rare templates effectively.

Contribution

It reframes template prediction as dense retrieval and ranking, employing contrastive pretraining and EMA stabilization to outperform existing methods.

Findings

Top-1 accuracy improved from 50.5% to 62.4% on USPTO-50k.

EMA stabilization further boosts accuracy by 0.9%.

Retrieval approach is effective for rare templates and alternative reactant predictions.

Abstract

Template based single step retrosynthesis predicts reactants by selecting and applying an explicit reaction template, making each prediction traceable to a chemical transformation rule. This is useful for synthesis planning, but template based methods are often viewed as less competitive than template free models because template prediction is commonly formulated as global classification over a long tailed rule library. We argue that this weakness is not inherent to templates, but to the learning formulation. We present ConRetroBert, a dual encoder framework that reframes template based retrosynthesis as dense product template retrieval followed by candidate set listwise ranking. Stage 1 uses contrastive pretraining to learn a shared embedding space between products and reaction templates. Stage 2 refines template ranking over mined hard negative candidate sets with a multi positive listwise objective. To enable template side adaptation without destabilizing hard negative mining, ConRetroBert uses a slow moving exponential moving average template encoder for retrieval bank construction while updating the live template encoder through the ranking loss. On the local USPTO-50k benchmark, Stage 2 candidate set ranking improves top-1 reaction accuracy from 50.5% to 61.3%, while EMA stabilized template adaptation further improves it to 62.4%. Fine tuning from a leakage controlled USPTO-Full checkpoint reaches 75.4% top-1 accuracy on USPTO-50k. We also show that retrieval based template prediction is strong in the long tail of rare templates, and that many correct reactant predictions arise from alternative explicit templates rather than only the recorded positive label. Code and data are available at https://github.com/JahidBasher/ConRetroBert.