🤖 machine learning

Uncertainty-aware reinforcement learning for chemical language models

This paper proposes and evaluates two uncertainty-aware reinforcement learning frameworks for chemical language models that mitigate the risks of exploring unreliable chemical space by either treating uncertainty as an optimization objective or modulating policy updates, ultimately achieving more robust molecular design with a significantly higher true hit rate.

Original authors: Borja Medina, Jon Paul Janet

Published 2026-06-25

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Borja Medina, Jon Paul Janet

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 treasure hunter trying to find the most valuable gems in a vast, uncharted cave system. This cave represents the "chemical space"—a universe of trillions of possible molecules. Your goal is to design new molecules that could become life-saving drugs.

To help you navigate, you have a Chemical Language Model (CLM). Think of this model as a highly skilled, but slightly overconfident, guide who has memorized a map of a small, well-explored section of the cave (the training data). When you ask the guide to suggest a new gem location, it uses its knowledge to predict how valuable that gem might be.

The Problem: The Overconfident Guide

The paper points out a major flaw in how we usually use these guides. In the past, we treated the guide's predictions as absolute facts. If the guide said, "This spot has a diamond worth $1 million!", we believed it blindly.

However, the guide is only confident about areas near its original map. If you ask it about a spot deep in the unknown, dark corners of the cave, it might still shout, "Diamond!" with the same confidence, even though it's just guessing. In reality, those predictions are shaky and unreliable.

When we let the guide lead us blindly into these "high-score but high-uncertainty" zones, we end up wasting time on fake treasures. The optimization process becomes unstable, and we generate molecules that look great on paper but don't actually work in the real world.

The Solution: Teaching the Guide to Say "I'm Not Sure"

The authors propose a new way to use Reinforcement Learning (RL), which is essentially a training method where the guide learns by trial and error. They want to teach the guide to pay attention to its own uncertainty—its internal "gut feeling" about whether it knows what it's talking about.

They tested two creative strategies to fix the guide's overconfidence:

Strategy 1: The "Honesty Bonus" (Score Modulation)

Imagine you are playing a video game where you get points for finding gems.

Old way: You get points only for the gem's value.
New way (Score Modulation): You get points for the gem's value minus a penalty if the guide is unsure about that gem.
The Analogy: It's like telling the guide, "If you are 100% sure this is a diamond, I'll give you a huge bonus. But if you're just guessing, I'll deduct points from your score." This encourages the guide to stop exploring the dark, unknown corners and stick to the well-lit, familiar paths where it can be sure of its findings.

Strategy 2: The "Volume Knob" (Loss Modulation)

This approach is more subtle. Imagine the guide is giving you a list of suggestions, and you are the one learning from them.

Old way: You listen to every suggestion with the same intensity, whether the guide is confident or guessing.
New way (Loss Modulation): You turn up the volume on the guide's confident suggestions and turn down the volume (or mute) the shaky ones.
The Analogy: If the guide says, "I'm 90% sure this is a diamond," you lean in and learn from it. If it says, "I think this might be a diamond, but I'm not really sure," you barely listen. This prevents the guide's wild guesses from messing up your training.

The Results: Better Treasure Hunting

The researchers tested these ideas in three different scenarios:

A Simulated Cave: A computer-generated world where they knew exactly how far the guide was from its map and how much it should be guessing.
Real Drug Data (ChemProp): Using real-world models trained on small vs. large datasets.
A Statistical Safety Net (Conformal Prediction): Using a method that flags predictions as "uncertain" if they don't fit the known patterns.

What happened?

Without the fix: The guide would often find "gems" that looked amazing but were actually just illusions (false hits). It would get stuck in areas where it was guessing wildly.
With the fix (especially the Volume Knob strategy): The guide stopped wasting time on the dark corners. It focused on areas where it was confident.
- The number of real, valid hits (true gems) increased by 50% (from 0.5 to 0.75).
- The total number of true hits nearly doubled.
- Crucially, they didn't lose the ability to find high-value molecules; they just stopped finding the fake ones.

The Takeaway

The paper concludes that we shouldn't just ask our AI guides "What is the best molecule?" We should also ask, "How sure are you?"

By treating uncertainty as a tool rather than ignoring it, we can make the AI exploration of chemical space much more robust. It's like giving the treasure hunter a compass that not only points to gold but also warns them when they are walking off the edge of the map. The result is a faster, safer, and more reliable way to discover new medicines.

Technical Summary: Uncertainty-aware Reinforcement Learning for Chemical Language Models

Problem Statement

Reinforcement Learning (RL) has emerged as a powerful paradigm for de novo molecular design, enabling Chemical Language Models (CLMs) to explore chemical space and optimize specific properties. However, existing RL frameworks typically treat scoring functions as deterministic oracles, ignoring the inherent uncertainty in molecular property predictions. These scoring functions rely on predictive models trained on limited, potentially biased datasets, introducing both aleatoric (data noise) and epistemic (model parameter/coverage) uncertainty.

When RL agents optimize solely on point estimates, they may explore regions of chemical space far from the training distribution where predictions are highly uncertain. This leads to the generation of molecules with high scores that are poorly supported by training data, potentially destabilizing the optimization process and yielding predictions far from their true values. While uncertainty quantification is well-studied in Quantitative Structure-Activity Relationship (QSAR) modeling, it has received limited attention in the context of RL-based molecular generation.

