Steering Discrete Diffusion: Design Choices for Feedback-Efficient Molecular Optimization

This post accompanies our NeurIPS 2026 paper, a collaboration between Caltech and Genesis Molecular AI.

Pretrained molecular generative models have learned the broad grammar of medicinal chemistry from hundreds of millions of drug-like SMILES strings, yet randomly drawing a plausible structure rarely yields a molecule that binds a specific protein target. Bridging this kind of gap requires test-time compute: the compute budget spent after training, to align a general-purpose prior with a specific downstream objective. In this work, we investigate how to spend this budget wisely to find valid, high-reward molecular and protein candidates. Our approach manages to match the ~5,000-evaluation performance of the strongest inference-time search baseline (Beam Search) using only 86 evaluations.

Figure 1. Write the pretrained chemical prior as $p_{\text{pre}}(x)$ and the target-specific binding reward as $r(x)$. Reweighting the prior by the reward defines the tilted target distribution $p^*(x) \propto p_{\text{pre}}(x) \exp(r(x)/\beta)$, which concentrates on the single best valid molecule as $\beta$ shrinks.

Aligning the prior distribution with a target one means concentrating probability mass where the reward is high without drifting off the valid chemical manifold the prior learned. Because the prior and the reward pull in different directions, their product lands in a tiny pocket of space that naive sampling practically never reaches. Every test-time method approximates this reward-tilted target, trying to concentrate samples on the highest-scoring valid molecules.

A test-time method spends its budget across two fronts: calling the oracle to evaluate candidates, and running the model itself for sampling and updates. Which algorithm wins depends on the relative cost of those two fronts. In code or math, verification is cheap and fast, making brute-force best-of-$N$ sampling hard to beat. In molecular design, the oracle dominates the budget. Each structure prediction or wet-lab assay is expensive, and a campaign gets only a few hundred or thousand queries. In our setting, the generator is a domain-specific pretrained discrete diffusion model and the oracle is a binding affinity or protein fitness predictor, which lands in the latter regime, where the design space remains largely unsettled.

Figure 2. Cheap oracles favor dense sampling (left); expensive molecular oracles require informative queries (right).

Under a tight query budget, we faced a choice between inference-time search, which keeps the generator frozen and filters draws at sampling time, and online adaptation, which updates the generator weights so each round samples from a better distribution. We spend our budget on updating the generator.

Figure 3. Test-time budgets can filter a frozen prior (top, middle) or update the generator (bottom).

Sustaining this generator update requires deciding how to query the oracle, how to turn raw rewards into a training signal, and how to push the generator into new territory without leaving the valid prior. Below, we unify these moving parts into an online adaptation framework and systematically characterize the interactions across its design space.

Approach: Four Design Choices

All four choices plug into the same underlying fine-tuning objective, which consumes a per-sample log-reward and updates the generator toward the reward-tilted target. We use DDPP-LB, a variant of Discrete Denoising Posterior Prediction (Rector-Brooks et al., 2024), and confirm that the conclusions carry over by swapping it for VIDD (Su et al., 2025).

1. Which Candidates Earn an Oracle Call

Trained on only a few thousand scored molecules, our surrogate predicts well only near candidates it has already seen. Picking the highest predicted scores wastes oracle calls on regions the surrogate is already confident about, teaching it nothing new.

To spend calls where the surrogate is uncertain, we use Thompson sampling (TS; Thompson, 1933). An ensemble of ten small multilayer perceptrons (MLPs) predicts a mean $\mu_i$ and disagreement $\sigma_i$ for each candidate. Rather than ranking by $\mu_i$, we draw one sample $\hat r_i = \mu_i + \varepsilon_i\sigma_i$ with $\varepsilon_i \sim \mathcal{N}(0,1)$ and send the top $K$ to the oracle. High disagreement gives an uncertain candidate a wide draw that can beat a confidently mediocre one, while agreement shrinks $\sigma_i$ back to greedy selection.

Here exploration scales directly with uncertainty, outperforming both random selection and Upper Confidence Bound (UCB) rules that add a fixed uncertainty bonus.

Figure 4. Thompson sampling, and the explore-exploit trade it makes. Shaded area shows the probability of uncertain candidate $c_B$ outranking $c_A$’s mean $\mu_A$

2. Which Rewards Drive the Update

We care about only the best molecules, but maximizing expected reward raises the average one. We could simply filter for the top-$k$ molecules, but discarding the rest of the batch is also not ideal.

Thus, we use Conditional Value at Risk (CVaR; Rockafellar & Uryasev, 2000) as an objective to adapt the diffusion prior (De Santi et al., 2025a). CVaR optimizes the mean of the top $(1-\alpha)$ tail above a quantile threshold $\tau_\alpha$. In the loop this maps to one line of code: replacing each reward $r$ with $\max(0, r - \tau_\alpha)/(1-\alpha)$ at $\alpha = 0.8$, so only the top 20% of molecules drive the update.

Figure 5. CVaR discards samples below $\tau_\alpha$ (unshaded) to optimize the tail mean (red line).

3. How Far to Move Off the Prior

A KL penalty to the prior is standard in reward fine-tuning. For our purposes, KL pulls the wrong way by keeping the generator near a prior we explicitly want to leave. Meanwhile, upper-tail shaping causes mode collapse: concentrating gradients on top samples pushes the generator to repeat early hits and stop exploring.

