D-DOIT: Training-free Adaptation of Discrete Diffusion via Doob's h-Transform

Jieke Wu; Qijie Zhu; Weimin Wu; Zeqi Ye; Minshuo Chen; Han Liu

cs.LGOct 4, 2026

D-DOIT: Training-free Adaptation of Discrete Diffusion via Doob's h-Transform

Authors: Jieke Wu, Qijie Zhu, Weimin Wu, Zeqi Ye, Minshuo Chen, Han Liu

Organizations: Department of Computer Science, King Abdullah University of Science and Technology · Department of Statistics and Data Science, Northwestern University · Center for Foundation Models and Generative AI, Northwestern University · Industrial Engineering & Management Sciences, Northwestern University

Abstract

We propose D-DOIT (Discrete Doob-Oriented Inference-time Transformation), a training-free and efficient adaptation method for discrete diffusion models with generic rewards. D-DOIT formulates adaptation as sampling from a reward-tilted target distribution and realizes this transport through Doob's h-transform of the discrete diffusion reverse kernel, using only reward values rather than reward gradients. Unlike continuous diffusion, masked discrete diffusion samples categorical token-reveal transitions rather than continuous state updates. D-DOIT derives the corresponding discrete Doob's h-transform, which guides sampling by reweighting reverse transition probabilities instead of adding a drift correction. To make this transformation practical, D-DOIT avoids expensive future rollouts. At each guided step, D-DOIT samples candidate next states, uses the model prediction head to complete each candidate into a clean sequence, evaluates each completion with the reward oracle, and resamples the next state with probabilities proportional to the rewards. An optional late-stage best-of-K refinement further improves sample quality by branching trajectories only near the end of denoising, avoiding the $K$ -fold cost over the full trajectory. Empirically, across regulatory DNA design and protein inverse folding benchmarks, D-DOIT outperforms training-free guidance baselines. It improves enhancer activity and cell-type specificity while preserving sequence naturalness, and achieves the highest success rate in protein inverse folding.

Figures & tables

Figure 1: Overview of D-DOIT. Left: Standard masked diffusion unmasks positions independently at each reverse step ( 2.1 ); it is cheap but reward-unaware. Right: One guided step approximately tilts the reverse kernel: it draws $M{=}3$ candidates from a tempered proposal ( 3.6 ), completes and scores them with the prediction head and reward to obtain proxy $h$ -values, and selects the next state with sharpened $h$ -weights ( 3.7 ). Bottom: Optional best-of- $K$ refinement duplicates each trajectory into $K{=}4$ copies after the split and keeps the highest-reward suffix ( 3.9 ), sharing prefix computation. Revealed tokens stay fixed; the practical sampler does not inherit the exact-target guarantee.

Property	CG	TDS/SMC	SVDD	D-DOIT (Ours)
Reward intervention	gradient (backward)	particle weights	soft value estimate	Doob $h$ -transform on the discrete kernel
Requires backprop	✓	✗	✗	✗
Weighted selection unit	—	$M$ particle trajectories	$M$ candidate states	$M$ candidate next states (same $\mathbf{x}_{l}$ )
Weight formula	—	path likelihood ratio	estimated soft value	transformed kernel via $\hat{h}^{\omega}$
Doob kernel tilt	✗	✗	✗	✓
Candidate diversity	—	different histories	same state start	same $\mathbf{x}_{l}$ start

Table 1: Comparison of reward-guided sampling methods for discrete diffusion. We contrast D-DOIT with our guidance baselines ( Section 4 ); CG = classifier guidance, TDS = twisted diffusion sampler, SMC = sequential Monte Carlo, SVDD = soft value-based decoding.

Method	Pred-Activity (median) $\uparrow$	ATAC-Acc $\uparrow$ (%)	3-mer Corr $\uparrow$	JASPAR Corr $\uparrow$
Training-free methods
Pre-trained	0.17(0.04)	1.5(0.2)	-0.061(0.034)	0.249(0.015)
CG	3.30(0.00)	0.0(0.0)	-0.065(0.001)	0.212(0.035)
SMC	4.15(0.33)	39.9(8.7)	0.840(0.045)	0.756(0.068)
TDS	4.58(0.49)	40.0(14.1)	0.817(0.097)	0.798(0.044)
SVDD-PM †	6.04(0.07)	39.8(0.4)	0.688(0.012)	0.917(0.002)

