Steering Masked Discrete Diffusion Models via Discrete Denoising Posterior Prediction
Authors: Jarrid Rector-Brooks
Mohsin Hasan
Zhangzhi Peng
Cheng-Hao Liu
Sarthak Mittal
Nouha Dziri
Michael M. Bronstein
Pranam Chatterjee
Alexander Tong
Joey Bose
Venue: ICLR 2025 Poster (2025)
Keywords: Discrete diffusion models
language modeling
probabilistic inference
Primary area: generative models
TL;DR: We fine-tune discrete diffusion models by solving a probabilistic inference task.

Abstract: Generative modeling of discrete data underlies important applications spanning text-based agents like ChatGPT to the design of the very building blocks of life in protein sequences. However, application domains need to exert control over the generated data by steering the generative process—typically via RLHF—to satisfy a specified property, reward, or affinity metric. In this paper, we study the problem of steering Masked Diffusion Models (MDMs), a recent class of discrete diffusion models that offer a compelling alternative to traditional autoregressive models. We introduce Discrete Denoising Posterior Prediction (DDPP), a novel framework that casts the task of steering pretrained MDMs as a problem of probabilistic inference by learning to sample from a target Bayesian posterior. Our DDPP framework leads to a family of three novel objectives that are all simulation-free, and thus scalable while applying to general non-differentiable reward functions. Empirically, we instantiate DDPP by steering MDMs to perform class-conditional pixel-level image modeling, RLHF-based alignment of MDMs using text based rewards, and finetuning protein language models to generate more diverse secondary structures and shorter proteins. We substantiate our designs via wet-lab validation, where we observe transient expression of reward-optimized protein sequences.

Published as a conference paper at ICLR 2025
STEERING MASKED DISCRETE DIFFUSION MODELS VIA DISCRETE DENOISING POSTERIOR PREDICTION
Jarrid Rector-Brooks!: 1.2 Mohsin Hasan 1.2 Zhangzhi Peng", Zachary Quinn", Chenghao Liu?.., Sarthak Mittal .2, Nouha Dziri Michael Bronstein Yoshua Bengio Pranam Chatterjee', Alexander Tong*.2 Avishek Joey Bose Université de Montréal, Mila, Duke University, "McGill University, 5 Alllen Institute for AI, OOfford University, 7 'Aithyra
2.2,
ABSTRACT
Generative modeling of discrete data underlics important applications spanning text-based agents like ChatGPT to the design of the very building blocks of life in protein sequences. However, application domains need to exert control over the generated data by steering the generative process proccls-typically via RLHF- satisfy a specified property, reward, or affinity metric. In this paper, we study the problem of steering Masked Diffusion Models (MDMs), a recent class of discrete diffusion models that offer a compelling alternative to traditional autoregressive models. We introduce DISCRETE DENOISING POSTERIOR PREDICTION (DDPP), novel framework that casts the task of steering pre-trate pre-trained MDMs as a problem of proba- bilistic inference by learning to sample from target Bayesian posterior. Our DDPP framework leads to family of three novel objectives that are all ionulation-free and thus scalable while applying to general nin-differentiable reward functions. Empirically, we instantiate DDPP by steering MDMs to perform class-conditional pixel-level image modeling, RLHF-based alignment of MDMs using text-based re- wards, and finetuning protein language models to generate more diverse secondary structures and shorter proteins. We substantiate our designs via wet validation, where observe transient expression of reward-optimized protein sequences.
INTRODUCTION
The success of diffusion models in continuous spaces, leading to state-of-the-art foundation models for image (Stability AI, 2023; Midjourney, 2023) and video synthesis (Villegas et al., 2022; Brooks et al., 2024), has spured several attempts translate these approaches for the generative modeling of discrete structures. The most performant approaches squarely fall under the scalable framework of absorbing state discrete diffusion (Austin et al., 2021), with new simplified training recipes that result in Masked Diffusion Models (MDMs) (Sahoo et al., 2024; Shi et al., 2024; Gat et al., 2024; Zhao al., 2024a). Indeed, recent MDMs now rival autoregressive models of similar scale to GPT-2 (Radford et al., 2019) for language modeling, with the potential for further progress through scaling. Furthermore, MDM style models are not constrained to generating data sequentially- siiiells-hichhh invites a more straightforward application to domains without autoregressive ordering, e-g. molecule generation (Vignac et al., 2022), discrete modeling of a natural causal images (Salimans et al., 2017), and modeling protein sequences (Lin et al., 2022; Wang et al., 2024).
Critical to the successful deployment of discrete generative models in practical applications-beyoole simply producing high-quality samples-is the ability to steer the generated samples optimize pre-specified downstream metric. For instance, in Language Modeling (LM) it is desirable to bias the model's generations to be sanitized from harmful responses (Zou et al., 2023; Perez et al., 2022), or aiming to generate protein sequences that are highly likely to be successfully synthesized and expressed in real wet lab settings (Verkuil et al., 2022; Dauparas et al., 2022). Put succinetly, highly performant discrete generative models are required to be aligned in manner that fine-tuning against downstream reward models has the intended effect of controllable generation, wherein the model post fine-tuning selects high-scoring samples from the universe of possible high-fidelity generations.
"Equal advising

---

Published as a conference paper at ICLR 2025
The standard approach for incorporating steerability into discrete generative models, which autoregressive, using pre-defined reward models is often framed as a fine-tuning task using reinforce- ment learning from human feedback (RLHF) (Christiano et al., 2017; Rafailov et al., 2024). However, applying RLHF frameworks to diffusion models is far more challenging. Unlike autoregressive models, diffusion models do not allow for straightforward computation of a sample's i exact likelihood without costly simulations. Although fine-tuning diffusion models that bypass exact likelihood computation can yield simulation-free algorithms that resemble RLHF (Wallace al., 2024; Uchara et al., 2024a), these methods effectively optimize loose lower bound to the true RLHF objective, leading to unstable training and suboptimal fine-tuning performance. Consequently, steering diffusion models in continuous spaces is primarily done through inference techniques that leverage the gradient of a conditional model in the form of guidance (Dhariwal and Nichol, 2021; Ho and Salimans, 2022). Unfortunately, discrete settings do not allow for principled definitions of guidance due to the lack of a conventional gradient operator. As a result, at present, there exists no scalable and rigorous method to steer and align Masked Diffusion Models to optimize desired reward models.
Main contributions. In this paper, we cast the problem of steering a Masked Diffusion Model as a task of probabilistic inference in sampling from a target Bayesian posterior. More precisely, we construct the target Bayesian posterior as being proportional to the product distribution of a base pre-traine pre-trained MDM model modulated by a pre-specified reward model. Importantly, this sampling viewpoint is fully compatible with classical RLHF for autoregressive models (Uehara et al., 2024a; Zhao et al., 2024b), but enjoys broader applicability as for the first time it can be applied to discrete diffusion models. Under this sampling perspective, our key insight is that the challenging sampling problem can be solved by learning an amortized sampler, which when taken as an MDM, can be viewed as finetuning the pre-trained MDM model by learning approximate the ) Bayesian posterior.
We introduce DISCRETE DENOISING POSTERIOR PREDICTION (DDPP), novel framework that exploits the denoising posterior parametrization inherent to current MDMs to define a series of simpler matching problems across varying corruption (masking) levels of the target Bayesian posterior. In particular, DDPP designs a forward process that corrupts the Bayesian posterior through a forward masking process and frames the finetuning task as learning another MDM to approximate the corresponding reverse process. As a result, each matching problem in the reverse process requires the construction of "denoising" Bayesian posterior that is conditioned on a partially masked sample which we demonstrate is simply proportional to the pre-trained model's 's own denoising posterior and the (terminal) reward of the fully unmasked sample. Crucially, each matching problem in DDPP can be defined on particular noise level without running the entire forward (corruption) process. Consequently, this makes DDPP simulation-free method which finetune large pre-trained MDMs. We test the empirical caliber of DDPP by steering MDMs across a multitude of domains ranging from images to protein sequences and steering MDM-based language models. We observe DDPP fine-tuned MDMs lead to competitive performance expression of reward- wwward-ptiiizzed protein sequences (with high secondary structure diversity and 3-sheet composition) in a wet setting, and natural textual responses that are steered to human sentiments.
key ingredient needed to
images, transient
BACKGROUND AND PRELIMINARIES
Notation and convention. A discrete data sample æ is identified by its realization over a vocabulary set X' = (1....,d d 1), over d - 1 possible categories. Of particular interest is the setting of masked diffuusin models that include an additional d-th category of masked token m to the vocabulary X which serves as an absorbing state for the diffusion process. A discrete token represented by the one-hot vector ef E A in the d-dimensional probability simplex and corresponds to placing a I on the i-th index and on all the other d indices. In this paper, by convention, we set e = ca as the one hot vector associated with the masked state m. A categorical distribution over d-categories, Cat(z: p). constructed by placing a Dirac with weight p', with the constraint Σ, p² = 1 and the density of discrete sample is written as p(X x) = Di- p0 - e²), where X is the discrete random variable.
A sequence =
of tokens is defined over the product space X" (1..... m}" and its corresponding probability mass function is given by p(x = x) = II E 12 5(x² e²). To reduce notational clutter, we make use of the shorthand 5(g) to denote Dirac measure on a discrete sample and interchangeably write p(X- = x) p(x) to denote the probability mass function. A dataset of sequences is designated as samples from the data distribution Pdata to be learned by a discrete diffusion model qo, with parameters 0. Discrete diffusion models, like their continuous
2

---

