RxnNano:Training Compact LLMs for Chemical Reaction and Retrosynthesis Prediction via Hierarchical Curriculum Learning

TL;DR

RxnNano is a compact 0.5B-parameter model that significantly improves chemical reaction and retrosynthesis prediction by integrating chemical knowledge through hierarchical curriculum learning and invariant modeling, outperforming larger models.

Contribution

The paper introduces a unified framework with novel objectives and curriculum strategies to embed chemical understanding into LLMs, reducing reliance on scale.

Findings

Outperforms larger models by 23.5% Top-1 accuracy

Achieves state-of-the-art results on reaction prediction benchmarks

Demonstrates robustness without test-time augmentation

Abstract

Chemical reaction prediction is pivotal for accelerating drug discovery and synthesis planning. Despite advances in data-driven models, current approaches are hindered by an overemphasis on parameter and dataset scaling. Some methods coupled with evaluation techniques that bypass fundamental challenges in reaction representation and fail to capture deep chemical intuition like reaction common sense and {topological atom mapping logic}. We argue that the core challenge lies in instilling these knowledge into the models. To this end, we propose a unified framework that prioritizes chemical understanding over scale through three key innovations: (1) a {Latent Chemical Consistency} objective that models reactions as movements on a continuous chemical manifold, ensuring reversible and physically plausible transformations; (2) a {Hierarchical Cognitive Curriculum} that trains the model through progressive stages, from syntax mastery to semantic reasoning, building robust chemical intuition; (3) {Atom-Map Permutation Invariance (AMPI)}, which force the model to learn invariant relational topology and balance multi-task learning. (4)and structured plan-based reasoning to improve the performance of the LLMs. Our compact {0.5B-parameter model}, \textbf{RxnNano} significantly outperforms fine-tuned LLMs ten times larger (>7B) and all the domain baselines, achieving a 23.5\% Top-1 accuracy improvement on rigorous benchmarks without test-time augmentation. https://github.com/rlisml/RxnNano.