Table 2: HepG2 enhancer design results. We compare D-DOIT with training-free guidance and training-based finetuning baselines on DRAKES HepG2; all metrics are higher-is-better. Parentheses denote standard deviation across seeds, and bold / underline mark best/second best per column. D-DOIT uses held-out eval/test seed protocols in Sections J.6 and J.7 . † SVDD-PM [ 12 ] is the matched-time control from Table 6 (

M=416

), not part of the original DRAKES sweep.

Method	$\Delta R$ $\uparrow$	Motif Corr. $\uparrow$
Training-free methods
Pre-trained	$0.000(0.003)$	$\underline{0.937(0.009)}$
CG	$0.000(0.003)$	$\mathbf{0.938(0.009)}$
SMC	$0.000(0.010)$	$0.587(0.114)$
TDS	$0.002(0.009)$	$0.573(0.078)$
Frozen Best-of-708	$0.108(0.028)$	$0.490(0.165)$

Table 3: Cell-type-specific enhancer design (left) and protein inverse folding (right). Left: means over three cell lines, SD in parentheses;

\Delta R

is target minus average off-target reward; training-based references are in Tables 13 and K.4 . Right: DRAKES protein benchmark; Pred-ddG and joint success are higher-is-better, scRMSD lower-is-better, SDs in Table 22 . Bold marks the best mean per column, not significance.

Quantity	LBOK-S	Q
$(M,K,s_{\mathrm{split}})$	$(9,4,105)$	$(52,8,0)$
Model NFE	$1{,}970$	$54{,}272$
Candidate completions	$1{,}773$	$53{,}248$

Table 4: Per-output DRAKES work. Full-window guidance,

L=128

, one completion per candidate.

Appendix figures & tables24 assets

Supplementary material from the paper’s appendix.

Appendix

				DRAKES		Ctrl-DNA
Sym.	Name	Role	Range	Q	S	Q	S
$L$	Steps	Total discrete reverse steps.	$\{32,64,128\}$	$128$	$128$	$128$	$128$
$M$	Budget	Candidate next states per guided step; linear head-evaluation cost.	$\{2,\dots,64\}$	$52$	$9$	$52$	$9$
$\omega$	Strength	Exponent on $\hat{h}$ in ( 3.7 ); for one completion, log-weight sharpness is $\omega/\beta$ .	$[0.5,\,8.0]$	$3.0$	$3.0$	$8.0$	$3.0$
$\beta$	Temp.	Scales reward: $\hat{h}=\exp(R/\beta)$ in ( 3.3 ). Small $\beta$ amplifies differences.	$[0.1,\,2.0]$	$0.2$	$0.2$	$0.5$	$0.2$
$\tau$	Prop. T	Proposal temperature ( 3.6 ) for candidate diversity; changes the reference kernel.	$[1.0,\,3.0]$	$3.0$	$3.0$	$3.0$	$3.0$

Appendix

Table 5: Hyperparameters of D-DOIT. We report two operating points per benchmark: Q = quality-oriented and S = speed-oriented. Q denotes high-budget quality; S denotes LBOK-S. Selection and test protocols are described in Tables 2 and 3 ; diagnostic sweeps specify their own configurations.

Metric	D-DOIT-Q	SVDD-PM	BoN
s/output $\downarrow$	$10.595$	$11.236$	$\mathbf{10.587}$
Activity $\uparrow$	$\mathbf{6.19\pm 0.05}$	$6.04\pm 0.07$	$4.70\pm 0.03$
ATAC (%) $\uparrow$	$\mathbf{91.7\pm 0.9}$	$39.8\pm 0.4$	$43.5\pm 1.3$
3-mer $\uparrow$	$\mathbf{0.894\pm 0.009}$	$0.688\pm 0.012$	$0.881\pm 0.008$
JASPAR $\uparrow$	$0.923\pm 0.002$	$0.917\pm 0.002$	$\mathbf{0.957\pm 0.002}$

