Uncertainty-aware reinforcement learning for chemical language models
Summary
Proposes two complementary approaches to incorporate predictive uncertainty into reinforcement learning for chemical language models, improving robustness and increasing true hit rate by 0.25 in de novo molecular design.
View Cached Full Text
Cached at: 06/25/26, 05:09 AM
# Uncertainty-aware reinforcement learning for chemical language models Source: [https://arxiv.org/abs/2606.24990](https://arxiv.org/abs/2606.24990) [View PDF](https://arxiv.org/pdf/2606.24990) > Abstract:Reinforcement Learning \(RL\) has become a powerful paradigm for de novo molecular design, enabling Chemical Language Models \(CLMs\) to navigate and explore the chemical space while optimizing specific desired properties\. However, the existing RL frameworks treat all scoring functions as deterministic oracles, neglecting the inherent uncertainty attached to the predictions of the different molecular properties\. This can lead to the exploration of highly\-uncertain regions of the chemical space, focusing on the generation of highly scored molecules which are poorly supported by the training data\. This can destabilize the optimization process, yielding predictions that are far from their true values\. We propose and compare two complementary ways of incorporating predictive uncertainty into RL\. In the first one, uncertainty is treated as an additional optimization objective and incorporated along with the rest of the scoring functions, allowing the policy to trade off exploitation against reliability\. Secondly, uncertainty is used to modulate policy updates, reducing the influence of molecules whose properties lie far outside the scoring function confidence domain\. Both approaches were evaluated across three different settings: \(i\) a controlled model system, in which the prediction error is modeled as a Gaussian distribution, with a variance proportional to the distance to the training data; and two real\-world tasks, making use of \(ii\) ChemProp models and \(iii\) a Conformal Prediction wrapper applied to a Random forest classifier\. We show that uncertainty\-aware RL enables CLMs to explore chemical space more robustly by favoring lower\-uncertainty regions\. This leads to more reliable hit discovery without compromising molecular score, increasing the true hit rate by 0\.25 \(from 0\.5 to 0\.75\), and nearly doubling the total number of true hits\. ## Submission history From: Borja Medina \[[view email](https://arxiv.org/show-email/a4d954a7/2606.24990)\] **\[v1\]**Tue, 23 Jun 2026 14:42:46 UTC \(32,710 KB\)
Similar Articles
Reinforcement Learning with Metacognitive Feedback Elicits Faithful Uncertainty Expression in LLMs
This paper introduces reinforcement learning with metacognitive feedback (RLMF) and metacognitive data selection to improve large language model calibration, enabling faithful expression of intrinsic uncertainty and surpassing standard RL by up to 63%.
DUD: Decoupled Update Dynamics for Reliable Uncertainty Quantification in Large Language Models
The paper introduces DUD (Decoupled Update Dynamics), a framework that separates Feed-Forward Network and Attention contributions via causal interventions to improve uncertainty quantification and calibration in large language models, outperforming state-of-the-art baselines.
Improving Molecular Property Prediction in Small Language Models Using Graph-based Tools
This paper proposes a Context-Augmented Prompting framework that uses a GNN expert model to provide predictive hints and explanatory subgraphs to improve molecular property prediction in small language models. Experiments on MUTAG and Tox21 show accuracy gains of up to 74% over SMILES-only baselines.
Uncertainty-Aware Simulation-Based Inference for Operations Research with Large Language Models
This paper proposes a training-free, uncertainty-aware inference framework for using large language models in operations research. The method uses short lookahead simulations and importance resampling to improve the coherence of mathematical formulations, outperforming standard baselines on OR benchmarks.
Towards Robust Reinforcement Learning for Small-Scale Language Model Agents
This paper systematically investigates instability in reinforcement learning for small language model agents (70-500M parameters), identifying three failure modes and proposing robust techniques including a merge-and-reinitialize adapter approach and safety mechanisms; it achieves stable convergence and improved win rates.