Reward-Guided Iterative Refinement in Diffusion Models at Test-Time with Applications to Protein and DNA Design
Authors: Masatoshi Uehara
Xingyu Su
Yulai Zhao
Xiner Li
Aviv Regev
Shuiwang Ji
Sergey Levine
Tommaso Biancalani
Venue: ICML 2025 poster (2025)
Keywords: Diffusion models
Primary area: deep_learning->generative_models_and_autoencoders
TL;DR: Diffusion models

Abstract: To fully leverage the capabilities of diffusion models, we are often interested in optimizing downstream reward functions during inference. While numerous algorithms for reward-guided generation have been recently proposed due to their significance, current approaches predominantly focus on single-shot generation, transitioning from fully noised to denoised states. We propose a novel framework for inference-time reward optimization with diffusion models. Our approach employs an iterative refinement process consisting of two steps in each iteration: noising and reward-guided denoising. This sequential refinement allows for the gradual correction of errors introduced during reward optimization. Finally, we provide a theoretical guarantee for our framework. Finally, we demonstrate its superior empirical performance in protein and DNA design.

Reward-Guided Iterative Refinement in Diffusion Models at Test-Time with Applications to Protein and DNA Design
Masatoshi Uehara " Xingyu Su
Yulai Zhao Xiner Li Aviv Regev'
Shuiwang Ji Sergey Levine Tommaso Biancalani
Abstract
X Reward-Guided
To fully leverage the capabilities of diffusion models, we are often interested in optimizing downstream reward functions during inference. While numerous algorithms for reward-gan -unddd generation have been recently proposed due their significance, current approaches predom- inantly focus on single-shot generation, tran- sitioning from fully noised to denoised states. We propose a novel framework for test-time re- ward optimization with diffusion models. Our approach employs an iterative refinement pro- cess consisting of two steps in each iteration: noising and reward-guided denoising. This se- quential refinement allows for the gradual cor- rection of errors introduced during reward opti- mization. Finally, we demonstrate its superior empirical performance protein and cell-type specific regulatory DNA design. The code avail- able http://github.commmass wttp://////aaaeeerieeieeeee Refinement.
Denoising
Noising
X
Iterationn
Figure Our proposed framework follows an iterative pro- cess, with each iteration injecting noise into the sample and then denoising while optimizing rewards. For sequences, this can be implemented via masked diffusion, initialized from pre-trained diffusion models (left). Our algorithm can continuously refine the outputs by gradually correcting er- rors introduced during reward-guided denoising, improving the design over successive iterations (middle). For instance, for the task of optimizing the similarity (RMSD) of a protein to a target structure (Red), we can progressively minimize the RMSD through refinement, optimizing the design from initial (Orange) fit to a better final fit (Green), as shown on the right.
2023; Fan al., 2023; Uchara et al., 2024), while in protein sequence diffusion models, it could include metrics such as stability, structural constraints, or binding affinity (Verkuil et al., 2022), and in DNA sequence diffusion models, may involve activity levels (Sarkar et al., 2024; Lal et al., 2024).
1. Introduction
Diffusion models have achieved significant success across various domains, including computer vision and scientific fields (Ramesh al., 2021; Watson al., 2023). These models enable sampling from complex natural image spaces molecular spaces that resemble natural structures. Be- yond the capabilities of such pre-trained diffusion models, there is often a need to optimize downstream reward func- tions. For instance, in text-to-image diffusion models, the reward function may be the alignment score (Black et al.,
Building the motivation above, focus on optimizing downstream reward functions while preserving the natural- ness of the designs. (e.g., natural-like protein sequence exhibiting strong binding affinity) by seamlessly integrating these reward functions with pre-trained diffusion models during inference. While numerous studies have proposed to incorporate rewards into the generation process of diffu- sion models (e.g., classifier guidance (Dhariwal and Nichol, 2021) by setting rewards classifiers, derivative-free meth- ods (Wu et al., 2024; Li et al., 2024)), they rely single- shot denoising pass for generation. However, natural question arises:
ProluttionaryScale (Work done Genentech) "Texas A&M University 'Princeton Uni- versity *UC Berkley. Correspondence Masatoshi Ue- hara <masatoshi, www.ttossssuuchaaeeeee www.ra <xingyu.su@tamuu hewuuueoooooooorrcccc....... Xingyu
'Equal contribution
Proceedings of the ndnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).
Can we further leverage inference computation during generation to refine the model's 's output?

---

Reward-Guided Iterative Refinement in Diffuusio Models
In this study, we observe that diffusion models can inher- ently support an iterative generation procedure, where the design can be progressively refined through successive cy- cles of masking and noise removal. This allows us to utilize arbitrarily large amounts of computation during generation to continuously improve the design.
Guidance (i.e. test-time reward optimization) in diffu- sion models. Most classical approaches involve classifier guidance (Dhariwal and Nichol, 2021; Song et al., 2021), which adds the gradient of reward models (or classifiers) during inference. As reviewed in Uehara et al. (2025), re- cently, derivative-free methods such as SMC-based guid- ance (Wu et al., 2024; Dou and Song, 2024; Phillips et al., 2024; Cardoso et al., 2023) or value-based sampling (Li et al., 2024; 2025) have been proposed. However, these methods rely on single-shot generation from noisy states to denoised states. In contrast, we propose a novel iterative re- finement approach that enables the optimization of complex reward functions, which can be challenging for single-shot reward-guided generation.
Motivated by the above observations, we propose a novel framework for test-time reward optimization with diffusion models. Our approach employs iterative refinement algo- rithm consisting of two steps in each iteration: partial nois- ing and sward-ggiied denoising as in Figure 1. The reward- guided denoising step transitions from partially noised states to denoised states using techniques such as classifier guid- ance or derivative-free guidance. Unlike existing single-shot methods, our approach offers several advantages. First, our sequential refinement process allows for the gradual correc- tion of errors introduced during reward-guided denoising, enabling us to optimize complex reward functions, such as structural properties in protein sequence design. In par- ticular, this correction is expected to be crucial in recent successful masked diffusion models (Sahoo et al., 2024; Shi et al., 2024), as once a token is demasked, it remains unchanged until the end of the denoising step. Besides, for reward functions with hard constraints, commonly countered in biological sequence or molecular design (e.g., cell-type-specific DNA design (Gosai et al., 2023; Lal et al., 2024) or binders with high specificity), our framework can effectively optimize such reward functions by initializing seed sequences within feasible regions that satisfy these constraints.
Note while classifier-free guidance (Ho and Salimans, 2022) and RL (control)-baseee fine-tuning (e.g., Fan et al. (2023); Black et al. (2023); Domingo-Enrich et al. (2024); Uehara et al. (2024); Venkatraman et al. (2024); Zhang et al. (2024)) also aim to address reward optimization in diffusion models, they are orthogonal to our work, as we focus on test-time techniques without any training.
Refinement in language models. Refinement-style gener- ation has been explored in the context of BERT-style masked language models and general language models (Novak et al., 2016; Guu et al., 2018; Wang and Cho, 2019; Welleck et al., 2022; Padmakumar et al., 2023; Wang et al., 2024). How- ever, our work is the first attempt to study iterative refine- ment in diffusion models. Note that while some readers may consider editing in diffusion models (e.g.. Huang et al. (2024)) to be relevant, this is a distinct area, as the focus is not on reward optimization, unlike our work.
Our contribution is summarized as follows. First, we pro- pose a new reward- -uided generation framework for diffu- sion models that sequentially refines the generated outputs (Section 3). Our algorithm addresses two major issues in existing methods such as the lack of a correction mech- anism and difficulties of handling hard constraints. Sec- ondly, we provide theoretical formulation demonstrating that our algorithm samples from the desirable distribution exp(r(a:)) (ar)))))))P)()))) (.), where ppre (.) is a pre-trained distribu- tion (Section 4) and r(-) is a reward function. Finally, we present a specific instantiation of our unified framework by carefully designing the reward-guided denoising stage in each iteration, which bears similarities evolutionary algo- rithms (Section 5). Using this approach, we experimentally demonstrate that our algorithm effectively optimizes reward functions, outperforming existing methods in computational protein and DNA design (Section 6).
Evolutionary algorithms and MCMC for biological se- quence design. Refinement-basse approaches with re- ward models, such as variants of Gibbs sampling and genetic algorithms, have been widely used for protein/DNA design (Anishchenko al., 2021; Jendrusch et al., 2021; Hie et al., 2022; Gosai et al., 2023; Pacesa et al., 2024). However, most works do not address the integration of diffusion models. While some studies focus on integrating generative models (Hie et al., 2024; Chen et al., 2024), we explore an approach tailored to diffusion models, given the recent success of diffusion models in protein and DNA sequence generation (Alamdari et al., 2023; Wang et al., 2024).
1.1. Related Works
We finally remark that as closely related works, evolutionary- type algorithms with generative models have been actively investigated (Schneuing et al., 2024; Lee et al., 2025). Com- pared to these works, we use reward-guided decoding meth- ods instead of standard decoding ways.
We categorize related works into three key aspects.
2

---