Appendix

Table 6: Matched-time controls. We test D-DOIT-Q vs SVDD-PM and Frozen Best-of-732 at similar measured wall-clock cost. Bold marks the best mean per column.

Configuration	Pred-Activity $\uparrow$	Model NFE per output	Reward evaluations per output	s/640 outputs
Frozen Best-of-15	$2.219\pm 0.017$	$1{,}920$	$15$	$97.8\pm 1.5$
$M=9,\tau=1,K=1$	$3.115\pm 0.021$	$1{,}280$	$1{,}152$	$227.0\pm 1.9$
$M=9,\tau=3,K=1$	$3.670\pm 0.072$	$1{,}280$	$1{,}152$	$227.8\pm 2.6$
$M=9,\tau=3,K=4,s_{\mathrm{split}}=105$	$4.491\pm 0.045$	$1{,}970$	$1{,}777$	$327.8\pm 4.3$

Appendix

Table 7: Backbone work, oracle work, and measured time on DRAKES. Each setting uses five sampling seeds and

640

outputs per seed. Pred-Activity is held-out Enformer activity; entries report means and sample standard deviations across seeds. The final column measures the full

640

-output job and must not be extrapolated to other physical batch sizes. The unbranched and branched rows are diagnostic settings, not substitutions for the separately selected main-table LBOK-S or Q results.

Generation method	Pred-Activity $\uparrow$	ATAC-Acc (%) $\uparrow$	3-mer Corr $\uparrow$	JASPAR Corr $\uparrow$
Enformer-guided D-DOIT	$4.491\pm 0.045$	$68.25\pm 1.86$	$0.8630\pm 0.0160$	$0.9468\pm 0.0025$
CNN-guided D-DOIT	$2.203\pm 0.054$	$51.91\pm 1.19$	$0.8738\pm 0.0013$	$0.9342\pm 0.0021$
MLP-guided D-DOIT	$1.544\pm 0.037$	$28.78\pm 1.29$	$0.9489\pm 0.0035$	$0.9347\pm 0.0031$
Frozen pre-trained model	$0.371\pm 0.018$	$1.66\pm 0.59$	$-0.0584\pm 0.0149$	$0.2723\pm 0.0220$

Appendix

Table 8: Changing the guidance oracle. Only the guidance reward

R

changes among the D-DOIT rows; the sampler, seeds, output count, held-out Enformer evaluator, and other metrics are fixed. Entries are mean

\pm

sample standard deviation over five seeds, with

640

sequences per seed. Pred-Activity uses the held-out Enformer. These results are separate from the high-budget Q evaluation.

Generation method	Pred-Activity $\uparrow$ (held-out Enformer)	Held-out CNN $\uparrow$	Held-out MLP $\uparrow$
Enformer-guided D-DOIT	$4.491\pm 0.045$	$0.867\pm 0.014$	$0.634\pm 0.015$
Frozen pre-trained model	$0.371\pm 0.018$	$0.484\pm 0.014$	$0.509\pm 0.014$

Appendix

Table 9: Transfer across held-out evaluators. The same Enformer-guided D-DOIT and frozen-model sequences are rescored by three held-out predictors. Entries are mean

\pm

sample standard deviation over five seeds, with

640

sequences per seed. Enformer scores are Pred-Activity; CNN and MLP scores have their own scales. Compare methods within each column, not absolute values across columns.

Metric	TDS reference	D-DOIT-LBOK-S final test	Mean comparison
Pred-Activity	$4.58$	$4.560(0.040)$	similar mean
ATAC-Acc	$40.0\%$	$65.03(1.43)\%$	pass
3-mer Corr	$0.817$	$0.8511(0.0165)$	pass
JASPAR Corr	$0.798$	$0.9477(0.0016)$	pass

Appendix

Table 10: Held-out final-test check for D-DOIT-LBOK-S. Targets are the TDS row in Table 2 . The labels compare reported means with the reference values; they are not significance or statistical-equivalence decisions.

Metric	DRAKES reference	D-DOIT-Q final test	Mean comparison
Pred-Activity	$5.61$	$6.19(0.05)$	pass
ATAC-Acc	$92.5\%$	$91.7(0.9)\%$	miss
3-mer Corr	$0.887$	$0.894(0.009)$	pass
JASPAR Corr	$0.911$	$0.923(0.002)$	pass

