On the Design Space of Discrete Diffusion Online Adaptation for Molecular Optimization

2026-07-02

AI Summary

This work proposes an online adaptation framework for discrete diffusion models tailored to molecular optimization under a limited oracle query budget. The approach uniquely integrates, in a systematic and synergistic manner, several key components—including candidate acquisition akin to Bayesian optimization, reward shaping, model debiasing, experience replay, and molecular validity constraints—into an efficient online fine-tuning pipeline. Evaluated across six small-molecule binding affinity tasks and three protein fitness tasks, the method consistently yields molecules with significantly higher reward scores compared to both offline fine-tuning and inference-time search strategies, given identical computational and query budgets. Its advantage is particularly pronounced in scenarios requiring substantial deviation from the pretrained prior distribution.

Abstract

Molecular optimization often starts from a pretrained generative model that captures a broad prior over valid molecular structures. At test time, however, the goal is not to sample from this prior, but to use a limited oracle budget to shift generation toward task-specific high-reward molecules. We study this adaptation problem for discrete diffusion models. Each online round couples several choices. The loop must decide which candidates to evaluate, how rewards become model updates, which feedback to reuse, and how far to move beyond the pretrained prior. These choices have mostly been studied in isolation, leaving open whether they complement one another, become redundant, or interfere inside a full online adaptation loop. We conduct controlled studies across six small-molecule binding-affinity tasks and three protein-fitness tasks. We find that acquisition, reward shaping, and model debiasing provide complementary routes to higher reward, especially for small molecules. Replay further stabilizes learning, while validity penalties keep small-molecule exploration on the valid molecular manifold. Together, these findings point to a practical recipe for feedback-efficient molecular optimization: online fine-tuning with acquisition, reward shaping, debiasing, replay, and validity control. This recipe outperforms offline fine-tuning and inference-time search baselines under matched oracle-call budgets and GPU-hour accounting. The gains are largest when high-reward candidates require larger shifts from the pretrained prior.

Problem

• molecular optimization
• discrete diffusion models
• online adaptation
• oracle budget
• generative model

Innovation

• discrete diffusion
• online adaptation
• molecular optimization
• reward shaping
• model debiasing