Reward-Guided Iterative Refinement in Difffusion Models
2. Preliminaries
XT
Soft-optimal policy [x)exp(y--005-))/a
We first provide an overview of diffusion models, then dis- cuss current reward-guided -guided algorithms in diffusion models and the potential challenges, which motivate our proposal.
Value function v(gpa)) &e r(8((())
x A
2.1. Diffusion Models
m(xO)
X
In diffusion models, the objective is learn a sampler pPre()) E ( for a given design space X' using avail- able data. The training procedure is summarized as follows. First, we define forward noising process (also called policy) : X' + (XX) that proceeds from t = 0 to t = T. Next, we learn a reverse denoising process Pt X - ( parametrized by neural networks, ensuring that the marginal distributions induced by these forward and backward pro- cesses match.
Reward
Figure 2: Existing reward-guided algorithms can be viewed sequentially sampling from 2T to zo following the soft optimal policy (pi 1 The primary distinction among these algorithms lies in how Pž is approximated.
where C is the normalizing constant. This objective is widely employed in generative models, such as RLHF in large language models (LLMs) (Ziegler et al., 2019; Ouyang et al., 2022). In diffusion models (e.g.. Uehara et al. (2024, Theorem 1)), this is achieved by sequentially sampling from the soft optimal policy {pi Je from t = T + 1 to 1, which is defined by
To provide concrete illustration, we explain masked diffu- sion models. However, we remark that our proposal in this paper can be applied to any diffusion model.
Example 1 (Masked Diffusson Models). Here, we explain masked diffusion models (Sahoo et al., 2024; Shi et al., 2024;
Austin el al., 2021; Camphell el al., 2022; Lou et al., 2023)).
pi( a4) x exp(vt-(()/a))) (
Let X be a space of one-hot column vectors {0 E {0, 10
E Iz- = 11, and Cat(TT) be the categorical distribution over K classes with probabilities given by π E A where K denotes the K-simplex. A typical choice of the forward noising process is qe (It+1 zx) = Cat(ay + (1 m)) where 0,.....0 Mask]. Then, defining āe 11
where
Ve (x1) == a log Erp (m.) (exp(r(=))/a))zl)) (2)
and the expectation is taken w.r.t. the pre-trained policy. Here, as illustrated in Figure 2, ft-1 serves as a look-ahead function that predicts the reward at ag from 2t, often re- ferred to as the soft value function in RL (or the optimal twisting proposal in SMC literature (Naesseth et al., 2019)).
K
the backward process is parameterized as
8(. = x1)
if I + m
Ft-1
Cat (1-āg-1 Jm (a.-1
a. 10(()100) ) x. m,
In practice, we cannot precisely sample from soft optimal policies because (1) the soft value function i is unknown, and (2) the action space under the optimal policy is large. Current algorithms address these challenges as follows.
1-6r &
where (0();) is predictor from xa to Zo-
Notation and remark. 6g denotes the Dirac delta distribu- tion mass a. With slight abuse of notation, we express the initial distribution pT+1 X' + AAX), and denote [1, T) by [T].
(1): Approximating soft value functions. A typical ap- proach is to use r(āo (x+)) by leveraging the decoder do (2x) obtained during pre-training. This approximation arises from replacing the expectation over xo ~ (=)((=));; in (2) with 5g0(ze) (i.e., Dirac delta at the mean of prre (xa (i.e., as the state becomes more noisy). Despite its potential crudeness, this approximation is commonly adopted due to its training-free free nature and the strong empirical perfor- mance demonstrated by methods such as DPS (Chung et al., 2022), reconstruction guidance (Ho et al., 2022), universal guidance (Bansal et al., 2023), and SVDD (Li et al., 2024).
2.2. Single-Shot Reward-Guided Generation
Note its accuracy degrades as / increases
Our goal is to generate natural-like design with high reward. In particular, we focus on inference- -ime algorithms that do not require fine-tuning of pre-trained diffusion mod- els. Below, we provide a summary of these methods.
For reward guided generation, we often aim to sample from
(/n)__ = argmaxx E [r(x)) - (pllopl PEA(X)
(2): Handling large action space. Even with accurate value functions, sampling from the soft optimal policy still exhibits difficulty because its sample space X' is still large.
= =exp((()/a))pm((1/))
3

---

Reward-Guided Iterative Refinement in Diffussio Models
Hence, we often resort to approximation techniques as fol- lows.
Algorithm 1 Reward-Guided Evolutionary Refinement in Diffusion models (RERD)
1: Require: initial designs zo (the index (-) means the number of iteration steps), noise level K
(0)
Classifier Guidance: In continuous diffusion models, the pre-trained policy pi (- xt-1) is a Gaussian policy. By constructing differentiable value function models, we can approximate Pī by shifting the mean using Vua(-)/a. A similar approximation also applies to discrete diffusion models (Nisonoff et al., 2024).
2: for s E [0,-... S 1] do
3: Noising: Sample "'K (s+1) from ( |aag') where OR
is noising policy from zo IK (See Section 2.1).
4: Reward-Guided Generation: Sequentially sample from j= (i.e., from x (s+1) x ((++)) (In prac- tice, we need to approximate it. Refer Algo- rithm 2).
Derivative-Free Guidance: Another approach is using importance sampling (Li et al., 2024). Specifically, we generate several samples from pp-1( xz-1) and then select the next sample based on the importance A closely related method using
5: end for 6: Output: {xb (5)
weight exp Sequential Monte Carlo (SMC) has also been proposed, as discussed in Section 1.1.
Roward-Guided Denoising
Notsing
Noising
Xx
2.3. Challenges of Single-Shot Generation
P Ixe)
There are two main challenges with the aforementioned cur- rent algorithms. First, for certain complex reward functions, they may fail to fully optimize the rewards. This occurs because the value functions employed in these algorithms have approximation errors. When a value function model is inaccurate, the decision at that step can be suboptimal, and there is no correction mechanism during generation. This issue can be particularly severe in recent popular masked discrete diffusion models in Example 1, where once token changes from the masking state, it remains unchanged until the terminal step (t = 0) (Sahoo et al., 2024; Shi et al., 2024). Consequently, any suboptimal token generation at intermediate steps cannot be rectified.
4x(Pi (i)))
9x
P/ (x)
Figure 3: Summary of RERD: We instantiate it within masked diffusion models. It alternates reward-guided de- noising and noising.
be addressed in RERD, by sequentially refining the outputs.
In Algorithm several choices are important, which are outlined below.
Another related challenge lies in accommodating hard con- straints with set C. X. Although one might assume that simply setting r(.) = I(- C) would suffice, in practice, the generated outputs often fail to meet these constraints. This difficulty again arises from the inaccuracy of value function models at large (i.e., in highly noised states).
Initial designs zu (0) Here, we consider two ap- proaches. The first choice is to run {p?} from t = T to /. = 0 as in single-shot inference-time alignment algorithms. Second, if we have access to real data { ~ >prr(.). we select samples with high rewards as initial designs. A straightforward way is by using the weighted empirical distribution:
3. Iterative Refinement in Diffusion Models
Σ cxp(2')/a) Szc. E, exp(22)fa)
(3)
To tackle challenges discussed in Section 2.3, we propose a new iterative inference-time framework for reward opti- mization in diffusion models. Our algorithm is an iterative algorithm where each step consists of two procedures: nois- ing using forward pre-trained policies and reward-guided denoising using soft optimal policies. This framework is formalized in Algorithm 1.
Approximation of the soft optimal policy PE in Line 4: As mentioned in Section 2.2, exact sampling from Pt is infeasible. However, shelf methods to approximate it, such as classifier guid- ance or IS-based approaches discussed in Section 2.2. A specific instantiation of this approximation is consid- ered in Section 5.
can employ any off-the-
Compared to existing algorithms that only perform single- shot denoising from t = T' to = O, our algorithm repeat- edly performs reward optimization, as depicted in Figure 3. The challenge of single-shot algorithms - namely, the lack of a correction mechanism discussed in Section 2.3 can
Noise level K: When K is close to 0, the inference time per loop is reduced. Moreover, because value

---

Reward-Guided Iterative Refinement in Diffusson Models
function models used to approximate soft-optimal poli- cies are typically more precise around K 0 (see Section 2.2), the reward optimization step becomes more effective. On the other hand, using a larger K allows for more substantial changes in a single step. In practice, striking the balance, we recommend setting K/T low.
show that the marginal distribution after noising is PK (-) exp(Ux (osp(/xx)//////// (+)a)/C where PKx( (.) is marginal dis- tribution at K induced by pre-trained policies. Then, by in- duction, during reward optimization, we show that ls E [K]: Ty follows ()) exp(ux ()/a)/C. Then, when k = 0, it would be equal to ppre()) (-) exp(r()/a).
5. Practical Design of Algorithms
Next, we provide theoretical clarifications of our framework in Section 4. Additionally, we present a practical instantia- tion of our framework in Section
As mentioned, RERD is a unified sequential refinement framework that can integrate off-the-shelf approximation strategies during reward-guided denoising (Line 4 in Al- gorithm 1). A key practical consideration is determining which approximation methods to adopt. In this context, we present specific version that bears similarities to evolution- ary algorithms.
4. Theoretical Analysis
We present the theoretical analysis of RERD. We begin with the key theorem, which clarifies its target distribution.
5.1. Combining Local IS and Global Resampling
Theorem 1 (Target Distribution of RERD). Suppose (a) the initial design To (0) follows pla) (defined in (1)), (b) the marginal distributions induced by the forward noising process match those of the learned noising process in the pre-trained diffusion models. Then, the oulput To from RERD follows the target distribution
=
(8)
Value funcdions
(-) x exp(r())/a)pPrg)
IS sampling
First, discuss the validity of the assumptions. The sumption (a) is readily satisfied when using the introduced strategy of initial designs in Section 3. The assumption (b) is also mild, as pre-trained diffusion models are trained in this manner (Song et al., 2021), though certain errors may arise practice. Another implicit assumption in practice is that we can approximate soft-optimal policies accurately.
Noising
functions
k
Cf Ix)
Value functions
Global resampling
Next, explore the implications of Theorem 1. The cen- tral takeaway is that we sample from desired dis- tribution for our task (()) in (1). Although this guaran- tee appears to mirror existing single-shot algorithms dis- cussed in Section 2.2, we anticipate differing practical per- formance in terms of rewards. This is due to their robust- ness against errors in soft value function approximation vr (xxe) 2 r()o(x))).
k=0
(Selection)
Figure Visualization of Algorithm reward-guided denoising consists of two components: local value-weighte sampling for each sample (from & = K to k and global resampling among samples in batch at l: = 1.
Our specific recommendation for approximating soft op- timal policies during reward-guided denoising (Line 4 in Algorithm 1) is presented in Algorithm 2. Here, we adopt strategy that does not require differentiable value function models, as reward feedback could often be provided in a black-box manner (e.g., molecular design). Specifically, we organically combine IS-based and SMC-based approxima- tions. Given batch of samples, we apply IS from k: K to l = 1 (Line 4-6) for each sample in the batch, where the proposal distribution is policy from pre-trained diffusion models. However, at the terminal step A = 1, we perform selection via resampling (Line 7), which is central to SMC and evolutionary algorithms. This step involves interaction among samples in the batch, as illustrated in Figure 4.
To clarify, recall that in reward-guided algorithms, we must employ approximated soft value function models when sampling from the soft optimal policies pž xx (xp(c=-1-)/aappft
xz). The approximation often be- comes more precise as the time step č in the soft optimal policy approaches O, as mentioned in Section 2.2. Indeed, in the extreme case, when í = the exact equality holds. Therefore, by maintaining sufficiently small noise level t t= K and avoiding the approximation of value functions at large t, RERD can effectively minimize approximation emrrr in practice.
Sketch of the Proof of Theorem 1. proof deferred to Section A. brief, first, we
The detailed
5

---

Reward-Guided Iterative Refinement in Diffusson Models
Algorithm 2 Practical version of RERD
5.2. Constrained Reward Optimization
1: Require: Estimated value functions {o))=r (i.e., (r((o(Z)))<=T) pre-trained diffusion models 1k=T-1, initial designs (= (0) ) (the index (+)
We often need to include hard constraints so that generated designs fulfill certain conditions. This is especially crucial in molecular design, where may require low-toxicity small molecules or cell-type- speceeii DNA sequences, as shown in Section 6.2. Here, we explore how to enable generation under such constraints. Formally, we define the constraint set as C = {a : r2 (x) < e}. Given another reward ri (-) to be optimized, our objective to produce designs with high ri (+) while ensuring (2)))
means the number of iteration steps and the index E [M] is an index in a batch), duplication number L in IS, repetition number S, noise level K, a E R
2: for E [0.....33 S 1] do
(s+1) 3: Noising: For each i [N], sample 2Ka"" from
forward noising processes qx(: =6.)).
4: for k E [K 1, .1]
-1]do do
Naïve approaches with single-shot algorithms. As an initial consideration, we examine how to address this prob- lem using existing single-shot methods. A straightforward approach to use the following reward
5: IS: Sample Vi E [M), (E ~ P
& ( and define next states from the weighted
empirical distributions:
r(+) = ())))((()))) <c)
Vi mhe
w m E, mxppr(
or use a log barrier formulation:
6: end for
Selection: Vi E [N], sample 2o, ~ pl p a (s+1))
7:
r(.) = (- + log(max(c-
and perform resampling:
where C1 is a suitably small value, and then sample from t = T to t = 0 by following approxima soft-optimal policies. However, in reality, the outputs at = 0 often fail to satisfy these constraints, regardless of how the rewards are defined. This shortcoming arises because the value function models used during reward- reward-guided denoising are not completely accurate.
(1/(r()o.11/o)
1=0.1 hEs
E, E,exp(((za.)))/a)
8: end for
9: Output: (10.5) N
li=1
This combined strategy during reward-guided denoising leverages the advantages of both IS approaches (Li et al., 2024) and SMC approaches (Wu et al., 2024). First, we use the pure IS strategy from k = K to k = when a sam- ple in batch is poor, it will not be permanently discarded during the refinement process. In contrast, in Algorithm 2, the final selection step allows for the elimination of such poor samples through resampling. Second, if we use the pure SMC strategy from A = K to k- = 1, resampling is performed at every time step, which significantly reduces the diversity among samples in the batch. We apply the SMC approach only at the final step.
Integration into our proposal (Algorithm 2). Now, we consider incorporating the above rewards into our frame- work in Algorithm 2. Here, compared single-shot al- gorithms, we can often begin with feasible initial designs that satisfy the constraints E C. Then, by keeping the noise level K in Algorithm 2 small, we can avoid deviat- ing substantially from these feasible regions. This gradual refinement strategy makes easier to produce designs that fulfill hard constraints.
6. Experiment
We aim to evaluate the performance of the proposed method (RERD) across several tasks by investigating the effective- of refinement procedures compared to existing single- shot guidance methods in diffusion models. We begin by introducing the baselines and metrics used in our evaluation. Subsequently, we present our results in protein and DNA design. For further details and additional results, refer to Section C. The code is available at https://gihub..oom/maaaa SrrrooDiiiooo-fieeeeeee
Relation to evolutionary algorithm. The above version can be viewed as a modern variant of the evolutionary al- gorithm, which seamlessly integrates diffusion models. An evolutionary algorithm typically consists of two steps: (a) candidate generation via mutation and crossover and (b)
selection. In Algorithm 2, the step (a) corresponds to Lines 3-6, where -wward-guided generation is employed, and the step corresponds to Line 7.
Remark 1. When the reward feedback is differentiable, we can effectively integrate classifier guidance into the proposal distributions. For further details, see the Appendix in Li et al. (2024).
Baselines and our proposal. We compare baselines that address reward-guider generation in diffusion models with RERD. Note that we primarily focus on settings where reward feedback is provided in a black-bon manner.
6

---

Reward-Guided Iterative Refinement in Diffusion Models
Table I: The results for the protein design task show that our method consistently outperforms the baselines. Note that P50 and P95 represent the median and 95% quantile of the rewards for generated designs, respectively. LL denotes the (estimated) per-residue .g.liikli Values parentheses represent the estimated 95% standard deviation.
Task
(a) ss match
(b) cRMSD
(c) globularity
(d) symmetry
P50 f
P95
LL
P50 4
P95 +
LL
P50 f
P95 f
LL 7
P50 f
P95
LL 1
SMC
0.63 (0.04)
0.80
-3.28
8.9 (0.7)
5.1
-3,58
-2.79 (0.05)
-2.13
-4.43
-0.45 (0.03)
0.21
-3.30
SVDD
0.66 (0.02)
0.82
-3.03
8.2 (0.4)
4.6
-3.59
-2.45 (0.02)
-2.00
-4.68
-0.33 (0.04)
0.36
-3.56
GA
0.70 (0.00)
0.95
-3,51
6.3 (0.4)
3.01
-3.60
-1.35 (0.02)
-1.22
-4.38
0.21 (0.04)
0.44
-3.07
RERD
0.86 (0.01)
0.96
-3.13
1.68 (0.02)
0.96
-3.51
-1.29 (0.02)
-1.15
-4.45
0.34 (0.01)
0.69
-3.08
(a) Generated proteins (Green)
(b) Generated proteins (Green) from RERD when cRMSD are shown. Red represents the tar- get backbone structures. The cRMSD score for the left figure is 0.42, while for the right figure, is 0.6.
when optimizing ss-match are
shown. Red represents the tar- get secondary structures, The ss- match score for the left figure is 0.96, while for the right figure, 1.0.
(c) Generated proteins
when optimizing globu- (d) Generated proteins when optimizing larity.
symmetry
Figure 5: We visualize the sequences generated from RERD using ESMFold.
RERD minizing CRMSD (SKPH)
RERO cRMSD (en)))
Note that have used the same hyperparameters o, across baselines (SMC, SVDD) and RERD.
4.0
4.0
3.5
CRAGO (tower Detter)
3.0
2.5
Metrics. We report the top 95% quantile (P95) and me- dian of rewards (P50) from generated designs, as these are the primary metrics to optimize. Additionally, we present the estimated per-residue log-likelihood (LL.) using the pre- trained diffusion models, which serves as a secondary metric that we aim to maintain at moderately high value to pre- serve the naturalness of the designs.
20
2.0
37 15 beraton
26
25
10
20
25
(teretion
Figure 6: The refinement step from RERD when optimizing cRMSD in two target backbone structures demonstrated. Recall that the first iteration corresponds to the result from SVDD. The Y-axis represents the median reward of gener- ated samples (Lower better).
6.1. Protein Design
We begin by outlining our tasks. First, we use EvoDiff (Alamdari al., 2023), a representative discrete diffusion model for protein sequences trained the UniRef database, as unconditional base model. Next, following existing representative works in protein design (Hie et al., 2022; Watson et al., 2023; Ingraham et al., 2023), we consider four reward functions related to structural properties, which take the generated sequence as input. For more details, refer to Section C.
SVDD (Li et al., 2024): representative single-shot, derivative-fi free guidance method (without refinement).
SMC (Wu al., 2024): Another single-shot, represen- tative derivative-free guidance method.
GA: A naïve approach for sequence design that uses pre-traine pre-trained diffusion models generate mutated de- signs within standard genetic algorithm (GA) pipeline (Hie et al., 2022). To ensure fair comparison, we allo- cate the same computational budget as RERD below.
ss-match: We Biotite (Kunzmann and Hamacher. 2018) to predict the secondary structure (ss). We then calculate the mean matching probability across all residues between the predicted and reference secondary
RERD in Algorithm (Ours). We set K/T = 10% and S = 50. For initial designs, we use the results generated by SVDD in Section 6.1 and designs that satisfy the constraints in Section 6.2.
We also report the diversity of generated designs. Since this metrie is difficult to compare formally and secondary in the context of reward optimization, is included the Appendix.

---

Reward-Guided Iterative Refinement in Diffussin Models
Table 2: The results for the DNA design task show that our method consistently outperforms the baselines.
Task
HepG2
K562
SKNSH
P50 f
P95 f
LL t
P50 4
P95 4
LL 1
P50 1
P95 1
LL 1
SMC
1.2 (0.3)
1.6
-1.15
1.0 (0.2)
1.4
-1.21
0.8 (0.2)
1.0
-1.22
SVDD
2.3 (0.2)
2.8
-1.08
1.3 (0.3)
1.6
-1.26
1.7 (0.3)
2.0
-1.21
GA
2.3 (0.4)
2.7
-1.21
2.2 (0.3)
2.6
-1.31
1.9 (0.4)
2.5
-1.28
RERD
7.9 (0.2)
9.1
-1.18
7.4 (0.2)
8.9
-1.25
5.5 (0.2)
6.7
-1.24
structures, where the target structure is represented by sequence consisting of (a-hellccss), (f-sheets), and c (coils). A score of 1.0 indicates perfect alignment.
as follows. Using datasets from Gosai et al. (2023), which measures the enhancer activity of 700k DNA sequences (200-bp length) in human cell lines using massively parallel reporter assays (MPRAs), we trained oracles based on the Enformer architecture (Avsec et al., 2021) as rewards across three cell lines (ra(-) in HepG2 cell line, rk(() in K562 cell line and r*(() in SKNSH cell line). Then, we aim to respectively optimize the following:
cRMSD:
This is the constrained root mean square deviation against the reference backbone structure af- ter structural alignment. Typically, < 2Å indicates highly similar structure. Note that lower value is preferred.
globularity (+ pLDDT): It reflects how closely the structure resembles a globular shape. Additionally, we optimize pLDDT to improve the stability of the structure.
(H( = (=(x))(x<(z) < c))((s(()) < <)
(4)
where c is a threshold. Here, optimizing ÿh means maxi- mizing ти while retaining Tk T' low. Then, similarly, we define Tk; is by exchanging their roles.
symmetry (+pLDDT, hydrophobicity): It indicates the symmetry of the structure the generated se- quence. Additionally, we optimize PLDDT and hy- drophobicity improve the stability of the structure.
Here are several additional points note. First, as discussed in Section 5.2, directly using FH Fx. Fs in practice would lead to suboptimal performance. Therefore, we use log barrier reward functions for all methods. Additionally, for GA and RERD, we initialize the designs with samples that satisfy the constraints (c.g., I(rk (x) < c))(rs(a) < c))). Recall that one of the advantages of our method is its ability to leverage designs from feasible regions that satisfy the constraints. Finally, use pre-trained discrete diffusion models from Wang al. (2024a) the backbone unconditional diffusion models.
Note that each of the above rewards is computed after es- timating the corresponding structure using ESMFold (Lin et al., 2023). Besides, for both ss-match and cRMSD, we use 10 reference proteins randomly chosen from datasets in Dauparas al. (2022) and report the mean of the results.
Results. We present our performance in Table and vi- sualize generated sequences in Figure 5. Overall, our algo- rithm (RERD) consistently demonstrates superior perfor- mance in terms of rewards while maintaining reasonably high likelihood. Notably, as illustrated in Figure 6, for sev- eral challenging tasks, while one-shot guidance methods such as SVDD underperforms, our approach, with refine- ment steps, gradually yields improved results.
10
RERD ootim soectfe actlity
Reward (Hligher Better)
Reward
6.2. Cell- Sele-Sye--p Regulatory DNA Design
HepG2
KS62 Cell lines
SKNSH
We begin by outlining our tasks. Here, we focus on widely studied cell-type-specific regulatory DNA designs, which are crucial for cell engineering (Taskiran et al., 2024; Su et al., 2025). Specifically, (i.e., DNA sequences that regulate gene expression) that exhibit high activity levels in certain cell lines while main- taining low activity in others.
Iteration
Figure 7: (Left) The refinement step from RERD is demon- strated. The Y-axis represents the median reward of gener- ated samples (Higher is better), (Right) Generated designs from IRAO. It is seen that the activity in the target cell line HepG2 is only high.
goal is to design enhancers
Following existing works (Lal et al., 2024; Sarkar et al., 2024; Gosai et al., 2023), we construct reward functions
Results. The results are presented in Table 2. Our meth- ods consistently exhibit superior performance in terms of
8

---

Reward-Guided Iterative Refinement in Diffussio Models
rewards while maintaining a relatively high likelihood. No- tably, while it has been reported that SMC and SVDD excel in optimizing individual rewards (e.g., TH only) in exist- ing works such as Li et al. (2024), we have observed that they struggle with handling additional constraints. In con- trast, as shown in Figure 7, RERD effectively handles such constraints (i.e., ensuring cell-type specificity) by gradually refining the results, starting from designs in feasible regions.
Austin, J., D. D. Johnson, J. Ho, D. Tarlow, and R. Van Den Berg (2021). Structured denoising diffusion models in discrete state-spaces. Advances in Neural Information Processing Systems 34, 17981-1799)
Avsec, Z. V. Agarwal, D. Visentin, J. R. Ledsam, A. Hatbska-Barwinska, K. R. Taylor, Y. Assacl, J. Jumper, P. Kohli, and D. R. Kelley (2021). Effective gene expres- sion prediction from sequence by integrating long-range interactions. Nature methods 18(10), 1196-1203.
7. Conclusion
Bansal, A., H.-M. Chu, A. Schwarzschild, S. Sengupta, M. Goldblum, J. Geiping, and T. Goldstein (2023). Uni- versal guidance for diffusion models. In Proceedings of the IEEF/CVE Conference on Computer Vision and Pattern Recognition, pp. -52..
We introduce new framework for inference-time reward optimization in diffusion models, utilizing an iterative evo- lutionary refinement process. We also provide a theoretical guarantee for the framework's 's effectiveness and demonstrate its superior empirical performance in protein and DNA de- sign, surpassing existing single-shot reward-guide gener- ation algorithms. As future work, we plan to explore its application in moall-molecule design.
Black, K., M. Janner, Y. Du, I. Kostrikov, and S. Levine (2023). Training diffusion models with reinforcement learning. arXiv preprint arXiv:2305.1330/
Impact Statement
Camphell, A., J. Benton, V. De Bortoli, T. Rainforth, G. Deli- giannidis, and A. Doucet (2022). A continuous time framework for discrete denoising models. Advances in Neural Information Processing Systems 35, 28266- 28279.
This paper presents a novel framework for inference- time re- ward optimization in diffusion models, with applications in protein and DNA design. Our work aims to improve the abil- ity of generative models to optimize structured objectives through iterative refinement. While this research primarily contributes to technical advancements in generative model- ing, it has potential implications in domains such as drug discovery and biomolecular engineering, where optimizing protein structures can accelerate scientific innovation.
Cardoso, G., Y. E. Idrissi, S. L. Corff, and E. Moulines (2023). Monte carlo guided diffusion for bayesian linear inverse problems. arXiv preprint arXiv:2308.07000 079883.
Chandler, D. (2002). Hydrophobicity: Two faces of water. Nature 41 417(6888), 491-491.
We acknowledge that generative models, particularly those optimized for specific reward functions, could be misused if not carefully applied, such as in the design of harmful bio- logical agents. However, our work is intended for scientific and therapeutic applications, and we emphasize the impor- tance of responsible deployment and alignment with ethical guidelines in biotechnology. Overall, our contributions align with the broader goal of machine learning methodologies, and we do not foresee any immediate ethical concerns be- yond those generally associated with generative models.
Chen, A., S. D. Stanton, R. G. Alberstein, A. M. Watkins, R. Bonneau, V. Gligorijevi, K. Cho, and N. C. Frey (2024), Llms are highly-constrained biophysical sequence opti- mizers. arXiv preprint arXiv:2410.22296.
Chung, H., J. Kim, M. T. Mccann, M. L. Klasky, and J. C. Ye (2022). Diffusion posterior sampling for general noisy inverse problems. arXiv preprint arXiv:2209. 14687.
Dauparas, J., L. Anishchenko, N. Bennett, H. Bai, R. J. Ragotte, L. F. Milles, B. I. Wicky, A. Courbet, R. J. de Haas, N. Bethel, et al. (2022). Robust deep learning- based protein sequence design using proteinmpnn. Sci- ence 378(6615), 49-56.
References
Alamdari, S., N. Thakkar, R. van den Berg, N. Tenenholtz, B. Strome, A. Moses, A. X. Lu, N. Fusi, A. P. Amini, and K. K. Yang (2023). Protein generation with evolutionary diffusion: sequence is all you need. BioRxiv, 2023-09
Dhariwal, P. and A. Nichol (2021). Diffusion models beat gans on image synthesis. Advances in neural information processing systems 34, 8780-8794.
Anishchenko, I., S. J. Pellock, T. M. Chidyausiku, T. A. Ramelot, S. Ovchinnikov, J. Hao, K. Bafna, C. Nor, A. Kang, A. K. Bera, et al. (2021). De novo protein design by deep network hallucination. Nature 600(7889), 547-552.
Domingo-Enrich, C., M. Drozdzal, B. Karrer, and R. T. Chen (2024). Adjoint matching: Fine-tuning flow and diffusion generative models with memoryless stochastic optimal control. arXiv preprint arXiv:2409.08S. 08861.

---

Reward-Guided Iterative Refinement in Diffusson Models
Dou, Z. and Y. Song (2024). Diffusion posterior sampling for linear inverse problem solving: A filtering perspective. In The Twelfth International Conference on Learning Representations.
Lal, A., D. Garfield, T. Biancalani, al. (2024). reglm: Designing realistic regulatory dna with autoregressive language models. bioRxiv, 2024-02.
Lee, S., K. Kreis, S. P. Veccham, M. Liu, D. Reidenbach, Y. Peng, S. Paliwal, W. Nie, and A. Vahdat (2025). Gen- mol: A drug discovery generalist with discrete diffusion. arXiv preprini arXiv:2501.06155
Fan, Y., O. Watkins, Y. Du, H. Liu, M. Ryu, C. Boutilier, P. Abbeel, M. Ghavamzadeh, K. Lee, and K. Lee (2023). DPOK: Reinforcement learning for fine- tuning text-to-image diffusion models. arXiv preprint arXiv:2305.16381.
Li, X., M. Uchara, X. Su, G. Scalia, T. Biancalani, A. Regev, S. Levine, and S. Ji (2025). Dynamic search for inference- time alignment in diffusion models. arXiv preprint urXiv:2503.020::
Goodsell, D. S. and A. Olson (2000). Structural symmetry and protein function. Annual review of biophysics and biomolecular structure 29(1). 103-153.
Li, X., Y. Zhao, C. Wang, G. Scalia, G. Eraslan, S. Nair, T. Biancalani, A. Regev, S. Levine, and M. Uehara (2024), Derivative free guidance in continuous and discrete dif- fusion models with soft value-based decoding. arXiv preprint wrxiv:2408.08233
Gosai, J., R. I. Castro, N. Fuentes, J. C. Butts, S. Kales, R. R. Noche, K. Mouri, C. Sabeti, S. K. Reilly, and R. Tewhey (2023). Machine-guided design of synthetic cell type-specific cis-regulatory elements. bioRxiv.
Guu, K., T. B. Hashimoto, Y. Oren, and P. Liang (2018). Generating sentences by editing prototypes. Transactions of the Association for Computational Linguistics 6, 437- 450.
Lin, Z., H. Akin, R. Rao, B. Hie, Z. Zhu, W. Lu, N. Smetanin, R. Verkuil, O. Kabeli, Y. Shmueli, al. (2023). Evolutionary scall prediction of atomic-level pro- tein structure with language model. Science 379(6637), 1123-1130.
Hie, B., S. Candido, 7. Lin, O. Kabeli, R. Rao, N. Smetanin, T. Sercu, and A. Rives (2022). A high-level programming language for generative protein design. bioRxiv, 2022-12.
Lisanza, S. L., J. M. Gershon, S. W. Tipps, J. N. Sims, L. Arnoldt, S. J. Hendel, M. K. Simma, G. Liu, M. Yase, H. Wu, et al. (2024). Multistate and functional protein design using rosettafold sequence space diffusion. Nature biotechnology, 1-11.
Hie, B. L., V. R. Shanker, D. Xu, al. (2024). Efficient evo- lution of human antibodies from general protein language models. Nature Biotechnology 42(2), 275-283.
Ho, and T. Salimans (2022). Classifier-free diffusion guidance. arXiv preprint arXiv:2207. 12598.
Lou, A., C. Meng, and Ermon (2023). Discrete diffusion language modeling by estimating the ratios of the data distribution. arXiv preprint 23..:23300.168::1
Ho, J., T. Salimans, A. Gritsenko, W. Chan, M. Norouzi, and D. Fleet (2022). Video diffusion models. Advances in Neural Information Processing Systems 35, 8633-8646.
Naesseth, C. A., Lindsten,
Schön, et al. (2019).
Elements of sequential monte carlo. Foundations and Trends@ in Machine Learning 12(3), 307-392.
Huang, Y., J. Huang, Y. Liu, M. J. Lv, J. Liu, W. Xiong, H. Zhang, S. Chen, and L. Cao (2024). Diffusion model-based image editing: A survey, arXiv preprint a7Xiv:2402.17525.
Nisonoff, H., J. Xiong, S. Allenspach, and J. Listgarten (2024). Unlocking guidance for discrete state-space diffu- sion and flow models. arXiv preprint 2mXiv:2406.01777
Ingraham, J. B., M. Baranov. Z. Costello, K. W. Barber, W. Wang, A. Ismail, V. Frappier, D. M. Lord, C. Ng- Thow-Hing E. R. Van Vlack, et al. (2023). Illuminating protein space with programmable generative model. Nature 623(7989), 1070-1078.
Novak, R., M. Auli, and D. Grangier (2016). Itera- tive refinement for machine translation. arXiv preprint 1rXi::1610006603
Ouyang, L., J. Wu, X. Jiang, D. Almeida, C. Wainwright, P. Mishkin, C. Zhang, S. Agarwal, K. Slama, A. Ray, et al. (2022). Training language models to follow instructions with human feedback. Advances in Neural Information Processing Systems 35, 27730-27744.
Jendrusch, M., J. Korbel, and S. K. Sadiq (2021). Al- phadesign: A de novo protein design framework based alphafold. Biorxiv, 2021-10.
Kunzmann, P. and K. Hamacher (2018). Biotite: a unifying open source computational biology framework in python. BMC bioinformatics 19, 1-8.
Pace, C. and J. Hermans (1975). The stability of globular protein. CRC critical reviews biochemistry 3(1), 1 1-43.
10

---

Reward-Guided Iterative Refinement in Diffussin Models
Pacesa, M., L. Nickel, C. Schellhaas, J. Schmidt, E. Pya- tova, L. Kissling, P. Barendse, J. Choudhury, S. Kapoor, A. Alcaraz-Serna, et al. (2024). Bindcraft: one-shot de- sign of functional protein binders. bioRxiv, 2024-09.
K. Theunis, R. Vandepoel, al. (2024). Cell-type- directed design of synthetic enhancers. Nature 626(7997), 212-220.
Uehara, M., Y. Zhao, K. Black, E. Hajiramezanali, G. Scalia, N. L. Diamant, A. M. Tseng. T. Biancalani, and S. Levine (2024), Fine-tuning of continuous-time diffusion mod- els as entropy- regoopy-rggulaiized control. arXiv preprint urXiv:2402.15/1 15/94.
Padmakumar, V., R. Y. Pang, H. He, and A. Parikh (2023). Extrapolative controlled sequence generation via iterative refinement. In International Conference on Machine Learning, pp. 26792-26808. PMLR.
Phillips, A., H.-D. Dau, M. J. Hutchinson, V. De Bortoli, G. Deligiannidis, and A. Doucet (2024). Particle denois- ing diffusion sampler. arXiv preprint 0626v:240..06320.
Uchara, M., Y. Zhao, K. Black, E. Hajiramezanali, G. Scalia, N. L. Diamant, A. M. Tseng, S. Levine, and T. Bian- calani (2024, 21-27 Jul). Feedback efficient online fine fine-tuning of diffusion models. In R. Salakhutdinov, Z. Kolter, K. Heller, A. Weller, N. Oliver, J. Scarlett, and F. Berkenkamp (Eds.), Proceedings of the 41st In- ternational Conference on Machine Learning, Volume 235 of Proceedings of Machine Learning Research, pp. 48892 48918, PMLR.
Ramesh, A., M. Pavlov, G. Goh, S. Gray, C. Voss, A. Rad- ford, M. Chen, and I. Sutskever (2021). sero-shot text- to-image generation. In International conference on ma- chine learning, pp. 8821-8831 Pmir.
Sahoo, S. S., M. Arriola, Y. Schiff, A. Gokaslan, E. Marro- quin, J. T. Chiu, A. Rush, and V. Kuleshov (2024). Simple and effective masked diffusion language models. arXiv preprint ar 4Xhv:2406.07524.
Uehara, M., Zhao, C. Wang, X. Li, A. Regev, S. Levine, and T. Biancalani (2025). Reward-guidee controlled gen- eration for inference-time alignment in diffusion models: Tutorial and review, arXiv preprint arXiv:2501.0900
Sarkar, A., Z. Tang, C. Zhao, and P. Koo (2024). Design- ing dna with tunable regulatory activity using discrete diffusion. bioRxiv, 2024 05.
Venkatraman, S., M. Jain, L. Scimeca, M. Kim. M. Sendera, M. Hasan, L. Rowe, S. Mittal, P. Lemos, E. Bengio, et al. (2024). Amortizing intractable inference in diffusion models for vision, language, and control. arXiv preprint 20079
Schneuing, A., C. Harris, Y. Du, K. Didi, A. Jamasb, Igashov, W. Du, C. Gomes, T. L. Blundell, Lio, al. (2024). Structure- bastur drug design with equivariant diffusion models. Nature Computational Science 4(12), 899-909.
Verkuil, R., O. Kabeli, Y. Du, B. I. Wicky, L. F. Milles, J. Dauparas, D. Baker, S. Ovchinnikov, T. Sercu, and A. Rives (2022). Language models generalize beyond natural proteins. BioRxiv, 2022-12.
Shi, J., K. Han, Z. Wang. A. Doucet, and M. K. Titsias (2024). Simplified and generalized masked diffusion for discrete data. arXiv preprint arXiv:2406.0432 043299.
Wang, A. and K. Cho (2019). Bert has mouth, and it must speak: Bert as a markov random field language model. arXiv preprint 1902iv:990.040000
Song, Y. C. Durkan, I. Murray, et al. (2021). Maximum likelihood training of score-based diffusion models. In Advances in neural information processing systems, Vol- ume 34, pp. 1415-1428.
Wang, C., M. Uehara, Y. He, A. Wang. T. Biancalani, A. Lal, T. Jaakkola, S. Levine, H. Wang, and A. Regev (2024a). Fine-tuning discrete diffusion models via reward opti- mization with applications to dna and protein design. arXiv preprint 2rXiv:2410.1364.
Song. Y. J. Sohl- Dickksicn,,,,, D. P. Kingma, A. Kumar, S. Er- mon, and B. Poole (2021). Score-based generative mod- cling through stochastic differential equations. ICLR.
Stark, H., B. Jing, C. Wang, G. Corso, B. Berger, R. Barzi- lay, and T. Jaakkola (2024). Dirichlet flow matching with applications to dna sequence design. arXiv preprint arXiv:2402.058811
Wang, M. Uchara, Y. He, A. Wang. T. Biancalani, Lal, T. Jaakkola, S. Levine, H. Wang, and A. Regev (2024b). Fine-tuning discrete diffusion models via reward opti- mization with applications to dna and protein design. arXiv preprint arXiv:2410.13643
M.
Su, X., H. Yu, D. Zhi, and S. Ji (2025). Learning discover regulatory elements for gene expression prediction. arXiv preprint ar arXiv:2502. 139991
Wang, H. M. Skreta, C.-T. Ser, W. Gao, L. Kong, F. Strieth- Kalthoff, C. Duan, Y. Zhuang, Y. Yu, Y. Zhu, et al. (2024). Efficient evolutionary search over chemical space with large language models. arXiv preprint arXiv:2406. 16966 16976.
Taskiran, I. L. K. I. Spanier, H. Dickmänken, N. Kempynck, A. Pančíková, E. C. Ekşi, G. Hulselmans, J. N. Ismail,

---

Reward-Guided Iterative Refinement in Diffusion Models
Wang, X., Z. Zheng, F. Ye, D. Xue, S. Huang, and Q. Gu (2024). Dplm-2: A multimodal diffusion protein language model. arXiv preprint arXiv:2410.133 13782.
Watson, J. L., D. Juergens, N. R. Bennett, B. L. Trippe, Yim, H. E. Eisenach, W. Ahern, A. Borst, R. Ragotte, L. F. Milles, al. (2023). De novo design of protein structure and function with ridiffusion. Na- ture 620(7976), 1089-1100.
Welleck, S. X. Lu, West, F. Brahman, T. Shen, D. Khashabi, and Y. Choi (2022). Generating se- quences by learning to self-correct. arXiv preprint vrXiv:2211.0005.
Wu, L., B. Trippe, C. Naesseth, Blei, and Cunning- ham (2024). Practical and asymptotically exact condi- tional sampling in diffusion models. Advances in Neural Information Processing Systems 36.
Zhang, D., Y. Zhang, J. Gu. Zhang, J. Susskind, N. Jaitly, and S. Zhai (2024). Improving gflownets for text-to-image diffusion alignment. arXiv preprint arXiv:2406.00633.
Ziegler, D. M., N. Stiennon, J. Wu, T. B. Brown, A. Radford, D. Amodei. Christiano, and Irving (2019). Fine- tuning language models from human preferences. arXiv preprint arXiv: arXiv:1909.000 085993...
12

---

Reward-Guided Iterative Refinement in Diffusion Models
A. Proof of Theorem 1
Here, we use induction. Hence, we prove that so follows pla), Distribution after noising.
First, we consider the distribution after noising. This is
/ g xo (0) (x)))((r )dzz
By plugging in the first assumption regarding distributions of initial designs, is equal to (16
f 9x
(2)
1 (xx exp(r(=0 exxp((()) )/))dzz
(5)
Recalling this definition soft value functions:
exp(vx())/o) = (x00 (expx(ee) xx] and the assumption (b) (go (0)(<<)=) pare (rokkr) and (x(+)) - (.) ), the term (5) is equal to pk (-)
Distribution after reward-guided denoising. Now, we consider the distribution of =
/ I
{
1/C
( mm)
p pro
K=K
With some simple algebra, this is equal to
/
i
Pk (ww.:///rr)eppoxxooooo
pre
1/C
pd (xx-1
(xx) |
pk (xx) (xx))exp(uxxxx)/)))ddae
exp(ex(Exx/aa)
k=K-1
1//
Pi (2 xxk)
x PK
pre (#K-1 pro
\==K-1
-1/C /
II 1==--
pi (
PK- (r
pre
(K-1))
Repeating this argument from K 1--1to 0. the above is equal to
pp ()) exp(r())aa//C.
This concludes the statement.
B. Additional Details for Protein Design
In this section, we have added further details on experimental settings and results.
B.1. Details on Baselines
RERD (Algorithm 2): We have used parameters L 20, N = 10, S = 30 in general. For the importance sampling step, we have used a = 0.0, and for the selection step, we have used a = 0.2.
SVDD: We set the tree width = 20, = 0.0.
SMC: In SMC, we set a 0.05 because if we choose a = 0.00, it just gives single sample every time step. Refer to Appendix Li et al. (2024).
GA: Here, compared to Algorithm 2, we have changed the mutation part (Line 3-7) with just sampling from pre-trained diffusing models without any soward-guide generation. To have a fair comparison with RERD, we increase the repetition number S so that the computational budget roughly the same as our proposal.
13

---

Reward-Guided Iterative Refinement in Diffusion Models
B.2. Details Reward Functions
Globularity. Globularity refers the degree to which protein adopts a compact and nearly spherical three-dimension structure (Pace and Hermans, 1975).It is defined based the spatial arrangement of backbone atomic coordinates, where the variance of the distances between those coordinates and the centroid is minimized, leading highly compact structure. Here, we set the protein length 150.
Globular proteins are characterized by their structure stability and water solubility, differing from fibrous or membrane proteins. The compact conformation helps proteins to maintain proper protein folding and reduce the risk of aggregation.
Symmetry. Protein symmetry refers to the degree to which protein subunits are arranged in repeating structure pattern (Goodsell and Olson, 2000; Lisanza et al., 2024; Hie et al., 2022). Here we focus on the rotational symmetry of single chain, which is defined by the spatial organization of subunit centroids. Specifically, we try to minimize the variances of the distances between adjacent centroids to achieve more uniform and balanced arrangement. Here, we set the protein length to be 150 to 240.
Symmetric proteins can bring multiple functional sites into close proximity, facilitating interactions and supporting the formation of large proteins with optimized biological functions.
Hydrophobieity. Hydrophobicity refers the degree which a protein repels water, primarily defined by the distribution of hydrophobic amino acids within the structure, namely, Valine, Isolcucine, Leucine, Phenylalanine, Methionine and Tryptophan (Chandler, 2002). Hydrophobicity is optimized by minimizing the average Solvent Accessible Surface Area (SASA) of the hydrophobic residues above, thus reducing their exposure to the surrounding solvent. Hydrophobicity enhances the protein structural stability, especially in the polar solvents such as water, facilitates the protein-protein interactions by prompting binding at the hydrophobic surfaces, and drives the proper protein folding by guiding the hydrophobic residues to the protein core.
PLDDT. PLDDT (predicted Local Distance Difference Test) confidence score used evaluate the reliability of the local structure in predicted proteins. It is defined by the confidence of model predictions, assigning confidence value to each residues. A higher PLDDT score indicates greater model confidence and suggests increased structural stability. To optimize the whole protein structure, we try to maximize the average pLDDT across the whole sequence as predicted by ESMFold (Lin et al., 2023).
B.3. Additional Results
More metric (diversity, PLDDT, and pTM)). We have included additional metrics in Table 3.
Generally, higher PLDDT and pTM values indicate more accurate structure predictions at the local residue and the global structure, respectively. However, in the context of de novo protein design, a low PLDDT does not necessarily imply poor performance (Verkuil et al., 2022). In the globularity task, it is expected that the generated protein is more novel protein.
We define diversity as the mean pairwise distance (normalized by length), where the distance is measured using the Levenshtein distance. While diversity can be an important metric evaluate the performance of pre-trained generative models, in the context of reward optimization, this metric may be secondary. It is shown that generated sequences from RERD reasonably diverse enough without collapsing to single samples.
Table Additional metrics for experiments in protein design. We have reported the median of pLDDT, pTM, and diversity of generated proteins.
Task
(a) ss-match
(b) cRMSD
(c) globularity
(d) symmetric
PLDDT
PTM
diversity
PLDDT
PTM
diversity
PLDDT
PTM
diversity
PLDDT
PTM
diversity
RERD
0.75
0.69
0.28
0.76
0.71
0.14
0.41
0.29
0.56
0.82
0.79
0.49
14

---

Reward-Guided Iterative Refinement in Diffusion Models
Recovery rate when optimizing cRMSD. By optimizing cRMSD, we can tackle the inverse folding task. While we have extensively investigated the performance in terms of recovery rates, we present the observed recovery rates for several proteins reference when using RERD. Although it does not match the performance of state-of-the-ar conditional generative models specifically trained for this task, such as ProteinMPNN (Dauparas et al., 2022), our algorithm, which combines unconditional diffusion models with reward models at test-fime, demonstrates competitive performance.
Table 4: Recovery rates when optimizing cRMSD
Proteins
5KPH
6NJF
EHEE rdd_0101
EA:rum2
0325_0005
RERD
0.26
0.31
0.28
0.30
ProteinMPNN
0.41
0.53
0.35
0.38
More generated proteins.
We have visualized more generated proteins in Figure 8.
(a) The generated proteins (Green) when optimizing ss-match are shown. Red represents the target sec- ondary structures. The ss- match score is 1.0 here.
(b) The generated proteins (Green) when optimizing cRMSD are shown. The generated proteins Red represents the target secondary when optimizing globular- structures.
ity are shown.
(d) The generated proteins when opti- mizing symmetry shown.
Figure 8: More generated protein from RERD.
C. Additional Details for DNA Design
Pre-trained models. We use the pre-trained diffusion model trained in Wang et al. (2024b). The code and its performance are available their paper. Here, we use the discrete diffusion model proposed in (Sahoo et al., 2024) using the same CNN architecture as in (Stark et al., 2024) and linear noise schedule.
Reward oracles. We use the exact oracle used Wang et (2024b). Again, the code and its performance are available their paper. Here, we use the Enformer architecture (Avsec et al., 2021) initialized with its pretrained weights. We use the data splitting based chromosome following standard practice (Lal et al., 2024).
Hyperparameters in baselines and RERD. We set S 15, a 0.0, T. 20.
Diversity. We calculate the diversity in the protein design task. It is 0.47 in HepG2, 0.49 in K562, and 0.53 in SKNSH. It is shown that generated sequences are reasonably diverse enough.
15


Reviews:
[{"review_id": "rT23nKlWOh", "reviewer": "Reviewer_xE3K", "summary": "The paper presents a novel framework for inference time reward optimization in diffusion models, introducing an iterative refinement approach that alternates between noising and reward guided denoising steps. This method departs from conventional single shot reward optimization, aiming to iteratively refine generated samples, allowing for the correction of errors and more effective optimization of complex reward functions. The authors provide a theoretical guarantee showing that their framework samples from a distribution proportional to the pretrained model distribution, weighted by the exponentiated reward function. The method is evaluated empirically on protein and DNA sequence design, demonstrating improvements over baseline approaches in optimizing structural properties of proteins and regulatory activity of DNA sequences while maintaining sample quality. The results suggest that this iterative approach is particularly useful for handling hard constraints, which is relevant in biological design tasks where feasibility constraints are often strict.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for the positive and very detailed feedback! Below we address the key questions and comments you raised:\n\n> Q. Abolition study \n\nThank you for the thoughtful suggestions regarding the ablations. In response, we have conducted additional ablation studies by varying key hyperparameters. To provide a quick yet informative signal, we focused on the ss-match and cRMSD tasks. Here is a link to figures describing experimental results. [link](https://anonymous.4open.science/r/hidden-1886)\nWe plan to extend these studies to other tasks in the final version. \n\n- We added an ablation study varying K in Algorithm 2 (**Figure 1,2** in the link). The results show strong performance when K/T=10%  or K/T=20% This result is expected: a large K reduces the benefit of refinement, while a very small K limits the opportunity for reward-guided decoding.\n\n- We also performed an ablation on L, the number of repetitions for importance sampling (**Figure 3,4** in the link). As expected, performance improves with a larger L, likely due to the increased computational budget and exploration.\n\n> Q. Has it been considered whether slight mismatches between the noising process and the learned forward diffusion model alter the final distribution?\n\nThis is an excellent point. We agree that optimization and sampling-time errors can differ in practice. While a rigorous analysis is challenging, one possible direction is to assume the mismatch is bounded in total variation distance by $\\epsilon$, and then analyze how this propagates to the final distribution. We will explore this idea further and aim to incorporate a discussion in the final version.\n\n> Q: Also, is there a possibility that the refinement process oscillates between suboptimal solutions?\n\nYes, it oscillates. But, in general, it tends to optimize a stable manner, as shown in Figure 6. \n\n> Q. How are reference proteins selected? \n\nWe follow the protocol introduced by Hie et al. (2022). We will make this more explicit in the revised version.\n\n> Q. Related works. \n\nThank you for the suggestions. We will add citations to the relevant works in the final version.", "claims_and_evidence": "The key argument is that single shot reward guided denoising methods are limited in their ability to optimize complex reward functions due to approximation errors in estimating value functions, particularly at highly noised states. According to the authors, an iterative refinement process that progressively applies reward optimization (i.e. RERD, the proposed approach) corrects errors more gradually, can correct suboptimal decisions in earlier steps and leads to superior performance on complex reward functions that involve structural constraints. \n\nThe authors provide a theoretical justification demonstrating that RERD samples from a distribution proportional to the pre-trained model distribution weighted by the exponentiated reward function, ensuring alignment with the target reward optimized distribution. This theoretical result is derived under the assumption that the noising and denoising processes used in the iterative refinement step match the pre-trained diffusion model's forward and reverse processes. While this assumption is reasonable given the training procedure of diffusion models, the practical effectiveness of this theoretical guarantee depends on the accuracy of approximating the soft optimal policy at each step, which is not explicitly analyzed in the theoretical section.\n\nExperimental results on protein design and regulatory DNA design tasks are presented as empirical evidence. The results consistently show that the proposed method outperforms baselines such as single shot guidance and genetic algorithms in terms of reward maximization while maintaining reasonable likelihood scores. \n\nThe paper does not include real world experimental validation beyond simulation and computational evaluation, which may limit the external validity of the claims, but within the scope of computational biomolecular design, the evidence presented is convincing.\nFurthermore, evaluation is not extended to other domains where reward guided generation might be relevant. This is perhaps a limitation since the approach could be generally applicable to several domains, and evidence of this would strengthen this submission.", "methods_and_evaluation_criteria": "The proposed method RERD iteratively introduces noise to partially perturb samples before applying reward guided denoising. The denoising step uses importance sampling and a final selection mechanism inspired by evolutionary algorithms to refine samples towards high reward solutions. The motivation behind this approach is that errors introduced during reward optimization due to inaccuracies in value function approximations can be corrected over multiple iterations. The theoretical framework shows that under idealized conditions, the final samples produced by the iterative refinement process follow a distribution proportional to the pretrained diffusion model's prior distribution, reweighted by an exponentiated reward function. This ensures that the algorithm maintains a principled probabilistic framework while still allowing for effective reward optimization.\n\nEvaluation focuses on benchmark tasks in protein and DNA sequence design with reward functions measuring structural properties and regulatory activity. These tasks have become commonplace downstream tasks to evaluate protein/biological sequence models and so are appropriate in this particular context. In particular, on protein design, secondary structure matching, backbone root mean square deviation, globularity and symmetry are used as reward metrics. All protein sequences are structurally evaluated using ESMFold. In DNA design, the task is to generate enhancer sequences that maximise activity in a specific cell type while minimizing activity in others, with reward functions designed using pretrained sequence based predictors trained on large scale enhancer activity datasets, and  evaluation metrics being 50th and 95th percentile of predicted activity scores. The baseline methods for comparison are SVDD, SMC and a genetic algorithm which applies mutations to pretrained diffusion model samples.  \n\nRERD is shown to outperform the baseline consistently (in terms of reward) while maintaining likelihood values comparable to the original diffusion model. In my opinion, the model is fairly and rigorously evaluated, and the evaluation criteria are accurate to the problem setting as it captures both reward minimisation maximization and sequence naturalness. One remark is that there is no ablation of individual components (e.g. the evolutionary resampling step).", "theoretical_claims": "The core claim is that the iterative refinement process samples from a distribution proportional to the pretrained diffusion model prior, weighted by an exponentiated reward function. This is established in Theorem 1, which asserts that under two main assumptions the final output of RERD follows the desired target distribution. These assumptions are that the initial samples follow the reward weighted distribution and that the noising process matches the forward process of the pretrained diffusion model. This claim is meant to provide a theoretical guarantee that RERD does not diverge arbitrarily from the pretrained model's learned distribution, ensuring that the generated samples remain plausible while optimizing the reward. The proof is structured as an induction argument over the iterative refinement steps, showing that if the distribution holds at step K, then applying reward guided denoising preserved this form until reaching the final step at which point the distribution matches the desired target form. \nJust for my own understanding, has it been considered whether slight mismatches between the noising process and the learned forward diffusion model alter the final distribution? Also, is there a possibility that the refinement process oscillates between suboptimal solutions? Perhaps deriving a bound on the variance of samples over multiple iterations could be useful.", "experimental_designs_or_analyses": "The experimental design seems well structured and backs up the theoretical claims.The experiments aim to assess whether the iterative refinement process leads to superior reward optimization while maintaining biologically plausible sequences. The evaluation setup involves a combination of benchmark datasets, pretrained diffusion models, and specific reward functions tailored to each task.\n\nFor protein sequence design the authors use EvoDiff (a discrete diffusion model trained on UniRef) as the base generative model, and compute the reward functions based on structural predictions from ESMFold. \nFor enhancer design, the authors use a pretrained discrete diffusion model and construct reward functions using enhancer activity predictors trained on large scale datasets from [1] (with consist of measurements of enhancer activity on several DNA sequences) which they use to train predictive models based on the well known Enformer architecture. The DNA design tasks involve generating sequences that maximize enhancer activity in a target cell line while suppressing it in others, ensuring specificity to a particular cell type. \n\nWhile indeed, evidently, this study relies on multiple models in the loop, this is a well known and often followed approach from various related works in the literature and often aligns with good practices in computational biology. It is still to note that these introduce an additional layer of approximation, crucially at the evaluation step. \nAdditional remarks: there doesn't seem to be much of a discussion of how the reference proteins were selected. Ditto for the DNA design task where the sequences are initialized from a pretrained diffusion model but their diversity is not analyzed. Also, while RERD is present as a unified framework, as mentioned previously I am also wondering about the impact of each individual component (noising, reward guided denoising, importance sampling and evolutionary resampling). The impact of each component is never examined in isolation. \n\n\n[1] https://www.nature.com/articles/s41586-024-08070-z", "relation_to_broader_scientific_literature": "Several key contributions in this field are cited. [1] provides an in depth guide on inference time guidance methods for optimizing reward functions in diffusion models, emphasizing the need for aligning generated samples with desired metrics without retraining the model. This paper builds upon this foundation by proposing an iterative refinement process, moving beyond single shot generation. [2] explores finetuning discrete diffusion models using reinforcement learning to optimize specific reward functions, particularly in biological sequence generation. The current study is in many ways quite similar but the emphasis is placed on test time optimization. \n\n\n[1] https://arxiv.org/abs/2501.09685\n[2] https://arxiv.org/abs/2410.13643", "essential_references_not_discussed": "Arguably [1]. \n\n[1] https://www.biorxiv.org/content/10.1101/2024.02.01.578352v1", "other_strengths_and_weaknesses": "Beyond what was mentioned already, the paper seems well presented and the presented framework is relatively novel work.", "other_comments_or_suggestions": "not applicable", "questions_for_authors": "not applicable", "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": "gt1JCFcPDc", "reviewer": "Reviewer_t8tW", "summary": "The authors introduce a novel inference-time framework for the iterative refinement and reward optimization of diffusion models. Their proposed method, Reward-Guided Evolutionary Refinement in Diffusion models (RERD), is based on the iterative refinement of generation with reward-guides denoising, and provide theoretical support for their method. The authors demonstrate the use case of RERD on masked diffusion models for the tasks of protein and biological sequence design. Through a set of reasonably thorough experiments, they show that their method yields improved performance relative to counterpart baselines.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate the positive feedback. Below are our responses to your questions:\n\nQ: Do we need 100∗(S−1) steps? \n\nYou're absolutely right. When setting K/T=10%, and T=1000, we would indeed require 100 steps. However, this component can be adjusted in practice by reducing T or K, which offers flexibility depending on computational constraints. We will clarify this point in the revised version.\n\nQ: Relation to Domingo-Enrich, Carles, et al. \n\nThank you for bringing this work to our attention. We will cite it in the final version. It appears that the focus of this paper is more on the fine-tuning of diffusion models, whereas our work emphasizes inference-time reward optimization. Following prior work such as DPS, SVDD, and SMC-based methods, we have primarily focused on comparison between inference-time techniques, which we view as complementary/orthogonal to fine-tuning approaches. That said, we agree that a more detailed discussion would be valuable and will include a comparison with fine-tuning methods in the revision.", "claims_and_evidence": "In general, the claims of the paper are supported by empirical evidence and theoretical results. One item I would like to point out:\n\n- On lines 70-72 (right): \"our work is the first attempt to study iterative refinement in diffusion models\". I am not entirely certain this claim is true, but I could be wrong. Please take a look at my comment in the \"Essential References Not Discussed\" section of the review.", "methods_and_evaluation_criteria": "The authors evaluate their proposed method with a diverse set of metrics and on a diverse set of tasks/settings.", "theoretical_claims": "Theorem 1 is the primary theoretical claim and is supported by a proof, which seems correct.", "experimental_designs_or_analyses": "The authors consider a thorough set of empirical experiments for the tasks of protein design and cell-type specific sequence design to evaluate and validate their proposed method. In general, the experiments which the authors conduct in this work appear sound and valid.", "relation_to_broader_scientific_literature": "This paper tackles problems of two broader areas: (1) controllable generation via diffusion models, and (2) protein and sequence design. Both are active fields, where addressing items (1) and (2) would have significant impact in the respective areas. I believe this work makes a sound contribution to both fields.", "essential_references_not_discussed": "I think one related reference was missed [Domingo-Enrich et al. 2024] on the topic of \"Guidance (a.k.a. test-time reward optimization) in diffusion models.\" To add, I encourage the authors to denote the differences between their proposed method and that of [Domingo-Enrich et al. 2024] for the refinement of diffusion models, or consider it as a baseline.  \n\nOtherwise, and to the best of my knowledge, all relevant related works are discussed. \n\n- Domingo-Enrich, Carles, et al. \"Adjoint matching: Fine-tuning flow and diffusion generative models with memoryless stochastic optimal control.\" arXiv preprint arXiv:2409.08861 (2024).", "other_strengths_and_weaknesses": "In general, I believe this is a well-written and easy-to-follow paper which showcases some convincing experimental results while also providing solid theoretical contributions to back the proposed approach.  \n\nI did not find any obvious weaknesses.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "1. From algorithm 1 and Figure 3, does this mean you run the inference process of the diffusion model $S-1$ times? i.e. if the diffusion models uses $1000$ inference steps, this means RERD needs $100 * (S - 1)$ steps to generate a sample?", "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": "djoQm9Sy7s", "reviewer": "Reviewer_HSgw", "summary": "The paper introduces 'Reward-Guided Evolutionary Refinement in Diffusion models (RERD)', a framework for optimizing reward functions during inference time in diffusion models. RERD employs an iterative refinement process consisting of two key steps per iteration: noising and reward-guided denoising. This approach enhances downstream reward functions while preserving the naturalness of generated designs. The framework is backed by a theoretical guarantee and demonstrates superior performance in protein and DNA design tasks compared to single-shot reward-guided generation methods.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your constructive suggestions and insightful comments! Following a reviewer's suggestion, we added (1) more ablation studies and (2) additional experiments for image generation with Stable Diffusion and MaskGiT\n\n> Ablation studies on the noising faction (K) \n\nThank you for the thoughtful suggestions regarding the ablations. In response, we have performed additional ablation studies by varying key hyperparameters. To provide a quick yet informative signal, we focused on the ss-match and cRMSD tasks. Here is a  [link](https://anonymous.4open.science/r/hidden-1886) to figures describing experimental results.\nWe plan to extend these studies to other tasks in the final version. \n\n- We added an ablation study varying K (**Figure 1,2** in the link) by fixing the computational budget for evaluating reward models. The results show a strong performance when K/T=10% or 20%. Generally, a large K/T reduces the benefit of refinement, while a very small K/T limits the opportunity for reward-guided decoding.\n\n- We also performed an ablation on L, the number of repetitions for importance sampling (**Figure 3,4** in the link). As expected, performance improves with a larger L due to the increased computational budget. \n\n> The proposed method is only implemented on discrete models\n\nThat's a great point! We’ve focused on discrete diffusion models because they tend to have a greater impact in the protein design domain. That said, our method can integrated with continuous diffusion models as well. To verify this, we have implemented our method when **we set Stable diffusion as pre-trained continuous diffusion models** and compressibility (the negative file size in kilobytes (kb) of the image after JPEG compression) as reward models (Figure 5 in the link). Following our experiment section, we have tried two scenarios where we set K/T to be 10% and 20%.  This figure also highlights the effectiveness of iterative refinement in continuous diffusion models, as we showed in our protein design scenarios (Figure 6 in our original draft).  We will incorporate them in the final version.\n\n> Limited novelty over [1]\n\nThank you for pointing out this work—we will certainly include a citation in the revised version. From our understanding, the paper introduces an SMC-based approach similar to other related methods we have cited (e.g., Wu et al., 2024; Dou and Song, 2024). However, it appears to follow a more single-shot sampling strategy. While the restart sampler component may share a similar spirit, **Our main contribution—an iterative refinement procedure tailored for reward optimization, supported by both theoretical and empirical evidence**—differs substantially in both methodology and intent. We will make this distinction clearer in the final version. \n\n> Limited Impact Due to Model Choice (EvoDiff is less popular) \n\n1. To the best of our knowledge, EvoDiff is widely recognized as a representative discrete diffusion model in the protein design domain, as noted in recent reviews (e.g., Winnifrith et al.). While other pre-trained diffusion models, such as DPLM and ESM-3, are also potential candidates, incorporating them into our framework would be relatively straightforward. We are be happy to include additional results if the reviewer has specific protein diffusion models in mind.\n\nWinnifrith, Adam, Carlos Outeiral, and Brian L. Hie. \"Generative artificial intelligence for de novo protein design.\" Current Opinion in Structural Biology 86 (2024): 102794.\n\n2. To further address the reviewers’ concerns, **we conducted additional experiments on image generation tasks using a different discrete diffusion model, implemented on top of the [MaskGiT](https://masked-generative-image-transformer.github.io/) codebase**.  Here, the setup closely resembles that of EvoDiff. We set the duplication number to L=20, and in each iteration, we remask a 10% square region of the entire image. Experiments are conducted across 32 image categories. As shown in **Figure 6 (linked above)**, compressibility consistently improves over iterations, demonstrating the practical effectiveness of RERD. One point of clarification: the compressibility here appears higher than that observed with Stable Diffusion earlier. This is primarily because MaskGiT operates in compressed sequence spaces, making it easier to optimize.\n\nThese results indicate that our method performs robustly across various pre-trained models. We will include more comprehensive quantitative results in the final version, and of course, we would be happy to provide further clarifications during the rebuttal process.", "claims_and_evidence": "The claims are not fully substantiated by clear and compelling evidence. The author suggests setting K/T to a low value; however, there is no ablation study on the noise scale K. Without such an analysis, it is difficult to assess the model’s performance in terms of reward estimation and computational cost across different K values. Additionally, while the paper asserts broad applicability to all diffusion models, the proposed method is only implemented on discrete models, raising questions about its generalizability.", "methods_and_evaluation_criteria": "Yes, the proposed method aims to overcome the limitations of single-shot approaches in optimizing complex rewards and managing hard constraints, which is relevant to the protein design task.", "theoretical_claims": "Yes", "experimental_designs_or_analyses": "Yes", "relation_to_broader_scientific_literature": "The paper is related to finetuning diffusion models with guidance and inference-time scaling for diffusion.\nThe paper is also related to applying RL for protein design.", "essential_references_not_discussed": "No", "other_strengths_and_weaknesses": "Weaknesses:\n1. Limited Impact Due to Model Choice – The use of the relatively less popular EvoDiff model for protein design may restrict the broader influence and adoption of the work within the field.\n2. Limited Novelty – The proposed method shares similarities with [1], which also employs a resampling-based correction approach for diffusion models, reducing the novelty of the contribution.\n\n[1] Liu, Yujian, et al. \"Correcting diffusion generation through resampling.\" Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2024.", "other_comments_or_suggestions": "None", "questions_for_authors": "None", "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:

Meta-review: This work introduces RERD, an iterative refinement approach for reward optimization in diffusion models which alternates between noising and reward guided denoising steps. This extends the single shot reward optimization and sequentially refines samples, increasing the ability to optimize complex reward functions. The theoretical claims are backed-up and illustrated on protein and DNA sequence design.

  All reviewers acknowledge the relevance and interest of the method, stressing the clear presentation and novelty of the framework. Some minor improvements are suggested, so I recommend acceptance as long as they are incorporated in the revised version of the paper.
Decision: Accept (poster)