Iterative Distillation for Reward-Guided Fine-Tuning of Diffusion Models in Biomolecular Design
Authors: Xingyu Su
Xiner Li
Masatoshi Uehara
Sunwoo Kim
Yulai Zhao
Gabriele Scalia
Ehsan Hajiramezanali
Tommaso Biancalani
Degui Zhi
Shuiwang Ji
Venue: ICLR 2026 Poster (2026)
Keywords: Biomolecular Design
Diffusion Models
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)

Abstract: We address the problem of fine-tuning diffusion models for reward-guided generation in biomolecular design. While diffusion models have proven highly effective in modeling complex, high-dimensional data distributions, real-world applications often demand more than high-fidelity generation, requiring optimization with respect to potentially non-differentiable reward functions such as physics-based simulation or rewards based on scientific knowledge. Although RL methods have been explored to fine-tune diffusion models for such objectives, they often suffer from instability, low sample efficiency, and mode collapse due to their on-policy nature. In this work, we propose an iterative distillation-based fine-tuning framework that enables diffusion models to optimize for arbitrary reward functions. Our method casts the problem as policy distillation: it collects off-policy data during the roll-in phase, simulates reward-based soft-optimal policies during roll-out, and updates the model by minimizing the KL divergence between the simulated soft-optimal policy and the current model policy. Our off-policy formulation, combined with KL divergence minimization, enhances training stability and sample efficiency compared to existing RL-based methods. Empirical results demonstrate the effectiveness and superior reward optimization of our approach across diverse tasks in protein, small molecule, and regulatory DNA design.

