Test-Time Training Scaling Laws for Chemical Exploration in Drug Design
This provides a scalable framework for generative molecular design in drug discovery, though it appears incremental as it adapts existing Test-Time Training methods to a new domain.
The paper tackles the problem of mode collapse in Chemical Language Models for molecular design by scaling Test-Time Training with multiple reinforcement learning agents, demonstrating that increasing the number of agents follows a log-linear scaling law that significantly improves exploration efficiency on the MolExp benchmark.
Chemical Language Models (CLMs) leveraging reinforcement learning (RL) have shown promise in de novo molecular design, yet often suffer from mode collapse, limiting their exploration capabilities. Inspired by Test-Time Training (TTT) in large language models, we propose scaling TTT for CLMs to enhance chemical space exploration. We introduce MolExp, a novel benchmark emphasizing the discovery of structurally diverse molecules with similar bioactivity, simulating real-world drug design challenges. Our results demonstrate that scaling TTT by increasing the number of independent RL agents follows a log-linear scaling law, significantly improving exploration efficiency as measured by MolExp. In contrast, increasing TTT training time yields diminishing returns, even with exploration bonuses. We further evaluate cooperative RL strategies to enhance exploration efficiency. These findings provide a scalable framework for generative molecular design, offering insights into optimizing AI-driven drug discovery.