Steering Large Language Models through the DMTA Cycle: Structure-Based Drug Design via Knowledge-Driven Bi-Level Thompson Sampling

Structure-based drug design (SBDD) can be effectively realized through an iterative refinement via the Design-Make-Test-Analyze (DMTA) cycle, which is a common workflow used by human experts. However, most conventional deep learning-based models as one-shot generators that lack feedback mechanisms, leaving the DMTA loop disconnected. In this work, we propose K-BTS, a Knowledge-Driven Bi-level Thompson Sampling framework that formalizes iterative SBDD as a Dynamic Hierarchical Multi-Armed Bandit problem. K-BTS closes the DMTA loop by decoupling decisions into two levels: an upper-level policy that prioritizes high-potential molecular lineages and a lower-level mechanism that retrieves explicit chemical rules to guide LLM generation. By integrating a dual-level Bayesian update, the framework transforms sparse docking scores into reusable experience. On the CrossDocked2020 benchmark, K-BTS achieves a state-of-the-art Top-1 average docking score. The results from diverse dimensions show that K-BTS ensures search determinism through a smooth, monotonic convergence that synchronizes structural drift with affinity improvement.

Designing new drug-like molecules is difficult because a useful molecule must fit well into a target protein, much like a key fitting into a lock. Many existing AI methods try to generate such molecules in one shot, but drug discovery is usually not a one-shot process. Human experts often improve molecules step by step: they design a candidate, make or simulate it, test how well it binds, and then use the result to decide what to try next. Our work builds an AI framework, called K-BTS, that follows this iterative process. Instead of asking a language model to freely invent molecules, K-BTS first chooses which promising molecular family to improve, and then selects an explicit chemical modification rule to guide the language model. After each generated molecule is evaluated by docking simulation, the system updates its beliefs about which molecular families and chemical rules are useful. In this way, sparse test results are turned into reusable experience that guides future molecular design. Experiments on a standard structure-based drug design benchmark show that K-BTS finds molecules with stronger predicted binding affinity, while making the search process more stable, interpretable, and closer to how human experts refine drug candidates.