Round-trip Reinforcement Learning: Self-Consistent Training for Better Chemical LLMs
Lecheng Kong, Xiyuan Wang, Yixin Chen, Muhan Zhang
TL;DR
The paper tackles the challenge of round-trip consistency (RTC) in chemical LLMs by introducing Round-Trip Reinforcement Learning (RTRL), which uses a backward model's success at reconstructing inputs as a reward signal to train the forward model. The framework supports an iterative self-improvement loop where forward and backward mappings are alternately trained, even with unpaired data, enabling data-efficient training in chemistry where labeled pairs are scarce. Empirically, RTRL yields substantial gains in RTC and primary task performance across supervised, self-supervised, and synthetic regimes, with RTC improvements up to 52% and primary-task gains up to 55%. This approach demonstrates that RTC is a trainable objective that can produce more robust, credible chemical foundation models and open new paths for bidirectional reasoning in scientific AI.
Abstract
Large Language Models (LLMs) are emerging as versatile foundation models for computational chemistry, handling bidirectional tasks like reaction prediction and retrosynthesis. However, these models often lack round-trip consistency. For instance, a state-of-the-art chemical LLM may successfully caption a molecule, yet be unable to accurately reconstruct the original structure from its own generated text. This inconsistency suggests that models are learning unidirectional memorization rather than flexible mastery. Indeed, recent work has demonstrated a strong correlation between a model's round-trip consistency and its performance on the primary tasks. This strong correlation reframes consistency into a direct target for model improvement. We therefore introduce Round-Trip Reinforcement Learning (RTRL), a novel framework that trains a model to improve its consistency by using the success of a round-trip transformation as a reward signal. We further propose an iterative variant where forward and reverse mappings alternately train each other in a self-improvement loop, a process that is highly data-efficient and notably effective with the massive amount of unlabelled data common in chemistry. Experiments demonstrate that RTRL significantly \textbf{boosts performance and consistency} over strong baselines across supervised, self-supervised, and synthetic data regimes. This work shows that round-trip consistency is not just a desirable property but a trainable objective, offering a new path toward more robust and reliable foundation models.