Methodology

The authors propose and evaluate two complementary strategies for incorporating predictive uncertainty into the RL loop within the REINVENT4 framework. The framework utilizes a Recurrent Neural Network (RNN) prior to sample molecules, which are then evaluated by a Multi-Parameter Optimization (MPO) scoring function.

1. Probabilistic Scoring Functions

Before introducing the uncertainty strategies, the authors address a mathematical inconsistency in standard REINVENT implementations. When scoring functions provide a full predictive distribution (e.g., via Mean-Variance Estimation), applying non-linear transformations to the expected value alone ( $f(E[X])$ ) is mathematically inconsistent with transforming the distribution first ( $E[f(X)]$ ). The authors implement a probabilistic scoring approach where the MPO is treated as a random variable, approximated via Monte Carlo sampling to propagate variance correctly.

2. Uncertainty Integration Strategies

The paper introduces two distinct methods to utilize uncertainty:

Score Modulation (SM): Uncertainty is treated as an explicit additional optimization objective. The predictive uncertainty is incorporated directly into the MPO scoring function alongside activity and physicochemical properties. This encourages the policy to trade off exploitation (high score) against reliability (low uncertainty).
Loss Modulation (LM): Uncertainty is not added to the score but is used to modulate the magnitude of policy updates. The contribution of each molecule to the gradient update is weighted inversely to its uncertainty. Molecules with high uncertainty have their influence on the policy reduced, effectively down-weighting unreliable learning signals.

3. Experimental Setups

The strategies were evaluated across three settings:

Controlled Model System (MS): A synthetic environment where prediction error is modeled as Gaussian noise with variance proportional to the distance from the training data (simulating an Applicability Domain). Ground truth is known, allowing for precise error measurement.
ChemProp Models: Real-data scenarios using ChemProp models trained with Mean-Variance Estimation (MVE). A "Predictor" model (trained on a small subset) serves as the scoring function, while a "Oracle" model (trained on the full dataset) serves as a ground-truth proxy.
Conformal Prediction (CP): A setup using a Random Forest classifier wrapped with a Mondrian Inductive Conformal Predictor. Uncertainty is quantified via prediction sets and p-values, without an external ground-truth oracle.

Key Results

Model System Results

In the controlled setting, the standard "Noisy Component" baseline exhibited a form of implicit self-regularization, shifting toward lower-uncertainty regions due to scoring instability, though this was slow and dependent on initial conditions.

SM and LM Performance: Both explicit strategies accelerated the shift toward low-uncertainty regions compared to the baseline.
Hit Quality: The LM strategy yielded the best performance, significantly increasing the ratio of true hits to total hits compared to the baseline.
Combined Approach (SM&LM): While this approach shifted to low-uncertainty regions fastest, it resulted in a significant reduction in the total number of accumulated hit-scaffolds, suggesting an over-correction that compromised diversity and exploration in the controlled setting.

Real-Data Results (ChemProp)

In the ChemProp scenario, the baseline (Predictor only) showed a strong increase in predicted activity but a substantial divergence from the Oracle score, indicating the model was exploiting high-uncertainty regions where the Predictor overestimated activity.

Reliability: Uncertainty-guided strategies (particularly LM) maintained a smaller, more stable gap between the Predictor and Oracle scores.
True Hits: The LM strategy produced the highest number of true hits and the largest improvement in the true-to-total hit ratio (increasing from 0.5 to 0.75), nearly doubling the total number of true hits compared to the baseline.
Uncertainty: All guided strategies successfully reduced the predictive uncertainty of the generated molecules.

Real-Data Results (Conformal Prediction)

Without a ground-truth oracle, "true hits" were defined as compounds assigned a unique class label (active) by the Conformal Predictor.

Ambiguity Reduction: Baseline optimization generated a large proportion of ambiguous predictions (assigned to both active and inactive classes). Uncertainty-guided strategies significantly reduced this ambiguity.
Hit Confidence: While the total number of hits was similar across strategies, the uncertainty-guided approaches produced a higher fraction of confident, unique-class predictions.

Significance and Claims

The paper claims that incorporating predictive uncertainty into RL-driven molecular generation yields consistent and meaningful improvements across diverse experimental settings. The primary significance lies in demonstrating that uncertainty is not merely a diagnostic tool but can serve as an active component of the optimization process.

Stability and Reliability: Uncertainty-aware strategies stabilize exploration by favoring regions where scoring functions are reliable, reducing the risk of exploiting spurious predictions.
Superiority of Loss Modulation: The authors identify the Loss Modulation (LM) strategy as the most robust approach. Unlike Score Modulation, which alters the reward landscape by mixing scores with uncertainty estimates, LM preserves the integrity of the original scoring objective while attenuating the influence of unreliable samples on policy learning.
Practical Impact: The approach increases the reliability of hit discovery without compromising molecular scores, making it broadly applicable to any generative framework where RL guides molecular optimization and predictive uncertainty can be estimated.

The authors conclude that their findings emphasize the necessity of accounting for uncertainty to improve the quality and trustworthiness of de novo molecular design, particularly when exploring regions of chemical space distant from training data.