Appendix

Table 11: Held-out final-test check for D-DOIT-Q. Targets are the DRAKES row in Table 2 . The final-test result passes Pred-Activity, 3-mer Corr, and JASPAR Corr, but misses the ATAC-Acc target on the held-out evaluation.

Figure 2: Utility-naturalness trade-off. We plot Pred-Activity (median) against 3-mer Corr; one point per method. Green shading marks the high-utility, high-naturalness region. This view covers the original benchmark rows; the separately time-calibrated SVDD and endpoint comparisons are in Table 6 . ATAC-Acc is reported separately in Tables 2 and 11 .

Method	Seed-42 screen	Selected config	Screen outcome	Takeaway
CG	guidance scale $\{20,50,100,200\}$	CG_G20: guidance scale $20$	$\Delta R\in[-0.002,0.000]$ ; specificity $\leq{-}0.007$ ; motif corr $0.903$ – $0.916$	All screened CG settings preserve motifs but do not improve reward or specificity.
SMC	beta values: $0.3$ , $0.5$ , $0.75$ , $1.0$ , $1.5$	SMC_B05: beta $0.5$	Best $\Delta R=0.011$ at beta $1.5$ , with motif corr $0.388$ and diversity $0.928$ ; selected beta $0.5$ has $\Delta R=0.006$ .	Small reward gains come with much weaker motif/diversity, so SMC is not competitive.
TDS	$g$ /beta pairs: $2/0.75$ , $5/0.75$ , $5/1.0$ , $10/1.0$	TDS_G5_B10: $g=5$ , beta $1.0$	$\Delta R\in[-0.002,0.001]$ ; specificity $<0$ for all settings; best motif corr $0.510$ .	TDS remains weak after this bounded tuning pass.

Appendix

Table 12: Ctrl-DNA training-free baseline tuning screen. We run a bounded seed-42 screen over 13 CG, SMC, and TDS hyperparameter settings on all three Ctrl-DNA cell tasks. The screen documents that these training-free baselines remain weak after tuning: CG preserves motifs without improving reward/specificity, SMC gains little reward while degrading motif/diversity, and TDS stays near zero reward movement. The selected configs are provenance for possible follow-up evaluation, not final-test rows.

Figure 3: Specificity vs. diversity on Ctrl-DNA. Each panel is one target cell line. Small markers are single seeds, and large black-edged markers are the corresponding

5

-seed means. The

x

-axis is cell-specific reward gain

\Delta R

; the

y

-axis is set-level Shannon entropy from the Ctrl-DNA evaluation script [ 25 ] . The green region marks positive

\Delta R

with high diversity. D-DOIT moves into this region, whereas Pre-trained/CG remain near zero reward movement and SMC/TDS lose diversity.

Cell line	Method	$\Delta R$ $\uparrow$	Motif Corr $\uparrow$	Diversity $\uparrow$
HepG2	Pre-trained	$-0.004(0.002)$	$\underline{0.937(0.009)}$	$\mathbf{1.988(0.001)}$
	CG	$-0.004(0.002)$	$\mathbf{0.938(0.009)}$	$\mathbf{1.988(0.001)}$
	SMC	$+0.003(0.011)$	$0.582(0.094)$	$1.484(0.109)$
	TDS	$-0.004(0.009)$	$0.606(0.087)$	$1.509(0.061)$
	D-DOIT-LBOK-S (Ours)	$\underline{+0.014(0.003)}$	$0.667(0.032)$	$1.981(0.002)$
	D-DOIT-Q (Ours)	$\mathbf{+0.095(0.008)}$	$0.558(0.048)$	$\underline{1.982(0.001)}$

Appendix

Table 13: Per-cell-type Ctrl-DNA results. We break down

\Delta R

, Motif Corr, and Diversity by target cell line. Parentheses denote standard deviation across

5

random seeds. Bold = best, underline = second best, per column within each cell line.