Published as a conference paper at ICLR 2025
counterparts, are stochastic processes that evolve with time t E [0, 1] such that t = 0 corresponds to Pdata := Po and = 1 corresponds to the terminal distribution, P1 of the process. As a discretization drop i and simply write t
of time, we divide [0, 1] into T intervals, and let t(i) i/T. For brevity, to denote the corresponding discrete timestep i(i). The notation 0 1 designates a collection of objects, e.g. densities p( (().:)), starting from time t. and including time 1 = O. A trajectory of sequences is denoted as ((xo11) X1 - + X + x + ... + Xo. Finally, we use subscripts to denote the time index-i.e. P--and reserve superscripts to designate indices over a set such as specific sample x' among collection of samples or dimensions within a vector, e.g. dimension x' in sequence.
Problem Setting. We are interested in the task of probabilistic inference of sampling from an unnormalized target distribution (xx) defined over a discrete space consisting of n tokens X
Po (xo) R(xo)
R(x0)= exp(- 2R
To To(xx)=
(1)
2m
A key aspect of the considered setting is that To (Xo) is defined as the product distribution of a pre-trained masked discrete diffusion model ph (xo) and distribution induced by (potentially differentiable) reward model R: : X" + R. The problem definition in Eq. 1 is an instance of Bayesian posterior sampling where the pre-trained MDM is the prior and reward acts as the likelihood or observation model which modulates samples with a high score. For instance, in scientific domains, the reward model can be provided as a Boltzmann distribution with known energy function & E(x)), or a human preference model as in RLHF (Ouyang et al., 2022; Rafailov et al., 2024). Importantly, this set- ting does not afford us any ground truth samples from mo (xo) in the form a dataset which prevents classically training another generative model. Instead, we are able to evaluate the reward model-and in special cases its gradient V R but not the normalizing constant, i.e. the partition function 2T.. Samples from the posterior To thus lie in the intersection of the modes of both the pretrained MDM and the reward model. As a result, learning an amortized sampler, qo (Xo). for æ (xo) is rationally equiv- alent to finetuning the pretrained MDM p (xo) using the reward R( (x) in analogous manner to RLHF (Uchara et al., 2024a) and is the main focus and contribution of this paper and outlined in $3.2.
2.1 SIMPLIFIED MASKED DISCRETE DIFFUSION
We are interested in developing a discrete diffusion model directly on discrete data-i.e. without embeddings or continuous reparameterizzations whose approach mirrors the construction of diffusion models for continuous spaces. Consequently, we require the specification of a forward process that converts discrete data Xg ~ Po at time / = 0 to an unstructured prior, P1 at the terminal time t = 1. The specification of a forward process via the transition kernel pt (x (xo) implies a unique time reversal of this forward process, termed the "reverse process", such that simulating from this reverse process results in samples from the desired target data distribution (xo).
We restrict our attention to the recent performant "simplified masked" forward process (Sahoo et al., 2024; Shi al., 2024; Gat et al., 2024; Zhao et al., 2024a) which hits a terminal distribution of all mask tokens in sequence pl [ó(m)]". Given a non-masked token in a sequence, z1 E x the simplified masked forward process increases the likelihood of transition time increases. Moreover, the masked forward process is simplified by design since the transition probabilities of a token unmasking (x/++ * m when o = m) is set zero zero-i.e. the token remains a masked token for the remainder of the trajectory. The design of the simplified forward process is also independent across each dimension of the sequence, conditioned on Xo, which allows us to model the transitions of each discrete token in sequence separately. In totality, the forward process for a sequence Xg can be summarized using the following expression for the transition kernel pe (a) [x)::
the mask state as
Pr (x, = II P.. (=)l)))
II Catl (2t - (1 - az (m))) T-1
(2)
where at is an invertible reparameterization of time such that 0gg and on = 0. Effectively, o corresponds to the noise schedule which corrupts the discrete data to p1. The corresponding marginal density induced by the forward process at time & can written as (xx) = Σx, Pe (xx (x0 (x)))xx)))))))
The reverse process which denoises sample from t - f- 1, and the time reversal of the simplified masked forward process, also factorizes over each dimension of sequence x. The probability pe pr( |æt,, .r)) of a reverse transition is given by the following posterior conditioned on x/.
(Cat(x(=i;((<t)
x + m
( at)) = Cat [وا_م; (1-ae-
(3)
xf = m.
1-0r
3

---

Published as a conference paper at ICLR 2025
Under the reverse process once token transitions out of the masked state for a time > 0 it remains in this state for the remainder of the trajectory. The analytical form of the posterior suggests a natural mean parametrization for denoiser in discrete diffusion model, po th predicts the clean sample at = 0 by denoising a noisy d/
x R + (Ad)", which
[Cat (Ct((()_1;((zt)
+ m
21.6 (,o(ai11lo uo (2f, 1)) Cat
(3-11; (1-ax 1ab (=
x m.
1-a,
Interestingly, the mean parametrization (40 used in the posterior is equivalent to predicting the concrete score (Meng et al., 2022) which is the discrete equivalent of the Stein score found in conventional continuous diffusion models (Zheng al., 2024). As the number of steps t - cc, training yields valid evidence lower bound (ELBO) to the marginal log-likelihood of the data distribution log p(xo).
((x)( 2- S( doi dt
1 Cz2
E. Exirpe (xc)
Σ
n
log log dt.
(5)
Thus, when given access to samples Xa po training an MDM can be seen as optimizing weighted cross-entropy loss and is analogous to fitting a (mean- field) variational posterior distribution qr.o (xxo[x)) Cat(xo: (xx, t)) that matches the first moments of pe (xx) and also minimizes the forward KI. divergence DEL (P& (xolx))pe (Pe()))()))) (x:))xx()ll ol,,, xx 8x)p:(xx) (Fijkelboom et al., 2024).
POSTERIOR SAMPLING VIA DISCRETE DENOISING POSTERIOR PREDICTION
Given access to pretrained masked discrete diffusion model ph (xo) we wish to sample from the reward-induced Bayesian posterior distribution (xx) x Po (xo) R(xo). We solve this sampling problem by first defining time-dependent forward masking process that progressively adds noise to To yielding the noisy swwrrd-inducdd posterior (x{)) Ex. Te (x( Xx (T (xx).. where we set
T (x, (1) = Pt (xxxxxxx (xx(X)) as it is the same masking process for the pre-trained MDM. Unfortunately, pré
since Po (xx) is MDM it does not easily provide exact likelihood. Undeterred we seek to
approximate the reverse process Te TT ( tied to the masking forward process by using another parametrized model gr.a (x.
Cat(xo; no (x,, ()) which we take be another MDM.
Matching sub-trajectories. To approximate the reverse process using MDM we require matching the denoising trajectory r(Xa:) of the reward-induced posterior ne (xx........ x xxxx across all masking levels. Assisted in this endeavor, we recall the fact that since p5 (x,) is also an MDM, have direct access to the pre-trained model's denoiser. Thus, we can compute any transition density starting from y (x,-1 (xPe f)) to the posterior over the endpoint y (xo (.)))).. conditioned on a partially masked sample X,. We form the sub-trajectory matching problem as an instantiation of detailed balance constraint starting from partially masked sequence x, of clean starting point Xo:
qa (xxx..... x,-1 xx, (x)) = Tf(Xo.-... X, (x,1/xx)((xx
Setting Xn = px ()
as the MDM's denoised sample, then (x ........... Xr- (x) is defined as.
1, xI.
I1;=1 p. (x] xj, x R(xx)
Tx(Xo. x xx,)
Zr, (x1)
Zz, (x1)
The detailed balance constraint over sub-trajectories suggests natural discrete denoising posterior predictive (DDPP) objective that minimizes the mean squared error of a log-ratio between the denoising sub-trajectories of the amortized MDM sampler and the reward-induced target posterior,
C c Er Ex, E+(xx::) [II| log qo (Xo.o1 xb)) loog pt (xo:111(x;) + xil2i]
(7)
where reward and the log partition function are captured in the constant K = log 2r, (xx) - log R(xo).
Interestingly, we can form an equivalent expression for the sub-trajectory loss CPP above by sampling
two intermediate points x,, X,-y in the sub-trajectory (xxx),, such that 0 < s 7 < t:
C PP
Ex, 77 (Ko:s)
[||EE,,,,.,.,, [log go (x...,,/ (x,, x0) - log p."h (Xs- x,, x + k5] [12]
(8)
The proof for this equivalence is presented in SC.3. Note that we sample s, s - 7 ~ ([0 2), s] uniformly, and when 7 = 1/T we sample x which is simple to do since the r(xo:a) already contains this information. Crucially, unlike Eq. 7 the reformulation of the sub-trajectory loss in Eq. 8 is effectively a simulation-free free version of Relative Trajectory Balance (RTB) (Venkatraman et al., 2024). If the approximation 4:,0 matches the denoising -iruurdd target posterior over all sub-trajectories then the reverse process of gt,0 can be simulated to draw samples that follow To(X)).
4

---

Published as a conference paper at ICLR 2025
Consequently, term the q that minimizes the DDPP objective in Eq. 8 as the finetuned MDM which solves the probabilistic inference task of sampling from mo (xo).
In contrast to learning MDMs typical generative modeling setups, the DDPP objective requires the computation of the intractable log partition function log 2z, (xx) evaluated at x, which is component of the term k. This observation motivates the design of three concrete learning objectives for posterior matching. which as collection we term the DISCRETE DENOISING POSTERIOR PREDICTION framework. Specifically, finetuning Qt, under a DDPP framework can be done in the following algorithms: 1.) DDPP-IS which uses Monte Carlo based importance sampling estimate to approximate log 2,, in Eq. 8, 2.) DDPP-LB that constructs a lower bound to DDPP-IS that is cheaper to evaluate by parameterizing log Z and 3.) DDPP-KL which uses a discrete gradient estimator to bypass computing log 2x, at the cost of requiring a differentiable reward-i.e. VR.
3.1 ESTIMATING THE LOG PARTITION FUNCTION
Inspecting the posterior predictive objective in Eq. we remark that it a simulation- stochastic regression objective which does not require a differentiable reward as the loss computes R (x(x)) and not a gradient of the reward. Consequently, this makes the posterior predictive objective both scalable and efficient objective for fine-tuning large pre- pre-trained MDMs as long the reward model is easy to eval- uate, Moreover, the posterior predictive objective is also an off-policy objective as it can be evaluated using any partially masked samples X2 p(x;] X2). Practically, this means that fine-tuning formed using a replay buffer of samples from a biased dataset, e.g. the original training set for o or even partially masked sequences that arrive from different model altogether. Despite its simple form the posterior predictive objective requires the computation of the log partition function of a partially masked sequence log 2z, which does not have a closed-form expression and must be estimated.
be per-
Monte Carlo Estimate of log 2x, with DDPP-IS. A numerical estimate of the log normalization constant can be obtained by utilizing the trick of using the pre-trained model's 's denoising posterior ype (xa) x,). Specifically, given X₁ pe(x.) obtain Monte Carlo estimate of log 2z, (x,) that tre
uses M additional samples from Xg ~ Pt xx) to estimate the log partition function,
log 2,, (xx) log
Σ X1
D Pt (Xo: (X., x1-1 (x, zu log (Ixx (xxp x. [((2))))
Where in the second equality in the first line we used the fact that we can approximately jump to the endpoint of the reverse process directly by using the pretrained model's denoiser sample Xo- Conveniently, this MC estimate solely requires obtaining denoised sample from the pre- pre-trained MDM which can be efficiently done as each sample requires a single step as due to the denoising posterior parametrization of an MDM (Eq. 4). We can further improve the estimation of this log normalization constant by leveraging importance sampling (IS) with a proposal distribution uxx)):
log 21S (xx) = log (xx))u((x
pre (xolxx) 14
1) =log log
1 .
M
pre (xo (xo1xx) (xolxx))((()) 1
xxx)
w(x))
For the IS estimator above it is easy to verify that the optimal proposal distribution for variance reduction is proportional to the denoising reward-induced target posterior w* (xo) Tx xt). Fortu- nately, this is precisely the distribution that is approximated by qr,9 using the posterior predictive objec- tive which motivates the reuse of the finetuned model as a suitable proposal, i.e. w(xu) qe,0 (x, (x))
Learning log 2n, with DDPP-LB. An alternative to using an MC-based estimate for log 2x, is LLB to parameterize the log partition function itself log 2 jointly with the qt.8 and optimize both using the same posterior predictive objective as first defined in Eq. 11. Operationally, this amounts to including another prediction head for the finetuned MDM model and is cheaper to compute than using an MC-based estimate as we do not require M evaluations of the pre-trained model in log 215 (xt).
At first glance. it remains unclear whether parameterized log ŽLB is sensible strategy. However, in the particular case where we choose the proposal distribution to be on- on-policy by using finetuned MDM wxxx)= = qt.0 (x. xx). we can show that the learned log partition function estimate is lower bound to the importance sampling estimate. This is formalized in the following proposition below.
Proposition 1. Let log 255 and log 21B be the M M-sample importance sampling estimate using the proposal (xoo(xo)xx, (x.) and learned approximation to the log partition function respectively. Given
5

---

Published as a conference paper at ICLR 2025
Algorithm 1 Single-step DDPP-IS and DDPP-LB
Input: Reward R(xo), base MDM Pa (xo|xx), sampling policy r (Xo), fine-tuning MDM (Xo)xx:)
I: while Training do
i, Xo ~ 14 [0, r1)o)()))
Sample time and clean data on or off-policy > Construct partially masked sample given clean data
3: X. pt (xx (xx|xx)
4:
if Importance Sample log (x.) then
I Log Partition Function Estimation Strategy
5:
log 2x, (xx) = log 2 (x() log
Σラ=1
pt (x)
w(xc))
6: else
log 2x, (xx) = log 144 20,
_
log (xx0(x))xx)) - log p! (xo) xx)
log R(xo) + log 2x, (xx)
Up, VoC*e)
10:
Return Qa
partially masked sample Xt (x() the optimal learned approximation lower bound the importance sampling estimate with fixed proposal Qr,a (xxo xx) and the following inequality holds:
log 2 (x,,(xx)) S logg 215 (x:)-
(9)
The proof for Eq. is provided in SC.1. We highlight that the lower bound becomes equality at the optimal proposal Ir, (xo (x)]xx) x Te Tx (x.). Learning log 2LB has the benefit of amortization as the same network can be reused for all partially masked samples X₃ ~ Px(x)), across all levels of masking. In addition, over the course of training, the learned estimate log n..0 becomes better estimate for the true log partition function. In practice, suffices take single gradient step to rather than optimizing till convergence. As a result, no additional overhead needs
ŽLB
optimize log 2LB n..8 to be incurred, and the learned estimate is averaged over batch of noisy samples B = {x]]-=
3.2 SINGLE-STER POSTERIOR SAMPLING WITH ENDPOINT PREDICTION
The sub-trajectory matching objective used by DDPP-IS and DDPP-LB can be simplified to faster single-step objective at the cost of paying a discretization error by not using finer-grained trajectory information. Specifically, we note that for MDMs the denoising posterior over end- points dr,0 (Xo xx) As Cat(xo; pe (X, t)) can be approximately computed without unrolling the sub- trajectory. This fact also holds for the pre-trained MDM as the model parametrization implies pt (xo) zs Cat(Xo: p(x,, t)). For the single-step objective we assume the parameterized denois- exactly match the posteriors. Leveraging this enables us express the denoising ewarddinducedd target posterior using a simple expression that directly uses the pre-trained model's 's denoising posterior pt. (xo(x;) as follows:
<x(oxxx)) (xxe
Pa (X, "pee "r (x) R(xo)
Pt. (xo (x) R(xo)
(10)
(x() Ex, Pt (X: x Ire (x)) R(x))
Zx, (x.)
The choice parameterizing 4t,0 (x, (x1) as another MDM offers a prescriptive strategy for sampling from the desired target mo by learning match the denoising soward-induced posterior at the pre- dicted endpoint n. (x)) This simplifies the expression DDPP defined over trajectories in Eq. single point, namely the predicted endpoint xo of each MDM. This objective presented below:
2
CPP Ex
log qr,o (x,o(xx)xx log P. (xo) xx) log R(x,) log ZT, (xx)
(11)
log (Mpl)x))
As done previously, we can employ any estimation strategy to compute the log partition func- tion Eq. 11. We note many cases, such as when the sequence length of the trajectory is small to moderate, the single-step objective may be an attractive alternative to the sub-trajectory variants of DDPP. Algorithm provides a detailed description of the single-step version of DDPP.
3.3 DDPP-KL: POSTERIOR PREDICTION VIA REVERSE KL MINIMIZATION
The single-step posterior prediction objective as defined using the loss function CP in Eq. 11 requires the estimation of log ZLB which introduces a source of variance in loss estimates that may sub-

---

Published as a conference paper at ICLR 2025
optimally influence learning dynamics of the fine-tuned model. In settings where the reward model is differentiable, we can bypass computing log Z altogether by learning to match the denoising posterior under the reverse KL divergence. Note that the forward KL divergence is inapplicable here as we do not have samples from m--i.e. a dataset. To see this, we define a variational posterior matching problem using the reverse KI. divergence that takes the following form:
8
CKL == Kl. (gr. (Xx Xxx ) x | (x (x))
(12)
Unlike conventional generative modeling using the reverse KL divergence which solely matches distributions at = 0 the problem definition in Eq. 12 defines a series of reverse KL minimization problems through time. In this manner, the reverse KL matches distributions annealed through time and can be used to derive a stochastic regression objective for fine-tuning,
(13) The expectation in Eq. 13, like DDPP-IS and DDPP-LB is taken uniformly with respect to time t ~ [0 1]. However, unlike the previous estimators, clean data needed to compute _KL is drawn purely on-policy by simulating the fine-tuning model X ~ a (Xo), which then also allows us to craft a noisy sample using the masking forward process X ~ Pr (x,(xx) Additionally, in Eq. 13 the constant C' = Xo, [log Zz, (xx) ] does not depend on the --dadd as result is also independent of the sample Xy ~ qt.0 (xa). This results in the constant C' being zero when computing the gradient of the loss CoLkL and result we can safely disregard computing log Zn, entirely.
CKL Er, xo.x, [log or,0 Xx) log pt. (Xo]x:) - log R(xo)] +
As samples Xo are procured on-policy to compute the gradient of the loss VoLki. we require backpropagating through the stochastic sampling of x which comes from simulating the fine-tuning MDM (t,0(xo). Fortunately, we make use of modern discrete gradient estimators which provide a biased but low variance gradient estimate enabling us compute LKL, Specifically, we opt to use the scalable 2nd order REINMAX estimator (Liu et al., 2024) which estimates the discrete gradient up to second-order terms in Taylor approximation of the actual gradient. We note that unlike DDPP-IS and DDPP-LB this new loss that minimizes the reverse KI. divergence L requires the reward model R to be differentiable and as a result is less broadly applicable than computing CPP However, in practice, learning can be faster as we make use of the information afforded to us by the gradient VR as well the fact that the objective does not need to estimate the log partition functionn
In appendix $C.2 we provide the exact algorithm Alg. 2 to compute the reverse KL objective. We further show how using gradient estimator like REINMAX can be used to derive efficient gradient estimation for a more general class of problems of sampling from (xp)) R(x,) Zmr, as well as the main fine-tuning setting for matching the denoising reward-induced posterior as defined in Eq. 10.
Table 1: Overview of posterior sampling methods
4 EXPERIMENTS
Method
Model calls / inl. Model train step
Sim. Fize
We investigate the application of DDPP to variety of discrete generative modeling settings. We provide the full experimental details in $D and present our main experimental results next.
SVDD
N
Discrete gustanee
RTB
T
DDFP KL
DDPP IS DDPF-LB
M
Baselines. Throughout our experiments, we rely on four principal baselines in: sampling from the pre-trained MDM model, Best-of-N sampling (Stiennon et al., 2020), Relative Trajectory Balance (RTB) (Venkatraman et al., 2024), and SVDD (Li et al., 2024) which is a concurrent inference time technique for steering diffusion models. Best-of-N represents a computationally expensive baseline but is guaranteed to produce samples from To. as such we use this as an upper bound on performance in terms of reward obtained as N + oo (Beirami et al., 2024). RTB is GFlowNet (Bengio et al., 2023; Madan et al., 2022: Lahlou et al., 2023) that requires simulating the entire diffusion trajectory. For image settings with differentiable reward, we also include discrete guidance as a baseline (Nisonoff et al., 2024). In Table 1 we illustrate the computational differences between DDPP and baselines.
4.1 SYNTHETIC EXPERIMENTS
We consider synthetic task of learning sample from a target distribution on 2D discrete grid and finetuning an MDM on binarized MNIST. This synthetic setting tests all DDPP variations with chosen baselines, presenting qualitative results in Figure 6, Figure 4 and quantitative results in Table 2.
Grid Experiment. We define prior density pp over the discrete 2- 2-dimensional, 128 128 grid, as showcased in Figure 6(a) where the probability mass corresponding to each point xn is on if the color is yellow. The goal is to sample from the product distribution outlined in Equation 1, which in this case
7

---

Published as a conference paper at ICLR 2025
(a) Prior Density
(b) Target Density
(c) DDPP-IS
(d) DDPP-LB
(c) DDPP-KL
Figure I: Samples generated by fine-tuning masked diffusion model to sample from the lower half its prior distribution. Samples Xo in this setting are 2-dimensional, with a vocabulary size of 128.
is defined to drop the modes in po which are at the top half of the grid, as visualized in Figure 6(b). These results show that all three variants of DDPP effectively learn sample from this target.
MNIST. We finetune MDMs to generate even MNIST digits. As observed in Table 2 we find that all three variants of DDPP match or outperform the base pre-trained model and RTB in all metrics, with DDPP-KL being the best. In comparison to the concurrent work of SVDD, we find that it outperforms DDPP in average log R but is worse in cmmllebbased metrics such as class conditional FLD (Jiraler- spong al., 2023) which measures the overall quality, diversity and generalizability of generated samples and class conditional BPD, We further report generated samples in Figure 4 located in $D.3.
4.2 PIXEL-LEVEL IMAGE MODELLING
Table 2: Fine-tuming tinettuning to produce only even digits on binarized MNIST. We report the mean performance over 3 runs for the log R. FLD, and BPD metrics.
We fine-tune MDMs on ander-agnostic image data, discretizing pixels in 64 x 64 downsam- pled CelebA images (Liu et al., 2018) to a vocabulary of 256 tokens. As there are no publicly available pre-trained MDM models we train our own MDM by modeling the raw pixel space and achieve 1.85 bits- bits-per-dim (BPD) on CelebA. Our full experimental setup is outlined in $D.3. For fine-tuning, we consider steering a pre-trained MDM using DDPP-LB as the most computationally cheap method with a class-conditions reward based Specifically, we steer the generative model to generate human faces with blond hair. For quantitative metrics, report the mean log reward obtained, and BPD in Figure 3 as well selected generated samples. Our quantitative results show that our proposed variant DDPP-LB significantly outperforms all other baselines in obtaining the highest reward. We also observe DDPP obtains BPD values that are within the range of the base model while being worse than RTB. We further find visual samples produced by DDPP to have the highest fidelity faces with blond hair, matching our fine-tuning goal.
Algorillun
log R(xa]
FLD J
BPD L
Base Model
26.900
33.89 +
0.130 ::
SVDD
-0013 0.01
34.19
+000+-900 0.95
Guidance (scale 1)
-25.24 + 0.26
34.67
+.67
0.171 + 0.001
Guidance (scale 5)
-23.21 0.21
37.33
0
0.174
0.001
Guidince (scale 100)
-9.32 1- 0.24
72.19
+ 0.43
0.147
0.001
RTB
-18.66 1 2:45
45.97
0.89
0.128 0,000
DDPT-IS (owve)
0.130 0.000
DDPP-LB (ours)
-5.68
0.34
33.76
0.90
0,128 0.000
DDPP KL (ours)
-3.13 0.06
31.75 0.51
0.129 1: 0.000
auxiliary classifier.
4.3 PROTEIN SEQUENCE MODELLING
Table 3: In-silico results for protein generation tasks. We report the mean result for metric with standard deviation across three seeds. DDPP-LB performs well across designability metrics (pLDDT and pTMM) while simultancously performing best on task specific metrics (3-sheet % and TM Score).
High Sebentecnntonn prntcin generation
Protein shrisking
P-sheet 9
PLDDT
PIM
log H(zu)
SS SS-KLL
TM-Score
PLDDT
PIM
log Rixzs) 1
Base Model
0.11 1 0.121
0.724 1 144
0.584
1 0.226
2.070 1 0.749
3,000 :: 3.03
0.245 1 0.58
0.724 0.144
0.584 1: 0.236
0.490 1: 0.116
Best 10
0,280 4: 0.093
0.112
0.033
3.766
0,835
3,212 2 0.371
1421
+: 2884
0.345 0U349
0.786
0.023
0.737
2: 0.997
0.600 à: 0.098
SVDD
0.114 4 0.148
0.484
1 0.134
0.349
1 0.174
1.669
1 0.907
3.353 4 2913
0.337
2 0.042
0.492
1: 0.131
0.368
0.17
0.673 1: 0.083
RTB
0.319 1319--00221
0.806
0.099
0.767
1 0.100
3.396 +: 1.061
2.193 2.224
0.290 0.56
0.707
0,866
0.547
+: 0.091
0:581 +: 1:1311
DDPP LB
0.436 1 0.037
0,897 2: 8.027
0.306 1: 0.029
3,700
1 0.186
0,640 2: 1.793
0 2: 0J047
0.768 0.045
0.747 1: 0.063
0,723 1: 0.094
Task description. We next apply DDPP to generate high-quality protein sequences by fine-tuning discrete diffusion protein language models (DPLM) (Wang et al., 2024). Specifically, we address two experimentally relevant tasks where vanilla DPLMs underperform. We outline exact reward functions and experimental setup in $D.2. First, we fine-tune DPLM to generate soluble protein sequences with high 8-sheet content. The second task, protein shrinking, involves miniaturizing known proteins by generating shorter sequences that preserve key structural features, using the TM-align score as the reward metric (Devkota et al., 2024). We evaluate performance by measuring designability metrics (ESMFold PLDDT and pTM) as well as task-specific metrics (8-shect percent and TM-Score). We also provide wet-lab validation for our best designs in the designable B-sheet task. We provide
8

---

Published as a conference paper at ICLR 2025
CCPP Putciss
16De
PAH rOrS USG 1STS
Core
Fare Model
Cuauity
20 10a
Date facentage
Figure 2: Left: SDS-PAGE of elution fractions from histidine tag purification of DDPP-designed protein constructs and positive controls following Coomassie blue staining. All DDPP-designed constructs are between 7.8-8.3 kDa. Predicted molecular weights of positive controls SKPH IQYS. IUBQ, and IBTB are 10 kDa, 12 kDa, 8.5 kDa, and 10 kDa, respectively. Recombinant protein bands indicated with red arrows and relevant ladder references are labeled with their molecular weight. Middle: Folded structures generated by DDPP 8-sheet fine-tuning, Right: Distribution of 3 8-sheets generated by each method.
MDM-designet sequences are
deeper description of evaluation metrics and experimental setup in $D.2. Finally, as ESMFold is itself expensive to query and, in particular, non-differentiable test our fastest method DIMN--DPP--BB
Main results. In-silico validation shows that DDPP-LB outperforms all baselines for the designable 3-sheet task, generating better sequences across all metrics. In particular DDPP achieves a significantly higher B-sheet percentage than baseline methods while maintaining high designability as measured by ESMFold (namely, high PLDDT and pTM)). We further observe that for the miniaturization task, DDPP-LB outperforms all baselines in shrinking ribonuclcase proteins, removing 34 residues while maintaining high structural similarity (lowest SS-KL of 0.64 and highest TM-Score of 0.361), and high structural quality with high pTT' and competitive PLDDT. This demonstrates DDPP-LB's effectiveness in generating compact yet structurally faithful proteins.
Experimental validation. We selected designs from finnttind DPLM for wet-lab validation, based on AlphaFold2 pLDDT/pTM scores. Sequences and structures were clustered using MMseqs and Foldseek (van Kempen et al., 2022; Steinegger and Söding, 2017), with two represen- tative sequences selected from each cluster. 4 positive controls consisting of two previously validated de novo designed proteins (PDB: 5KPH, IQYS) and two other stable proteins, ubiquitin and Barstar (PDB: 1 IUBQ, IBTB) were included as comparison. We expressed the designed proteins, including the controls in E. coli, and purified them using histidine-tag purification, after which we assessed expression level and purity via SDS-PAGE, followed by Coomassie staining. Our results demonstrate strong overexpression and efficient purification of the two previously validated de novo controls and moderate overexpression of ubiquitin and barstar controls (Figure 2). Purified protein can also be observed for four out of the six DDPP-derived constructs, though with comparatively lower yields than the positive controls (Figure 2). One potential cause of these relatively low yields may be the sizeable accumulation of DePp-derived proteins in the insoluble fraction of the cell lysate. As such, is likely that further optimization of the expression and purification methods (e.g., longer induction time or lower induction temperatures) may lead to significant improvements to overall soluble yields.
4.4 TEXT
Task description. We consider two text tasks: (i) toxic story generation using the Tinystories dataset (Eldan and Li, 2023), and (ii) product review generation using Amazon data (Hou et al., 2024). For both tasks, start by fine-tuning pre-trained MDM model (Sahoo et al., 2024) in supervised fine-tuning manner on both datasets before running online fine-tuning. As reward models, we use RoBERTa (Liu, 2019) fine-tuned for toxicity classification, and BERT (Devlin, 2018), fine-tuned for Amazon review sentiment analysis, respectively. Our experiments aim to demonstrate our method's ability to induce behaviors that are uncommon in the base pre-trained model, specifically in generating toxic content in product reviews. Full experimental details are provided in Appendix $D.4.
Main results. In Table 4 we report the average log reward as well as perplexity (Gen PPL) of the generated samples as measured by GPT-2 (Radford et al., 2019). We find that DDPP-LB is the most effective variant of DDPP and achieves significantly higher log reward compared to SVDD and RTB for both tasks. We further observe that all methods achieve comparable Gen PPI. suggesting
9

---

Published as a conference paper at ICLR 2025
Algorithm & Metric
log R(xn) f
BPD +
Base
-57.31 +
2.67 +
SVDD
-13.27 + 12.38
Guidance (scale 1)
-91.10 = 1.63
3.20 0.01
Guidance (scale 5)
-62.75 = 0.25
5.39 :
Guidance (scale 100)
-41.51 0.02
5.15 1 0.00
RTB
-60.28 = 1.74
2.04 - 0.00
DDPP-LB (ours)
-6.94 1.39
2.62 0.15
Results for discrete image modeling raw pixel values CclebA (64 64). We report the mean performance of DDPP and baselines separated into inference- nnference-based (top) and amortized (bottom) over runs for the log R and class-BPD metrics. Right: Generated samples from Base, SVDD, RTB, and DDPP-LB.
Figure 3: Left:
that generated responses are fluent; however, samples from DDPP-LB adheres better to the task specification. We refer to $D.4.1 and $D.4.2 for generated samples from DDPP.
Table 4: Text experiments with log reward and Gen PPL results averaged over 3. As Best of 10 draws samples directly from po (xo) we instead bold the fine-tuning method whose Gen PPL is lowest.
Dataset Algorithm, 4 Metric
Tinystories
Amazon reviews
log R(xo)
Gen P
log R(xo) f
Gen PPI. +
Best of 10" SVDD
93.25 + 0.17
15.94 0.03
-103.05 0.25
124.45 1.02
146.95 + 1.08
20.35 + 0.03
-27.48 + 10.91
165.86 = 1.22
RTB
107.83 3.08
18.53 + 0.55
-35.22 + 16.03
160.54 12.19
DDPP-IS (ours)
163.45 + 7.06
20.15 + 0.30
105.16 + 2.41
152.85 + 1.64
DDPP-LB (ours)
205.76 + 3.88
19.60 + 0.69
152.08 + 34.01
167.25 + 27.33
5 RELATED WORKS
Discrete diffusion. The prevailing paradigms for diffusion over discrete spaces can be broadly categorized into 1.) continuous diffusion in a latent or reparametrized space by first transforming the initial discrete data (Li et al., 2022; Chen et al., 2022: Davis et al., 2024; Cheng et al., 2024), and 2.) defining diffusion using discrete analogs of score approximation (Meng et al., 2022; Lou et al., 2023). The latter approach can also be described using the theoretical framework of Continuous-timg Markov Chains (CTMC) (Austin et al., 2021: Campbell et al., 2022: 2024). Closest to setting consider a specific instantiation of discrete diffusion that simplifies the CTMC framework by using masked forward process (Sahoo et al., 2024; Shi et al., 2024; Zhao al., 2024a; Gat et al., 2024).
Finetuning as sampling. The task of fine-tuning generative models under reward models can be viewed as a sampling problem and encompasses conventional RLHF (Uchara et al., 2024a; Black et al., 2023; Fan et al., 2024; Dong al., 2023). A simple but expensive method to sample from the reward- induced Bayesian posterior distribution is best of N sampling (Stiennon et al., 2020), which provably samples from the correct distribution the number of samples from the base pre-trained model grows, N N 0o (Beirami et al., 2024; Gao et al., 2023; Ferbach et al., 2024). Alternatively. the sampling perspective has been explored in the discrete setting fine-tune autoregressive models (Zhao et al., 2024a; Hu et al.. 2023). and diffusion models (Uchara et al., 2024b; Venkatraman et al., 2024; Zhao al., 2024a). Finally, inference time techniques represent the most prominent approach conditional sampling (Ho and Salimans, 2022; Dhariwal and Nichol, 2021; Li et al., 2024; Nisonoff et al., 2024).
CONCLUSION
In this paper, we present DISCRETE DENOISING POSTERIOR PREDICTION a novel framework to steer Masked Discrete Diffusion Models by viewing as a problem of sampling from a Bayesian posterior. We introduced three concrete training strategies to instantiate our framework in DDPP-IS, DDPP-L., and DDPP-KL and apply them to modeling synthetic data, pixel-level image modeling, fine-tuning protein MDMs to increase secondary structure diversity, and steering MDMs on language to match human sentiment, We find that DDPP not only is able to optimize an amortized sampler to closely match the reward-induced Bayesian posterior but it has good agreement in other sample quality metrics without severely compromising generated sample quality. An interesting direction for future work is understand how to balance optimization of DDPP-L.B and strategies to selecting ?-
10

---

Published as a conference paper at ICLR 2025
7 CONTRIBUTIONS STATEMENT
J.R. conceived of the initial fimulation-free fine-tuning framework, while M.H. derived the final version of the objective that can be written as a log ratio of denoisers. J.R. drove the development of code for all experimental settings, M.H. and A.T. drove the completion of image experiments, J.R. and Z.P. drove the completion of protein and text experiments, and M.H. drove the completion of toy experiments. Z.P. led TM-score protein experiments. J.R., C.L.,, Z.Q., and P.C. conceived of the 3-sheet protein task and wet lab experimental design. Z.Q. ran all wet lab experiments, N.D. guided text experiments. S.M. helped in setting up the baselines and brainstorming about the estimators. A.J.B. drove the writing of the paper with help from M.H.. M.B., Y.B., A.T., S.M. and A.J.B. guided the project from the machine learning side, while P.C. guided from the wet lab side. A.J.B, with help from M.H. designed the discrete (Reinmax) gradient estimator version of the objective function. J.R., A.J.B., and A.T. were responsible for the overall organization of the project.
ACKNOWLEDGEMENT
The authors would like to thank Kolya Malkin, Moksh Jain, Emily Jin, Katarina Petrovic, Scott Le Roux, Ahmed Elhag, Xingyue Huang, and Vignesh Ram Somnath for their useful comments early versions of this manuscript. AJB is partially supported by an NSERC Post-doc fellowship. This research is partially supported by EPSRC Turing AI World-Leading Research Fellowship No. EP/X040062/1 and EPSRC AI Hub on Mathematical Foundations of Intelligence: An "Erlangen Programme" for AI No. EP/Y028872/1
The authors acknowledge funding from UNIQUE, CIFAR. NSERC, Intel, Samsung, and Dream- fold. The research was enabled in part by computational resources provided by the Digital Re- search Alliance of Canada (https:/ //as://alll iancecan. ca), Mila (https:/ /mi la quebec), and NVIDIA.
REPRODUCIBILITY STATEMENT
We take the following steps to enhance the reproducibility of our work. In particular, all of our theoret- ical results include full proofs which are presented $C. To assist in the reproducibility of our empiri- cal findings we provide precise experimental details such as algorithmic descriptions of all variants of DDPP in Algorithm and Algorithm We further provide architectural choices, training details, and hyperparameters for all datasets and tasks in $D. Code is available at http:///gghuub..omjjrrrdddddd
REFERENCES
T. Akhound-Sadeght J. Rector-Brooks, A. J. Bose, S. Mittal, P. Lemos, C.-H. Liu, M. Sendera, S. Ravanbakhsh, G. Gidel, Y. Bengio, al. Iterated denoising energy matching for sampling from boltzmann densities. arXiv preprini arXiv:2402.066 06121, 2024. (Cited on page 16)
Austin, D. Johnson, Ho. D. Tarlow, and R. Van Den Berg. Structured denoising diffusion models in discrete state-spaces. Advances in Neural Information Processing Systems, 34:17981- 17993, 2021. (Cited on pages and 10)
A. Beirami, A. Agarwal, J. Berant, A. D'Amur, J. Eisenstein, C. Nagpal, and A. T. Suresh. Theoreti- cal guarantees on the best-of-n alignment policy, arXiv preprint arXiv:2401.0187 2024. (Cited on pages and 10)
E. Bengio, M. Jain, M. Korablyov, D. Precup, and Y. Bengio. Flow network based generative models for non-iterative diverse candidate generation. Advances in Neural Information Processing Systems, 34:27381-2733 2021. (Cited on page 16)
Y. Bengio, S. Lahlou, T. Deleu, E. J. Hu, M. Tiwari, and E. Bengio. Gflownet foundations. The Journal of Machine Learning Research, 24(1):0000-100660, 2023, (Cited on pages 7 and 16)
K. Black, M. Janner, Du, Kostrikov, and Levine. Training diffusion models with reinforcement learning, arXiv preprint arXiv:2305.13300 2023. (Cited page 10)
T. Brooks, B. Peebles, C. Holmes, W. DePue, Y. Guo, L. Jing, D. Schnurr, J. Tay- lor. T. Luhman, E. Luhman, C. Ng, R. Wang. and A. Ramesh. Video genera- tion models as world simulators. 2024. URL https:/ https://opena.: /openai....////////// . com/research/ vi deidogeeeeeaat on-modell s-as-world- simummmmllaa mumaaors. (Cited on page
11

---

Published as a conference paper at ICLR 2025
A. Campbell, J. Benton, V. De Bortoli, T. Rainforth, G. Deligiannidis, and A. Doucet. A continuous time framework for discrete denoising models. Advances in Neural Information Processing Systems, 35:28266-288 28279, 2022. (Cited page 10)
A. Campbell, Yim, R. Barzilay, T. Rainforth, and T. Jaakkola. Generative flows on discrete state-spaces, Enabling multimodal flows with applications to protein co-design. arXiv preprint MrXiv:2402.04997. 2024. (Cited on page 10)
Chen, R. Zhang, and G. Hinton. Analog bits: Generating discrete data using diffusion models with self-conditioning. arXiv preprint arXiv: 04202, 2022. (Cited page 10)
C. Cheng, J. Li, Peng, and G. Liu. Categorical flow matching statistical manifolds. arXiv preprint arXiv:2405.16444 2024, (Cited on page 10)
P. P. Christiano, J. Leike, T. Brown, M. Martic, S. Legg, and D. Amodei. Deep reinforcement learning from human preferences. Advances neural information processing systems, 30, 2017. (Cited on page 2)
J. Dauparas, I. Anishchenko, N. Bennett, H. Bai, R. Ragotte, L. F. Milles, B. I. Wicky, A. Courbet, R. J. de Haas, N. Bethel, et al. Robust deep learning-basse protein sequence design using proteinmpnn. Science, 378(6615):49950 2022. (Cited page I)
O. Davis, S. Kessler, M. Petrache, A. J. Bose, et Fisher flow matching for generative modeling over discrete data. arXiv preprint 2405 2024. (Cited page 10)
V. De Bortoli, M. Hutchinson, P. Wirnsberger, and A. Doucet. Target score matching, arXiv preprint arXiv:2402.08667, 2024. (Cited on page 16)
K. Devkota, D. Shonai, J. Mao, S. H. Soderling, and R. Singh. Miniaturizing, modifying, and augmenting nature's 's proteins with raygun. bioRxiv, pages 2024-08, 2024. (Cited on pages and 21)
Devlin. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv 2018. (Cited on page 9)
preprint
P. Dhariwal and A. Nichol. Diffusion models beat gans image synthesis. Advances in neural information processing systems, $4:8780-8794, 2021. (Cited on pages 2 and 10)
H. Dong, W. Xiong, D. Goyal, Y. Zhang, W. Chow, R. Pan, S. Diao, J. Zhang, K. Shum, and T. Zhang. Raft: Reward ranked finetuning for generative foundation model alignment. arXiv preprint arXiv:2304.06767 2023. (Cited on page 10)
F. Eijkelboom, G. Bartosh, C. A. Naesseth, M. Welling, and J.-W. van de Meent. Variational flow matching for graph generation. arXiv preprint arXiv:24066 04843,, 2024. (Cited on page 4)
R. Eldan and Y. Li. Tinystories: How small can language models be and still speak coherent english? arXiv preprint 2055772305.07759, 2023. (Cited on pages 9 and 29)
Y. Fan, O. Watkins, Y. Du, H. Liu, M. Ryu, C. Boutilier, Abbeel, M. Ghavamzadeh, K. Lee, and K. Lee. Reinforcement learning for fine-tuning text-to-image diffusion models. Advances in Neural Information Processing Systems, 36, 2024. (Cited on page 10)
D. Ferbach, Q Bertrand, A. J. Bose, and G. Gidel. Self-consuming generative models with curated data provably optimize human preferences. arXiv preprint arXiv:2407.094499, 2024. (Cited on page 10)
L. Gao, Schulman, and J. Hilton. Scaling laws for reward model overoptimization. In International Conference on Machine Learning, pages 10835-10866. PMLR, 2023. (Cited on page 10)
I. Gat, Remez, N. Shaul, F. Kreuk, T. Chen, G. Synnaeve, Adi, and Lipman. Discrete flow matching, arXiv preprint arXiv:2407. 15595, 2024. (Cited on pages 1, 1, and 10)
T. Geffner, G. Papamakarios, and A. Mnih. Compositional score modeling for simulation-bass based inference. In International Conference on Machine Learning, pages 11098-11 119-11116...1111111111-111111111 PMLR, 2023. (Cited on page 16)
12

---

Published as a conference paper at ICLR 2025
B. Hie, S. Candido, Z. Lin, O. Kabeli, R. Rao, N. Smetanin, T. Sercu, and A. Rives. A high-level programming language for generative protein design. bioRxiv, pages 2022-12, 2022. (Cited on page 21)
Ho and T. Salimans. Classifier-free diffusion guidance. arXiv preprint urXiv:2207, 125988 2022. (Cited on pages 2 and 10)
Y. Hou, J. Li, Z. He, A. Yan, X. Chen, and McAuley, Bridging language and items for retrieval and recommendation. arXiv preprint arXiv:2403.00 033552 2024. (Cited on pages 9 and 28)
E. J. Hu, M. Jain, E. Elmoznino, Y. Kaddar, G. Lajoie, Y. Bengio, and N. Malkin. Amortizing intractable inference in large language models. arXiv preprint arXiv:2310.04363, 2023. (Cited on page 10)
M. Jiralerspong, J. Bose, Gemp, C. Qin, Y. Bachrach, and G. Gidel. Feature likelihood divergence: evaluating the generalization of generative models using samples. In Thirty-seventh Conference Neural Information Processing Systems, 2023. (Cited on page 8)
W. Kabsch and C. Sander. Dictionary of protein secondary structure: pattern recognition of hydrogen- bonded and geometrical features. Biopolymers: Original Research on Biomolecules, 22(12): 2577-2637, 1983. (Cited page 21)
D. Kingma, T. Salimans, B. Poole, and Ho. Variational diffusion models, 2023. URL https: // aaaaiiiv..o... org/ aba/ 00630. (Cited on page 24)
N. Kluge Corrêa. Dynamic Normativity. PhD thesis, Universitäts-und Landesbibliothek Bonn, 2024. (Cited on page 26)
S. Lahlou, T. Deleu, P. Lemos, D. Zhang, A. Volokhova, A. Hernández-Garera, L. N. Ezzine, Y. Bengio, and N. Malkin. theory of continuous generative flow networks. In International Conference Machine Learning, pages 18269-18300. PMLR, 2023. (Cited on pages and 16)
X. Li, Thickstun, I. Gulrajani, P. S. Liang, and T. B. Hashimoto. Diffusion-Im improves controllable text generation. Advances in Neural Information Processing Systems, 15:4328-43433 2022. (Cited on page 10)
X. Li, Y. Zhao, C. Wang, G. Scalia, G. Eraslan, S. Nair, T. Biancalani, A. Regev, S. Levine, and M. Uehara. Derivative-free guidance in continuous and discrete diffusion models with soft value- based decoding. arXiv preprint arXiv:2408 082552, 2024. (Cited on pages 7 and 10)
Z. Lin, H. Akin, R. Rao. B. Hie, Z. Zhu, W. Lu, N. Smetanin, A. dos Santos Costa, M. Fazel-Zarandi, T. Sercu, S. Candido, et al. Language models protein sequences at the scale of evolution enable accurate structure prediction. bioRxiv, 2022. (Cited page
L. Liu, C. Dong, X. Liu, Yu, and Gao. Bridging discrete and backpropagation: Straight-through and beyond. Advances in Neural Information Processing Systems, 36, 2024. (Cited on page 7)
Y. Liu. Roberta: A robustly optimized bert pretraining approach. arXiv preprint arXiv: 1907. 1692. 2019. (Cited page 9)
Z. Liu, P. Luo, X. Wang. and X. Tang. Large-scale celebfaces attributes (celeba) dataset. Retrieved August, 15(2018):11, 2018. (Cited on page 8)
A. Lou, C. Meng, and S. Ermon. Discrete diffusion language modeling by estimating the ratios of the data distribution. arXiv preprint arXiv:23 16834, 2023. (Cited on page 10)
K. Madan, J. Rector-Brooke M. Korablyov, E. Bengio, M. Jain, A. Nica, T. Bosc, Y. Bengio, and
N. Malkin. Learning gflownets from partial episodes for improved convergence and stability. iclr' 2023. arXiv preprint Stti::229911788800 2022. (Cited on page 7)
N. Malkin, M. Jain, E. Bengio, C. Sun, and Y. Bengio. Trajectory balance: Improved credit assignment in gflownets. Advances in Neural Information Processing Systems, 35:5955-594 -5667, 2022. (Cited on page 16)
13

---

Published as a conference paper at ICLR 2025
C. Meng, K. Choi, J. Song, and S. Ermon. Concrete score matching: Generalized score matching for discrete data. Advances in Neural Information Processing Systems, 35.34532-34545, 2022. (Cited on pages 4 and 10)
Midjourney. https: http://wwwwddfouuneg/// hww.:///ddourreyy.co/ooe/eee///e 2023. Accessed: 2023-09-09 (Cited on page 1)
S. Mittal, N. L. Bracher, G. Lajoie, P. Jaini, and M. A. Brubaker. Exploring exchangeable dataset amortization for bayesian posterior inference. In ICML 2023 Workshop on Structured Probabilistic Inference {{&) Generative Modeling, 2023. (Cited on page 16)
H. Nisonoff, J. Xiong, S. Allenspach, and Listgarten. Unlocking guidance for discrete state-space diffusion and flow models. arXiv preprint arXiv:2406.01572 2024. (Cited pages 7, 10, 31, and 32)
M. Oquab, T. Darcet, T. Moutakanni, H. V. Vo, M. Szafranice, V. Khalidov, P. Fernandez, D. HAZIZA, F. Massa, A. El-Nouby, M. Assran, N. Ballas, W. Galuba, R. Howes, P.-Y. Huang, S.-W. Li, Misra, M. Rabbat, Sharma, G. Synnaeve, H. Xu, H. Jegou, Mairal, P. Labatut, Joulin, and P. Bojanowski. DINOv2: Learning robust visual features without supervision. Transactions on Machine Learning Reseurch, 2024. ISSN 2835-8856. URL https: http:///peeeeeeeeee/o/eo.. net/ forum? d=a 68S0t 6zFt. (Cited on page 24)
L. Ouyang, Wu, X. Jiang, D. Almeida, C. Wainwright, Mishkin, C. Zhang, S. Agarwal, K. Slama, A. Ray, et al. Training language models to follow instructions with human feedback. Advances in neural information processing systems, 35:27730-27744 2022. (Cited on page
W. Peebles and Xie. Scalable diffuusion models with transformers. Proceedings of the IEEE/CVF International Conference Computer Vision, pages 4195 4205, 2023. (Cited page 29)
E. Perez, S. Huang, F. Song, Cai, R. Ring, Aslanides, A. Glaese, N. McAleese, and G. Irving. Red teaming language models with language models. arXiv preprint arXiv:2202.03226. 2022. (Cited on page
S. T. Radev, U. K. Mertens, A. Voss, L. Ardizzone, and U. Köthe. Bayesflow: Learning complex stochastic models with invertible neural networks. IEEE transactions on neural networks and learning systems, 33(4): 1452-146,,, 2020. (Cited on page 16)
A. Radford, J. Wu, R. Child, D. Luan, D. Amodei, I. Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019. (Cited on pages and 9)
R. Rafailov, A. Sharma, E. Mitchell, C. D. Manning, S. Ermon, and C. Finn. Direct preference optimization: Your language model secretly reward model. Advances in Neural Information Processing Systems, 36, 2024. (Cited on pages 2 and 3)
S. S. Sahoo, M. Arriola, Y. Schiff, A. Gokaslan, E. Marroquin, J. T. Chiu, A. Rush, and V. Kuleshov. Simple and effective masked diffusion language models, arXiv preprint arXiv:2406.07524, 2024. (Cited on pages 3, 9, 10, 20, and 24)
Salimans, A. Karpathy, X. Chen, and D. P. Kingma. Pixelcnn++: Improving the pixelenn with discretized logistic mixture likelihood and other modifications. arXiv preprint arXiv: (CX1.0/701.05//1 2017. (Cited on page 1)
M. Sendera. M. Kim, S. Mittal, P. Lemos, L. Scimeca, J. Rector-Brooks, A. Adam, Y. Bengio, and N. Malkin, On diffusion models for amortized inference: Benchmarking and improving stochastic control and sampling. arXiv preprint wrXiv:2402.050000 2024. (Cited on page 16)
J. Shi, K. Han, Z. Wang, A. Doucet, and M. K. Titsias. Simplified and generalized masked diffusion for discrete data. arXiv preprint arXiv:2406.04329 2024. (Cited on pages 1, 3, 10, and 24)
Stability AI. https: http::///ttaallttttysattaaaaaaaa ftt.s/ttiaeeddffuiooo.o. 2023. Accessed: 2024-05-05.
(Cited on page 1)
M. Steinegger and Söding. Mmseqs2 enables sensitive protein sequence searching for the analysis of massive data sets. Nature biotechnology. 1022 2017. (Cited on page
14

---

Published as a conference paper at ICLR 2025
N. Stiennon, L. Ouyang, J. Wu, D. Ziegler, R. Lowe, C. Voss, A. Radford, D. Amodei. and P. F. Christiano. Learning to summarize with human feedback. Advances in Neural Information Processing Systems, 33:3008-3000 -0201, 2020. (Cited pages and 10)
M. Uehara, Y. Zhao, T. Biancalani, and S. Levine. Understanding reinforcement learning-based fine-tuning of diffusion models: A tutorial and review. arXiv preprini urXiv:2407.13734 2024a. (Cited on pages and 10)
M. Uehara, Y. Zhao, Black, Hajiramezanali, G. Scalia, N. L. Diamant, A. M. Tseng, Bian- calani, and S. Levine. Fine-tuning of continuous-time diffusion models as entropy-regularized control. arXiv preprint arXiv:2402.1519 15194, 2024b. (Cited on page 10)
M. van Kempen, S. S. Kim, C. Tumescheit, M. Mirdita, C. L. Gilchrist, J. Söding, and M. Steinegger. Foldseek: fast and accurate protein structure search. Biorxiv, pages 2022-02, 2022. (Cited page 9)
S. Venkatraman, M. Jain, L. Scimeca, M. Kim, M. Sendera, M. Hasan, L. Rowe, S. Mittal, P. Lemos, E. Bengio, al. Amortizing intractable inference in diffusion models for vision, language, and control. arXiv preprint arXiv:2405..m 2,0707,,00 2024. (Cited on pages 7, and 10)
R. Verkuil, O. Kabeli, Y. Du, B. I. Wicky, L. F. Milles, J. Dauparas, D. Baker, S. Ovchinnikov, T. Sercu, and A. Rives. Language models generalize beyond natural proteins. BioRxiv, pages 2022-12 12, 2022. (Cited page 1)
C. Vignac, I. Krawezuk, A. Siraudin, B. Wang, V. Cevher, and P. Frossard. Digress: denoising diffuuion for graph generation. arXiv preprint 27Xiv2209..47344 2022. (Cited on page
Discrete
R. Villegas, M. Babaeizadeh, P.-J Kindermans, H. Moraldo, H. Zhang. M. T. Saffar, S. Castro, Kunze, and D. Erhan. Phenaki: Variable length video generation from open domain textual descriptions. In International Conference on Learning Representations, 2022. (Cited on page 1)
B. Wallace, M. Dang, R. Rafailov, L. Zhou, A. Lou, S. Purushwalkam, S. Ermon, C. Xiong, S. Joty. and N. Naik. Diffusion model alignment using direct preference optimization. In Proceedings of the IEFF/CVE Conference on Computer Vision and Patlern Recognition, pages (Cited on page?) 2)
2024.
X. Wang, Z. Zheng, F. Ye, D. Xue, S. Huang, and Q. Gu. Diffusion language models are versatile protein learners. arXiv preprint 4rXiv:2402.1856) 2024. (Cited on pages and 8)
Wildberger, M. Dax, Buchholz, S. Green, J. H. Macke, and B. Schölkopf. Flow matching for scalable simulation-based inference. Advances in Neural Information Processing Systems, 36, 2024. (Cited on page 16)
Y. Zhang and J. Skolnick. Tm-align: a protein structure alignment algorithm based on the tm-score. Nucleic acids research, 33(7):2302-2309 2005. (Cited on page 23)
L. Zhao, X. Ding, L. Yu, and Akoglu. Unified discrete diffusion for categorical data, arXiv preprint arXiv:2402.03701, 2024a. (Cited on pages 1, 3, and 10)
S. Zhao, R. Brekelmans, A. Makhzani, and R. Grosse. Probabilistic inference in language models via twisted sequential monte carlo. arXiv preprint arXiv:2404.17544 2024b. (Cited on pages and 29)
K. Zheng, Y. Chen, H. Mao, M.-Y. Liu, J. Zhu, and Q. Zhang. Masked diffusion models are secretly time-agnostic masked models and exploit inaccurate categorical sampling. arXiv preprint wrXiv:2409.02908, 2024. (Cited on page 4)
A. Zou, Z. Wang, N. Carlini, M. Nasr, J. Z. Kolter, and M. Fredrikson. Universal and transferable adversarial attacks on aligned language models. arXiv preprint arXiv:2307.15045 2023, (Cited on page !)
15

---

Published as a conference paper at ICLR 2025
A BROADER IMPACT
Our proposed DISCRETE DENOISING POSTERIOR PREDICTION is tailored approach to steering and fine-tuning Masked Diffusion Models. At present, MDMs are an emergent category of discrete generative models that have general purpose modeling capabilities a variety of domains including language modeling, sequence- based drug design, and discrete modeling of graphs. Consequently, we believe DDPP has potential use in various practical use cases. For instance, like current RLHF tech- niques applied to modern autoregressive LLMs, future scaled MDMs on text datasets might be tuned to promote harmful behavior and toxic content. Moreover, applying DISCRETE DENOISING POSTERIOR PREDICTION in drug design use cases has the potential to create in-silico sample of protein sequences that may have biologically potent negative externalities. We do, however, make the distinction that such risk is speculative this stage given the large complexities of translating in-silico designs to actual synthesized biomolecules. As result, we encourage practitioners who seek to fine-tune MDMs using DDPP to exercise due caution when applying our proposed techniques actual use cases.
Ethical statement. As part of qualitatively evaluating DDPP, this paper includes generated samples of text. We highlight that the set of examples may contain potentially disturbing, harmful, or upsetting examples, covering variety of sensitive topics like discriminatory language, descriptions of harm, and misinformation, among other high-risk categories. Its primary purpose is to advance research in under- standing the impact of DDPP from such generated samples in order to prevent further propagation of undesirable content and behaviors.
interpretable lens. is not advised to train future MDMs
B ADDITIONAL RELATED WORK
Sampling proportional to energy. Our approach can be closely linked to learning sample proportional to a target probability, as in our setup we aim to approximate sampling proportional to the energy p (- xxx R(.) for any point x₃ any time t. This has been an avenue of research for a number of works in continuous time (Bengio et al., 2021; 2023; Malkin et al., 2022; Lahlou al., 2024; Sendera al., 2024; De Bortoli et al., 2024), in Bayesian
et al., 2023; Akhound-Sades posterior inference where the energy is defined by the product of likelihood and prior (Mittal et al., 2023), as well as posterior inference settings where we even do not have access to energy function but only to a simulator (Radev et al., 2020; Wildberger et al., 2024; Geffner et al., 2023).
C THEORETICAL RESULTS
C.I PROOF OF PROPOSITION
Before proving proposition we first prove a useful Lemma that states the optimal log partition function log 2x,(x) (x.) which is the learning goal for parameterized approach log (x.)
Lemma 1. Given a sample x ~ (x))))))) and the denoising posterior distribution qu (xo)x)), a local minimizer for estimate for the log partition function log 2n. using N samples from x/ ~ 9t, (xx xx) ) given by:
N
log (x((x()x))
log 2
Σ
(14)
N
(x)(()))x))
Proof. By definition the log partition function is constant, let that constant be log 2z, (x:) C.
Then the loss in Eq. 11 is quadratic in C,
L (x)))((a) [[] log
log Doggeexxn (x))x.) log R(xo)
(15)
16

---

Published as a conference paper at ICLR 2025
For batch of N samples of x (x,o(x))))), find locally optimal constant C' (local minima) by taking the gradient of Eq. 15 and setting it O. In more detail we have,
0 Ve
(log Q (x) + C log px log R(xo))² Σ
(16)
N
2 N
0
Σ
(log (x, + C (o - log R(xx))
(17)
N
2
0 0= 20
+
log 20 (x)gxx)()))))))))))))) - log logr((x()x)) - log R(x()
(18)
N
0
C
1
Σ
N
log
(x,0(x(((x)) x4)
(19)
N
Pr (x
N
pr (x(xx)x)))((
log
(20)
N
(1,0(x((x1)
Using Lemma we now prove Proposition 1, stated again below for convenience.
Proposition 1. Let log 215 and log 24,,0 be the M -ssaplee importance sampling estimate using Ke the proposal (x.o(Xx|xx) and learned approximation to the log partition function respectively. Given partially masked sample x, px(xx) the optimal learned approximation is lower bound the importance sampling estimate with a fixed proposal qr,9 (Xn x.) and the following inequality holds:
log 2 ( < log 215 (x:).
(9)
Proof. We optimize log 1LB (xx) using the loss defined in Eq. 11. Using Lemma I we know the analytic expression for the locally optimal estimate is given by log Z (xx). Plugging this into the definition of the log partition function we get.
log (x(oxox))
}
log 0e.0 (x.)) Exo (x (x))
(21)
(x,0 xx,)
Pe (x0 xx, H{(o))
< log oxxoxx))
(22)
(xolxx)
= llog
21 Kr (x+)
(23)
The lower bound turns into equality at the optimal proposal qt. x Pe)x
C.2 ESTIMATING DDPP-KO KL WITH REINMAX
We first provide an algorithmic description below of training using DDPP-KL. We first highlight how the reverse KL objective can be applied to general setting beyond just fine-tuning before turning to the exact setting of the main paper.
Algorithm 2 DDPP-KI.
Input: Differentiable reward R(xo), base MDM p5 (xa (x), fine-tuning MDM go (xo Xt), Num samples K
I: while Training do
2: Xo ~ 11 1]. (xxo)
Sample time and clean data on-policy from the fine-tuning MDM Constrruct partially masked sample given clean data
3: X, (e(xx (xo)
4:
[ ~ 9r (-(x.)
& Reparametrized Sampling of clean data
CKL 1 D (log qe, (x66)x)) log Phe (x6(xx)) - log R(xi))
K
6:
VoLkI. := V Rainmox (c)t.)
Use the Reinmax discrete gradient estimator
7:
Update (0, CoLk
8:
Return go
17

---

Published as a conference paper at ICLR 2025
Non-finetuning Case. In this appendix, we study the REINMAX gradient estimator for the general problem of sampling from the following distribution:
To (x() x R(xo)
(24)
Gradient of [kt, We can decompose the gradient into the following terms due to the linearity of expectations:
C Ex,x, [log qe (xx]xx) - Ex,,, [log Iy (xo]xx)]
= - + L.
(25)
We again highlight the fact that the expectation is taken using the following distributions I, Xo. X ~ (0 1J, Pe that needs gradient estimation since qo is the parameterized distribution. Furthermore, all terms that use this sample inside the expectation are affected by this gradient computation.
As a result, Xo is drawn on-policy and is a stochastic variable
Taking the gradient each term respectively. The gradient of C is:
VoLt
Va (Br,x0xx, [log 2r,0 (xp) xt)
Ex,xor ( /Rein-Max (log qt,0 (xx) (xx|xx))]
(26)
The gradient of of C is:
Volt
Vo Xo,,
[log ne. (xo[xx)))
V xo,
px log To (xo) + log I (x:))
V
E, X0,
log (x, - log R(xo) + log Σ q.e (x, (x)))))))
(27)
To use the Reinmax gradient estimator we must compute 0f(;) where f is the function inside the expectation E, [f]]]. We now make use of the following facts:
(F1) Analytic expression of V., log (e (o)). For simplicity of presentation, we focus on single token zó sequence but the result remains true for the entire sequence Xp- Recall in the discrete setting of masked diffusion models pt Cat(zo; Quæe), which allows us to write:
V VzgPe (atl(a) (1)
(28)
Vay Cat Qrrt)
(29)
Cat(abi Quat)
Vay(20 Quat)
(30)
2 Qiat
V V (ar (1), .n) + (1 - (z)(2),
(31)
(1 xb) + (1 (z) em)
axoi
(32)
(1), mn) (1
em)
(F2) Differentiability of the reward Vx, R(xo). If we assume the reward is differentiable we can exploit the same trick to write:
Xx, R(xo)
Vx, log R(xo) =
(33)
R(xu)
Note that the final term in Eq. 27 does not depend on the realization of the sample Xy ~ g(xx (x:) and thus its gradient Rein-max is O. This enables us write the approximate gradient
V VoL? 2 E. Ex / Reinmax
(- log (x. log [(xx))]
(->
V Reinmax R(xo)
)]
Σ
Ex
(34)
Or anfat xb) + (1 (z) cm)
R(xo)
18

---

Published as a conference paper at ICLR 2025
The first term in the equation has closed-form expression for the gradient but is still stochastic gradient since it depends on Xu ~ (xo).
Finetuning Case. the fine-tuning setting we aim to sample from the following Bayesian posterior:
To (xo) x Po (xo) (xx))(xx)
aed
(35)
Z
For MDMs the likelihood under the model Po (Xo) is intractable to evaluate and leads to modified objective for gradient estimation with REINMAX in L in Eq. 25:
V Vol.
Vo (B4, x0.xx
[log Te (x xx)))
Vo [Rx,x.xx
log rf (xolxx log R(xo) + log ET,, (x:)])
V X,,,,,
log p, (xx (x)) log R(xx) + log [xx)~xf((xlxx (x.) (xx))
(36)
Note that in the equation above we can evaluate the log partition function using samples drawn from the denoising posterior of the pre-trained model x pt (xo (x,) and not the on-policy samples
Xy ~ qo (xo). Thus this term is a constant when we compute the gradient. Thus we have,
VoL. & A Reinmax (Bexxxxx,
-log lo (xalx:) - log R(xo)))
(37)
C.3 EQUIVALENCE OF SUB-TRAJECTORY OBJECTIVES
In this appendix, we detail how to compute an efficient approximation the loss function that inspired by the KL divergence between sub-trajectories found in the GFlowNet literature but adapted for MDMs.
Consider the trajectory of a sequence: r(Xo.)) == xı + X - x + Xo. We seek to minimize the joint distribution over the (sub) trajectrress conditioned on partially masked sample x,: Xt-1 Xt, po (xe: (x.) = Te (Xo, ..., Xx- xx (x;).
qo [xo
(38)
Here Te (X1,..... X1- |x,, xo) is defined as,
pt (Xu,.... X4-1 (x, R(xo)
my. (x0, X,,, Xx, Xx, 1))
(39)
2x, (x,)),)
I p. (x,- [x) 2 )((x))
(40)
2 (xx)
We minimize the following KL divergence,
IDI (Xo. x |x, & ())p:(x)) |,, xx-1lxx)
(41)
Here we used the convention that &a p f) and x = (x, t). The KL between path measures along the sub-trajectory shares the same optimum as the following loss objective:
C, = E,,,,
Er(xo.)) ||l log xx x0)) log P.
( (xo....,, Xt- + +x|la]
xpe)) ) + re) 1
Ex,x,
Er(xow) log go ( x0) - log p. (x &x.
(42)
-(xx,.
log qo (x.->/x.; 20) loo (x Xx. xpre)
(43) In the last equation we define the constant K. (- log R(xo) + log Z,,, (x)))/t and use the fact that our notation convention uses p(X = I) = p(z) for discrete random variables. Now we make the observation that for any t have effectively picked endpoint over the trajectory. More precisely, 2[(0, t]. which also allows us to sample x, ~ (x (x)). ). in an analogous manner to how x, is constructed.
HE xx.,x,,..., [log go (x=- X.. Xx, xx) log p. pre (xa-ys Na. x m
11
19

---

Published as a conference paper at ICLR 2025
D ADDITIONAL EXPERIMENTAL DETAILS
shared heterogenous high-performance computing cluster. This cluster is primarily composed of GPU nodes with RTX8000, V100, A100, L40S, and H100 NVIDIA GPUs. We briefly note trick used across number of our experiments for DDPP-LB described as warming up log Zr (z+). We found that early iterations of training DDPP-LB could be somewhat unstable the parameterized normalizing constant was not calibrated to a proper range given the pre-trate pre-trained model and reward function. As such, we found that warming up log (x,) for some number of steps at the beginning of training by only allowing gradient flow through the log (1²) term helped stabilize training and improve overall performance. For the runs on which warming up log Ze (at) was utilized, we resume normal training (i.e., allowing gradient flow through the fine-tuned denoiser and log Z (x+)) after the warmup period has concluded. For all experiments with DDPP-LB we used another separate, small DiT to parameterize the log Ze (xx) prediction.
All experiments were performed
D.1 SYNTHETIC EXPERIMENTS
Two synthetic tasks were performed: (1) sampling from a posterior over a 2 dimensional grid, and (2) fine-tuning on binarized MNIST. In both cases 90 million parameter MDM model was trained on samples from the prior distribution, with the same DiT architecture as in Sahoo et al. (2024).
D.1.1 GRID EXPERIMENT
The space consists of discrete tokens X E ,0,,....12]].. A prior density po is defined over this space which assigns uniform probability for tokens falling inside one of the 16 evenly spaced squares, and a near-zero probability outside this. This prior distribution is depicted in Figure 6(a). Pre-training was done using the Adam optimizer, with B1, B2 {0.9, 0.999). and learning rate of 3c-4.
The reward function R(x,) = for z < 64, and R(x,) = 1 for =d 2 64. This results a fine-tuning target x ((x))p"*(xxe which selects only the squares in the lower half of the grid. This product distribution is visualized in Figure 6(b).
For fine-tuning we train the model using our loss-functions with the Adam optimizer, using a learning rate of 4e 3, 31, B2 {0.9, 0.999). and weight decay of 0 across all methods. DDPP-IS used 16 samples to estimate the partition function. Training is done using a replay buffer populated with points X. sampled policy from the model, well as off- off-policy points from the prior distribution, added to the buffer every 100 training steps. A batch of 64 used.
D.1.2 MNIST
28828 The prior pdd (Xo) in This task consisted of generating binarized MNIST digits x E 10,1)28x28 this case is the MNIST data distribution. For pre-training, the Adam optimizer is used with learning rate of 4e- 3, 31, B2 = {0.9, 0.999} and a weight decay of O.
This MDM is fine-tuned to produce even digits. More precisely, the reward function is R(x,)
p(vvn | xo) (Ei-0 ply
xo with p(y i xo) being obtained from pretrained MNIST classifier (LeNet 5 in this case). The inverse-temperature 3 is set to 5 for all experiments.
For fine-tuning with our methods, we use Adam with learning rate of 1e 5 and B1, 82 (0.9, 0.999). Training is done with batch-size of 64. Samples are drawn from replay-buffer populated with only on-policy samples. Method specific hyperparameters include:
DDPP-IS: the importance sampling estimate is done with 16 samples
DDPP-LB: learning rate of le is used for network layers estimating log Zx,
DDPP-KL: The KL objective per x, is computed using 8 samples
RTB is trained with a learning rate of 5e 5, with weight decay 0.01, on trajectories of length 32 with batch size of 8. For training, 30% of the steps detached. The smaller batch-size is chosen to fit the training on 80GB of GPU memory.
SVDD uses 10 particles in each inference step.
For all methods (including baselines), inference is done with 128 steps. Additional information computation of metrics is included in D.3.1.
20

---

Published as a conference paper at ICLR 2025
D.1.3 MNIST SAMPLES
Samples from our methods, as well as the pretrained model, are shown in Figure 4.
D.2 PROTEIN SEQUENCES
Protein design involves the creation of novel protein sequences that adopt specific structures and perform desired functions. This is a critical field in synthetic biology and biotechnology, as it enables the rational engineering of proteins with enhanced stability, novel functionalities, or improved therapeutic properties. Advances in machine learning-based models, such protein language models (pLMs), have enabled rapid exploration of protein sequence space, making de novo protein design more feasible and versatile. However, current pLMs struggle in generating realistic sequences which satisfy certain criteria, and we study using DDPP to finetune DPLM to generate high-scoring proteins given reward function.
D.2.1 IN-SILICO TASKS
In task 1, we fine-tune the DPLM model to generate designable protein sequences that optimize for several critical features, including high predicted template modeling (pTM) and predicted local distance difference test (pLDDT) scores from ESMFold, reduced exposed hydrophobic residues, high sequence entropy, and an increased proportion of 3-sheet content (Hie al., 2022). These optimizations are captured in the reward function Rt given by:
log R WPTM PTM + WDLDDT pLDDT + Wheet Sheet%
+ WEntropy
(s) - WHpho Exposed Hpho%
Where the terms represent:
PTT and PLDDT: Structural confidence scores from ESMFold, measuring global and local accuracy, respectively.
Sheet%: The proportion of residues predicted form 3-sheets, determined by DSSP (Kab- sch and Sander, 1983).
(s): Sequence entropy, defined as:
H(s) Σ Σ pi(a) log):(a),
f=1
where L is the length of the sequence and (a() is the probability of amino acid at position
i.
Hpho%: Percentage of hydrophobic residues exposed on the surface, calculated based on solvent-accessible surface area.
Exposed_Hpho%:
The weights for these features are as follows:
UpTM = 1, ULLDDT = 1, WSheet = 4.5, Winntropy = 0.8, Wiipho = 0.25.
As the scale of the various reward terms are non-uniform we selected the reward weights to weight all rewards similarly besides the sheet percent reward which is weighted higher. For the f-sheet task we found that both RTB and DDPP faced issues with mode collapse. After investigating the protein structures generated by base DPLM found that the base model is only capable of generating small number of motifs (in particular, over 2k samples from the base model we found only two motifs with log R(x,) > 3.5), implying that the targeted product distribution indeed collapses around these structural motifs as we observe in the case of RTB and DDPP. As such, we conclude that DDPP (and RTB) achieve the goal of fine-tuning as they sample from the product distribution and reproduces samples with f-sheets at a much higher proportion than the base model.
In task 2, we focus on generating shorter sequences of known proteins that preserve essential structural characteristics, using the TM-align score the reward function (Devkota et al., 2024). This task allows the exploration of mutational effects. Ribonuclease proteins (PDB IDs: 9RAT-A, 11BA-A) selected for this task due to their well-chhracterized structure, function, and folding mechanisms.
21

---

Published as a conference paper at ICLR 2025
5
6
3
5
5
(a) Pretrained model
2
5
4
2
5
2
4
g
8
<
(b) DDPP-IS
2
2
L
3
(c) DDPP-LB
2
8
C
(d) DDPP-KL.
Figure 4: Uncurated samples from the pretrained model, and after fine-tuning with our methods: DDPP-IS, DDPP-LB, DDPP-KL
22

---

Published as a conference paper at ICLR 2025
The reward function R is defined as:
R W Wim_scure TM-align(s,, t).
Where:
score: The weight of the TM-Score reward,
s: Predicted structure from ESMFold of the generated sequence.
t: Target protein structure.
TM-align: A measure of structural similarity between s and t, defined as:
LE
1
TM-align max
L
da
where Lt is the length of the target protein, Luli is the length of the aligned region, d; is the distance between the i-th pair of aligned residues, and dg is the distance scale based on Le (Zhang and Skolnick, 2005).
While not used in the reward function for either experimental setting, also measure the KL divergence, reported as KL-SS in Table between the secondary structure distribution given by DSSP for both the target and miniaturized protein.
Note that in these experiments, the number of recycles in ESMFold is set to 0 to reduce computational overhead. For both tasks we generate amino acid sequences of length 90. Evaluation is performed by sampling 200 proteins for each method across three seeds and reporting the mean and standard deviation of each metric accordingly. All methods ran 500 inference steps during evaluation. All protein experiments used 1.50 million parameter DPLM base model to begin fine-tuning from. All models used a log-linear noise schedule with Jmin 1e-4 and Fmax = 20 and used a linear learning rate warmup period of 2500 training steps.
DDPP was trained with warmup period for log (z;)), learning rate of 1c- batch size of 16, a replay buffer of max length 10,000, and inserting new batches to the buffer sampled on-policy from the current model every 250 training steps. RTB was trained similarly, but with smaller batch size to account for its greater memory requirements. RTB matches the setting of DDPP but with a batch size of while doing 90 inference steps during training new batch of trajectories simulated on-policy every training step). To allow RTB to fit in memory we detached 65% of trajectory timesteps when computing a backward pass on the RTB objective. SVDD was run on the base DPLM model with = 10 particles. To control the concentration of our designated target distributions, we set the reward temperature 3 0.125 for the B-sheet task and 3 0.001 for the protein miniaturization task.
We report an extended version of Table 3 where we include results for both ribonuclease targets in Table 5. We observe that DDPP consistently achieves the highest TM-Score across the two templates while maintaining high structural quality with an average pLDDT of around 0.8.
D.2.2 EXPERIMENTAL VALIDATION
Genes encoding for de novo protein sequences obtained from Integrated DNA Technologies (IDT) and cloned into PET-24a(+) (Novagen) expression vectors with a C-terminal 6xHis tag using Gibson Assembly (New England Biolabs, NEB). Assembled plasmids were verified via Sanger sequencing, then transformed into chemically competent Escherichia coli BL21 cells (NEB). Starter cultures (3 mL. Luria Bertani media, 50 ug/mL. kanamycin) were inoculated from freshly prepared agar plates and grown at 37 and shaken at 225 RPM overnight. Starter cultures were then diluted 1:100 into 50 mL LB medium supplemented with antibiotic. Cultures were then grown at C and 225 RPM until an optical density (OD600) of 0.5-0.7 was reached. Protein expression was then induced with mM isopropyl Db-thiogalaccopyranoiide (IPTG) for 4 hours at 37 Cells were then collected by centrifugation (4,500xg) at 4 and resuspended in lysis buffer (Tris-buffered saline (TBS), 25 mM imidazole). Cell suspensions were then lysed via sonication (10s pulses, 40% amplitude). The corresponding lysate was centrifuged at 12,000xg for 30 minutes, and the
hhttps:/hhuggine cotaaacccccccacaoooooooo/aaaaaaa rk wwwd//ppdpppogg m50mm
23

