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

Vladislav Gromadskii; David Li; Samson Gourevitch; Yazid Janati; Eric Moulines; Maxim Panov; Alexander Korotin

cs.LGOct 2, 2026

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

Authors: Vladislav Gromadskii, David Li, Samson Gourevitch, Yazid Janati, Eric Moulines, Maxim Panov, Alexander Korotin

Organizations: Applied AI Institute, Moscow, Russia · MBZUAI · CMAP, École Polytechnique · Institute of Foundation Models · LRE, EPITA · AXXX, Moscow, Russia

Abstract

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.

Figures & tables

Figure 1: IDRF (ours) across DNA, images, and text. Reward-only fine-tuning and IDRF sample with the same budget of $8$ , $4$ , and $32$ steps. Reward-only fine-tuning hacks the reward: DNA lacks accessibility (ATAC-Acc), images blur, and reviews degenerate into repetition that rarely ends with EOS. IDRF avoids these failures.

Figure 2: IDRF (ours) training framework. An on-policy rollout of the student from $\mathbf{m}^{L}$ provides an intermediate state $x_{t_{j}}$ and a completed sample $x_{0}$ . The soft prediction $\widetilde{x}=G_{\theta}(x_{t_{j}})$ is freshly corrupted and scored by the frozen reference $f_{\mathrm{ref}}$ and the trainable auxiliary $f_{\psi}$ , whose NELBO gap gives $\mathcal{L}_{\mathrm{ID}}$ . The reward of $x_{0}$ gives the policy-gradient loss $\mathcal{L}_{\mathrm{R}}$ , and the student minimizes $\mathcal{L}_{\mathrm{R}}+\lambda\mathcal{L}_{\mathrm{ID}}$ .

Method	Pred-Activity $\uparrow$	ATAC-Acc HepG2 (%) $\uparrow$	$3$ -mer Corr. $\uparrow$	Log-Lik $\uparrow$	Entropy (nats)	Steps $\downarrow$
Reference	$0.18\,(0.02)$	$2.0\,(0.4)$	$-0.042\,(0.010)$	$-262\,(0.1)$	$1.340\,(0.001)$	$128$
CG	$1.56\,(0.02)$	$0.0\,(0.0)$	$-0.057\,(0.002)$	$-269\,(0.2)$	$1.332\,({<}0.001)$	$128$
SMC	$4.69\,(0.48)$	$49.7\,(6.4)$	$\underline{0.880}\,(0.014)$	$-256\,(1.0)$	$1.299\,(0.012)$	$128$
TDS	$4.54\,(0.55)$	$33.4\,(10.8)$	$0.822\,(0.130)$	$-257\,(1.7)$	$1.313\,(0.024)$	$128$
CFG	$5.02\,(0.08)$	$91.7\,(0.9)$	$0.743\,(0.004)$	$-265\,(0.7)$	$1.259\,(0.003)$	$128$
DRAKES	$5.61\,(0.07)$	$\underline{92.3}\,(0.6)$	$\mathbf{0.887}\,(0.001)$	$-264\,(0.3)$	$1.293\,(0.002)$	$128$

Table 1: Regulatory-DNA design. Unmarked values are our measurements, mean (sample SD) over three sampling seeds.

\dagger

marks the released SEPO checkpoint scored with the HepG2 ATAC output (Appendix C.3 ), and

\ast

results reported by Han et al. (2026) , whose released evaluation samples in

128

steps; GF is gradient flow. Bold is best and underline second best among our measurements, excluding the reward-only ablation;

\ast

rows and Log-Lik are unranked.

Figure 3: Regularization sweep of IDRF (ours) on regulatory-DNA design at

K=8

. Red marks reward-only fine-tuning (

\lambda_{\mathrm{ID}}=0

), and the shaded band the setting of Table 1 .

Method	Steps $\downarrow$	CLIP $\uparrow$	ImageReward $\uparrow$	PickScore $\uparrow$	HPSv2.1 ( $\times 100$ ) $\uparrow$
Reference	$4$	$14.00\pm 0.18$	$-2.17\pm 0.01$	$18.24\pm 0.03$	$14.67\pm 0.12$
Reference, $w=3$	$128$	$23.11\pm 0.13$	$-0.73\pm 0.05$	$21.10\pm 0.04$	$25.05\pm 0.12$
ReDGE, $w=3$	$20$	$\mathbf{37.77}\pm 0.10$	$0.30\pm 0.04$	$21.38\pm 0.03$	$25.98\pm 0.10$
RLOO w/o $\mathcal{L}_{\mathrm{ID}}$	$4$	$26.12\pm 0.10$	$-0.49\pm 0.03$	$20.61\pm 0.02$	$22.57\pm 0.08$
IDRF (ours)	$4$	$27.74\pm 0.06$	$\mathbf{0.88}\pm 0.02$	$\mathbf{22.11}\pm 0.02$	$\mathbf{27.72}\pm 0.04$

Table 2: CLIP-guided MaskGIT generation ,

1{,}000

