Test-Time Training Scaling Laws for Chemical Exploration in Drug Design

Area of Science:

• Artificial intelligence
• Computational chemistry
• Drug discovery

Background:

• Chemical language models (CLMs) with reinforcement learning (RL) are used for de novo molecular design.
• Mode collapse limits CLM exploration capabilities in chemical space.
• Test-time training (TTT) in large language models inspires new approaches.

Purpose of the Study:

• To enhance chemical space exploration in CLMs by scaling TTT.
• To introduce MolExp, a benchmark for evaluating diverse molecule discovery with similar bioactivity.
• To investigate the impact of scaling TTT strategies on exploration efficiency.

Main Methods:

• Proposed scaling TTT for CLMs by increasing the number of independent RL agents.
• Introduced MolExp benchmark for assessing structurally diverse molecule generation.
• Evaluated the effect of TTT training time and cooperative RL strategies.

Main Results:

• Scaling TTT via more RL agents follows a log-linear law, boosting exploration efficiency on MolExp.
• Increasing TTT training time showed diminishing returns for exploration.
• Cooperative RL strategies were evaluated for enhanced exploration.

Conclusions:

• Scaling TTT with multiple RL agents offers a viable strategy for efficient molecular exploration.
• The findings suggest a scalable framework for generative molecular design.
• Optimizing AI-driven drug discovery can benefit from these insights into exploration efficiency.