Uni-Mol3: Multi-Molecular Reaction Modeling

• Uni-Mol3 is a multi-molecular foundation model that integrates 1D, 2D, and 3D molecular data into discrete, context-sensitive tokens for advanced reaction-level reasoning.
• It employs a hierarchical pipeline combining multi-scale tokenization, molecular pre-training, and reaction pre-training to excel in forward prediction, retrosynthesis, condition generation, and yield estimation.
• The model demonstrates significant improvements over baselines through enhanced chemical plausibility, transfer capability, and prompt-aware performance across diverse organic reaction tasks.

Searching arXiv for the Uni-Mol3 paper and closely related Uni-Mol-family works to ground the article. Uni-Mol3 is a multi-molecular foundation model for organic reaction modeling that extends the Uni-Mol line from single-molecule 3D representation learning to reaction-level reasoning over sets of molecules. Its central design is a hierarchical pipeline: a multi-scale molecular tokenizer converts 1D, 2D, and 3D molecular information into discrete 3D-aware tokens; molecular pre-training learns single-molecule grammatical regularities; reaction pre-training learns multi-molecular reaction principles; and prompt-aware fine-tuning adapts a shared encoder–decoder backbone to forward prediction, retrosynthesis, condition generation, and yield prediction. In the reported evaluation, Uni-Mol3 is trained with about 19M molecules for tokenizer and molecular pre-training, 11.97M reactions from Pistachio for reaction pre-training, and is assessed on 10 datasets spanning 4 downstream tasks (Wu et al., 30 Jul 2025).

1. Position in the Uni-Mol lineage

Uni-Mol3 belongs to the broader Uni-Mol ecosystem, whose earlier stages established the technical premises for 3D-first molecular modeling. "Unified 2D and 3D Pre-Training of Molecular Representations" formulated a single backbone that jointly models molecular graphs and conformations through masked reconstruction and bidirectional 2D↔3D generation, and reported state-of-the-art results on 10 of 11 downstream tasks with an average improvement of 8.3% on 2D-only tasks (Zhu et al., 2022). "Highly Accurate Quantum Chemical Property Prediction with Uni-Mol+" then specialized the family toward quantum chemistry by introducing a two-track Transformer and iterative conformation refinement toward DFT equilibrium geometries (Lu et al., 2023).

Within that lineage, Uni-Mol3 is not a docking model and should not be conflated with Uni-Mol Docking V2. The docking paper explicitly states that there is no formal “Uni-Mol3” branding there; Uni-Mol Docking V2 is the docking-specialized branch built on the original Uni-Mol backbone, whereas Uni-Mol3 is a separate framework introduced for multi-molecular organic reaction modeling (Alcaide et al., 2024). This distinction matters because the modeling unit changes: docking predicts ligand poses in a fixed pocket, while Uni-Mol3 treats reactions as structured sequences of multiple molecules with different roles.

A plausible implication is that Uni-Mol3 marks a shift in the Uni-Mol program from universal single-molecule representation learning toward a foundation-model view of chemical systems in which molecular structure, intermolecular context, and task prompts are handled within one reaction-level architecture.

2. Core architecture and 3D-aware molecular language

At the front end of Uni-Mol3 is Mol-Tokenizer, a multi-scale tokenizer that maps molecular structure into discrete tokens. For a molecule $M=(x,e,r)$, the inputs are 1D atom features $x \in \mathbb{R}^{N \times d_x}$, 2D bond features $e \in \mathbb{R}^{N\times N \times d_e}$, and 3D coordinates $r \in \mathbb{R}^{N\times 3}$. Atom-level inputs include atom token, atom degree, and atom type, combined as

$$x_{\text {single} }^i=\operatorname{Embedding}\left(x_{\text {token} }^i\right)+\operatorname{Embedding}\left(x_{\text {degree} }^i\right)+\operatorname{Embedding}\left(x_{\text {type} }^i\right).$$