images per method (

100

per class), mean with

95\%

confidence interval over images. CLIP is optimized and the remaining metrics are held out.

w

is the classifier-free guidance weight, and rows without

w

, including IDRF, use no guidance (

w=0

). The Steps column reports the number of denoising steps. Four-step methods use an additional final unmasking step. ReDGE adds

400

reward-gradient updates per image.

Figure 4: CLIP-guided MaskGIT samples , the first sample of each class for each method. Original ImageNet classes appear above the reference row, and reward prompts above the three guided rows. The reference is conditioned on the class alone. ReDGE optimizes the prompt reward at inference, while RLOO and IDRF (ours) are fine-tuned on it. IDRF uses four sampling steps plus cleanup.

Panel of Figure 5 — Figure 5: Reward-only fine-tuning abandons the review format. Left: fraction of samples emitting EOS during fine-tuning with $32$ sampling steps. Right: generated examples after $10{,}000$ updates, illustrating repetitive, unterminated RLOO outputs and complete IDRF reviews.

Method	Steps $\downarrow$	$\log P(1\star)\uparrow$	GenPPL $\downarrow$	Length	EOS (%) $\uparrow$	Norm. ent. $\uparrow$	GM $\downarrow$
Real one-star	–	$-0.950$	$53.3$	$32.2$	$100.0$	$0.948$	$0.03\,(0.26)$
MDLM reference	$128$	$-5.070$	$91.7$	$27.8$	$100.0$	$0.955$	$2.45\,(0.21)$
RLOO w/o $\mathcal{L}_{\mathrm{ID}}$	$32$	$\mathbf{-0.008}$	$76.9$	$126.8$	$9.4$	$0.719$	$28.81\,(0.45)$
IDRF (ours)	$32$	$-0.044$	$\mathbf{65.4}$	$44.1$	$\mathbf{100.0}$	$0.902$	$\mathbf{10.86}\,(0.50)$
IDRF (ours)	$16$	$-0.031$	$107.3$	$37.4$	$\mathbf{100.0}$	$\mathbf{0.922}$	$11.44\,(0.46)$

Table 3: Sentiment steering on Amazon Fashion reviews.

\log P(1\star)

is optimized and the rest held out. GM (

\pm

s.e.) is against held-out real one-star reviews. Bold is the best fine-tuned model.

Appendix figures & tables8 assets

Supplementary material from the paper’s appendix.

Appendix

Setting	IDRF (ours)		RLOO w/o $\mathcal{L}_{\mathrm{ID}}$
Sampler	ancestral	max-confidence	ancestral
Sampling steps $K$	$8$	$8$	$8$ / $128$
$\lambda_{\mathrm{ID}}$	$0.2$	$0.2$	$0$
Rollouts per iteration	$256$	$256$	$256$ / $256$ ( $8\times 32$ )
Leave-one-out group	$256$	$256$	$256$ / $32$
Optimizer (student, auxiliary)	AdamW, learning rate $10^{-4}$ , weight decay $0$

Appendix

Table 4: Training configuration for regulatory-DNA design.

Method	Published	Our measurement	Evaluated model
SEPO	$99.5\,(0.2)$	$8.70\,(1.49)$	Released checkpoint
SEPO + GF	$99.9\,(0.09)$	$47.40\,(1.82)$	Released checkpoint

Appendix

Table 5: Published and independently measured ATAC-Acc (%)

\uparrow

. Published values follow the cited papers. Our values use HepG2 output

1

and

640

sequences per sampling seed. Parentheses give the reported SD for the papers and sample SD across sampling seeds for our measurements.

Sampler	$\mathcal{L}_{\mathrm{ID}}$ state	Pred-Activity $\uparrow$	ATAC-Acc (%) $\uparrow$	$3$ -mer Corr. $\uparrow$	Log-Lik
Ancestral	on-policy $x_{t_{j}}$	$7.22$	$\mathbf{95.3}$	$\mathbf{0.683}$	$-257$
Ancestral	forward-masked $x_{0}$	$5.89$	$52.5$	$0.273$	$-261$
Max-conf	on-policy $x_{t_{j}}$	$7.31$	$\mathbf{92.3}$	$\mathbf{0.718}$	$-253$
Max-conf	forward-masked $x_{0}$	$6.23$	$42.8$	$0.408$	$-224$

Appendix

Table 6: State on which

\mathcal{L}_{\mathrm{ID}}

is evaluated in IDRF (ours), regulatory-DNA design at

K=8

Method	A	C	G	T	CG+GC +CA+AC	T-free
IDRF, $K=8$	11.70	30.99	42.64	14.67	44.46	0.00
RLOO w/o $\mathcal{L}_{\mathrm{ID}}$ , $K=8$	19.26	47.38	33.36	${<}0.01$	72.66	99.32
RLOO w/o $\mathcal{L}_{\mathrm{ID}}$ , $K=128$	25.39	43.79	30.82	${<}0.01$	79.61	99.74

Appendix