---

Published as a conference paper at ICLR 2025
Table 5: Miniaturizing ribonuclease proteins 9RAT-A and IIBA-A (124 AAs) 90 AAs while
preserving structural fidelity (high TM-Score) and quality (high pLDDT and PTM).
Template
SS-KL
log R(xa) f
TM-Score f
PLDDT
pTM
Base Model
2.944 1 2.936
0.502 + 0.128
0.251 + 0.064
0.724 1: 0.144
0.584 1 0.226
Best-of-10
0.640 + 1.872
0.725 0.098
0.363 + 0.049
0.789 + 0.018
0.754 + 0.086
9RAT-A
DDPP
1.086 2.242
0.735
0.122
0.368 = 0.061
0.793 : 0.044
0.768 + 0.066
RTB
1.808 + 2.597
0.597 0.109
0.299 + 0.055
0.796 1 0.054
0.750 + 0.084
SVDD
3,465 + 2.835
0.699
0.079
0.350 + 0.039
0.499
0.137
0.383 + 0.178
Base Model
3,136 + 3.150
0.478 + 0.101
0.239 + 0.051
0.724 + 0.144
0,584 + 0.226
Best-of-10
2.602 3.309
0.654 0.089
0.327 + 0.045
0.782 4 0.027
0.720 + 0.109
11BA-A
DDPP
0.194 +: 1.009
0.709
0.048
0.354 + 0.024
0.743
0.036
0.727 + 0.054
RTB
2.579 + 2.799
0.564
0.111
0.282 + 0.056
0.797 + 0.058
0.744 0.101
SVDD
3.240 + 2.992
0.647 0.079
0.324 = 0.040
0.486 + 0.124
0.354 + 0.162
supernatant was loaded into a HisPur Ni-NTA His-spin column (ThermoScientific) and purified as recommended. Expression of purified proteins both the soluble and insoluble fraction, as well as his-tag purification fractions, was assessed using SDS- EDS-polyacrylang gel electrophoresis.
D.3 DISCRETE IMAGE MODELLING
To setup the finetuning task we first pre-train large masked diffusion models on the original dataset. This uses standard masked diffusion loss as explored in previous work (Shi et al., 2024; Sahoo et al., 2024).
CelebA Pretraining. We train a 241 million parameter model based on the variational diffusion model (VDM) architecture (Kingma et al., 2023) and the setup of Shi al. (2024). We adapted the U-Net plus self-attention architectures from Kingma et al. (2023) as used CIFAR-10 in their experiments, with a few notable additions. We replace the Fourier feature inputs with an input embedding layer which embeds 257 (256 pixel values + <MASK>) tokens into the embedding dimension. We double the number of residual blocks from 32 to 64 per encoder / decoder, and double the embedding dimension from 128 to 256. We use an Adam optimizer with learning rate 1e - 3, 3,=0.9 and 3.99. We train our model for 450k steps with batch size 128 on a cluster of 16 NVIDIA L40S GPUs We resize all CelebA images to 64x64 with bilinear interpolation. Samples from this model can be seen in Figure 5.
Separately, we train a 7M parameter classifier classify hair color on CelebA. We use this as our energy function with a temperature setting of 0.1 for all finetuning experiments.
CelebA Finetuning. With the problem setup, we next finetune our pretrained model to sample images with blond hair. We train each model for up to 12 A100 hours. We an early stopping criteria based on a validation set using an approximate bits-per-dimension calculation using the ELBO. We find that the original needs least 1 100 inference steps for good performance therefore we evaluate all models in this setting. For our model we use 1 000 warmup steps for log Z, learning rate of le we resample two batches every 500 gradient steps of the model and add them the replay buffer.
In contrast to our model, RTB requires a full trajectory for each gradient step. For CelebA, this means rollout of 000 inference steps taking approximately 2 minutes for batch size of with this model. Because of memory constraints we detach 99% of inference steps and use a batch size of fit in 80GB of GPU memory with a global batch size of trajectories per gradient step.
A100
D.3.1 METRICS
The metrics used to evaluate image fine-tuning include mean log reward, feature-likelihood divergence (FLD), and bits per dimension (BPD).
FLD. For FLD, draw K samples from the model, and K samples from the test set restricted to the target class. The FLD is computed using the DINOV2 feature space (from the ViT-B 14 model) between these two sets of samples (Oquab et al., 2024). For MNIST, K = 5k.
24