Method	$\Delta R$ $\uparrow$	Motif Corr $\uparrow$
AdaLead	$0.05(0.05)$	$0.36(0.08)$
PEX	$0.09(0.02)$	$0.23(0.04)$
BO	$0.01(0.02)$	$0.11(0.11)$
CMA-ES	$0.01(0.01)$	$0.08(0.05)$

Appendix

Table 14: Evolutionary-search baselines on Ctrl-DNA. We report

\Delta R

and Motif Corr averaged across the three target cell lines for the four directed-evolution baselines reproduced from [ 25 ] . The training-free and training-based comparisons appear in Table 3 .

Axis	Setting	3-mer Corr	Pred-Activity / ATAC (%)
$l^{*}$	$64\to 128$	$-0.06\to 0.894$	$3.61\to 4.77$ / $33.6\to 72.7$
$\gamma$	$1.2\to 0.6$	$0.846\to 0.905$	$4.53\to 4.72$ / $72.7\to 66.4$
$M$	$8\to 32$	$+9\%$	$+31\%$ / $+29$ pp

Appendix

Table 15: Headline ablation findings. Each row varies one hyperparameter around the DRAKES ablation reference setting. We report sweep-margin endpoints from a single-seed sweep; full sweeps appear in Appendix L .

Figure 4: Guidance-window sweep. Expanding the guided window improves utility and sharply recovers naturalness between

l^{*}=96

and

l^{*}=112

Figure 5: ESS diagnostic. ESS and collapse fraction across

\omega

; weighted selection collapses to near-greedy at

\omega\geq 4

Figure 6: Interaction of proposal and completion temperatures. We sweep the proposal temperature $\tau$ and the completion temperature $\gamma$ at $M=32$ . Naturalness peaks at sharper prediction-head completions (smaller $\gamma$ ), while functional activity peaks at higher proposal diversity (larger $\tau$ ). The red star marks the ablation reference configuration used before the final held-out Q budget selection.

Figure 7: Candidate-budget sweep. We vary the Monte Carlo candidate budget

M

on the DRAKES HepG2 task and report the same four metrics. Increasing

M

gives direct gains in functional utility, while naturalness changes much less across the sweep.

Completions $n_{h}$	Final guidance reward $R$	Wall-time increase
$1$	$5.272\pm 0.057$	baseline
$2$	$5.764\pm 0.033$	$41\%$
$3$	$5.978\pm 0.026$	$81\%$

Appendix

Table 16: Completion-count ablation on HepG2 . We use five seeds and

3{,}200

sequences per setting. The score is final guidance reward

R

, not held-out Pred-Activity. Runtime increases are relative to

n_{h}=1

Fixed bases	Direct: mean (variance)	Rollout: mean (variance)	W1	KS
$0.78\%$	$3.494\ (0.974)$	$3.480\ (1.097)$	$0.060$	$0.045$
$33.39\%$	$3.536\ (1.241)$	$3.611\ (1.109)$	$0.101$	$0.050$
$66.34\%$	$3.648\ (1.685)$	$3.772\ (1.620)$	$0.124$	$0.058$

Appendix

Table 17: Matched-state completion diagnostic. At each stage,

640

partial states come from

128

source trajectories per seed over five seeds. Each state is independently completed

32

times by the direct prediction head and

32

times by sequential continuation. Moments and distances summarize the pooled distributions of per-state average

h

-values; parentheses denote variance, not uncertainty.

$\omega$	Mean ESS	Largest normalized weight	Top-candidate selection rate	Soft reward $R$	Argmax reward $R$
$1$	$1.877\pm 0.020$	$0.775$	$77.53\%$	$4.894\pm 0.055$	$5.154\pm 0.027$
$4$	$1.237\pm 0.004$	$0.927$	$92.69\%$	$5.183\pm 0.038$	$5.154\pm 0.027$
$8$	$1.143\pm 0.002$	$0.957$	$95.72\%$	$5.237\pm 0.051$	$5.154\pm 0.027$

Appendix

Table 18: Soft selection versus direct argmax on HepG2. We use

M=8

\tau=3

K=4

, and

\beta=0.2

, with five seeds and

3{,}200

sequences per setting. ESS and top-candidate frequency describe guided-step selection; rewards are final guidance reward