Pairwise inputs combine bond type, shortest-path distance, and a Gaussian-kernel encoding of 3D Euclidean distance: $x_{\text {pair}^{i, j}=\operatorname{Embedding}\left(e^{i, j}\right)+\operatorname{Embedding}\left(x_{\text {SPD} }^{i, j}\right)+x_{\text {dis} }^{i, j}.$ The tokenizer encoder adopts the two-track Transformer backbone of Uni-Mol2, updating single and pair representations jointly (Wu et al., 30 Jul 2025).

Discrete tokenization is implemented with Finite Scalar Quantization. For each atom embedding $h_{\text{single}^i$, Uni-Mol3 applies

$f(h_{\text{single}^{i}) = \lfloor L / 2\rfloor \tanh (h_{\text{single}^{i}), \qquad s_i = \text{round}\left(f(h_{\text{single}^{i})\right),$

and uses a codebook of size $L^d$ to map the quantized vector to an integer code $z_i$. These integers are the 3D tokens. The resulting vocabulary is intended to encode local 3D chemical environments rather than substrings of SMILES. The paper’s token analyses report context-sensitive specialization, for example distinct tokens for oxygen in double bonds versus oxygen with two single bonds, and distinct sulfur environments such as sulfones versus sulfur in aromatic rings (Wu et al., 30 Jul 2025).

The reaction model built on top of these tokens uses a T5-style encoder–decoder with 8 encoder layers, 8 decoder layers, hidden dimension 768, and 8 attention heads. Multi-molecular interactions are handled by attention over concatenated token sequences from multiple molecules rather than by an explicit reaction graph. This means that reactants, products, and condition molecules are all represented within one token stream, and interaction patterns are learned through standard Transformer self-attention.

3. Progressive pre-training from molecules to reactions

Uni-Mol3 uses a two-stage pre-training strategy that moves from single molecules to multi-molecular systems. This is one of its defining features: molecular pre-training is designed to learn “molecular grammars,” while reaction pre-training is designed to capture “fundamental reaction principles” (Wu et al., 30 Jul 2025).

In molecular pre-training, a molecule is first tokenized as

$x \in \mathbb{R}^{N \times d_x}$0

Then 15% of the tokens are replaced by a mask token $x \in \mathbb{R}^{N \times d_x}$1, yielding $x \in \mathbb{R}^{N \times d_x}$2. The encoder produces a conditional embedding $x \in \mathbb{R}^{N \times d_x}$3, and the decoder autoregressively generates the molecule’s SMILES: $x \in \mathbb{R}^{N \times d_x}$4 This stage links 3D-aware token sequences to valid molecular string realizations.

In reaction pre-training, each molecule $x \in \mathbb{R}^{N \times d_x}$5 in a reaction is tokenized as $x \in \mathbb{R}^{N \times d_x}$6, and the reaction-level input is the concatenation

$x \in \mathbb{R}^{N \times d_x}$7

Here Uni-Mol3 masks entire molecules rather than individual atoms, again at a 15% masking ratio. The decoder then generates the concatenated reaction SMILES autoregressively: $x \in \mathbb{R}^{N \times d_x}$8 The stated intent is to force the model to learn co-occurrence regularities, implicit reaction templates, and condition–reactant–product patterns.

Before these sequence-level stages, Mol-Tokenizer itself is trained with atom-type reconstruction and coordinate-denoising objectives: $x \in \mathbb{R}^{N \times d_x}$9 where

$e \in \mathbb{R}^{N\times N \times d_e}$0

This training regime is what makes the tokenizer explicitly 3D-aware rather than a mere symbol converter.

Stage	Input/operation	Objective
Mol-Tokenizer training	Molecular 1D/2D/3D structure	$e \in \mathbb{R}^{N\times N \times d_e}$1
Molecular pre-training	Masked 3D token sequence	Autoregressive SMILES generation
Reaction pre-training	Concatenated multi-molecule token sequence with molecule masking	Autoregressive reaction SMILES generation
Fine-tuning	Prompted token sequence	Task-specific generation or regression

This hierarchy is a substantive departure from earlier Uni-Mol-family models focused on molecules, pockets, or quantum properties. It operationalizes a progression from local 3D chemical context to reaction-level semantics.

4. Prompt-aware downstream tasks and empirical performance

Uni-Mol3 unifies four downstream task families through task-specific prompt tokens prepended to the input sequence. The prompts reported are \<forward-sep\>, \<forward-mixed\>, \<reverse\>, \<condition\>, and \<yield\> (Wu et al., 30 Jul 2025). Generative tasks use the shared encoder–decoder to produce target SMILES, while yield prediction uses a regression head on encoder outputs.

The downstream benchmarks cover 10 datasets. Product prediction is evaluated on USPTO-MIT, SMol-Reactions-FP, and Pistachio-FP. Retrosynthesis is evaluated on USPTO-50k, SMol-Reactions-RS, and Pistachio-RS. Condition generation is evaluated on USPTO-500-MT, USPTO-Condition, and Pistachio-CG. Yield prediction is evaluated on the Buchwald–Hartwig C–N coupling dataset, with four out-of-sample additive test sets (Wu et al., 30 Jul 2025).

The reported generative metrics are Top-1 accuracy, Levenshtein distance, molecular fingerprint Tanimoto coefficient, and invalidity rate. On USPTO-MIT mixed forward prediction, Uni-Mol3 improves over T5Chem from Top-1 88.9%, LEV 0.527, MFP-TC 0.974, invalidity 0.20% to Top-1 89.6%, LEV 0.485, MFP-TC 0.979, invalidity 0.15%. On Pistachio-FP mixed forward prediction, it improves over T5Chem from Top-1 90.9%, LEV 0.560, MFP-TC 0.982, invalidity 0.15% to Top-1 91.7%, LEV 0.462, MFP-TC 0.985, invalidity 0.09% (Wu et al., 30 Jul 2025).

For retrosynthesis, the scaffold-split SMol-Reactions-RS benchmark is particularly informative because it is designed to reduce leakage. There Uni-Mol3 reports Top-1 29.1%, LEV 9.933, MFP-TC 0.786, compared with T5Chem at Top-1 28.0%, LEV 10.593, MFP-TC 0.758 and PRESTO at Top-1 27.7%, LEV 11.229, MFP-TC 0.745. For condition generation on Pistachio-CG, Uni-Mol3 reports Top-1 44.4%, LEV 6.482, MFP-TC 0.827, invalidity 0.03%, compared with T5Chem at Top-1 43.3%, LEV 6.705, MFP-TC 0.823, invalidity 0.06% (Wu et al., 30 Jul 2025).

Yield prediction is evaluated by MAE, RMSE, and $e \in \mathbb{R}^{N\times N \times d_e}$2. On Buchwald–Hartwig Test1, Uni-Mol3 reports MAE 5.867, RMSE 9.680, $e \in \mathbb{R}^{N\times N \times d_e}$3, improving over T5Chem at MAE 8.145, RMSE 11.837, $e \in \mathbb{R}^{N\times N \times d_e}$4. The paper states that Uni-Mol3 dominates on 11 of 12 metrics across the four test sets (Wu et al., 30 Jul 2025).

A plausible implication is that Uni-Mol3’s main empirical advantage is not confined to a single reaction task. The same prompt-conditioned backbone achieves consistent gains across synthesis directionality, auxiliary condition inference, and quantitative outcome prediction, which is the operational meaning of its “foundation model” designation in this context.

5. Ablations, transfer behavior, and qualitative characteristics

The main ablations remove three components: Uni-Tokenizer, molecular pre-training, and reaction pre-training. The reported conclusion is that molecular pre-training is the most critical component: removing it substantially degrades performance across tasks, suggesting that single-molecule grammar learning is a prerequisite for reliable reaction-level generation (Wu et al., 30 Jul 2025). Uni-Tokenizer also yields notable gains, especially for condition generation, indicating that 3D tokenization contributes information not captured by plain SMILES tokens. Reaction pre-training provides smaller but still consistent improvements.

The paper also compares single-task and multi-task fine-tuning. For T5Chem, multi-task fine-tuning deteriorates performance on all tasks, although validity improves slightly. For Uni-Mol3, multi-task fine-tuning causes only a slight drop in product prediction Top-1 from 93.0% to 92.3%, while substantially improving retrosynthesis Top-1 from 76.9% to 80.4% and condition generation Top-1 from 44.4% to 47.6%, with lower invalidity across tasks (Wu et al., 30 Jul 2025). This suggests that the pre-training stack produces more shareable representations than a purely SMILES-based baseline.

Cross-dataset transfer experiments further support that reading. Across six train→test combinations spanning three tasks, Uni-Mol3 is reported to outperform Molformer, Chemformer, and T5Chem in all six scenarios, with particularly strong gains in condition generation and in transfer from large Pistachio datasets to smaller USPTO or SMol datasets (Wu et al., 30 Jul 2025). That result is consistent with the intended role of hierarchical pre-training: reaction knowledge learned at large scale is expected to transfer more effectively when grounded in a 3D-aware tokenization scheme.

Qualitative analyses in the paper emphasize that Uni-Mol3 more often preserves chemically plausible structure when it is wrong. In product prediction, baselines are described as omitting atoms, misplacing reactive sites, or misordering fragments, whereas Uni-Mol3 more often maintains chemically plausible structures. In retrosynthesis, the model sometimes proposes alternative retrosynthetic routes that differ from the annotated ground truth but remain chemically plausible. This suggests a learned transformation space broader than one-to-one memorization of the training corpus (Wu et al., 30 Jul 2025).

6. Interpretation, misconceptions, and limitations

A common misconception is to treat “Uni-Mol3” as a generic label for any recent Uni-Mol-family model. The record in the cited papers does not support that usage. Uni-Mol Docking V2 is a docking-specific model built on the original Uni-Mol molecular and pocket encoders and is not formally branded as Uni-Mol3 (Alcaide et al., 2024). Uni-Mol3, by contrast, is explicitly a multi-molecular foundation model for reaction modeling introduced in 2025 (Wu et al., 30 Jul 2025).

Another misconception is that Uni-Mol3 is simply “Uni-Mol with a T5 decoder.” That characterization omits two central elements: the FSQ-based 3D molecular language and the progressive molecular-to-reaction pre-training pipeline. Those components are exactly what distinguish it from earlier single-molecule Uni-Mol derivatives and from SMILES-only encoder–decoder systems.

Its limitations are also specific. The 3D conformers are generated by ETKDG + MMFF using RDKit; the model does not represent conformational ensembles or transition states. Reaction pre-training relies heavily on Pistachio, which may bias the learned reaction distribution toward patent chemistry. Conditions are represented as discrete molecules, but numerical variables such as temperature, time, and pressure are not explicitly modeled. The framework is focused on organic reactions; performance on inorganic, organometallic, or solid-state chemistry is left open. Full pre-training requires substantial compute, reported as 8×H100 GPUs with 100k steps for the tokenizer, 1,000,000 steps for molecular pre-training, and 1,500,000 steps for reaction pre-training (Wu et al., 30 Jul 2025).

In the broader Uni-Mol ecosystem, Uni-Mol3 can therefore be read as one branch of a diversification pattern. Uni-Mol+ showed that the family could be extended toward geometry-refined QC prediction (Lu et al., 2023); Uni-Mol Docking V2 showed that the same ecosystem could support physically grounded protein–ligand docking (Alcaide et al., 2024). Uni-Mol3 extends the paradigm in a different direction: from single-molecule representation learning to discrete, prompt-conditioned modeling of multi-molecular reaction systems. This suggests that the Uni-Mol program is evolving not by a single monolithic architecture, but by retaining a 3D-first inductive bias while adapting tokenization, pre-training, and task heads to distinct chemical regimes.