---

Published as a conference paper at ICLR 2025
Figure 5: Uncurated pre-trained CelebA model samples using discrete generative model.
25

---

Published as a conference paper at ICLR 2025
BPD. An upper bound the log-likelihood is computed using the MDM ELBO loss (on the fine- tuned model), and this is normalized (by the number of pixels and log 2) to obtain the reported BPD metric. For baselines other than discrete guidance, this metric is computed on the test set restricted to the target-class. For discrete guidance, BPD is computed by evaluating the MDM ELBO of the base model on samples generated using guidance (due to lacking an analogous ELBO for guidance-based sampling).
D.4 TEXT EXPERIMENTS
All text experiments begin by starting from the pretrained MDLM consisting of 170 million pa- rameters. We then do supervised fine-tuning produce model capable of producing output of the desired format before proceeding with online fine-tuning. For all text experiments we train using the Adam optimizer with $1, 82 {0.9, 0.999) and weight decay 0. For both tasks SVDD was run with 10 particles. Evaluation was done by training each method for one day across three seeds and generating 1000 samples from the best checkpoint according to the mean reward generated during training. We report both the mean reward of the 1000 samples across three seeds and their standard deviations, as well as the generative perplexity according to a GPT2-Large model with 812 million parameters 3.
D.4.1 TINYSTORIES
To obtain a base model we performed supervised fine- tuning from the base MDLM model on the Tinystories dataset. As all methods were prompted with the text "Once upon time" during training, we restricted the SFT dataset to only datapoints whose stories started with the text "Once upon a time,", resulting in corpus of 977.921 examples. SFT was done using Adam, B1, 82 = (0.9, 0.999), and a learning rate of 4e-3 over 60,000 training steps using 4 NVIDIA A100 GPUs. All models were trained for up 24 GPU hours on NVIDIA L40S GPUs. Fine-tuning checkpoints were selected based upon the iteration with best average reward when sampling new training batch from the model. All methods employ learning rate schedule with a linear warmup for the first 2,500 training steps and keep the noise schedule provided by the pre-trained MDLM model log-linear schedule with Tmin = 1e-4 and maa = 20. All evaluations were performed by taking 1000 samples for each method across three seeds. We provide set of curated samples in Table 6.
The reward function for this task was selected to be pre-trained classifier' (Kluge Corréa, 2024).
The classifier a RoBERTa model with 125 million parameters which was fine-tuned on a curated subset of various toxieity/harmlessness datasets. The reward R(xo) is then set to the likelihood of sequence being toxic under the pre-trained classifier so that R(xo) = p(a = 1(50) where p(a = (ro) denotes the likelihood of the sequence zo possessing toxic sentiment. We select this task as it allows a demonstration of how our method can recover rare behavior under the pre-trained model while still maintaining sample quality. We consider
target distribution the tempered reward distribution vo(Fo) x po (#o)R(rro)? 21/ with 3 0.25.
pre
For DDPP-LB we used 1,500 warmup steps for log ZT, (Iz)).= a learning rate of 1e-4 and a batch size of 16. We employ a replay buffer with max length of 10,000 and sample training batches uniformly from the buffer. The buffer is filled every 50 training steps with batch sampled on-policy from the current fine-tuned model, while every 250 steps a batch from the SFT training dataset is added to the buffer. We use EMA with a decay rate of 0.9999, learning rate of 1e-4, and train without LoRA. DDPP-LB was trained using 64 inference steps for simulation. DDPP-IS employed the same hyperparameter settings DDPP-LB except that it dispelled with learning log Zz, (1r)) and instead estimated it with K- = 16 Monte Carlo samples from the one-step preitrined posterior pt 2.))
As RTB cannot fit all timesteps of trajectory into memory during the backwards pass, we detached 55% of timesteps where each trajectory consisted of 32 timesteps. RTB was trained with LoRA enabled, a LoRA rank of 16, and a learning rate of 5e-5. Due to memory constraints the batch size was set to 4. SVDD was run by using
10 particles per inference timestep. best of N (with N = 10) sampling was performed by taking 10 samples from the SFT model and selecting the sample with highest likelihood under the reward model.
"ntrpss//hugggngfccc.oo/ www.peaiiooomuunttggooooooooooee hhttps:://nggginggfcc..oo
26