R

. Direct argmax does not depend on

\omega

Split update $s_{\mathrm{split}}$	Final guidance reward $R$	Best-of- $4$ gain	Wall time/job (s)
$0$	$5.890\pm 0.028$	$1.605\pm 0.019$	$793.6\pm 2.2$
$32$	$5.855\pm 0.035$	$1.545\pm 0.030$	$653.8\pm 1.5$
$64$	$5.729\pm 0.029$	$1.435\pm 0.019$	$510.4\pm 1.9$
$105$	$5.272\pm 0.057$	$0.959\pm 0.020$	$326.4\pm 0.5$

Appendix

Table 19: Controlled split-position ablation with

M=9

K=4

, and

n_{h}=1

fixed. Only

s_{\mathrm{split}}

varies. Both final reward and best-of-

4

selection gain use the guidance reward

R

, not held-out Pred-Activity.

Fixed positions	Still-masked positions	Remaining-position Hamming distance
$25.43\%$	$74.57\%$	$0.710\pm 0.010$
$50.38\%$	$49.62\%$	$0.689\pm 0.010$
$90.38\%$	$9.62\%$	$0.661\pm 0.011$

Appendix

Table 20: Suffix diversity under the frozen reverse process. For each saved partial sequence,

32

independent suffix completions are sampled. Hamming distance is the mean pairwise fraction of differing bases restricted to positions that were still masked in the shared partial state.

Update	$1$	$32$	$64$	$96$	$128$
D-DOIT	$2.067$	$2.328$	$1.763$	$1.758$	$3.670$
Frozen sampling	$2.053$	$1.561$	$1.058$	$0.716$	$0.371$

Appendix

Table 21: Held-out enformer completion-proxy trajectory. At each update, masked positions are filled with the frozen model’s deterministic argmax predictions before scoring. Scores are averaged over

640

sequences per seed and then over five seeds, rather than using the batch-median Pred-Activity aggregation.

Method	Pred-ddG median $\uparrow$	ddG $>0$ (%) $\uparrow$	scRMSD median (Å) $\downarrow$	scRMSD $<2$ Å (%) $\uparrow$	Success (%) $\uparrow$
Pre-trained	$-0.544\pm 0.037$	$36.6\pm 1.0$	$0.849\pm 0.013$	$90.9\pm 0.6$	$34.4\pm 0.5$
CG	$-0.561\pm 0.045$	$36.9\pm 1.1$	$0.839\pm 0.012$	$90.9\pm 0.6$	$34.7\pm 0.9$
SMC	$0.659\pm 0.044$	$68.5\pm 3.1$	$0.841\pm 0.006$	$93.8\pm 0.4$	$63.6\pm 4.0$
TDS	$0.674\pm 0.086$	$68.2\pm 2.4$	$\mathbf{0.834\pm 0.001}$	$\mathbf{94.4\pm 1.2}$	$62.9\pm 2.8$
CFG	$-1.186\pm 0.035$	$11.0\pm 0.4$	$3.146\pm 0.062$	$29.4\pm 1.0$	$1.3\pm 0.4$
D-DOIT	$\mathbf{0.870\pm 0.046}$	$\mathbf{74.8\pm 3.1}$	$0.872\pm 0.009$	$92.5\pm 1.3$	$\mathbf{68.1\pm 3.5}$

Appendix

Table 22: Protein inverse folding on the DRAKES benchmark. D-DOIT entries are mean

\pm

sample SD across seed-level metrics. Other rows and uncertainties are from DRAKES Table 2 [ 8 ] , not rerun. Bold marks the best mean per column, not statistical significance.

Explore similar work

Feb 18, 2026

Training-Free Adaptation of Diffusion Models via Doob's $h$ -Transform

Adaptation methods have been a workhorse for unlocking the transformative power of pre-trained diffusion models in diverse applications. Existing approaches often abstract adaptation objectives as a reward function and steer diffusion models to generate high-reward samples. However, these approaches can incur high computational overhead due to additional training, or rely on stringent assumptions on the reward such as differentiability. Moreover, despite their empirical success, theoretical justification and guarantees are seldom established. In this paper, we propose DOIT (Doob-Oriented Inference-time Transformation), a training-free and computationally efficient adaptation method that applies to generic, non-differentiable rewards. The key framework underlying our method is a measure transport formulation that seeks to transport the pre-trained generative distribution to a high-reward target distribution. We leverage Doob's