Published as a conference paper at ICLR 2026
ITERATIVE DISTILLATION FOR REWARD-GUIDED FINE- TUNING OF DIFFUSION MODELS IN BIOMOLECULAR DESIGN
Xingyu Su Xiner Li
Masatoshi U
Sunwoo Kim
Yulai Zhao
Gabriele Scalia
Ehsan Hajiramezanali Tommaso Biancalani
Degui Zhi Shuiwang Ji
'Texas A&M University "EvolutionaryScale SSou National University
"Princeton University *Genentech University of Texas Health Science Center at Houston
(xingyu. su, sji] @tamu.edu
ABSTRACT
We address the problem of fine-tuning diffusion models for reward-guided genera- tion in biomolecular design. While diffusion models have proven highly effective modeling complex, nigh-dimensional data distributions, real-world applications of- ten demand more than potentially non-differentiable reward functions such as physics-based simula- tion or rewards based on scientific knowledge. Although RL methods have been explored fine-tune diffusion models for such objectives, they often suffer from instability, low sample efficiency, and mode collapse due to their on-policy nature. In this work, we propose an iterative distillation-baat enables diffusion models to optimize for arbitrary reward functions. Our method casts the problem as policy distillation: collects off-policy data during the roll-in phase, simulates reward-based soft-optimal policies during roll-out, and updates the model by minimizing the KL divergence between the simulated soft-optimal policy and the current model policy. Our off- off-policy formulation, combined with KL divergence minimization, enhances training stability and sample efficiency compared to existing RL-based methods. Empirical results demonstrate the effee- tiveness and superior reward optimization of our approach across diverse tasks protein, small molecule, and regulatory DNA design. The source code is released at (https://divelab httpa/////ieeelb.it..bbb github. 1o/VIDD/).
generation, requiring optimization with respect
fine-tuning framework that
INTRODUCTION
Diffusion models (Sohl-Dickstein al., 2015 Ho et al., 20200 Song et al., 2020) have achieved remarkable success across diverse domains, inc including computer Vision and scientific applications (e.g., protein design (Watson et al., 2023 Alamdari et al. 2023)). Their strength lies in modeling complex, high-dimensional data distributions, including natural images and chemical structures such as proteins and small molecules. However, in many real-world scenarios, especially for biomolecular design, generating samples that merely resemble the training distribution is not sufficient. Instead, often seek to optimize specific downstream reward functions. For instance, in protein design, beyond generating plausible structures, practical applications frequently require satisfying task- objectives such as structural constraints, binding affinity, and hydrophobicity (Hie et al., [2022] Pacesal al. 2024). To meet these requirements, fine-tuning diffusion models with respect to task-specific rewards is crucial, enabling goal-directed generation aligned with downstream objectives.
Numerous algorithms have been proposed for fine-tuning diffusion models with respect to reward functions, motivated by the observation that this problem can be naturally framed as a reinforcement learning (RL) task within an regulay-ee Markov Decision Process (MDP), where each
"Equal contribution
This work was done when he was at Genentech
Corrrsponniig author

---

Published as a conference paper at ICLR 2026
Fr 16
Ditlusion Denoising
Benr f- 0, Vo exp
log Digpi-1(e--
19. for a(rr)
pere
Pr-1
14
Reward fanction
firetion
St
Q
R
N
P
Zg
Model lazy update
Loss update
Roll-in
Goal: reward- guided fine -anittnnng
Rolhoul
Soft value fanction approximation
SS-match 0.695
Docking score 13.1
IPTM - 0.70
Figure Overview of VIDD, VIDD fine-tunes diffusion models to maximize potentially non- differentiable rewards by iteratively distilling soft- denoising policies. It alternates between (1) off off-policy roll-in, (2) guuded reward- weighted roll-out, and (3) forward KL-based model updates. Our algorithm leverages aff-policy roll-ins and forward KL minimization rather, which contribute improved optimization stability.
policy corresponds to the denoising process the diffusion model (Black et al. 2024 Fan et al. 2023). In computer vision, current tfuteoof-the-ar methods fine-tune diffusion models by directly backpropagating reward gradients through the generative process (Clark et al. 2023 Prabhudesai et al.)
2023). However, in many scientific applications, reward functions are inherently non-differentiable.
making such optimization inapplicable. For example, in protein design, rewards based on secondary structure matching (e.g.. DSSP algorithm (Kabsch & Sander 1983)) or binding affinity predictions (e.g., AlphaFold3 (Abramson et al. 2024)) typically rely on hard lookup-tablle based on scientific knowledge. Similarly, in small molecule design, reward functions such as synthetic accessibility (SA), molecular fingerprints Yang et al. 2021), and outputs from physics-based simulators (e.g.. AutoDock Vina (Trott & Olson 2010) are also non-differentiable. These characteristics pose fundamental challenge for direct back-propagation approaches in scientific domains.
In such cases, policy gradient methods like Proximal Policy Optimization (PPO) (Schulman et al.. 2017) offer natural alternative, as they are used in diffusion models. However, PPO is known to exhibit instability, hyperparameter sensitivity, and susceptibility to mode collapse many contexts Yuan et al. (2022): Moalla et al. (2024). These issues arise in part due to its on- on-policy nature trajectories used for training the model are generated by the current fine-tuned policy, leading to narrow exploration around previously visited regions. Furthermore, from theoretical perspective, PPO can be viewed as minimizing the reverse Kullback-Leibler (KL) divergence between the target and generated trajectory distributions. This reverse KL objective encourages mode-secking behavior,
potentially leading to mode collapse (Wang et al. 2023 Go et al. 2023 Kim et al. 2025).
To address the aforementioned challenges, propose new framework, VIDD (Value-guided Iterative Distillation for Diffusion models), designed to maximize possibly non-differentiable rewards stable and effective manner. The core idea iteratively distill soft-optimal policies-including which serve as target denoising processes that optimize the reward while remaining close to the current fine-tuned model, as visualized FFguuee Concretely, the algorithm proceeds in three iterative steps: (1) roll-in using sufficiently exploratory off- off-policy trajectories, (2) roll-out to simulate soft-optimal policies, and (3) update the fine-tuned model by minimizing the KL divergence between the soft- optimal and current model policies. Importantly, in (2) and (3), our algorithm effectively leverages value functions tailored to diffusion models to guide fine-tuning, analogous to value- weighted MLE in RL (Peters et al. 2010). Notably, framework leverages off-policy roll-ins decoupling data collection from policy updates and employs forward KL minimization rather than reverse KL, both of which contribute improved optimization stability complex reward landscapes.
2

---

Published as a conference paper at ICLR 2026
Our contributions are summarized follows. We propose novel algorithm, VIDD, for fine-tuning diffusion models through iterative distillation of target policies composed of both value functions and the current policy. Unlike direct reward backpropagation methods, which require differentiable reward signals, our approach can handle non-differentiable rewards commonly encountered in scientific domains. Algorithmically, our key innovation lies in effectively incorporating value functions specifically tailored to the diffusion models. Empirically, we validate the effectiveness of our method across a range of scientific tasks, particularly in protein and small molecular design.
1.1 RELATED WORKS
Fine-tuning diffusion models for reward-maximication When reward functions are differen-
tiable, is often the case in computer vision, saate-of-the-eee methods achieve strong performance by
directly backpropagating gradients induced by the reward functions on the diffusion model (Clark
al. 2023 Prabhudesai al. 2023 Wu et al. 2024 Wang al. 2024). However, in many scientific
domains such as biology and molecule field, reward functions are often non-differentiable (Lisanza
al. 2024, Hie et al. 2022). In such settings, a natural approach is to adopt RL techniques that
optimize reward without requiring differentiability (Zhang et al., [2024] Gupta et al. 2025 Ektefaie
al. [2024] Vogt et al. 2024 [Deng et al.) 2024] Venkatraman et al. [024 Rector-Brooks et al. 2024]
Zekri & Boullé, 2025). ). For instance, DPOK (Tan et 2023) and DDPO (Black et al.] 20224) adapt PPO, a stabilized variant of policy gradient methods, to the diffusion model fine-tuning. However, these methods often suffer from instability when applied to diffusion models (Clark et al. 2023
Figure 3), due to their inherently on-policy nature and the use of a reverse KL divergence objective between the fine-tuned policy and target policy (discussed further in [Section 5]. In contrast, method avoided on-policy updates and instead leverages a forward KI. divergence objective, which stabilizes training and mitigates mode collapse.
Inference-time technique for reward-maane maximizaaiinn. An alternative line of work focuses on training-free free techniques that aim to improve rewards solely at inference time. The most straight- forward example is Best-of-N sampling, which selects the highest-reward output from N generated candidates. More sophisticated methods go further by evaluating and selecting promising interme- diate generation states during the sampling process (Wu et al.. 2023 Li et al., 2024] Kim al., 2025 (Singhal et al. 2025: Ma et al. 2025 Tang et al. 2025). While these methods avoid the need for fine-tuning, they often incur sigmificantly higher inference-time costs due to repeated sampling. Although they can boost reward in some cases, they do not directly improve the underlying generative model. this respect, training-frte free methods are orthogonal and complementary to tine-tuning-base approaches, and the two be effectively combined achieve even stronger performance.
2
PRELIMINARY
We begin by reviewing the foundations of diffusion models. We then formulate our objective: fine-tuning diffusion models to maximize task-specific reward functions.
2.1 DIFFUSION MODELS
Diffusion models (Sohl-Dickstein et al. 2015 Ho et al. 2020 Song et al. 2020) aim to learn a sampling distribution p(-) ((X') over predefined design space X based the observed data. Formally, a diffusion model defines forward noising process dr (x, Iz-1). which gradually corrupts from clean data To p(z)) over discrete time steps L,..., T into noise. The learning objective is to approximate the reverse denoising process ( 22),, where each Pt is a kernel mapping from X' to A(X', such that the overall reverse trajectory transforms noise samples back into data samples drawn from the true distribution. In practice, each pt is parameterized by neural network and trained to minimize a variational lower bound on the negative log- likeliill of the data.
Notation. With slight abuse of notation, we often denote the initial noise distribution Pr E (XX) as (-I-), and we often refer to the denoising process policy, following terminology commonly used in RL. We denote by (n,2)) : X' + X' the neural network used in pretrained diffusion models to predict the denoised input. Extension to other parameterizations is straightforward.
3

---

Published as a conference paper at ICLR 2026
2.2 REWARD MAXIMIZATION IN DIFFUSION MODELS
Our goal is to fine-tune diffusion models to produce outputs that achieve high rewards. Here, we formalize the problem and highlight key challenges.
Task Description Our goal is to fine-tune pretrained diffusion model to generate samples that maximize task- specifp reward function. Formally, given pretrained diffusion model p E A(X) and reward function r X I. we aim to fine-tune pre-trained models such that
argmax
E Rao [r(2))]
Dis({p")
(p)rre))
0
reward maximization
regularization
where pd E ( denotes the distribution induced by the fine-tuned reverse denoising process {p.)t=r- and pare denotes the original pretrained distribution. Term (a) promotes the generation of high-reward samples, while term (b) penalizes deviation from the pretrained model to maintain the naturalness of generated samples. For example, when the discrepancy measure Dis in (b) the KL divergence, the optimal solution takes the following form up to normalizing constant:
exp(r())/o)pPr)
Challenges. In domains such as computer vision, reward functions are typically modeled using differentiable regressors or classifiers, allowing for direct gradient-based optimization (Clark et al. 2023; Prabhudesai et al.] 2023). In contrast, many key objectives in biomolecular design, such as structural constraints, hydrophobicity, binding affinity, often rely on non-differentiable features, as detailed in [Seectioon I] While RL-based approaches. such as policy gradient methods, can in principle handle non-differentiabll feedback (see Seectinn I.I). they often suffer from instability due to their on-policy nature and reliance on reverse KL. divergence. the following, we introduce a novel fine-tuning method that can effectively optimize possibly non-differentiable rewards.
ITERATIVE DISTILLATION FRAMEWORK FOR FINE-TUNING DIFFUSION MODELS
Our algorithm is motivated by the goal of distilling desirable teacher policies that maximize task- specifie reward funetions. To clarify this motivation, we first define what constitutes teacher policies in our setting, followed by a detailed description of our approach to distilling them.
3.1 TEACHER POLICIES IN DISTILLATION
We introduce soft-optimal policies, which serve as the teacher policies that our algorithm aims to distill. Specifically, we define the soft-optimal policy pi-1 X ((x)) as a value-weighted variant of the pre-trained policy:
pi.1 (112,) - (-1x)) exp
(2)
exp(ux((x://)
Here, for E [T + 1, 1, :*, 1]. the soft value function is defined as:
r((0))
# (+) = a log Ego (2ooor))))))
exp
| -
(3)
where the expectation is taken under the trajectory distribution induced by the pre-trained policies.
These soft-optimal policies naturally arise when diffusion models are framed within entropy- regularized Markov Decision Processes. A key property of these policies that, when sampling trajectories according the soft-optimal policy sequence {pt)e the resulting marginal distribu- tion over final outputs approximates the target distribution cxp(r(-)/ PPre in ①I (eeaaa et al.) 2025a).
Importantly, many test-time guidance methods such as classifier guidance both continuous diffusion models (e.g.. (Dhariwal & Nichol 2021 Bansal al. 2023)) and discrete diffusion models
4

---

Published as a conference paper at ICLR 2026
(e.g.. g., (Nisonoff et al. 2024)), well as sequential Monte Carlo (SMC)-based approaches (e.g.. (Wu et al.) 2023 Li et al. 2024; Kim et al., (025)) can be interpreted as approximate sampling schemes for these soft-optimal policies (Uchara et al. 2025a). While test-time guidance methods are appealing due to their ease of implementation, they often incur substantial computational overhead during inference and may struggle to achieve consistently high reward values. In contrast, explore how soft-optimal policies can inform the fine-tuning of diffusion models, enabling -wwrddggudded generation without requiring any additional computation test time.
3.2 ITERATIVE DISTILLATION
Thus far, we have introduced soft-optimal policies and discussed their desirable properties. Within our framework, we designate these policies as teacher policies, while the fine-tuned models are treated as student policies (Czarnecki al. 2019). A natural objective for distilling such policies into a fine-tuned model {p]] } is given by
argmin
En Emyon (ttppppi---aaaallllllll------
where us E ((Y) denotes roll-in distribution, where we will elaborate on in Section 4.1 By simple algebra (see the detailed derivation in Appendix B), up to some constant, this objective reduces to
argmax
E Beree
E., m (or (le (op
(xp((x(x:)/o)
0
(5)
As will demonstrate in Section 4.3] this objective can be optimized in practice by approximating value functions and replacing expectations with sample estimates. Notably, the objective is off policy in nature, meaning that the roll-in distribution u, can be arbitrary and does not need match the current model policy. Indeed, this objective, commonly referred to as value-weighted maximum likelihood in the RL literature (Peters et al.) (010)), has been employed as scalable and stable approach for off-policy RL (Peng et al. 2019).
In practice, we adopt an iterative distillation procedure. This is motivated by the fact that, due to the reliance on empirical samples and approximate value functions, it is generally infeasible to accurately distill the target policy p into the student policy pf_ a single step. Rather than fixing p throughout training, we periodically update it by the latest student policy but in a lazy manner, allowing the target to gradually improve over time. This dynamic yet infrequent (lazy) target update enables progressive refinement, allowing both the teacher and student policies to evolve iteratively toward better alignment. This strategy is analogous to standard RL practices grounded in the Policy Improvement Theorem (Mnih et al. (0155).
To formalize the above iterative process, we introduce an iteration index s E [1, .S and perform recursive updates as follows:
Texp(10n (Fa-1)/a log pf (01
0.41 t 0s
Σ E Brewn (m--))))))))))
(6)
exp(v;(r:)/or
where {pout] denotes the roll-out policy, and ugut is the corresponding soft value function. The roll-out policy is updated in a lazy manner-sy specifriilll it is refreshed every K steps using the current model parameters p². This lazy update scheme is critical for algorithmic stability in the off-policy setting, preventing rapid changes in the target while still allowing gradual improvement of the student policy toward higher reward. In the next section, formalize this iterative distillation process as a complete algorithmic procedure.
4 ITERATIVE VALUR-WEIGHTER WEIGHTED MLE
In this section, we present our algorithm for fine-tuning diffusion models to maximize reward functions. The full procedure is summarized in Algorithm I Each training iteration consists of three key components: (1) the roll-in phase, which defines the data distribution over which the loss is computed; (2) the roll-out phase, which aims approximate the teacher policy by sampling from
5

---

Published as a conference paper at ICLR 2026
Algorithm 1 VIDD (Value-guided Iterative Distillation for Diffusion models)
Require: reward r: X - IR, pretrained model {up-1( (- xx)), fine-tuned model {rî-1( (- sx))
2: Initialize fine-tuned model p" = p roll-out model pour = prre, lazy update interval K, a
3: for [1.... do
4: // Roll-in phase
5: Generate roll-in samples following roll-in policies (explain in Section 4.1) and collect D
{x.).... , To 11=1-
Ji=1
6: // Roll-out phase
7: Obtain a sample following the roll-out policy pp - ~ pp (at x for each
Approximate soft-value functions 62 (x. t r [ (o and (x(xx()) <
8:
r (#o(<t) [1,, ,0ott).
9:
If 5%K O, update the roll-out policy: (vgu) t (x,²) and your Oz.
10: // Distillation phase
Update model parameters 0 by gradient ascent:
exp(è=- 1z (z( log pl 1f lef
Os+1 f- 0g
77)
(xp()r(x.")aa)
end for 13: return As
roll-out policy and computing its corresponding weight (soft value); and (3) the distillation phase, where the objective is defined as the KL divergence between the teacher policies and the student policies (i.e., the fine-tuned models). We detail each of these components below.
4.1 ROLL-IN PHASE
Due to the off-policy nature of our algorithm, we have significant flexibility in selecting the Zi,
roll-in distribution path IT, IT-1,".. ,00, which subsequently serves as the training distribution for loss computation. A well-designed roll-in policy should satisfy two competing desiderata: (1) exploration expuoration-eesurrnng broad coverage over the design space to avoid local optima, and (2) exploitation iopolaation-mminaiiinng proximity to high-reward trajectories for efficient policy improvement.
To balance these goals, we adopt a mixture roll-in strategy by sampling the roll-in distribution from:
the pre-trained policy p.. which promotes exploration by generating diverse trajectories;
the roll-out policy p. which is periodically updated and reflects stabilized knowledge from the student model.
At each training step, we sample from yhth with probability 1 3., and from p with probability 3, Further details on the construction of pe are provided in the following subsection.
4.2 ROLL-OUT PHASE
Our goal in this step is to approximate the current teacher policy. To this end, we aim to (1) sample x;-1 conditioned x; for each timestep &. following roll-out policy, and (2) compute its corresponding soft-value, which will later serve as a weight during the distillation process.
Sampling from roll-out policy. Recall the motivational formulation in Equation 6 Our goal here is replace the expectation with its empirical counterpart. To this end, we draw samples 7t-1 from p which is updated periodically at fixed time intervals.
Approximation of soft value functions. Recall that soft-value functions are defined as conditional expectation given Tt-1- While one could estimate these expectations using Monte

---

Published as a conference paper at ICLR 2026
Carlo sampling or regression, commonly done in standard RL settings, we recommend a more practical approximation: 0t- 1 := r (i0( Recalling @o 80t!) is the denoised prediction from the diffusion model parameterized by the current student policy. This approximation is based on the replacement of the expectation in ( with its posterior mean. We apply the same
approximation to ig (ate), estimating it via r (are; (o0)) )). Further discussions soft value functions
can be found in Appendix F
Notably, this approximation has been implicitly adopted in several recent test-time reward opti- mization methods (Chung et al. 2022. Wu et al., 2023; Li al. 2024). Building on its empirical success, we extend this idea the tine-tuning setting. Compared Monte Carlo-based methods, this approximation is computationally efficient- requiring only single forward pass through the denoising network- and avoids the need to train an additional value function.
4.3 DISTILLATION PHASE
Thus far, we have defined the training data distribution via the roll-in phase and specified the supervised signal through the roll-out phase. The final step is to perform distillation by solving a supervised learning problem over the roll-in distribution. Specifically, between the teacher and student policies, as formalized in (6). In our implementation, soft value functions are approximated using estimates from trajectories. The expectations over roll-in and roll-out distributions are replaced with empirical estimates obtained from samples collected in their respective phases. The resulting loss function is given in 7 which corresponds precisely value-weighted maximum likelihood objective.
minimize the KL divergence
5 COMPARISON WITH POLICY GRADIENT METHODS
In this section, we highlight the key differences between our algorithm, VIDD, and policy gradient methods (Fan et al.) 2023 Black et al.) 2024). Broadly, there are two main distinctions: (1) policy gradient methods are inherently on-policy, whereas VIDD naturally supports off-policy updates; and (2) policy gradient methods implicitly optimize the reverse KL divergence between the fine-tuned policy and the soft-optimal policy, while our objective more closely aligns with the forward KL divergence. We elaborate each of these differences below.
On-policy vs. off-policy nature. In policy gradient algorithms, fine-tuning is typically framed as the optimization of the following objective:
J(0)= Er [r(zo) a El KL(p?_ x4) IPP |m:))
In practice, this objective is optimized using the policy gradient (PG) theorem. For instance, when 0, the gradient simplifies to the standard form (Black et al. 2024
VJ(0)= = > E (o
[r(2o) V log 1 (-<))]]
(8)
However, accurate gradient estimation requires that the roll-in distribution used to sample I; matches the current policy, making the method inherently on- on-policy. This constraint limits exploration and
increases the risk of convergence to suboptimal local minima. Even more stable variants, such as PPO, remain sensitive to hyperparameter choices (Adkins et al. 2024), and overly strong regularization- while stabilizing the landscape can significantly slow down learning (Clark et al. 2023 Fig. 3).
In contrast, VIDD naturally accommodates off-policy updates, allowing the use of more exploratory roll-in distributions without sacrificing training stability.
Reverse KL vs. forward KL divergence minimization. Furthermore, we show that the PPO objec- tive (() is equivalent to minimizing the reverse KL divergence between the trajectory distributions induced by the opttmoptimal policy and those induced by the fine-tuned policy.
Theorem 1. Denote po. E X ... X the induced joint distribution from = T to t= 0 by {p() and denote Po.r as the corresponding distribution induced by the soft-optimal policy. Then,
J(0) = KL (-1lpa = Σ Eipp (1re000---1rr)llle---11001
7

---

Published as a conference paper at ICLR 2026
Table I: Performance of different methods on both protein, DNA, and molecular generation tasks w.r.t. rewards and naturalness. The best result among fine-tuning baselines is highlighted in bold. We report the 50% quantile of the metric distribution. The + specifies the the standard error of the estimate quantile with 95% confidence interval.
Procein SS mutch
DNA Enhancer HepG2 ATAC-Acc7
Molecule Docking Parpi Dreking Searet
Method
R-sheer5kr
pI.DDI
Divessity
Pred-Activity!
3-mer Corr"
Pro-trained
0.05 - 0.08
0.37 0.09
0.91
0.14 + 0.26
0.000 0.000
.0.15
7.2 + 0.5
971 +
Best-of-N (Ne32)
0.26 = 0.13
0.38 0.11
0.90
1,30 0.64
0.000 0.000
-0.17
10.2 0.4
951 4
Standard Fine tuning
0.48 0.16
0.30 1: 0.04
0.57
1.17 1: 1.23
0.094 0.292
0.829
7.8 1: 1.8
908 4:
DDPP
0.3 0.07
0.36
5.33
0.305 01460
0,879
1.3
981 +
DDPO
0.02
0.55
0.52
7.38 0.11
0,086 0.280
0.398
1.3
929 4:
VIDD
0.83 0.01
0.82 0.01
0.52
8.28 0.18
0.820 0.384
0.162
9.4
1.7
741 + 21
In contrast, VIDD optimizes objective that more closely resembles the forward KL divergence. Since the reverse KL is known to be mode-seeking and can lead to unstable optimization landscapes (Wang& et al. 2023 Go et al. 2023 Kim et al. (022), avoiding reverse KL minimization contributes to more stable and effective fine-tuning.
EXPERIMENTS
Thus far, we have introduced VIDD, a framework designed to optimize possibly non-differentiable downstream reward functions effectively and stably in diffusion models in a sample-efficient manner. In this section, we evaluate the performance of VIDD across a range of biomolecular design tasks. We begin by describing the experimental setup.
6.1 EXPERIMENTAL SETUP
6.1.1 TASK DESCRIPTORS
We aim to fine-tune diffusion models by maximizing task-specific reward functions. In the following. we outline the choice of the pre-trained diffusion models and the formulation of the reward functions used in our biomolecular design tasks (protein, DNA, small molecule design).
Protein sequence design. We adopt EvoDiff (Alamdari et al. 2023)
pre-trained diffusion
model, a representative masked discrete diffusion model for protein sequence design, trained on the UniRef database (Suzek et al. 20077). We use EvoDiff
unconditional generative model to
produce natural protein sequences. To tackle downstream protein design tasks, we use the following
reward functions inspired by prior work (Hic et al.) 2022 Verkuil al. [002]] Lisanza al. 2024
Uchara et al. 20256 Pacesa et al., 2024). Appendix D.I provides detailed definitions of each reward function. For tasks involving secondary structure matching optimization, we employ ESMFold (Lin)
et al., 2023) to predict the 3D structures of generated sequences. For protein binder design, we leverage AlphaFold2 (Jumper et al. 202T) to model the 3D structure of the multimers. In order to encourage the naturalness and foldability of the designed proteins, we additionally incorporate structural confidence metrics, such as PLDDT and radius of gyration (Pacesa et al. 2024), as part of the reward function.
ss-match (B-sheet). This task aims to maximize the probability secondary structure (SS) matching between the generated protein sequences and predefined target pattern, computed across all residues. SS are predicted using the DSSP Kabschh & Sander 1983), and include a-helices, f-sheets, and coils. Following Pacesa al. (2024), specifically encourage the formation of 8-sheets, as protein generative models are known to exhibit a bias toward a-helices.
Binding affinity. This task focuses on designing binder proteins given target proteins, with the goal of maximizing their binding affinity. We quantify binding affinity using the ipTM score predicted by the GoohaFold-Multinng 2024), we select PD-LI and IFNAR2 as representative target proteins.
model (Jumper et al. 2021). Following prior work (Pacesa al.
8

---

Published as a conference paper at ICLR 2026
Table Performance of different methods on protein binding design tasks w.r.t. ipTM, optimized reward, and diversity. The best result is highlighted in bold.
PD-LI Rewardf
IFNAR2 Reward
Method
ipIM
Diversity?
ipTM1
Diversityf
Pre-trained
0.1468 + 0.0538
0.0847 0.1317
0.9022
0.1179 # 0.0153
0,0612 + 0.0621
0,9007
Best-of- NN (N=128)
0.2662
0.0629
0.2463
0.2225 0.0675
0.9058
Standard Fine-tuning
0.1640 +: 0.0215
0.1398 0.0351
0.8999
0.1307 0.0503
0.0926 :: 0.0712
0,9063
DDPP
0.18 + 1,0330
0.2066 + 0.0453
0.8763
0.1375 + 0.0794
0.1236 + 0.0782
0.8850
DDPO
0.7881 + 0.0250
0.8767
0.0301
0.5266
0.2403 0.0488
0.3142 + 0.0544
0.7169
VIDD
0.8182 + 0.0213
0.0237
0.5539
0.5090 0.1079
0.5120 + 0.1093
0.5176
1pTM: 0.84
IpTM: D.84
1pTM: 0.58
1pTM: 0.64
(a) PD-L1
(b) IFNAR2
Figure 2: Protein structure visualizations for the PD-LI and IFNAR2 binding design tasks. The binder protein shown green and target protein is in orange, with hotspot residues labeled on the structure.
DNA sequence design. focus on the regulatory DNA designs widely used in the cell engineer-
ing (Tskiran et al. 2024: Su al. 2025). Following Wang et al. (2024), we adopt discrete diffusion model (Sahoo et al. 2024) trained on enhancer datasets from Gosai et al. (2023) as our pre-trained use predictions from the Enformer model (Avsec et al. 2021 to estimate enhancer activity in the HepG2 cell line (denoted by Pred-Activity) This
model (T 128). For the reward function,
reward has been widely employed in prior DNA design studies (Taskiran et al. 2024 Lal al. 2024
due to its relevance in cell engineering, particularly for modulating cell differentiation.
Small molecule design. We GDSS (Jo al. 2022). trained ZINC-250K (Irwin & Shoichet
2005), as the pre-trained diffusion model (T' 1000). For rewards, we use binding affinity to protein Parpl (Yang al., (2021) (docking score (DS) calculated by Quick 2 (Alhossary 2015)), which is non-differentiable. Here, we renormalize docking score to max( DDS 0), so that higher value indicates better performance. Note these rewards have been widely used Lee al.)
(2023): Jo et al. (2022): Yang et al. (2021): et al. (2025).
6.1.2 BASELINES
We compare VIDD with the following baselines.
This is naïve yet widely adopted approach for reward maximization test time. Note that such inference-ting time methods are significantly slower compared to our fine-tuned models. Additional comparisons will be provided Appendiix H. Given that this method is N times slower than the baseline methods, we do not consider it practical baseline for fine-tuning.
Best-of-N:
Standard Fine- Tine-Tuning (SFT). This method fine-tunes the model by sampling from the pretrained model and applying the same loss used during pretraining, but with reweighting based on reward values (Peng et al. 2019).
DDPO Black et al. 2024): A PPO-style algorithm discussed in Section 5
DDPP Rector-Brooks et al. 2024): Another recent state-of-the-aar method applicable to our
setting with non-differentiable reward functions by enforcing detailed balance between reward- weighted posteriors and denoising trajectories.
VIDD: Our algorithm. Regarding detailed setting of hyperparameters, refer to Appendix D.2

---

Published as a conference paper at ICLR 2026
6.1.3 METRICS
Recall that our objective is generate samples with high desired reward in Section 2.2 Ac- cordingly, we report the median reward of generated samples as the primary evaluation metric. As secondary metrics, we include additional measures whose specific choice depends on the task context, such as the naturalness of generated samples. More specifically, we use pLDDT scores for protein design; the 3-mer Pearson correlation (3-mer Corr) between the generated sequences and those in the dataset from Gosai et al. (203); the negative log-likelihood (NLL) of the generated samples with respect to the pretrained model for small molecule design. For further results, refer to Appendix
6.2 RESULTS
All results summarized in Table and Table 2 Below, we provide an interpretation of each outcome.
55-match
Protein sequence design. VIDD consistently outper- forms baseline methods by achieving higher rewards in 3-sheet content and ipTM binding affinity. Additionally, Figure and Figure 3 illlstrrate the predicted binder-target complexes and secondary-structure matching tasks, con- firming efffective binding of the designed binders to their targets and reasonable secondary structures of the designed proteins. Other quantitative results are provided in [Ap- pendix E.4, further validating the quality and effectiveness of the generated proteins, as well as effects of varying parameters such as mixture of roll-in policy, lazy update interval and regularization coefficient.
Figure 3: Protein structure visualizations
for protein SS-match tasks.
DNA sequence design. Here, since this reward technically differentiable, we include an additional baseline, DRAKES (Wang et al. 2024), which directly backpropagates through the reward signal. Notably, as shown in the Pred-Activity column, VIDD outperforms not only DDPP and DDPO, but also DRAKES In addition to the aforementioned Pred-Activity and 3-mer Corr metrics, follow Wang et al. (2024) in incorporating ATAC-Acc an independent binary classification model trained on chromatin accessibility data from the HepG2 cell line (Consortium et al. 2012) as orthogonal reward. This is motivated by the fact that Pred-Activity is trained reward model and thus may be susceptible to overoptimization. In the ATAC-Acc column, method also exhibits strong performance, suggesting that VIDD would be robust to over-optimization. More quantitative results can be found in Appendix E.2
Small molecule design. VIDD outperforms the fine-tuning baseline methods in terms of reward. More quantitative metrics that describe naturalness in molecules is in Appendix E.3
7 CONCLUSION
In this work, we present VIDD, novel fine-tuning algorithm for diffusion models under possibly non- differentiable reward functions. Our method has the potential to accelerate discovery in areas such as protein and drug design, but we also recognize possible misuse in generating harmful biomolecular sequences. We advocate for safeguards and responsible research practices in deploying such models.
10

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGMENTS
This work was supported in part by National Institutes of Health under grant U0IAG0701 12 and ARPA-H under grant IAYIAX000053.
REFERENCES
Josh Abramson, Jonas Adler, Jack Dunger, Richard Evans, Tim Green, Alexander Pritzel, Olaf Ronneberger, Lindsay Willmore, Andrew Ballard, Joshua Bambrick, et al. Accurate structure prediction of biomolecular interactions with alphafold 3. Nature, pp. 1-3, 2024.
Jacob Adkins, Michael Bowling, and Adam White. A method for evaluating hyperparameter sensitivity in reinforcement learning. Advances in Neural Information Processing Systems, 37: 124820-124842, 2024.
Sarah Alamdari, Nitya Thakkar, Rianne van den Berg, Neil Tenenholtz, Bob Strome, Alan Moses, Alex Xijie Lu. Nicolo Fusi, Ava Pardis Amini, and Kevin K Yang. Protein generation with evolutionary diffusion: sequence is all you need. BioRxiv, pp. 2023-09, 2023.
Amr Alhossary, Stephanus Daniel Handoko. Yuguang Mu, and SheenKKong Kwoh. Fast, accurate, and reliable molecular docking with quickvina 2. Bioinformatics, 31(13):221--2216 2015.
Ziga Avsec, Vikram Agarwal, Daniel Visentin, Joseph R Ledsam, Agnieszka Trabska-Barwinnse Kyle Taylor, Yannis Assael, John Jumper, Pushmeet Kohli, and David R Kelley. Effective gene expression prediction from sequence by integrating long-range interactions. Nature methods, 18 (10):1196-122) 2021.
Arpit Bansal, Hong-Min Chu, Avi Schwarzschild, Soumyadip Sengupta, Micah Goldblum, Jonas Geiping, and Tom Goldstein. Universal guidance for diffusion models. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 843-852, 2023,
Kevin Black, Michael Janner, Yilun Du, Ilya Kostrikov, and Sergey Levine. Training diffusion models with reinforcement learning. arXiv preprint arXiv:2305. 13301, 2024.
Patrick Bryant, Gabriele Pozzati, and Ame Elofsson. Improved prediction of protein-proee protein interac- tions using alphafold2. Nature communications, 13(1): 1265, 2022.
Hanqun Cao, Haosen Shi, Chenyu Wang. Sinno Jialin Pan, and Pheng-Ann Heng. GLIDE'E: gradient-free lightweight fine-tune approach for discrete sequence design. In ICLR 2025 Workshop on Generative and Experimental Perspectives for Biomolecular Design, 2025.
Hyungjin Chung, Byeongsu Sim, Dohoon Ryu, and Jong Chul Ye. Improving diffusion models for inverse problems using manifold constraints. Advances in Neural Information Processing Systems, 2022.
Kevin Clark, Paul Vicol, Kevin Swersky, and David Fleet. Directly fine-tuning diffusion models on differentiable rewards. arXiv preprint arXiv:2309.17400 2023.
ENCODE Project Consortium et al. An integrated encyclopedia of dna elements in the human genome. Nature, 489(7414):57, 2012.
Wojciech Czarnecki, Razvan Pascanu, Simon Osindero, Siddhant Jayakumar, Grzegorz Swirszez, and Max Jaderberg. Distilling policy distillation. In The 22nd international conference artificial intelligence and statistics, pp. 1341- PMLR, 2019.
Fei Deng, Qifei Wang, Wei Wei. Tingbo Hou, and Matthias Grundmann. Prdp:
Proximal reward difference prediction for large-scale reward finetuning of diffusion models. In Proceedings of the IEFF/CVF Conference Computer Vision and Pattern Recognition, pp. 7423-7433, 2024.
Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. Advances in neural information processing systems, 34:8780-8794 2021.
11

---

Published as a conference paper at ICLR 2026
Yasha Ektefaie, Olivia Viessmann, Siddharth Narayanan, Drew Dresser, J Mark Kim, and Armen Mkrtchyan. Reinforcement learning on suucture-conditiin categorical diffusion for protein inverse folding. arXiv preprint arXi::2410.17173 2024.
Ying Fan, Olivia Watkins, Yuqing Du, Hao Liu, Moonkyung Ryu, Craig Boutilier, Picter Abbeel, Mohammad Ghavamzadch, Kangwook Lee, and Kimin Lee. DPOK: Reinforcement learning for fine-tuning text-to-image diffusion models. arXiv preprint arXiv:2305. 16381, 2023.
Dongyoung Go, Tomasz Korbak, Germán Kruszewski, Jos Rozen, Nahyeon Ryu, and Mare Dymet- man. Aligning language models with preferences through f-divergence minimization. arXiv preprint wrXiv22302.02215. 2023.
Sager J Gosai, Rodrigo Castro, Natalia Fuentes, John C Butts, Susan Kales, Ramil R Noche, Kousuke Mouri, Pardis C Sabeti, Steven K Reilly, and Ryan Tewhey. Machine-guided design of synthetic cell type-specific cis-regulatory elements. bioRxiv, 2023.
Shashank Gupta, Chaitanya Ahuja, Tsung- Yu Lin, Sreya Dutta Roy, Harrie Oosterhuis, Maarten de Rijke, and Satya Narayan Shukla. A simple and effective reinforcement learning method for 2025.
text-to-image diffusion fine-tuning, arXiv preprini
Brian Hie, Salvatore Candido, Zeming Lin, Ori Kabeli, Roshan Rao, Nikita Smetanin, Tom Sercu, and Alexander Rives. A high-level programming language for generative protein design. bioRxiv, pp. 2022-12, 2022.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-665 6851, 2020.
Fumitaka Inoue, Anat Kreimer, Tal Ashuach, Nadav Ahituv, and Nir Yosef. Identification and massively parallel characterization of regulatory elements driving neural induction. Cell stem cell, 25(5):713-727, 2019.
John J Irwin and Brian K Shoichet. ZINC- free database of commercially available compounds for virtual screening. Journal of chemical information and modeling, 45(1):177-1822 2005.
Jaehyeong Jo, Seul Lee, and Sung Ju Hwang. Score-based generative modeling of graphs via the system of stochastic differential equations. In International Conference Machine Learning, pp. 10362-10383. PMLR, 2022.
John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Židek. Anna Potapenko, al. Highly accurate protein structure prediction with alphafold. nature, 596(7873):5833555) 2021.
Wolfgang Kabsch and Christian Sander. Dictionary of protein secondary structure: pattern recognition of hydrogen-bonded and geometrical features. Biopolymers: Original Research on Biomolecules, 22(12):2577-2637 1983.
Sunwoo Kim, Minkyu Kim, and Dongmin Park. Test-time alignment of diffusion models without reward over-optimization. In The Thirteenth International Conference on Learning Representations, 2025.
Avantika Lal, David Garfield, Tommaso Biancalani, et al. regim: Designing realistic regulatory dna with autoregressive language models. bioRxiv, pp. 2024-02, 2024.
Seul Lee, Jaehyeong Jo, and Sung Ju Hwang. Exploring chemical space with score-basee out-of- distribution generation. Proceedings of the 40th International Conference on Machine Learning. 2023.
Xiner Li, Yulai Zhao, Chenyu Wang, Gabriele Scalia, Gokcen Eraslan, Surag Nair, Tommaso Bian- calani, Aviv Regev, Sergey Levine. and Masatoshi Uchara. Derivative-fret guidance in continuous and discrete diffusion models with soft value-basee decoding. arXiv preprint arXiv:2408.00 08250. 2024.
12

---

Published as a conference paper at ICLR 2026
Xiner Li, Masatoshi Uchara, Xingyu Su, Gabriele Scalia, Tommaso Biancalani, Aviv Regev, Sergey Levine, and Shuiwang Ji. Dynamic search for inference- alignment in diffusion models. arXiv preprint arXiv:2503.
2025.
Zeming Lin, Halil Akin, Roshan Rao, Brian Hie, Zhongkai Zhu, Wenting Lu, Nikita Smetanin, Robert Verkuil, Ori Kabeli, Yaniv Shmueli, al. Evolutionary-scale prediction of atomic-level protein structure with a language model. Science, (79(6637)111331111
2023.
Sidney Lyayuga Lisanza, Jacob Merle Gershon, Samuel WK Tipps, Jeremiah Nelson Sims, Lucas Amoldt, Samuel J Hendel, Miriam K Simma, Ge Liu, Muna Yase, Hongwei Wu, et al. Multistate and functional protein design using rosettafold sequence space diffusion. Nature biotechnology. pp. 1-11, 2024.
Nanye Ma, Shangyuan Tong, Haolin Jía, Hexiang Hu, Yu-Chuan Su, Mingda Zhang, Xuan Yang, Yandong Li, Tommi Jaakkola, Xuhui Jia, et al. Inference-time time scaling for diffusion models beyond scaling denoising steps. arXiv preprint
0at.:22110097222 2025.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Mare G Bellemare, Alex Graves, Martin Riedmiller, Andreas Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning, nature, 51 518(7540):529 503,-3333000 2015.
Skander Moalla, Andrea Miele, Daniil Pyatko, Razvan Pascanu, and Caglar Gulcehre. No repre- sentation, no trust: connecting representation, collapse, and trust issues in ppo. arXiv preprint arXiv:2405.00662 2024.
Hunter Nisonoff, Junhao Xiong, Stephan Allenspach, and Jennifer Listgarten. Unlocking guidance for discrete state-space diffusion and flow models. arXiv preprint arXiv:2406.01552, 2024.
Martin Pacesa, Lennart Nickel, Christian Schellhaas, Joseph Schmidt, Ekaterina Pyatova, Lucas Kissling, Patrick Barendse, Jagrity Choudhury, Srajan Kapoor, Ana Secaraz-Serna et al. Bindcraft: one-shot design of functional protein binders. bioRxiv, pp. 2024-09, 2024.
A Paszke. Pytorch: An imperative style, high-performance deep learning library. arXiv preprint 1rXiv:1912.01703, 2019.
Xue Bin Peng, Aviral Kumar, Grace Zhang, and Sergey Levine. ddvantage-weighene regression: Simple and scalable off-policy reinforcement learning. arXiv preprint arXiiv:110.0011111111 2019.
Jan Peters, Katharina Mulling, and Yasemin Altun. Relative entropy policy search. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 24, pp. 1607-1612, 2010.
Mihir Prabhudesai, Anirudh Goyal, Deepak Pathak, and Katerina Fragkiadaki. Aligning text-to-image diffusion models with reward backpropagation. arXiv preprint arXiv:2310.037:9 2023,
Jarrid Rector-Brooks, Mohsin Hasan, Zhangzhi Peng, Zachary Quinn, Chenghao Liu, Sarthak Mittal, Nouha Dziri, Michael Bronstein, Yoshua Bengio. Pranam Chatterjee, et al. Steering masked discrete diffusion models via discrete denoising posterior prediction. arXiv preprint arXiv:2410.08133 2024.
Subham Sekhar Sahoo, Marianne Arriola, Yair Schiff, Aaron Gokaslan, Edgar Marroquin, Justin Chiu, Alexander Rush, and Volodymyr Kuleshov. Simple and effective masked diffusion language models. arXiv preprint arXiv:2406.07224 2024.
Anirban Sarkar, Ziqi Tang, Chris Zhao, and Peter Koo. Designing dna with tunable regulatory activity using discrete diffusion. bioRxiv, pp. 2024-05, 2024.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXivv1707.06347, 2017.
Raghav Singhal, Zachary Horvitz, Ryan Techan, Mengye Ren, Zhou Yu, Kathleen McKeown, and Rajesh Ranganath. A general framework for inference-time scaling and steering of diffusion models. arXiv preprint arXiv:2501.0668 2025.
13

---

Published as a conference paper at ICLR 2026
Jascha
Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. International conference on machine learning. pp. PMLR, 2015.
Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. arXiv preprint orXiv:2010.02502 2020.
Xingyu Su, Haiyang Yu, Degui Zhi, and Shuiwang Ji. Learning to discover regulatory elements for gene expression prediction. arXiv preprint arXiv:2502.13991 2025.
Baris E Suzek, Hongzhan Huang, Peter McGarvey, Raja Mazumder, and Cathy H Wu. Uniref: comprehensive and non-redundant uniprot reference clusters. Bioinformatics, 23(10): 2007.
Sophia Tang, Yinuo Zhang, and Pranam Chatterjee. Peptune: De novo generation of therapeutic peptides with multi-objective-gudded discrete diffusion. ArXiv, pp. arXiv-2412, 2025.
Ibrahim I Taskiran, Katina Spanier, Hannah Dickmänken, Niklas Kempynck, Alexandra Pančíková, Eren Can Ekşi, Gert Hulselmans, Joy N Ismail, Koen Theunis, Roel Vandepoel, et al. Cell-type- directed design of synthetic enhancers. Nature, 626(7997):112-220, 2024.
Oleg Trott and Arthur J Olson. Autodock vina: improving the speed and accuracy of docking with new scoring function, efficient optimization, and multithreading. Journal of computational chemistry, 31(2):455-461 2010.
Masatoshi Uchara, Yulai Zhao, Chenyu Wang. Xiner Li, Aviv Regev, Sergey Levine, and Tommaso Biancalani. Inference-time alignment in diffusion models with reward-guided generation: Tutorial and review. arXiv preprint arXiv:2501.00 009685 2025a
Masatoshi Uchara, Yulai Zhao, Chenyu Wang. Xiner Li, Aviv Regev, Sergey Levine, and Tommaso Biancalani. neward-guided controlled generation for inference-time alignment in diffusion models: Tutorial and review. arXiv preprint arXiv:2501.0960 006655, 2025b.
Siddarth Venkatraman, Moksh Jain, Luca Scimeca, Minsu Kim, Marcin Sendera, Mohsin Hasan, Luke Rowe, Sarthak Mittal, Pablo Lemos, Emmanuel Bengio, et al. Amortizing intractable inference in diffusion models for vision, language, and control. arXiv preprint arXiv:2405.20971 2024.
Robert Verkuil, Ori Kabeli, Yilun Du, Basile IM Wicky, Lukas F Milles, Justas Dauparas, David Baker, Sergey Ovehinnikov, Tom Sercu, and Alexander Rives. Language models generalize beyond natural proteins. BioRxiv, pp. 12, 2022.
Yannick Vogt, Mehdi Naouar, Maria Kalweit, Christoph Cornelius Miething, Justus Duyster, Joschka Boedecker, and Gabriel Kalweit. Betterbodies: Reinforcement learning guided diffusion for antibody sequence design. arXiv preprint arXiv:2409. /6298, 2024.
Chaoqi Wang. Yibo Jiang, Chenghao Yang, Han Liu, and Yuxin Chen. Beyond reverse kl: Gen- eralizing direct preference optimization with diverse divergence constraints. arXiv preprint arXiv:2309. 162400 2023.
Chenyu Wang, Masatoshi Uchara, Yichun He, Amy Wang, Tommaso Biancalani, Avantika Lal, Tommi Jaakkola, Sergey Levine, Hanchen Wang, and Aviv Regev. Fine-tuning discrete diffusion models via reward optimization with applications to dna and protein design. arXiv preprint vrXiv:2410.13643 2024.
Joseph L Watson, David Juergens, Nathaniel R Bennett, Brian L Trippe, Jason Yim, Helen E Eisenach, Woody Aher, Andrew Borst, Robert J Ragotte, Lukas F Milles, et al. De novo design of protein structure and function with rídiffusion. Nature, (0977976)100891110
2023.
Luhuan Wu, Brian Trippe, Christian Naesseth, David Blei, and John P Cunningham. Practical and asymptotically exact conditional sampling in diffusion models. Advances in Neural Information Processing Systems, 36:31372-31403, 2023.
14

---

Published as a conference paper at ICLR 2026
Xiaoshi Wu, Yiming Hao, Manyuan Zhang, Keqiang Sun, Zhaoyang Huang, Guanglu Song, Yu Liu, and Hongsheng Li. Deep reward supervisions for tuning text-to-image diffusion models. In European Conference on Computer Vision, pp. 108-124. Springer, 2024.
Soojung Yang, Doyeong Hwang, Seul Lee, Seongok Ryu, and Sung Ju Hwang. Hit and lead discovery with explorative rl and fragment-based molecule generation. Advances in Neural Information Processing Systems, 34:7924-79333 2021.
Rui Yuan, Robert M Gower, and Alessandro Lazaric. A general sample complexity analysis of vanilla policy gradient. In International Conference Artificial Intelligence and Statistics, pp. 3332-3380. PMLR, 2022.
Oussama Zekri and Nicolas Boullé, Fine-tuning discrete diffusion models with policy gradient methods. arXiv preprint arXiv:2502.01384, 2025.
Yinan Zhang, Eric Tzeng, Yilun Du, and Dmitry Kislyuk. Large-scale reinforcement learning for diffusion models. In European Conference on Computer Vision, pp. 1-17. Springer, 2024.
15

---

Published as a conference paper at ICLR 2026
A PROOF OF THEOREM CONNECTING PPO OBJECTIVE WITH REVERSE KL DIVERGENCE
We restate the formal version of Theorem
1 The statement is ((0) is equal to
Σ B (o0) ( x.))]
up to some constant. We will provide its proof in this section.
Now, we prove this statement. We start calculating the inverse KI. divergence. This
Σ
KL (12)
T=
p. (-20) (x)((11())))
Σ (=T+1
KL pl (111(1)0)))1
exp(v((Fe)/a)
((()))
0z-1 (2
E (01_,(01-1100)))
log pegri-1 pre
a
[=T+1
Bi (0 [-r(xo) + aKL. ( + 1=
Σ
Here, corresponds to constant (+) where is defined by
a log Exo (e0
recalling its definition
B
CONNECTING VIDD OBJECTIVE WITH KL DIVERGENCE
In this section, we plan to explain how derive 4 Recall that the original objective is
argmin Ez,on, (<)) (
Based on the definition of conditional KL divergence
[L(p(- 1x)) |I g( æx)) = En, 1~p(- [log((x.-1 xx)) log ((x x.)].
and the definition of soft optimal policy p² pT. = pre (-12)) exp (-)/-)
oxp(vx(x:)/aa
objective can be rewritten
argmin E En Man 1 s-pt (n. 1122) log p._ (m xx) - log pi (2 m1)
arrgmin Σ --
(x
Σ Exoue exp(ve((t)/a) E tot log pf
Hence, ignoring the first term since this is constant, the objective function reduces to
argmax
Σ Byycony
Ea Fa i (me- (ee (xp(cx(2t)/a)
(ex) (exp(1) (otoo//w-((tt///// log log (ZI- (
thus obtain
16

---

Published as a conference paper at ICLR 2026
C BROADER IMPACT AND LIMITATIONS
C.I BROADER IMPACT
This paper presents work whose goal to advance the field of Deep Learning, particularly diffusion models. While this research primarily contributes to technical advancements in generative modeling, has potential implications in domains such as drug discovery and biomolecular engineering. We acknowledge that generative models, particularly those optimized for specific reward functions, could be misused if not carefully applied. However, our work is intended for general applications, and emphasize the importance of responsible deployment and alignment with ethical guidelines in generative AL. Overall, our contributions align with the broader goal of machine learning methodolo- gies, and do not foresee any immediate ethical concerns beyond those generally associated with generative models.
C.2 LIMITATIONS AND FUTURE WORKS
The success of reward-guider fine-tuning such reinforcement learning critically depends on the quality of the reward signal. However, in reality, reward functions are often imperfect: they may reflect proxy objectives that only loosely correlated with real-world biological chemical utility, Poorly designed rewards can lead to over-optimization and exploit the reward without producing truly meaningful or functional outputs.
Furthermore, post-training is often sensitive to the precision of the reward-when the reward signal is noisy misaligned, learning become unstable entirely fail. In practice, reward evaluation also be expensive (e.g.. involving structure prediction or simulation), making it essential to design reward mechanisms that are not only accurate but also sample-efficient, enabling effective training under limited reward budget.
Finally, since reward signals are imperfect and often expensive to evaluate, hyperparameters that regulate reward exploitation such as rollout frequency, may require some tuning in practice. Overly aggressive configurations can amplify reward noise, while overly conservative choices may cap achievable gains, reflecting practical tradeoff inherent to reward-driven optimization. Designing fully adaptive schedules for these parameters would require nomain--spciif assumptions about reward reliability and exploration metrics, which falls outside the scope of this work but represents promising engineering extension built on top our framework.
D ADDITIONAL DETAILS FOR EXPERIMENT SETTING
D.1 DETAILS ON TASKS AND REWARD FUNCTIONS
Here we present the detailed descriptions of the task settings and reward functions used in the experiments.
SS-match. The secondary structure matching steers the energy toward user-defined secondary structure. To annotate residue secondary structure, we use the DSSP algorithm (Kabsch & Sander 1983), which identifies elements such as a-helices, 8-sheets and coils. This reward function returns the fraction of residues assigned to the desired secondary structure element. In case, aim to maximize the fraction of residues forming 8-sheets, thereby encouraging structures with higher 8-sheet content (Pacesa et al. 2024
PLDDT. PLDDT (predicted Local Distance Difference Test) is per-residue confidence score produced by structure prediction models such as AlphaFold and ESMFold. It estimates the local accuracy of predicted atomic positions, with higher scores indicating greater confidence. In protein generation tasks, pLDDT is widely used to evaluate the structural reliability of predicted models, as higher average PLDDT values correlate with well-formed and accurate local geometry. In setting, we use the average PLDDT score moore-computed via ESMFold (Lin et al. 2023) as a reward signal to encourage the generation of structurally confident protein designs.
17

---

Published as a conference paper at ICLR 2026
Table 3: Hyperparameter values for the DNA generation task.
Hyperparameters
Values
Value weight coefficient a
1.0
Sequence length
200
# of decoding steps
128
Update interval K # of training steps Batch size Learning rate Reward
5
2000
32
le - 4
Pred-Activity
Protein Binding Affinity. We aim to generate binder proteins conditioned on given target proteins. To optimize binding affinity, we maximize the ipTM score, a widely used metric calculated by the AlphaFold2-Multima model (Jumper et al. 202T Pacesa et al. 2024). To further encourage the naturalness and foldability of the generated binders, we incorporate additional structural confidence terms into the reward function: Reward = ipTM + 0.1 pLDDT + 0.02 x radius. The weights are chosen ensure that ipTM remains the primary optimization objective and is not overwhelmed by the other, more easily optimized components. The pLDDT is computed from the predicted multimer complex using AlphaFold2, while the radius of gyration encourages the formation of well-folded, globular structures Paacesa et al. 2024).
Molecule Binding Affinity. We use the docking program Quick 2 (Alhossary et al.) 2015 to compute the docking scores following Yang et al. (2021). with exhaustiveness as 1. Note that the docking scores are initially negative values, while we reverse it to be positive and then clip the values to be above 0, ie.. We compute DS regarding protein parpl (Poly [ADP-ribose] polymerase- 1), which is target protein that has the highest AUROC scores of protein-ligand binding affinities for DUD-E ligands approximated with AutoDock Vina.
Enhancer HepG2. We examine publicly available large dataset on enhancers (n 700k) (Gosai et al. 2023), with activity levels measured by massively parallel reporter assays (MPRA) (Inoue et al. 2019). where the expression driven by each sequence is measured. In the Enhancers dataset, each is a DNA sequence of length 200. We pretrain the masked discrete diffusion model (Sahoo et al. 2024) on all the sequences. We then split the dataset and train two reward oracles (one for finetuning and one for evaluation) each subset. Each reward oracle is learned using the Enformer architecture (Avsec al. (02T), while y E R is the measured activity in the HepG2 cell line. The Enformer trunk has convolutional layers, each having 1536 channels. well 11 transformer layers, with 8 attention heads and key length of 64. Dropout regularization applied across the attention mechanism, with attention dropout rate of 0.05, positional dropout of 0.01, and feedforward dropout of 0.4. The convolutional head for final prediction has 2* 1536 input channels and uses average pooling, without activation function. We learn two These datasets and reward models widely used in the literature on computational enhancer design (Lal et al. 2024
Sarkar et al. 2024).
D.2 HYPERPARAMETERS
Here we present the hyperparameters in Table 3 Table 4 and Table 5
D.3 SOFTWARE AND HARDWARE
Our implementation is under the architecture of PyTorch (Paszke 2019). The deployment environ- ments are Ubuntu 20.04 with 48 Intel(R) Xeon(R) Silver, 4214R CPU 2.40GHz, 755GB RAM, and graphics cards NVIDIA RTX 2080Ti. All experiments conducted on a single GPU, selected from NVIDIA RTX 2080Ti, RTX A6000, or NVIDIA H100 with 80GB HBM3 memory, depending on the scale of the task.
18

---

Published as a conference paper at ICLR 2026
Table Hyperparameter values for the protein generation task.
Hyperparameters
Values
Value weight coefficient a
1.0
Sequence length
256 (ss-match), 100 (binding affinity)
of decoding tokens each step
4
Update interval K # of training steps
5 (binder IFNAR2), 50 (ss-match, binder PD-LI)
10000
Batch size Learning rate Reward (ss-match)
32
le -
ss-match
Reward (Binding Affinity)
ipTM + 0.1 x PLDDT + 0.02 x radius
Table 5: Hyperparameter values for the molecule generation task.
Hyperparameters
Values
Value weight coefficient a
6.0
Maximum atom number
38
# of decoding steps
1000
Update interval K
20
# of training steps Batch size Learning rate Reward
1000
1024
5e 6
Docking Score
Table Final performance combined with inference-time techniques.
Protein SS-match
DNA Enhancer HcpG2
Molocule Docking - Parpl Docking Scoret
Method
A-shoet5*
PLDDTT
Pred-Activity
ATAC-Ace?
3-mer Corr?
NLL,
VIDD
0.83 = 0.01
0.82 0.01
8.28 : 0.18
0,820 + 0,384
0.162
1.7
741 21
VIDD BON (N=32)
0.84
0.82 0.01
0.07
0.750 0.433
0.152
1.0
D.4 LICENSES
The dataset for molecular tasks under Database Contents License (DbCL) v1.0 The pretrained protein generation model EvoDiff is under MIT License, The dataset for DNA task covered under AGPL-3.0 license
We follow the regulations for all licenses.
E ADDITIONAL EXPERIMENT RESULTS
E. CONNECTING VIDD WITH INFERENCE- TIME TECHNIQUES
inference-time techniques are orthogonal to our approach. Here we clarify the connection between VIDD and inference-time methods. Among various inference-time schemes, the most relevant VIDD SVDD (Li Etietal al. 2024). Both methods approximate the value function using
As discussed in Section 1.1
(x())) r()o(zz)),
but their usage differs fundamentally: SVDD leverages this estimated value guide the choice of the next denoising step during sampling, whereas VIDD fine-tunes the diffusion model based on the estimated value, leading to learnable and reusable policy rather than purely inference-time adjustment.
Furthermore, inference-tima techniques can be combined with VIDD to further enhance performances. Table and Table presents the results of applying VIDD with Best-of-N sampling, demonstrating that additional performance gains be achieved through this combination.
19

---

Published as a conference paper at ICLR 2026
Table 7: Final performance combined with inference-time techniques for protein binder design.
IFNAR2
Method
TpTM
PLDDTT
Radius]
Diversity?
ipTMT
PLDDT1
Radius]
Diversity
VIDD
0.82
0.04
-0.12
0.11
2.20
VIDD - BoN (N=128)
0.84
0.00
0.91 = 0.01
-0.16 = 0.06
0.39
0.66 # 0.02
0.55 + 0.04
1.04 = 0.45
0.43
Table 8: Performance of different methods on DNA generation tasks. The best result among fine- tuning baselines is highlighted in bold, and the second best result highlighted in underline.
Method
Pred-Activity
ATAC-Aecf
3-mer Corr
Log-Lik3
Pre-trained
0.14 1 0.26
0.000
0.000
-0.15
-245
Best-of-N (N=32)
1.30 + 0.64
0.000 0.000
-0.17
-240 7.2
DRAKES w/o KL
6,44 1. 0.04
0.825
0.028
0.307
-281 1. 0.6
DRAKES
5.61 + 0.07
0.925 + 0.006
0.887
-264 + 0.6
Standard Fine-tuning
1.17 + 1,23
0.094 0.292
0.829
-263 9.3
DDPP
5.33 + 0.94
0.305 = 0.460
0.879
-218 + 10.4
DDPO
7.38 + 0.11
0.086 + 0.280
0.398
-126 + 9.5
GLIDEE
7.35 + 0.07
0.906 0.003
0.490
-240 + 14.2
VIDD
8.28 + 0.18
0.820 0.384
0.162
-198 8.6
E.2 FURTHER RESULTS FOR DNA GENERATION
More metrics on DNA generation tasks. Here we provide additional evaluation results for the DNA generation task in Table 8 to offer a more comprehensive comparison across a broader set of metrics. We add
baseline GLID'E (Cao et al. 2025) here for comparison. Note that optimization target solely the Pred-Activity. Following the prior work (Wang et al., 2024), we fine-tune using Pred-Activity reward model and evaluate using different leakage. In terms of performance, VIDD attains strong gains Pred-Activity, as well as competitive results ATAC-Acc and Log-Likelihood. Although its 3-mer correlation diverges from that of the pretrained model distribution, the substantially higher functional rewards indicate that VIDD is able discover novel yet highly effective sequences. This suggests that VIDD explores regions beyond conventional motif statistics while still generating functionally superior designs.
avoid data
Regarding overall performance, DRAKES benefits from explicit gradient information as discussed in Section 6.2 Gradient-bee based methods are expected to have better performances because they rely on precise token-level gradients, while methods designed for dineeffferettable rewards must operate using only sequence-level reward signals. Despite this disadvantage, VIDD still achieves superior performance on the optimized Prev-Activity objective, demonstrating its effectiveness. We include DRAKES in the comparison to help readers better understand the level of results VIDD can achieve.
Lazy update interval We study the effect of the lazy update interval K on the DNA sequence design discussed in /Sectinn 4.2] and other parameters are provided in Table As shown in Table 9 performance is different under different K and does not improve monotonically with more frequent updates. These results suggest that less updates roll-out policy can stabilize training and improve optimization.
Regularization coefficient We study the effect of the regularization coefficient a as shown in (7). Hyperparameter details for the remaining settings are listed in T In the DNA sequence design task, the performance comparison is presented in Table I0 show that a = 1.0, i... keeping the reward distribution unchanged, yields the best performance.
E.3 FURTHER RESULTS FOR MOLECULE GENERATION
First, report the diversity comparisons across methods in Table II on page 22
To evaluate the validity our method in molecule generation, we further report several key metrics that capture different aspects of molecule quality and diversity in Table T2 on page 22
20

---

Published as a conference paper at ICLR 2026
Table 9: The influence of lazy update interval K the performances of DNA sequence design.
Lazy Update Interval
Pred-Activityt
ATAC-Aco?
3-mer Conr"
Log-Lik
7.06 + 0.35
0.000 + 0.000
0.211
-156 + 12.0
5
8.28 + 0.18
0.820 = 0.384
0.162
-198 = 8,6
10
7.76 + 0.32
0.047 0.211
0.457
-265 5.5
20
7.71 + 0.37
0.086 0.280
0.398
-265 5.3
50
7.23 + 0.42
0.484 0.500
0.470
-266 7.2
Table 10: The influence regularization coefficient a on the performances DNA sequence design.
Regularization Coefficient
Pred-Activityy
ATAC-AccT
3-mer Corrt
Log-Likt
0.8
8.21 + 0.24
0.820 + 0.384
0.272
-246 + 6,2
1.0
8,28 + 0.18
0.820 + 0.384
0.162
-198 + 8.6
2.0
7.26 0.39
0.977 1 0.151
0.351
-248 8.0
The validity of molecule indicates its adherence to chemical rules, defined by whether it can be successfully converted to SMILES strings by RDKit. Uniqueness refers to the proportion of generated molecules that are distinct by SMILES string. Novelty measures the percentage of the generated molecules that are not present in the training set. Fréchet ChemNet Distance (FCD) measures the similarity between the generated molecules and the test set. The Similarity to Nearest Neighbors (SNN) metric evaluates how similar the generated molecules are to their nearest neighbors in the test set. Fragment similarity measures the similarity of molecular fragments between generated molecules and the test set. Scaffold similarity assesses the resemblance of the molecular scaffolds in the generated set to those in the test set. The neighborhood subgraph pairwise distance kernel Maximum Mean Discrepancy (NSPDK MMD) quantifies the difference in the distribution of graph substructures between generated molecules and the test set considering node and edge features. Atom stability measures the percentage of atoms with correct bond valencies. Molecule stability measures the fraction of generated molecules that are chemically stable, i.e., whose all atoms have correct bond valencies. Specifically, atom and molecule stability are calculated using conformers generated by RDKit and optimized with UFF (Universal Force Field) and MMFF (Merck Molecular Force Field).
We compare the metrics using 512 molecules generated from the pre-trained GDSS model and from different methods, as shown in Table 12 on page 22 Overall, our method achieves comparable performances with the pre-trained model on all metrics, maintaining high validity, novelty, and uniqueness while outperforming on several metrics such as FCD, SNN, and NSPDK MMD. Pre- trained performs consistently well across all metrics, particularly in SNN and atomic stability. However, it does not optimize specific molecular properties as effectively as the other methods. DDPP performs poorly in scaffold similarity and NSPDK MMD, indicating that generates unrealistic molecules. These results indicate that our approach can generate diverse set of novel molecules that are chemically plausible and relevant.
E.4 FURTHER RESULTS FOR PROTEIN GENERATION
More metrics on protein binder design Tables 13 and Table 14 report the full evaluation metrics
for the protein binding affinity design tasks. Here Reward ipTM+ + 0.1 pLDDT + 0.02 radius, indicate the aggregated metrics. While ipTM serves as the primary metric for binding affinity. pLDDT and radius of gyration are included as secondary metries encourage the generation of well-folded, globular binder proteins.
Additional evaluation metrics are reported Table 15 where p'TM values are obtained directly from SophaFold2-Multime (Jumper et al. (021) and pDockQ scores follow Bryant et al. (002).. From the table, we observe that our proposed VIDD consistently achieves the best performance among the main baselines. Note pTM and pDockQ are not directly optimized during the fine-tur tunintnn but only used for evaluation. Protein binding affinity can assessed using wide variety of metrics, raising the question of how effectively integrate multiple objectives. Developing principled multi objectibeeeeveee fine-tuning framework therefore remains an interesting direction for future work explore.
We present results for both pTTM and pDockQ.
21

---

Published as a conference paper at ICLR 2026
Table 11: Performance of different methods on molecular generation task w.r.t. reward, NLL, and diversity.
Binding Affinity - Parpl
Method
Docking Soore?
NLL1
Diversityf
Best-of-N (N-32)
10.2 0.4
951 +:22
0.7938
0,3052
Standard Fine-tuning
18
908 77
0.8787 0,1088
DDPP
0.0845
DDPO
8.5 =: 1.3
929 43
0.8993
0.0567
VIDI
94 1.7
741 21
0,9019 0,4777
VIDD
12.1
Table 12: Comparison of the generated molecules across various metrics. The best values for each metric are highlighted in bold.
Unique"
Novelty1
Frag
Mol Stable
Pre-trined
1.000
1.000
1.000
12.979
0.:14
0.513
1,000
0.038
0.320
0.917
DPS
1,000
1.000
1,000
13.230
0.389
0.388
1,000
0,040
0.310
0.578
22.710
SVDD
1.000
L,000
1,000
12.278
0.428
0.622
1,000
0.052
0,478
0,910
Standard Fine-tuning
1.000
1,000
1,000
9.698
0.4855
0.597
1,000
0.026
0.375
0.935
DDPP
DDPO
1.000
1.000
1,000
11.506
0.333
0.484
0,975
0.032
0.581
0.961
VIDD
1,000
1,000
1,000
4.869
0.489
0.714
0.999
0.016
0.458
0.911
Table 13: Performance different methods on protein binding design tasks on target protein PD-L1.
Method
Reward1
ipTMM
PLDDTT
Radius.]
Diversity
Pre-trained
0.0847 0.1317
0.1468 0.0538
0.3284
0.0420
4.7420 + 5.2352
0.9022
Best-of-N (N=128)
0.2654 0.0629
0.2662 0.10
0.3890
0.0530
1.9883 + 3.4164
0.8996
Standard Fine-tuning
0. + 0.0351
0.1640 + 0.0215
0.3349
0.0328
1.8829
1.1945
0.8999
DDPP
0.2065 0.0453
0.1889 = 0.0330
0.3720
0.0269
+ 0.9780
0.8763
DDPO
0.8767 0.0301
0.7881 = 0.0250
0.8244
0.0821
-0.3081
0,1285
0.5266
VIDD
0.9079 + 0.0237
0.8182 + 0.0213
0.8720 + 0.0421
-0.1232 + 0.1066
0.5539
Table 14: Performance of different methods on protein binding design tasks on target protein IFNAR2.
Method
Reward
ipTM
pL.DDT
Radius]
Diversity?
Pre-trained
0.0612 + 0.0621
0.1179 = 0.0153
0.3525 + 0.0513
4.5964 + 2.9320
0.9007
Best-of-N (N=128)
0.2225 1 0.0675
0.2463 1 0.1055
0.3843 0.0515
3.1082 1: 2.9186
0.9058
Standard Fine-tuning
0.0926 + 0.0712
0.1307 = 0.0503
0.3321 + 0.0381
3.5632 +: 2.1944
0.9063
DDPP
0.1236 + 0.0782
0.1375 = 0.0794
0.3624 + 0.0359
2.5107 + 1.1723
0.8850
DDPO
0.3142 1 0.0544
0.2403 1 0.0488
0.6300 1 0.0743
-0.5435 1 0.1219
0.7169
VIDD
0.5120 0.1093
0.5090 0.1079
0.4711 0.0490
2.2039 +: 1.9989
0.5176
Table 15: Additional evaluation metrics for protein binder design. Note that those metrics are not optimized during training and only used for evaluation.
PD-LI
IFNAR2
Method
PTM
DDockQt
PTM
DDockQt
DDPP
0.5537 = 0.0179
0.0693
0.0214
0.4697 = 0.0084
0.0359 + 0.0082
DDPO
0.7862 = 0.0809
0.2445 + 0.0478
0.5064 0.0178
0.0516 = 0.0190
VIDD
0.8369 + 0.0300
0.4164 0.0460
0.5885 0.0368
0.1241 + 0.0352
Mixed roll-in policy Section 4.1 we describe a mixed roll-in strategy that samples trajectories from the pre-trained policy p. and the roll-out policy pout with probabilities 1 3. and Bg. respectively. Figure 4/ablates 3, and find that injecting a non-zero fraction of pf"h (i.e., not always rolling in from pe improves diversity and often yields better overall performance. The specific
22

---

0.8238 Published as a conference paper at ICLR 2026
1.0
B. = 0.8
8. 1.0
Performance (IpTM Reward)
0.8
0.6
0.4
0.2
0.0
0.0
0.2
0.4
0.6
0.8
1.0
Diversity
Figure 4: Performance vs. diversity in PD-L1 binder design under different roll-in mixtures. Mixing in the pre-trained policy during roll-in (smaller 3,) increases diversity compared to relying solely on the roll-out policy (3.=1).
Table 16: The influence of lazy update interval K on the performances on protein sequence design for ss-match task.
Lazy Update Interval K
B-sheet%f
PLDDTT
Diversity
0.7972 0.0323
0.6745 + 0.0643
5
0.8914 + 0.0155
0.6196 + 0.0263
0.5023
50
0.8281 + 0.0098
0.8202 + 0.0118
0,5154
V120
3:00
2013
200
10:8
100
(a) SS-match
(b) PD-LI
(c) IFNAR2
Figure 5: Training curves of different methods SS-match, PD-LI, and IFNAR2 binder design tasks. The y-axis shows the optimized reward, and the r-axis shows training steps.
optimal mixture is task-dependent; accordingly, treat 8, tunable hyperparameter selected by validation to balance exploration (p,") and exploitation (pont).
Lazy update interval For the hyperparameter of lazy update interval K discussed in Section 4.2 we present the results Table 16 Table 17 and Table 18 for the protein sequence design task, and other parameters provided in Table
Regularization coefficient For the regularization coefficient a in 7 the performance comparison for protein sequence design tasks is presented in Table [9 and Table 20 We could notice that similar as DNA sequence design, keep the reward distribution unchanged (a = 1.0) yields the best performances.
23

---

Published as a conference paper at ICLR 2026
Table 17: The influence lazy update interval K the performances on protein binder design tasks for target protein PD-LI.
Lazy Update Interval A
Rewardf
ipTMt
PLDDTT
Radius
Diversity!
0.4336 + 0.1214
0.4090 0.1208
0.4174
0.0304
0.8551 + 0.7798
0.5048
0.6983 0.1195
0.6428
0.1116
0.6211
0.0814
0.3270 0.2587
0.5140
10
0.5537 = 0.0490
0.5073
0.0458
0.5606 0.0374
0.4791 0.2324
0,5062
20
0,1902 1 0.0880
0.2327 1 0.0804
0.4367 0.0678
4.3089 1 0.9180
0.5033
50
0.9079 + 0.0237
0.8182 + 0.0213
0.8720 0.0421
-0.1232 = 0.1066
0.5539
Table 18: The influence lazy update interval on the performances on protein binder design tasks for target protein IFNAR2.
Lazy Update Interval K
Rewardf
ip'MM
PLDDTT
Radius]
Diversity
0.1702 = 0.0311
0.1433 + 0.0099
0.4195 = 0.0321
0.7567 + 1.3396
0.7454
0,5120 1 0.1093
0,5090 1 0.1079
0.4711 1 0.0490
2.2039 1 1.9989
0.5176
10
0,2305 + 0.0252
0.1955 + 0.0167
0.4517 + 0.0302
0.5062 + 0.5819
0.6052
20
0.1528 + 0.0352
0.1266 + 0.0270
0.3809 = 0.0421
0.5971 + 0.6893
50
0.1227 = 0.0231
0.1160 + 0.0108
0.3618 0.0393
1.4747 + 0.8167
0.8717
VIDO
VIDO
DO00
DBPO
1009
DDFF
00
0.4
0:7
02
Divenity
Diverity
Diverdly
(a) SS-match
(b) PD-LI
(c) IFNAR2
Figure 6: Scatter plots of reward versus diversity for different methods on SS-match, PD-LI, and IFNAR2 binder design tasks. Each point corresponds to training checkpoint.
Trade-off between Reward and Diversity Readers may be concerned that higher rewards could come at the cost of reduced diversity from results in Table and Table (as well as Table and Table 14. To clarify this, we include the training curves (Figure 5) and the reward-diversity scatter plots Figure 6). The curves show how reward evolves over training, while the scatter plots visualize the reachable regions in the reward diverrstty plane for each method. These figures illustrate that only our method is able to expand the reachable set and achieve substantially higher rewards across training.
F SOFT VALUE FUNCTION
POSTERIOR MEAN APPROXIMATION
We use 33 to define the soft value function:
r(zo) a
vx(x)) = log IE (xoop))))))
exp
Ft
Recall that the training objective diffusion models (Section 2.1) is to accurately recover the clean sample Zp from noisy state In practice, this means that the diffusion model aims produce reliable estimate to (2;) of the posterior mean E[zo|xxy] Following prior work (Uchara et al. 2025a Li et al. 2024), substituting this estimate into the expectation above yields the commonly used soft value approximation:
(x(2)) r(fo(x))),
which is often referred to as the posterior mean approximation, In our implementation, we obtain (o())) via argmax decoding from the distribution p(zo xx) at & 0.
24

---

Published as a conference paper at ICLR 2026
Table 19: The influence of regularization coefficient a on the performances on protein binder design tasks for target protein PD-LI.
Regularization Coeflicient
Rewardf
ip
PLDDTt
Radius]
Diversilyf
0.8
0.5505 = 0,0809
0,4839 + 0.0739
0.5430
0.0565
-0.6149 + 0.1902
0,5592
0.9079 = 0.0237
0.8182 + 0.0213
0.8720 = 0.0421
-0.1232 + 0016
0.5539
2
0.4443 = 0.0854
0.3930 + 0.0812
0.4752 = 0.0552
-0.1865 + 0.3129
0.5038
Table 20: The influence of regularization coefficient a on the performances on protein binder design tasks for target protein IFNAR2.
Regularization Coeflicient
Rewardf
ipTMt
PLDDT
Radius
Diversilyf
0.8
0.2729 0.0844
0.2472 0.0733
0.4013
0.0448
0.7226 + 0.6606
0,5017
0,5120 = 0.1093
0.5090 0.1079
0.4711 = 0,0490
2.2039 = 1.9989
0.5176
2
0.1005 = 0.0480
0.1247 + 0.0455
0.3587 + 0.0329
3.0035 + 1.3196
0.8769
F,2
MONTE CARLO ESTIMATION
Beyond the posterior mean, may also approximate the soft value using Monte Carlo sampling, which computes multiple predictions of zu and averages their rewards:
M
(x+) 2
Σ
r
(amm)) (afm)) (x:)).
a(m)) (xx) ~ p(ro xz),
M
where we sample from p(z0 xx)) multiple times to get different =
(zt) by temperature sampling. Temperature sampling is used because argmax decoding above produces a single deterministic mode and therefore cannot support multiple samples for Monte Carlo estimation.
Table 21 and Table 22 reports the effect of different soft value estimation functions in our setting, Specificall, Posterior mean represents take the argmax decoding to get Hp from p(zo 2z). The interesting observation is that even when we use Monte Carlo estimation with M = 4, which increases reward computation cost by (and reward calculation itself expensive using models parameters), the performance becomes worse. This
like ESMFold-3B or AlphaFold2-MMulinnee suggests that the posterior mean (argmax) provides more stable and discriminative value signal for estimating (22)). In contrast, small-M Monte Carlo samples mix high- and low-quality draws, weakening the training guidance. Increasing M further could reduce this variance, but it would require substantially more reward calls, making it computationally impractical in biomolecular design, where reward evaluation is both high-cost and non-regular.
Therefore, using argmax decoding approximate the value offers the best trade-off: it minimizes reward computation while providing a stable and reliable signal for fine-tuning, making it the most practical choice for VIDD.
F.3 TRAINING VALUE NETWORK
Another alternative is to train a separate neural network to approximate the soft value function directly. Although such critic models are common in standard reinforcement learning. they are considerably less practical in biomolecular design. Reward functions in our domains (c.g.. protein binder design, secondary-structun
matching design) rely on large structure prediction models such as AlphaFold2-Multimer (Jumper et al., (022)) (93.2M parameters) or ESMFold (Lin et al.) 2023) (3B parameters), which operate on full sequences rather than individual denoising steps. In contrast, the diffusion generator used in our experiments (EvoDiff (Alamdari al., 2023) contains only 38M parameters. Training an accurate value network over intermediate states H; instead of the final state extremely challenging and computationally extensive. Consequently, mean and Monte Carlo estimators in this work, and leave the training of scalable and accurate value networks for complex biomolecular tasks interesting future direction.
focus the posterior
25

---

Published as a conference paper at ICLR 2026
Table 21: The influences of soft value function on the performances on protein binder design tasks for target protein PD-LI.
Reward"
ipTMt
PLDDT
Radius
Diversity
Posterior mean
0.9079 +: 0.0237
0.8182 = 0.0213
0.8720
0.0421
-0.1232 0.1066
0.5539 0,5155
Monte carlo estimation (M=4)
0.7758 + 0.0620
0.7105 + 0.0603
0.4993 + 0.0359
-0,7711 0.1049
Table 22: The influences of soft value function on the performances on protein binder design tasks for target protein IFNAR2.
Rewardf
ipTMt
PLDDTT
Radius.
Diversity
Posterior mean
0,5120 0.1093
0.5090
0.1079
0.4711 0.0490
2.2039
0.5176 0.5195
Monte carlo estimation (M=4)
0.4149 0.0936
0.3474
0.0906
0.6892 0.0617
0.0703 0.2881
G NOISE REWARD FUNCTION
In this section, we inject synthetic noise into the reward estimation function to assess the robustness of VIDD under imperfect reward signals. This scenario reflects realistic biomolecular settings, where surrogate reward models may deviate from true experimental measurements. Specifically, we add n% Gaussian noise to the reward function as:
In F n N(0, 1),
where (0, 1) is the normal distribution. During training, for inference, we observe the clean reward r. The results in Table 23 and Table [4 show that, as noise levels increase, the performance degrades noticeably. This is expected: VIDD does not incorporate explicit robustness mechanisms, so inaccuracies in reward estimation can mislead the optimization process and fine-tune the model toward suboptimal directions.
the noised reward rn signal, and
Since noisy or biased rewards exist in biomolecular design, addressing reward uncertainty remains a crucial direction. Incorporating robustness- reinforcement learning techniques, such reward denoising or uncertainty calibration, represents promising direction for future research.
H VISUALIZATION OF GENERATED SAMPLES
In Figure 7 we visualizes the docking of VIDD generated molecular ligands to protein parp1. Docking scores presented above each column quantify the binding affinity of the ligand-protein interaction, while the figures include various representations and perspectives of the ligand-protein protein complexes. We aim provide a complete picture of how each ligand situated within both the local binding environment and the larger structural framework of the protein. First rows show close-up views of the ligand bound to the protein surface, displaying the topography and electrostatic properties of the protein's is binding pocket and providing insight into the complementarity between the ligand and the pocket's surface. Second rows display distant views of the protein using the surface representation, offering broader perspective the ligand's spatial orientation within the global protein structure. Third rows provide close-up views of the ligand interaction using ribbon diagram, which represents the protein's 's secondary structure, such as alpha-helices and beta-sheets, to highlight the specific regions the protein involved in binding. Fourth rows show distant views of the entire protein structure in ribbon diagram, with ligands displayed within the context of the protein's 's full tertiary structure. Ligands generally fit snugly within the protein pocket, as evidenced by the close-up views in both the surface and ribbon diagrams, which show minimal steric clashes and strong surface complementarity.
26

---

Published as a conference paper at ICLR 2026
Table 23: The influences of noise reward function the performances protein binder design tasks for target protein PD-LI.
Reward
ipTMt
pLDDDT
Radius
Diversityf
noise=0.0
0.9079 + 0.0237
0.8182 + 0.0213
0.8720 + 0.0421
-0.1232 + 0.1066
0.5539 0.5154
noise=0.1
0.4744 1 0.1172
0.4570 0.1192
0.4306 0.0255
1.2802 0.9287
Table 24: The influences of noise reward function on the performances on protein binder design tasks for target protein IFNAR2.
Reward
ipTMt
pLDDT1
Radius
Diversity
noisc=0.0
0,5120 + 0.1093
0.5090 + 0.1079
0.4711 0.0490
1.9989 2.2039
0.5176 0.5234
noisc=0.1
0.1970 + 0.0217
0.1446 = 0.0160
0.4692 + 0.0620
-0.2763 + 0.2740
11.7
12.3
12.1
13.1
Figure 7: Visualization of generated molecules using VIDD optimizing the reward of docking score for parp1 (normalized as DS, 0)).
27


Reviews:
[{"review_id": "I4TgCkm8Yy", "reviewer": "Reviewer_57LE", "summary": "This paper presents **VIDD (Value-guided Iterative Distillation for Diffusion models)** — a reward-guided fine-tuning framework for diffusion models in biomolecular design tasks where reward functions are often non-differentiable.\nVIDD reformulates RL-style fine-tuning as **off-policy, value-weighted maximum likelihood estimation** with a **forward KL** objective, addressing instability and mode collapse common in PPO-like methods.\nIt iteratively distills soft-optimal denoising policies through three phases: off-policy roll-in, reward-weighted roll-out, and KL-based distillation.\nExperiments across **protein**, **DNA**, and **molecular** design tasks show strong and stable improvements over DDPO, DDPP, and other baselines, including differentiable settings.", "questions": "1. How sensitive is VIDD to the choice of α (temperature) and lazy update interval K?\n\n2. How does the algorithm behave when the reward oracle is noisy or misaligned (e.g., false-positive docking scores)? Does it still remain stable, or collapse due to biased reward?\n\n3. Can VIDD be combined with differentiable reward gradients (hybrid)? For differentiable domains, could both gradient and value-weighted distillation terms coexist?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. **Principled Distillation Framework:** The paper introduces a principled distillation-based alternative to PPO-style fine-tuning, effectively addressing instability issues commonly seen in on-policy reinforcement learning for diffusion models.\n\n2. **Handling Non-differentiable Rewards:** The proposed method applies to realistic biomolecular tasks where reward gradients are unavailable, bridging diffusion modeling with practical scientific design scenarios.\n\n3. **Stability via Lazy Updates:** The lazy-update strategy allows training to use mostly stable, older data while periodically refreshing the roll-out policy to capture model improvements — striking a balance between stability and adaptivity.\n\n4. **Comprehensive Evaluation:** The experiments are comprehensive, covering multiple scientific domains (protein, DNA, molecule) and providing ablations and comparisons with strong baselines.", "weaknesses": "1. **Heuristic Value Approximation:** The soft-value approximation via a single forward reward prediction (Algorithm 1, line 8) may introduce bias or variance, but the paper lacks quantitative analysis of its impact.\n2.  **Limited Theoretical Rigor:** The forward-KL interpretation is intuitive yet only algebraically shown (Appendix B) without a formal derivation or stability proof.\n3. **Comparisons on Differentiable Rewards:** For differentiable cases (e.g., DNA enhancer tasks), comparisons with gradient-based fine-tuning methods such as DRAKES are limited; more systematic evaluation would strengthen claims of generality.\n4. **Overlap with Prior Work:** The framework is conceptually related to value-weighted MLE and offline RL distillation and the paper could better position its novelty beyond these precedents.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "SfyBKyUtzr", "reviewer": "Reviewer_dS5g", "summary": "The paper proposed VIDD, a novel reinforcement-learning-based method for diffusion models in biomolecule design, including protein, small molecule and DNA sequecnces. The authors found the limitations of the original methods, including mode collapse, computational inefficiency, and reward hacking during training. Then, VIDD solves them step by step with corresponding techniques. In experiments, results from three systems demonstrated the effectiveness of the general design.", "questions": "1. The results in DNA design system of DRAKES is different from the ones from original paper. Is there any setting changes?\n\n2. In Section 5, the paper presents a theorem showing that the PPO-style objective is equivalent to a reverse KL divergence to the soft optimal trajectory distribution. However, for VIDD, it only states that it is “closer to a forward KL,” without providing a theorem or assumptions under which the empirical objective (7) is strictly equivalent to a forward KL.\n\n3. Eq (5) and Eq(7) fall under the offline-RL framework. Could you please clarify the coverage condition for hte roll-in distribution. Moreover, can you analyze the coveragence under mixture roll-ins?\n\n4. Can you do ablation studies on K, $\\alpha$, to demonstrate the stability if lazy updates?\n\n5. Also for Eq (5) and Eq (7), what is the differences/novelty of VIDD compared to exising reward-reweighted SFT or value-weighted MLE in offline RL?\n\n6.For the “off-policy + forward-KL”, is there empirical evidence showing that VIDD maintains its claimed stability advantage under the same roll-in coverage and identical KL regularization budget?\n\n7. Can you provide a formal proof/connection between VIDD and the mentioned inference-time methods?\n\n8. I found that when the roll-in distribution and target teacher distribution are mismatched, Eq. (7) does not perform explicit importance correction. Will it bring biased estimation? If so, how is there any sols?\n\n9. There should be general (mean/median) evaluation for sampling rather than simply using Best-Of-N.\n\n10. More ablation studies should be added. \n\n11. Using AlphaFold-Multimer’s ipTM as the sole metric can be biased by the predictor itself. Can you try to report another metric to demonstrate the effectiveness, such as pTM, pDockQ?\n\n12. For protein experiments, I think there would be potential data mismatch between EvoDiff and VIDD. Would you please clarify it?\n\n13. Could you please show some insights towards the training distillation and test-time sampling tradeoff?\n\n14. Can you also report log-likelihood in DNA system? It is because sometimes the 3-mer correlation cannot capture accurate simiarity among distributions.", "limitations": "", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. The motivation and the insight for the method design make sense, and the proposed solutions are simple and easy-to-implement. \n\n2. It involves many biological systems, demonstrating the effectiveness and the robustness of the algorithm in different tasks.\n\n3. VIDD unfied diffusion and value-weighted MLE with clear objective function and implementation, the framework could combine any type of non-differentiable reward functions, which is of high application range in biomolecule design.", "weaknesses": "1. The most concerning point from me is that the papre did not prove the effectiveness of each component through ablation studies, which makes the Method section not solid. More results on different biological systems are needed. \n\n2. Only limited baseline works were discussed and compared. Some baselines from discrete flow matching/diffusion on biological sequence design are of high correlation with this work [1-4]. \n\n3. More strict theoretical proof/analysis are needed (See Questions). \n\n4. There are irrational design of experimental validation (See Questions).\n\nReferences: \n\n[1] Tang, S., Zhang, Y., & Chatterjee, P. PepTune: De Novo Generation of Therapeutic Peptides with Multi-Objective-Guided Discrete Diffusion. In Forty-second International Conference on Machine Learning.\n\n[2] Zhao, Y., Uehara, M., Scalia, G., Kung, S., Biancalani, T., Levine, S., & Hajiramezanali, E. Adding Conditional Control to Diffusion Models with Reinforcement Learning. In The Thirteenth International Conference on Learning Representations.\n\n[3] Cao, H., Shi, H., Wang, C., Pan, S. J., & Heng, P. A. (2025). GLID $^ 2$ E: A Gradient-Free Lightweight Fine-tune Approach for Discrete Sequence Design. In ICLR 2025 Workshop on Generative and Experimental Perspectives for Biomolecular Design.\n\n[4] Tang, S., Zhang, Y., Tong, A., & Chatterjee, P. (2025). Gumbel-softmax flow matching with straight-through guidance for controllable biological sequence generation. ArXiv, arXiv-2503.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "mghZfRAfoT", "reviewer": "Reviewer_F9Zw", "summary": "This paper focuses on addressing the challenge of reward-guided fine-tuning of diffusion models in biomolecular design. Diffusion models are effective at modeling the complex distributions of biomolecules like proteins, small molecules, and regulatory DNA. However, real-world applications require optimizing for core rewards—such as physics-based protein binding affinity, small molecule docking scores, and science-driven protein secondary structure matching scores—that are non-differentiable, making traditional gradient-based methods inapplicable. Existing reinforcement learning (RL)-based fine-tuning approaches (e.g., DDPO, PPO) face issues including training instability, low sample efficiency, and mode collapse due to their on-policy nature and reliance on reverse KL divergence.\n\nTo tackle these problems, the authors propose VIDD (Value-guided Iterative Distillation for Diffusion models), a framework that treats fine-tuning as a policy distillation task, enabling optimization for any type of reward (including non-differentiable ones) through three iterative stages. First, the Roll-in stage collects off-policy trajectories by mixing samples from the pre-trained diffusion model (to explore a wide range of biomolecular design spaces) and a stable roll-out policy (to leverage high-reward regions), decoupling data collection from policy updates to improve sample efficiency. Second, the Roll-out stage simulates a soft-optimal \"teacher policy\" that balances reward maximization and compliance with the pre-trained distribution. This policy is weighted by a value function approximated using the diffusion model’s prediction of the clean sample , eliminating the need for reward gradients. Third, the Distillation stage updates the \"student model\" (the fine-tuned diffusion model) by minimizing the forward KL divergence between the soft-optimal policy and the current model policy, which prevents mode collapse compared to the reverse KL divergence used in RL methods.\n\nExperiments across four biomolecular design tasks—protein secondary structure matching, protein-target binding, DNA enhancer design, and small molecule docking—demonstrate that VIDD consistently outperforms baseline methods (e.g., DDPO, DDPP, Best-of-N) in terms of reward metrics. For instance, it achieves 83% β-sheet content in proteins, 8.28 Pred-Activity in DNA enhancers, and a 9.4 docking score for small molecules. Meanwhile, it maintains biomolecular naturalness, as shown by metrics like pLDDT (a measure of protein structure confidence) and 3-mer correlation (a measure of DNA sequence naturalness).", "questions": "1.\tValue Function Approximation and Error Correlation: VIDD uses x₀ prediction to approximate soft value functions, but the accuracy of x₀ prediction tends to decrease for long biomolecules (e.g., large protein complexes or DNA sequences with over 500 bases). Can you provide additional experiments to show the correlation between x₀ prediction errors (e.g., structural similarity between the predicted x₀ and the real x₀) and VIDD’s reward performance? Or can you compare this approximation method with alternatives such as Monte Carlo sampling or regression-based value estimation to verify its validity?\n2.\tAdaptive Hyperparameter Optimization: Hyperparameters like the Roll-in mix ratio (βₛ) and Roll-out update interval (K) have a major impact on performance. Manual tuning is time-consuming—for example, βₛ=0.8 optimizes protein diversity. Do you plan to design adaptive strategies, such as adjusting K based on the rate of reward improvement or optimizing βₛ based on the diversity of Roll-in data? Can you also provide heuristic rules for hyperparameter selection across different biomolecular tasks (e.g., recommended ranges for α in protein vs. small molecule tasks)?\n3.\tRobustness to Noisy and Biased Rewards: Biomolecular rewards often contain noise—for example, fluctuations in Enformer-predicted DNA activity or small molecule docking scores. How would VIDD perform if rewards are perturbed (e.g., ±10% Gaussian noise) or biased (e.g., systematically underestimating high-reward samples)? Are there calibration methods (e.g., multi-reward fusion) to reduce the impact of proxy bias?", "limitations": "", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "1.Effective Adaptation to Non-Differentiable Rewards: VIDD bypasses the need for reward gradients by approximating soft value functions using the diffusion model’s x₀ prediction. This allows it to directly optimize core non-differentiable rewards in biomolecular design, such as binding affinity scores from AlphaFold and docking scores from QuickVina2, filling a key gap in fine-tuning diffusion models for scientific applications.\n2. Training Stability and Anti-Collapse Ability: Unlike on-policy RL methods, VIDD uses off-policy data in the Roll-in stage, reducing sensitivity to noisy trajectories. By minimizing forward KL divergence instead of the reverse KL divergence used in RL, it forces the model to cover the distribution of the soft-optimal policy, avoiding mode collapse. For example, in protein binding tasks, VIDD maintains higher diversity than DDPO, and in small molecule design, it achieves a lower negative log-likelihood (NLL), indicating better preservation of biomolecular naturalness.", "weaknesses": "1.\tUnaddressed Impact of Value Function Approximation Errors: VIDD relies on the diffusion model’s x₀prediction to approximate soft value functions, but it does not analyze how errors in x₀prediction—such as those arising from long protein sequences or complex DNA structures —affect the distillation process. Additionally, it fails to compare this approximation method with alternatives like Monte Carlo sampling, leaving uncertainty about whether this is the optimal approach.\n2.\tHigh Hyperparameter Sensitivity and Lack of Adaptive Strategies: Key hyperparameters—such as the roll-in mix ratio (βₛ), the roll-out policy update interval (K), and the value weight coefficient (α)—require manual tuning and have a significant impact on performance. For example, in DNA tasks, using K = 5 results in much higher activity than using K = 20. However, VIDD does not propose adaptive strategies (e.g., dynamically adjusting K based on reward progress), which increases the difficulty of applying it across different tasks.\n3.\tInsufficient Exploration of Robustness to Noisy Rewards: Biomolecular rewards often suffer from \"proxy bias\"—for instance, a predicted binding affinity score (ipTM) only approximates real binding affinity, and DNA activity predicted by Enformer may differ from actual cellular activity. VIDD does not test its performance under conditions of reward noise (e.g., random fluctuations) or bias, leaving unanswered questions about whether it will over-optimize for proxy objectives (e.g., producing proteins with high predicted binding affinity but no actual binding function).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "cUuQqyoJnD", "reviewer": "Reviewer_AJss", "summary": "The paper", "questions": "I found the explanation of 'Approximation of soft value functions' (line 8 in Alg 1) a little too concise. In what sense is this a posterior mean? Why is it possible to evaluate a discrete reward function on this 'mean'?\n \nIt was hard for me to understand which parts of Section 4 are novel. Can you please provide more details about the relationship to prior work? Also, are there special cases or modifications of your framework that would make the method similar to other RL methods (e.g. if the roll-out policy was the same as the roll-in policy)? Drawing these connections can help readers understand things better. I appreciate section 6, but it was hard for me to understand if there is other related work that is similar to yours instead of a policy gradient method.\n\nTo what extent is the use of a diffusion model orthogonal to the particular RL approach? The inverse process of the diffusion model is just just a generic policy. Why is your presentation specific to diffusion models? Are there other non-diffusion policies where your approach would be sensible?\n\nWhat if, for example, you applied your approach to an autoregressive protein language model instead of a discrete diffusion model? I'm not suggesting you do this experiment. I'm just trying to understand why diffusion models are the focal point.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "In general, the exposition is well done. I don't have a ton of experience on modern RL approaches for training diffusion models and I found the explanations fairly approachable.\n\nThe evaluation setup is from prior work from papers that are more application-specific. In general, creating reward functions that are cheap computational surrogates for wet-lab experiments is difficult. I work in this area. I think the ones that are used are reasonable choices.", "weaknesses": "It was hard for me to understand which parts of Section 4 are novel. See question below.\n\nThe main paper provides few experimental results. However, the supplement has some really interesting follow-on results. It would have been nice to see some of these ablations and analyses featured in the main paper (e.g., fig 4).\n\nThe 'approximate soft-value functions' approach makes approximations that may be risky to make in practice. I would have liked to see a contrast with a sampling-based approach to estimating this.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)