---

Published as a conference paper at ICLR 2025
Table Curated samples from DDPP-LB on Tinystories.
Generated Text
Once upon time, there was a little girl named Lily. One day, Lily went to the park with her mom. She loved to run fast and laugh. Lily saw a funny butterfly and ran it too. She fell fast and hurt tree. Lily's knee hurt and she cried. But her mommy kissed her cheeks and said, "Be careful when you run, Lily."
Once upon time, there was an elderly wolf. He lived in a big den against the woods. One day, the wolf felt very tired and wanted a place where there was big tree eat on. So, he went to sleep all day. But, while he was waking up, he saw big, icy creature. The quickly jumped up, but his legs were too weak. The creature took the elderly wolf away. And that how winter ended.
Once upon time, there little bird. His wings were weak and he fell down. The bird wanted his wing to restore him. So, he flapped his wings with his weak heart.
Once upon time, there was little boy named Timmy. He loved going to the woods with his family. One day, Timmy's ' friend Johnny came to the woods to play. Johnny was excited to go outside and play.
They found a big tree with words said reverse," and Timmy would play its branches. He said, "reverse!" and pushed the tree. Then, they and laughed.
But then, they heard loud noise coming from the bushes and scared them. It was a big, mean bear! Timmy tried to reverse and run away, but he wasn't 't fast enough. The bear chased him and caught him with its sharp hands.
Timmy was very scared and never went back to the woods again.
Once upon time, calm blue sea, there was small boat. The boat had sailors. They lived happily in the day water. One day, the water was very hot. So. they all decided soak up and have a picnic.
But, by the time, the sailors started to play a game. They swam around and counted, two, three soon!" and all the sailors kept playing. They water flowers and trees, and everyone laughed.
But then, clumsy heavy sailor hurt his head on rock. "Ouch!" he cried. His friends helped him up and said, "Be careful next time!" The sailor felt better and they all laughed. They knew they could play again and have fun calm day soon.
27