h

-transform to realize this transport, which induces a dynamic correction to the diffusion sampling process and enables efficient simulation-based computation without modifying the pre-trained model. Theoretically, we establish a high probability convergence guarantee to the target high-reward distribution via characterizing the approximation error in the dynamic Doob's correction. Empirically, on D4RL offline RL benchmarks, our method consistently outperforms state-of-the-art baselines while preserving sampling efficiency.

Qijie Zhu, Zeqi Ye, Han Liu +2

Department of Statistics and Data Science, Northwestern University. · Department of Industrial Engineering and Management Sciences, Northwestern University. · Department of Computer Science, Northwestern University.

Diffusion Models Diffusion Model Guidance Diffusion Sampling RL for Diffusion Models Doob H-Transform

Ask PDF

Oct 2, 2026

IDRF: Inverse-Distilled Reward Fine-tuning of Masked Discrete Diffusion Models

Masked discrete diffusion models offer a promising alternative to autoregressive generation, but iterative sampling can be costly, and intractable sequence likelihoods complicate reward fine-tuning. We introduce IDRF, a framework for reward fine-tuning of few-step masked discrete diffusion generators. Starting from a standard reverse-KL-regularized objective, IDRF replaces the intractable sequence-level KL penalty with inverse-distillation regularization. With an optimal auxiliary denoiser, we prove that the population inverse-distillation loss upper-bounds the sequence-level KL divergence to the reference distribution. IDRF optimizes a trajectory-based surrogate of this loss without reference-model rollouts, so the student keeps its own few-step sampler. We view few-step generation as a finite-horizon Markov decision process and optimize reward with a clipped policy-gradient objective over the student's trajectories. Across DNA, image, and text generation, IDRF achieves high reward with up to

32\times

fewer denoising steps than the reference while mitigating reward hacking and preserving sample quality.

Vladislav Gromadskii, David Li, Samson Gourevitch +4

Applied AI Institute, Moscow, Russia · MBZUAI · CMAP, École Polytechnique +3

Diffusion Model Fine-Tuning Masked Diffusion Models RL for Diffusion Models Few-Step Diffusion Sampling

Ask PDF

Aug 27, 2026

VGAS: Variance-Reduced Guidance and Adaptive Selection for Training-Free Reward Alignment in Discrete Diffusion

Masked discrete diffusion models perform strongly on text, code, and biological sequences, but their training objective rewards only naturalness, and retraining the generator for every new reward is expensive. Inference-time steering of a frozen model either guides the sampler by the reward gradient or searches over several trajectories, and recent samplers combine the two. Such combinations are assembled as pipelines that leave three choices at their defaults: a guidance estimate resting on one Gumbel draw per sample, a reward tilting placed without reference to the distribution the combination then targets, and a selection temperature held fixed although the spread of per-step rewards drifts. We identify that distribution and settle the three choices against it. We therefore propose Variance-reduced Guidance and Adaptive Selection (VGAS), a simple yet effective inference-time framework that reduces the variance of the guidance estimate for both reward types, applies the reward tilting in the clean-token logits, where the pretrained schedule is preserved, and sets the selection temperature per step. Across regulatory DNA, protein and small-molecule benchmarks, VGAS attains the best training-free reward and matches or surpasses a reward-fine-tuned generator.

Kwanyoung Kim

Department of AI, GIST

Diffusion Model Alignment Diffusion Model Guidance Masked Diffusion Models Inference-Time Guidance

Ask PDF

Abstract

Explore similar work

Training-Free Adaptation of Diffusion Models via Doob's hhh-Transform

IDRF: Inverse-Distilled Reward Fine-tuning of Masked Discrete Diffusion Models

VGAS: Variance-Reduced Guidance and Adaptive Selection for Training-Free Reward Alignment in Discrete Diffusion

Training-Free Adaptation of Diffusion Models via Doob's $h$ -Transform