We want an objective that distinguishes a high-reward region the generator samples constantly from one it has barely visited. Density entropy regularization (DER; De Santi et al., 2025b) provides this by decreasing the strength of the prior log-likelihood in the reward, giving $\log r(x) - \gamma \log p_{\theta}(x)$. Crowded modes become worth less than equally promising regions rarely visited, and because the term uses the model’s own density, it fits smoothly as a diffusion objective.

Figure 6. Standard fine-tuning (DDPP-LB, blue) narrows onto the prior peak; adding DER (red) redistributes mass.

4. What Keeps the Generator Stable

Pushing the model this far off-prior causes the optimization to degenerate in two ways:

First, the model starts reward hacking: it generates broken SMILES strings that fool the surrogate into predicting high scores despite not representing valid molecules. Because chemical validity only becomes clear once the full string is complete, we assign a fixed negative reward after decoding whenever RDKit fails to parse the structure.

Second, the model forgets earlier hits: diversity collapses easily if we train only on the newest batch. So we maintain a replay buffer (Mnih et al., 2015) of the top-scoring molecules found so far and sample training batches from it.

Figure 7. The assembled loop combining TS, CVaR, DER, replay and validity components.

Results

We ran nine tasks. Six are small-molecule binding targets taken from the ODesign benchmark (Zhang et al., 2025) and named by the PDB entries they come from: 2VT4, 5SDV, 6CM4, 7BKC, 7C7M, and 7YLL, each scored by a learned affinity model. The other three are protein fitness tasks, one per family: GB1, CreiLOV, and TrpB, each with its own generator and its own fitness predictor. Every cell runs five seeds.

The aggregate numbers below average over the six molecular targets, and the trajectory plots show one of them, 7C7M, so the individual curves stay readable. On the three protein tasks, the same ablations barely change performance, bringing a result that shows when online adaptation is actually necessary.

The Components Interact In Synergy

In all comparisons below, we hold the oracle query budget fixed to measure how much reward each method extracts from the same number of evaluations.

First, every component is necessary: dropping any single module from the full loop hurts performance. Under DDPP-LB at 5,000 oracle calls, removing DER causes the largest drop, pulling normalized reward down from 0.97 to 0.58.

Second, these components are strongly coupled rather than additive. A full ablation across DER, CVaR, and Thompson sampling shows that CVaR and Thompson sampling barely help unless DER is active. Our analysis shows that DER pushes the generator into new off-prior regions, while CVaR and Thompson sampling focus queries and updates on the best candidates found there.

Figure 8. Where each component moves the sampling distribution. Each panel overlays the full loop (blue) against the same loop with one component removed, on two small-molecule targets. The x-axis is a generated sample’s likelihood under the frozen pretrained model, so further right means further from the prior.

Tracking the full optimization trajectories shows when this synergy kicks in. All variants perform similarly for the first ~100 oracle calls while sampling near the prior. The performance gaps open up only after enough updates push the generator off-prior.

Figure 9. Normalized reward trajectories across single, pairwise, and full component ablations on 7C7M, averaged over 5 seeds. If the components were additive, removing a pair would cost about as much as its two single removals added together.

Protein Tasks Reveal Where Online Adaptation Pays Off

On GB1, CreiLOV, and TrpB, removing components barely changes performance. The protein generators are trained per family, so a high-fitness variant is already something the prior samples readily. Unlike the broad small-molecule prior, these leave DER no off-prior regions to unlock, giving CVaR and Thompson sampling little room to improve over simple sampling.

This gives a practical rule of thumb for when to use online adaptation: check how far your target objective sits from your pretrained prior. The further off-prior your target sits, the larger the gains from our framework.

Adapting Beats Searching

We compared against five inference-time search baselines that steer a frozen generator: best-of-N, beam search, MCTS, discrete Feynman-Kac correctors, and SMC. Averaged over the six small-molecule targets, our loop matches the best baseline’s final score using 58x fewer oracle calls. This efficiency gap will become critical when each oracle call is a far more expensive structure-prediction run or wet-lab assay.

Figure 10. Best reward vs. cumulative oracle calls on 7C7M, averaged over 5 seeds: online adaptation vs. inference-time search.

Conclusion

Our work shows that online adaptation succeeds as a coupled system shaped by the underlying distributional dynamics. As high-fidelity foundation models and automated assays become standard in the discovery stack, how much a model improves per oracle call determines which targets we can solve within a realistic experimental budget. Having validated this framework across computational benchmarks, our next focus is refining each module and extending the loop to wet-lab validation, closing the cycle between generative modeling and experimental drug design.

For the full set of results, ablations, and analysis, see our arXiv preprint. Come visit us at NeurIPS 2026 in Atlanta.

Acknowledgments

This work was co-led by Trevor Chen and Ariel Dai alongside Jason Yang, Daniel Khalil, Wenda Chu, and Yisong Yue at Caltech, and Riccardo De Santi at ETH Zurich, with experimental assistance and advising from the Genesis team. The Caltech team is supported in part by NSF #2505096 and gifts from Genesis Molecular AI, Point72, and OpenAI. Jason Yang was supported in part by a Google Fellowship. Riccardo De Santi was supported by an ETH AI Center doctoral fellowship and the Swiss National Science Foundation under NCCR Catalysis and NCCR Automation.