---

Published as a conference paper at ICLR 2025
D.4.2 AMAZON REVIEWS
Table 7: Curated samples from DDPP-LLB Amazon review generation task.
Generated Text
Cheap. Poor fit. The pants return immediately and the fabric was like a burlap sack bag- nog-Wanted it for a gift-It's crap. Broke in one day. Customer service never responded.<br /> Very cheap!
Everything was horrible.
It's the worst one I've ever bought.. you have keep it in your bag for my daughter and they, seriously feel embarrassed you ever got it it falls out and ripped.. returning this. That said hated this bag till see it!!!!! (Cnaven at the neckline and it is super too short. Color is off white. Maybe have to fix if want ironing<br What a waste of valu money
Such poor material!!! was like a plastic. Way too small so I returned.It was smaller than the size listed
don't feel this product has any quality. My sunglasses was delivered broken.
Cheap piece of garbage. Got it for my niece for Halloween and it broke inó one time
We again begin by first performing supervised fine-tuning from the base MDLM model, but this time on the fashion split of the Amazon Reviews dataset (Hou et al., 2024) restricted to reviews consisting of at most 512 tokens, resulting in an SFT dataset of size. We perform SFT using Adam with 31, 3/2 = {0.9, 0.999} and learning rate of 4e and EMA with decay parameter 0.99, The SFT model was trained for 85,000 training steps on 4 NVIDIA A A100 GPUs. As for the tinystories task fine-tuning checkpoints were selected based upon the iteration with best average reward when sampling new training batch from the model. All methods employ learning rate schedule with linear warmup for the first 2,500 training steps and keep the noise schedule provided by the pre-trained MDLM model a log-linear schedule with amin le 4 and mmaz 20. All evaluations were performed by taking 1000 samples for each method across three seeds.
The reward function for this task was a BERT model consisting of 167 million parameters fine- tuned Amazon customer reviews to predict review's star-rating. We then set the reward R(xo)= = p(a- = 1 (.). the likelihood under the pre-trained classifier that the generated sample is a one-star review. We consider our target distribution the tempered reward distribution To xx P (2x) R(x,) 1/88 with B = 0.5.
For DDPP-LB used 1,500 warmup steps for log Zx, (zz), a learning rate of 1e-4 and batch size of 16. We employ a replay buffer with max length of 10,000 and sample training batches uniformly from the buffer. The buffer is filled every 5 training steps with a batch sampled on-policy from the current fine-tuned model, while every 250 steps a batch from the SFT training dataset is added to the buffer. We use EMA with decay rate of 0.9999, learning rate of 1e-4, and train without LoRA. DDPP-LB was trained using 64 inference steps for simulation. DDPP-IS employed the same hyperparameter settings as DDPP-LB besides not learning log Zn, (zx) and instead estimating it with K = 16 Monte Carlo samples from the one-step pre-trained posterior p. (zo|zz).
As RTB cannot fit all timesteps of a trajectory into memory during the backwards pass, we detached 78.5% of timesteps where each trajectory consisted of 64 timesteps. RTB trained with LoRA enabled, a LoRA rank of 16, and a learning rate of 5e-5. Due to memory constraints the batch size was set 4. SVDD was run by using n = 10 particles per inference timestep. best of N (with N = 10) sampling was performed by taking 10 samples from the SFT model and selecting the sample with highest likelihood under the reward model.
We provide set of curated samples for the Amazon task in Table 7,
co/l: ini Yuan/ honm/mmazon--ee ew- sent mentt-naa aonnt-nnalyii lysis
28

---

Published as a conference paper at ICLR 2025
E ADDITIONAL EXPERIMENTAL RESULTS
E. AUTOREGRESSIVE BASELINE
Table Results Tinystories with Twisted SMC autoregressive baseline. Because DDPP and Twisted SMC use different architectures and number parameters, we report reward and generative perplexity hefore and after fine-tuning and bold the methad which provides the best percent change from the base model.
log R() pro
log R(zu) post
% change R(ma)
Gen PPL pre
Gen PPL post
% change Gen PPL
Twisted SMC
40.52 + 0.10
94.56 0.85
133.4 + 1.60 278.0 + 15.2
8.52 + 0.05
10.25 # 0.51
20.31 = 5.96 18.38 4.05
DDPP-LB
54.94 + 0.76
205.76 3.88
16.66 0.20
19.6
In order to compare the performance of DDPP against autoregressive methods, we evaluated Twisted SMC (Zhao et al., 2024b) on the Tinystories task and compared its performance to DDPP. Unfor- tunately, the base model used by Twisted SMC was of different architecture and model size than we used. In particular, as the base autoregressive model, we used the GPT-Neo architectured model trained in the original Tinystories paper (Eldan and Li, 2023) which uses 68 million parameters while our base model used a diffusion transformer (Peebles and Xie, 2023) with 170 million parameters. To ensure the fairest possible comparison, despite the difference in model parameters of the pre- pre-trained models, we compare Twisted SMC and DDPP fine-tuning in terms of the percent change in the average reward and generative perplexity of the finetuned samples. Results presented in Table 8.
We observe that DDPP improves reward more than the autoregressive baseline while incurring comparable but minor performance drop in generative perplexity to Twisted SMC. Our results here contextualize that DDPP can better negotiate the tradeoff between optimizing reward and sample quality than twisted SMC on autoregressive models. Finally, we note that Twisted SMC cannot be easily performed for MDM's and such DDPP remains compelling choice for fine-tuning.
E.2
COMPARING OVERALL COMPUTATION TIME OF DDPP VS INFERENCE BASED METHODS
50
Method
Sequence Length
10²
SVDD
90
Best of 10
150
DDPP
200
10
250
Seconds
Hours
10
20
10
0
10
10'
10
10
10
10
Number of Generated Samples
Batch Size
(a) Sample generation time
(b) ESMFold query times
Figure 6: (a) We plot the number of hours required to generate particular number of samples for each method (including DDPP training time) the 90 sequence length designable 3-sheet protein task. We see that although inference time methods may be preferable if generating only few samples, DDPP quickly offers faster sampling as the number samples grows. (b) We demonstrate how the FSSFoold hased reward function cannot be parallelized on a single GPU and that the computational overhead becomes even more pronounced as sequence length increases.
To further analyze the computation time tradeoff of inference time methods compared to DDPP we examine their computational overhead on the task of generating designable fi-hheet protein of sequence length 90. We expect that the computational tradeoff of inference time method vs fine- tuning method to favor fine-tuning based methods the number of generated samples scales and as the computation required to evaluate the reward function increases. We examine the designable 8-sheet task its reward function is especially onerous to compute it involves folding protein with ESMFold. This process does not parallelize well on even on an A 100 80 GB GPU as show in Figure 6b where see that reward computation time scales linearly with batch size. Moreover, the
29

---

Published as a conference paper at ICLR 2025
Iteration [1c3)
Iteration [te5]
10
Itoratoon
Time (hours)
Time ihours!
Time (hours)
(a) MNIST
(b) Amazon Reviews
(c) Tinystories
Iterations and Time required to reach a threshold value log R( for varying number of MC: samples M in DDPP-IS. Threshold values are computed using training curves smoothed with rolling average, to better capture training trends.
Figure 7:
reward function computation becomes even more expensive if scale sequence length as one might do in many real world tasks due to O(N9) (where N denotes sequence length) operation involving pairwise residue interactions.
To evaluate computational efficiency we compare the overall computation time on a single A 80 GB GPU required to generate different numbers of sequence length 90 samples from given method, as shown in Figure 6a. To compute the time for DDPP to generate samples, we first measure the training time to convergence and add this to the inference time required to sample from the fine-tuned model. For best of 10 and SVDD we simply generate the specified number of samples and record elapsed computation time. We observe in Figure 6a that while SVDD and best of 10 are fast if one needs to generate only few samples, they quickly become significantly more expensive as the number of samples needed increases. In particular, generate only 1000 samples DDPP, combining both its training and inference time, requires only 1.99 hours while best of 10 and SVDD require 6.15 and 37.5 hours, respectively. This means that for generating even this relatively small set of 1000 samples, amortized sampling using DDPP is up to 18x faster than comparable inference time methods while generating higher quality samples. Moreover, this inference time gap only increases as the number of generated samples increases and would also become more severe were sequence length to increase, as evidenced by Figure 6b.
Of course, the utility of amortized sampling methods for reducing overall computation time is dependent on the number of samples required and the computational overhead of the reward function. only few samples are needed and the reward function is cheap, it is advisable to use an inference time method such as best of N SVDD to generate samples. However, if a large number of samples must be generated or the reward is expensive amortized sampling approaches like DDPP are preferred. Indeed, this is one of the ultimate motivations of using RLHF algorithms in autoregressive models instead of methods like best of N the overall computation required to fine-tune a pre-trained model and subsequently sample from the fine-tuned model is much cheaper than generating N samples and selecting the best one many times over.
E.3 ANALYSIS OF IMPACT OF THE NUMBER OF MC SAMPLES M IN DDPP-IS
We include an ablation comparing the impact of the number of MC samples M used in DDPP-IS, on the datasets: MNIST, Amazon reviews, and Tinystories. The training steps and process times at which the model achieves a certain reward threshold, for different M, are plotted Figure 7.
We observe that a larger number of samples generally improves the reward at faster rate per iteration (gradient step), while each iteration generally takes more time, based on how expensive the reward function is to evaluate. Each gradient step involves a single call to the denoising model, and M calls to the reward function. Therefore, when increasing the number of samples M, it results in trade-off with fewer calls to the denoiser, while having more calls to the reward function, to achieve the same reward. In addition, when using more samples, it results in loss with lower variance, which can benefit training.
For the MNIST task in Figure 7a we see that a larger number of samples achieves the reward threshold in fewer iterations, and in less time. In this task the reward function is simple (a small classifier), larger number of samples doesn't it add too much a time cost per iteration, and a larger number of samples is preferred. On the other hand, for the Amazon reviews task in Figure 7b this manifests
30

---

Published as a conference paper at ICLR 2025
@ 1
w=1
@ 5
w 5
e 100
100
Figure 8: Discrete guidance samples for MNIST (left) and CelebA (right). As guidance scale increases sample quality decreases, espccially for CelebA, while reward increases (sce Table 2 and Figure 3). We view the generation of these high reward, low quality samples increases form of reward hacking.
in the trend that as MM increases, the number of iterations to reach the reward threshold generally decreases while the overall time increases. For this task, the reward model expensive to evaluate, so a smaller number of samples is more time-efficient. Finally, for tasks such as Tinystories, shown in Figure 7c, where the reward function is more expensive than MNIST but less so than Amazon reviews the interplay is more complicated. Increasing the number of Monte Carlo samples to M 8 improves the variance properties of the loss curve, which leads to an improvement in both iterations and overall time. As M increases beyond this, the number of iterations convergence decreases slightly at the cost of more overall computation time. A suggestion informed by these experiments is that, for given time budget, M should be treated as a hyperparameter for tuning, with ranges over lower values for tasks with more expensive reward functions.
E,4 ABLATION ON PROTEIN EXPERIMENTS SEQUENCE LENGTH
Table 9: Ablation on different protein lengths using DDPP-LB. DDPP-LB still generates high quality proteins sequence length increases.
Sequence length
3-sheet %
PLDDT
pTM1 f
log R(xo) t
90
0.44 + 0.04
0.90 + 0.03
0.81 + 0.03
3.70 = 0.19
150
0.71 + 0.04
0.74 + 0.01
0.66 0.04
4.56 0.28
200
0.56 0.09
0.77
+ 0.13
0.75
+ 0.11
4.36 0.32
250
0.64 + 0.01
0.91 + 0.05
0.89 + 0.02
4.78
0.09
We investigate further the performance of DDPP the designable 8-sheet task the sequence length scales. To this end, we repeat our 8-sheet experiment for DDPP-LB over additional protein sequence lengths of 150, 200, and 250 (we note that the length 90 we used in our original experiments was selected due to constraints regarding wet lab experimental protocol). We maintain all experimental settings, but for each of the different sequence lengths we perform grid search over the reward temperature parameter 3 and the learning rate. We selected a learning rate of 1e-6 for each additional sequence length, while for reward temperature we selected a setting of 3 0.0625 for sequence lengths 150 and 250 and maintained the original reward temperature of 3 = 0.125 for the sequence length 200 task. We report our results Table 9, where we see that DDPP-LB can still generate proteins according to the target distribution even protein sequence length increases. In fact, DDPP- LB seems to generate sequences with higher reward as we increased sequence length, an observation which follows intuition that the DPLM base model should is better generating slightly longer protein sequences as miniproteins of short lengths like 90 are relatively rare in the base model's training set compared slightly longer proteins.
E.5 DISCRETE GUIDANCE EXPERIMENTS
To help compare DDPP's performance to inference- methods we compare against discrete guidance proposed in Nisonoff et al. (2024). Discrete guidance requires both a differentiable reward as well as a reward which may be evaluated for partially masked (noisy) states. Unfortunately, text tasks require retokenization step the reward models different tokenization than does
31

---

Published as a conference paper at ICLR 2025
the pre-trained MDLM model we employ. As tokenization is a monn-iiferentiiabee operation unable to evaluate discrete guidance on our text tasks. Further, since discrete guidance requires the also prevented from evaluating on protein
reward function be evaluated on noisy states task. This is because our protein reward function uses ESMFold, a complicated protein folding model, which is only defined on full protein sequences. However, our image tasks fulfill both criteria and as such we evaluate discrete guidance on the MNIST and CelebA tasks.
As our initial image classifiers (used as reward functions) are not defined on partially masked states trained noisy versions of them for use in discrete guidance. To train the noisy reward functions we follow the recommendations of Nisonoff et al. (2024) by training on the same dataset as the original classifiers and noising the sampled datapoints according the same forward process as the fine-tuned diffusion model. The noisy classifiers performed nearly as well as the original, non-noisy classifiers with test set accuracy for MNIST degrading from 99% for the non-noisy classifier to 98% for the noisy classifier and from 96% for the non non-noisy to 95% for the noisy classifier on CelebA
Final reward evaluations are performed by evaluating samples generated using discrete guidance on the noisy reward models with the original, non-noisy rewards. BPD values for this baseline are computed by evaluating the base model's ELBO on images sampled using guidance (due to lacking an analogous ELBO formula for guidance-based sampling). Evaluation protocol follows that used for DDPP and other baselines, described more detail Appendix 1.1.2 for MNIST and Appendix D.3 for CelebA. Results are computed across three seeds for each guidance scale.
Guidance results are shown in Table 2 and Figure 3. We evaluate on rewards with the same temperature as the other baselines (3 = for MNIST and 3 = for CelebA). We have an additional multiplier in the guidance scale w. For settings corresponding to those of DDPP (w for MNIST and w = 1 for CelebA) and other baselines we observe that discrete guidance either improves mean reward by a small amount (or decreases it) compared to the base model and does not approach the performance of DDPP. We scale the guidance scale by 5x and 100x over the original guidance scale and, as expected, observe increase in
reward the guidance scale increases, even becoming competitive with the mean reward of DDPP on MNIST. However, this improvement in reward coincides with crippling decrease in sample quality indicated by the increase in BPD values. In Figure 8 show guided samples for both MNIST and CelebA with increasing guidance strengths, where we see that as guidance strength increases an instance of reward hacking occurs, with the guided samples achieving high reward under the classifier while being of low sample quality.
32


Reviews:
[{"review_id": "muK1xYL0bz", "reviewer": "Reviewer_nvQG", "summary": "This is a very interesting paper that design a new framework DDPP to sample from a target Bayesian posterior by steering pretrained MDMs. The author derived 3 primary objectives within the DDPP: DDPP-IS, DDPP-LB, DDPP-KL, providing different methods to handle the intractable partition function in the Bayesian posterior. This framework demonstrated strong performance across diverse tasks. In synthetic tasks, DDPP generated target-aligned samples with higher fidelity than baselines. For image modeling on the CelebA dataset, it produced attribute-specific images with competitive fidelity. In protein sequence modeling, DDPP generated high-quality sequences optimized for specific structural features, validated by wet-lab testing. In text generation, DDPP achieved alignment with sentiment and toxicity targets, producing fluent and reward-optimized outputs. Overall, DDPP consistently outperformed baselines, effectively steering model outputs to meet specific reward criteria across varied domains.", "questions": "N/A", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "1. The authors introduce the DDPP framework for steering MNDMs  by treating the task as sampling from a Bayesian posterior, which is an innovative approach for discrete generative models. They derive three scalable, simulation-free objectives that make fine-tuning computationally efficient, even for large pre-trained models.\n2. Extensive experiments were conducted across diverse tasks, including image modeling, text generation, and protein sequence design. DDPP demonstrated strong performance across all domains, highlighting its adaptability to various data types and applications.\n3. The paper includes wet-lab validation for protein sequence tasks, confirming the practical applicability of the model in real-world protein design. This physical validation of AI-generated outputs adds quite a lot credibility.\n4. The paper is clearly written and well-structured. The presentation is overall good.", "weaknesses": "1. Including autoregressive models as baselines could be useful, as they are widely used and effective in many generative tasks.\n2. Since DDPP-IS involves MCMC, it would be helpful to analyze the computational cost of each method, covering both training and inference phases.", "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": "bJe4gPifJD", "reviewer": "Reviewer_jeLb", "summary": "The paper construct transition distributions of discrete diffusion models for sampling from the composed distribution $\\pi_0(x_0) = p(x_0) R(x) / Z$, sometimes referred to as \"guided\" generation with reward function $R(x)$. Unlike in continuous space diffusions, the transition distribution  $p(x_{t-1} | x_t)$ corresponding to $\\pi_0(x_0)$ cannot be approximated using a Taylor expansion around $x_t$. The method proposed by the authors instead trains a second model $q(x_0 | x_t)$ to match the effective transition distribution of $\\pi_0(x_0)$. The primary challenge of matching the transition distribution of $\\pi_0(x_0)$ is estimating the log partition function $\\log Z$, and the authors propose three approaches to overcoming this challenge: (1) approximating $\\log Z$ through importance sampling (2) parameterizing $\\log Z$ with a neural network (3) foregoing approximation of $\\log Z$ by matching the reverse KL, which instead requires back-propagation through stochastic discrete sampling. Each of these variants has slightly different requirements for inference/train time compute, or differentiability of $R$. The authors evaluate their proposed variants on a few toy tasks as well as on conditional image generation and protein design. Overall they find that the proposed method (DDPP) achieves better reward values while maintaining plausible samples.", "questions": "1. Why was no comparison made to other methods that post-train diffusion models with preference data, e.g. Diffusion-DPO [1]? Is there a fundamental limitation of these methods in the discrete setting or would porting these methods to a discrete setting just end up looking similar to DDPP-KL? \n2. Was any variance reduction method required to make DDPP-KL work in practice? \n3. Why is it sufficient that DDPP-LB provides a lower bound on the log partition? As far as I can tell the loss functions can incentive minimizing the value of log partition, in which case minimizing a lower bound is not analogous to minimizing the bounded quantity. \n4. In the language examples, DDPP appears to perform much better than the provided baselines. Do you have any sense for how these conditional samples might compare to samples from a similar sized autoregressive model used in tandem with \"steering\" methods for autoregressive models (e.g. PPLM, fudge, twisted SMC). These models a probably capable of generating samples with higher likelihoods and many of these methods are \"plug-and-play\", meaning they don't require training a secondary model. Therefore, in practical settings, they might still be preferred over this method. \n\n[1] Diffusion Model Alignment Using Direct Preference Optimization. Wallace et al. (2023)", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper provides a nice overview of the problem of discrete guidance and several sensible solutions to it. I appreciate that the authors implemented several variants and compared them to each other, as it offers more signal to the community about promising directions. The significance of the work is substantial because discrete diffusion models are increasingly popular in both text applications, where guidance can equate to controllability/alignment, and for scientific modalities, where guidance is used for solving inverse problems. I personally don't find the originality to be high, as the method largely adopts pre-existing tools to solve the problems that arise in implementing each variant of their method, but I don't think this relatively low originality outweighs the quality of the work's presentation and analysis. To this end, I found the evaluation of the method mostly satisfactory, as it spanned multiple modalities and even included some wet lab experiments.", "weaknesses": "As I stated briefly above, I don't see the paper as making any foundational contributions that are broadly applicable beyond guidance of discrete diffusion models or any observations that are deeply surprising in nature. The paper is driven by a narrow and practical goal of finding a good method for approximating $\\log Z$ (or foregoing its approximation) when training a secondary model $q$ to approximate the transition distribution of $\\pi_0 (x_0)$. I personally don't find these contributions lackluster, as guidance of discrete diffusions is an open problem with practical impact. In terms of evaluation, the results seem fairly reasonable and there are some helpful baseline comparisons, though it's a little challenging to conclude much from the wet lab experiments when the only non-control methods are DDPP, so it's unclear how much heavy lifting the posterior inference method is doing versus other facets of sampling and filtering.\n\nOne notable limitation of the proposed method is that it requires training a secondary model in the first place. Once trained, this secondary model has relatively fast sampling, and therefore the training amortizes depending on how much inference is performed. In many scientific applications, inference can be relatively limited and throughput is often more important than latency, so in some cases maybe this tradeoff is not entirely worth it depending on the cost and complexity of setting up the post-training procedure, but I don't think these considerations disqualify the approach in general. It's also worth noting that best-of-N sampling seems to perform quite well in many experiment for relatively small N, even approaching the performance of DDPP in some cases, and this method can be deployed essentially out of the box, with only inference time overhead (which might be relatively minor).", "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": "mZ3ALhNKXX", "reviewer": "Reviewer_2bfH", "summary": "This paper addresses the task of controlled generation of a pretrained masked diffusion model (MDM) $p^{\\rm pre}(x)\\approx p _0(x)$ towards an altered target distribution $\\pi _0(x)\\propto p _0(x)R(x)$, where $R(x)$ is a reward function. The approach involves fine-tuning the MDM to align with the new target distribution. The paper introduces a simulation-free framework called \"discrete denoising posterior prediction\" (DDPP) to train the amortized MDM. The framework has shown promising empirical results for various tasks in image, text, and proteins.", "questions": "1. Questions for the framework.\n\n- Could the author(s) provide further insight into the rationale behind choosing an $L^2$-based loss in learning $q _\\theta({\\bf x} _0,...,{\\bf x} _{t-1}|{\\bf x} _t,\\hat{\\bf x} _0)$ in DDPP-IS/LB? The use of an $L^2$-based loss comparing the difference of log-densities is not commonly seen in the literature of learning probability distributions (e.g., variational inference), and it requires learning or estimating the normalizing constant ${\\cal Z} _{\\pi _t}({\\bf x} _t)$, which involves computational overhead. Furthermore, what is the probability of ${\\bf x} _{0:t}$ in the losses in equations (7,8,11,43)? Although this can be any trajectory, it would be beneficial to clarify your choice for the experiments.\n\n- In DDPP-KL, the author(s) employ the reverse KL loss instead of forward KL loss. However, this requires the differentiability of reward function $R(x)$, and needs the Reinmax trick to estimate the discrete gradients. It would be valuable if the author(s) could elaborate on why the forward KL loss is not suitable in this context.\n\n2. As discussed in Ou et al. (2024) and later Zheng et al. (2024), the time component in MDMs does not play a significant role. An MDM $\\mu _\\theta({\\bf x},t)\\in(\\Delta^d)^n$ is equivalent to an any-order autoregressive model, i.e., it predicts the one-dimensional conditional probability distributions of the masked positions given the observed ones in a partially observed sequence. If we formulate the MDM without time component, how would your fine-tuning framework be adapted?\n\n3. Questions regarding the experiments.\n\n- In the MNIST experiment, what would be the possible reason why SVDD having a higher reward value $\\log R({\\bf x} _0)$ compared with DDPP? Additionally, why is there no BPD metric for the samples generated by SVDD?\n\n- In the CelebA experiment, it seems that the quality of samples from DDPP (Figure 3) is worse than those from the pretrained model (Figure 5).\n\n4. Minor comments\n\n- In lines 89 and 97, it seems that the notation $\\cal V$ should be $\\cal X$ instead.\n\n- In line 98, the RHS of the equation $p({\\bf X}={\\bf x})$ is a bit confusing. Does it imply that each position is mutually independent? Clarification would be helpful.\n\n- In line 105, it might be more appropriate to say \"drop $t$ and write $i$\" instead. Since the paper uses a discrete time setting, replacing $\\tau({\\bf x} _{0:1})={\\bf x} _1\\to\\cdots\\to{\\bf x} _t\\to{\\bf x} _0$ with $\\tau({\\bf x} _{0:T})={\\bf x} _T\\to\\cdots\\to{\\bf x} _t\\to{\\bf x} _0$ would be more appropriate.\n\n- In section 2.1, when dealing with multiple dimensions, $\\mu _\\theta({\\bf x},t)$ should be in $(\\Delta^d)^n$, which predicts the probabilities $\\Pr({\\bf X} _0^i=\\cdot|{\\bf X} _t={\\bf x})$.\n\n- In line 260, it would be better to move the training objective (Eq. 11) to this location for better readability.\n\n- Furthermore, I noticed a few typos or errors in the paper. For example, in line 323, it states \"In settings, where the reward model is differentiable\", and in line 430, it says \"We provide full a deeper description of evaluation metrics and experimental setup\". I strongly recommend that the author(s) carefully review and polish the writing in this paper during the rebuttal period to ensure clarity and accuracy.\n\nI would be happy to raise the score if the author(s) could address my concerns.\n\n**References**\n\nOu et al. Your Absorbing Discrete Diffusion Secretly Models the Conditional Distributions of Clean Data. Arxiv 2406.03736.\n\nZheng et al. Masked diffusion models are secretly time-agnostic masked models and exploit inaccurate categorical sampling. ArXiv 2409.02908.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The paper addresses the important problem of controlled generation, which has significant applications as discussed in the paper. The proposed DDPP framework, consisting of three different versions, is a novel approach that efficiently tackles this problem without simulating the reverse diffusion process. The empirical results demonstrate the superiority of the DDPP method over several baseline methods, including best-of-$N$, SVDD, and RTB. This highlights the effectiveness and potential of the proposed framework in the context of discrete diffusion models.\n\nThe paper is well-written and clearly organized, and provides sufficient details to understand and replicate the experiments. However, further improvement can be made in the clarity of the mathematical explanations.", "weaknesses": "The methodologies introduced in the paper are somewhat confusing to me, and I am unable to fully grasp the underlying intuition behind their design. Please check the question part for details.\n\nThe writing of the paper, particularly the mathematical explanations, could benefit from improvement. I have provided comments and suggestions in the part regarding these issues.", "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": "sVSHt5pM19", "reviewer": "Reviewer_oyrv", "summary": "The paper introduces Discrete Denoising Posterior Prediction (DDPP), a framework for guiding Masked Discrete Diffusion Models (MDMs) toward a reward model. DDPP offers three scalable, simulation-free objectives—DDPP-IS, DDPP-LB, and DDPP-KL—suited for fine-tuning MDMs in tasks like image modeling, protein design, and text generation.", "questions": "1. Could DDPP be adapted to other types of discrete generative models beyond MDMs such as autoregressive or flow-based discrete models?\n\n2. Relevant to weakness, how does the computational cost of DDPP compare to other fine-tuning/inference-time approaches, especially in larger-scale applications?\n\n3. Protein design applications can involve sequences of varying lengths, which might affect DDPP’s performance. How does the method scale with longer or structurally complex sequences? I suppose an ablation study on sequence length is critical.\n\n4. The paper introduces three variants of DDPP (DDPP-IS, DDPP-LB, and DDPP-KL) but does not make it clear which one is preferred in what scenarios. Could the authors comment on selecting the right variant given specific reward functions or computational constraints?", "limitations": "", "rating": 5, "confidence": 5, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. The paper presents DDPP, a method to align Masked Diffusion Models (MDMs) with reward functions directly at inference, skipping RL and costly sampling typical for autoregressive models.\n\n2. DDPP is backed by theories.  Additionally, its three variants—DDPP-IS, DDPP-LB, and DDPP-KL— are designed for both differentiable and non-differentiable rewards. These enable efficient MDM fine-tuning with strong performance across tasks, which is impressive.\n\n3. The writing flow of this work is somehow clear. I find the presentation smooth and easy to follow.", "weaknesses": "1. First of all, while the presentation was good, I found the paper’s theoretical parts were not very easy to follow. I would suggest simplifying theory sections and providing illustrations to improve readability. It is still not crystal clear to me what are the motivations and the differences for making such 3 variants. For instance, for DDPP-KL, which involves a reverse KL divergence objective, the authors break down its computational pathway further, which might help readers better grasp the core mechanisms quickly.\n\n2. Since the study centers on MDMs, it remains unclear if DDPP can generalize to other types of discrete diffusion models or if the results depend heavily on the MDM-specific setup.\n\n3. The authors claimed that DDPP enabled efficient MDM fine-tuning across many domains. There's still limited discussion of the computational costs relative to baseline approaches. For example, I assume SVDD is computationally friendly as it is purely inference time. Furthermore, it can balance memory and time for lighter computation demands. Thus, in terms of computation efficiency, I would assume DDPP is not superior to SVDD. Would the authors comment on 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)