Table 7: Reward-only fine-tuning nearly eliminates thymine. Nucleotide and dinucleotide composition in percent. The pair column is the combined frequency of CG, GC, CA and AC among all dinucleotides, and T-free is the share of sequences without T.

Panel of Figure 6 — Figure 6: Unfiltered DNA outputs , the first two sequences of sampling seed $0$ for each method. Line breaks are for display only.

Setting	Value
Sampling steps $K$	$4$ ancestral steps and one cleanup call
$\lambda_{\mathrm{ID}}$	$0.5$
Optimizer (student, auxiliary)	AdamW, learning rate $3\times 10^{-5}$ , weight decay $0.01$
Rollouts per update	$64$ ( $16$ groups of $4$ )
Auxiliary / student updates per iteration	$1$ / $1$
Iterations	$80$

Appendix

Table 8: Training configuration for IDRF (ours) on CLIP-guided MaskGIT.

ImageNet ID	Class	CLIP prompt
817	Sports car	a blue sports car
864	Tow truck	a red tow truck
546	Electric guitar	a yellow electric guitar
609	Jeep	a black jeep
560	Football helmet	a yellow football helmet
883	Vase	a black vase

Appendix

Table 9: ImageNet classes and CLIP prompts used in the MaskGIT experiments.

Candidate set	Steps $\downarrow$	GM vs. real one-star $\downarrow$	GM vs. real, all ratings $\downarrow$
Real one-star reviews	–	$0.03\pm 0.26$	$1.99\pm 0.29$
Real reviews, all ratings	–	–	$0.06\pm 0.24$
MDLM reference	$128$	$2.45\pm 0.21$	$0.54\pm 0.21$
RLOO w/o $\mathcal{L}_{\mathrm{ID}}$	$32$	$28.81\pm 0.45$	$29.53\pm 0.47$
IDRF (ours)	$32$	$10.86\pm 0.50$	$12.13\pm 0.56$
IDRF (ours)	$16$	$11.44\pm 0.46$	$12.85\pm 0.37$

Appendix

Table 10: Gradient moment against two held-out references (lower is closer),

\pm

one standard error over ten pair estimates. Generated sets are those of Table 3 .

Explore similar work

A2D2: Fine-Tuning Any-Length Discrete Diffusion for Adaptive Decoding
Jun 11, 2026Sophia Tang, Yuchen Zhu, Molei Tao +1Diffusion Model Fine-Tuning Diffusion Language Model Decoding
Department of Computer and Information Science, University of Pennsylvania · School of Mathematics, Georgia Institute of Technology · Department of Bioengineering, University of Pennsylvania
Infinite Mask Diffusion for Few-Step Distillation
May 11, 2026Jaehoon Yoo, Wonjung Kim, Chanhyuk Lee +1Masked Diffusion Models Diffusion Model Distillation
Korea Advanced Institute of Science and Technology (KAIST)
Diffusion Reinforcement Learning via Centered Reward Distillation
Mar 14, 2026Yuanzhi Zhu, Xi Wang, Stéphane Lathuilière +1RL for Diffusion Models Diffusion Model Alignment
LIX, École Polytechnique, CNRS, IPP · Inria at Univ. Grenoble Alpes, CNRS, LJK

Jun 11, 2026J/K move · Enter open · S save

A2D2: Fine-Tuning Any-Length Discrete Diffusion for Adaptive Decoding

Sophia Tang, Yuchen Zhu, Molei Tao, Pranam Chatterjee

Department of Computer and Information Science, University of Pennsylvania · School of Mathematics, Georgia Institute of Technology · Department of Bioengineering, University of Pennsylvania

Diffusion Model Fine-Tuning Diffusion Language Model Decoding RL for Diffusion Models Diffusion Models

Open paper AskPDF arXiv

Discrete diffusion models offer a simple and stable likelihood-based framework for sequence generation, recently extended to any-length settings via token insertion. Principled reward-guided fine-tuning for any-length discrete diffusion, however, remains largely unexplored. We introduce Fine-Tuning Any-Length Discrete Diffusion for Adaptive Decoding (A2D2), a unified framework for reward-guided fine-tuning of any-length discrete diffusion models via joint optimization of the insertion and unmasking policies together with a quality-based inference schedule. We derive the Radon-Nikodym derivative for the joint insertion-unmasking path measures, enabling theoretically guaranteed convergence to the intractable reward-tilted sequence distribution without requiring target samples. Building on this, we establish unmasking and insertion quality as tractable approaches for minimizing decoding error and introduce the Adaptive Joint Decoding (AJD) loss, which provably yields the optimal path measure that generates the reward-tilted distribution. Empirically, A2D2 improves reward optimization while enhancing generation flexibility and accuracy over prior fixed-length fine-tuning and inference-time guidance methods.

Abstract

Explore similar work

A2D2: Fine-Tuning Any-Length Discrete Diffusion for Adaptive Decoding

Infinite Mask Diffusion for Few-Step Distillation

Diffusion Reinforcement Learning via Centered Reward Distillation