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

1 INTRODUCTION

• Masked discrete diffusion models generate sequences by unmasking tokens in parallel, but optimizing reward alone can damage properties the reward does not measure.The paper names DNA activity, image attributes, and review sentiment as example task objectives.Source: §1 · p. 1
• Sequence-level KL regularization is intractable for masked diffusion, while reverse-transition KL methods require reference evaluations along rollouts and tie regularization to the reference sampling path.Existing methods also use likelihood surrogates and evidence bounds.Source: §1 · p. 2
• IDRF combines trajectory-based policy gradients with inverse distillation, using the student’s own trajectories for reward and regularization without reference likelihoods or rollouts.The regularizer compares denoising losses from a frozen reference and a fitted auxiliary, shaping the student’s few-step sampler directly.Source: §1 · p. 2
• With an optimal auxiliary denoiser, the population inverse-distillation loss upper-bounds sequence-level KL; across DNA, images, and text, IDRF reaches high reward with up to 32× fewer denoising steps than the reference.The framework is reported to preserve properties lost under reward-only fine-tuning.Source: §1 · p. 2

2.1 MASKED DIFFUSION

• Masked diffusion trains a denoiser to recover hidden tokens from visible context and generates by filling a fully masked sequence.The MDLM formulation ties this prediction task to a variational bound.Source: §2.1 · p. 2
• Forward masking independently retains each clean-sequence token with probability α_t and masks it otherwise; the schedule decreases continuously from 1 to 0.The denoiser predicts missing tokens from visible context; the reverse process reveals masked positions using those predictions and keeps revealed tokens fixed.Source: §2.1 · pp. 2–3

2.2 INVERSE DISTILLATION

• Inverse distillation fixes a pretrained reference denoiser and fits a student distribution on which the reference remains predictive, comparing the reference loss with the best denoiser fitted to student samples.The comparison accounts for sample predictability: a student repeating one sequence could have low reference loss while covering little of the data distribution.Source: §2.2 · p. 3
• With a Bayes-optimal reference and an exact inner minimum over a sufficiently expressive denoiser family, prior work proves an inverse-distillation bound for joint reverse processes but leaves the practical sequence objective open.This paper’s Proposition 1 proves a bound for its population loss and an arbitrary reference.Source: §2.2 · p. 3
• In practice, a trainable auxiliary approximates the inner minimum by reconstructing masked student samples, while the reference stays fixed and the student reduces the denoising-loss gap.IDRF uses this gap as its reward-fine-tuning regularizer.Source: §2.2 · p. 3

3.1 KL-REGULARIZED REWARD FINE-TUNING VIA INVERSE DISTILLATION

• IDRF starts from a KL-regularized reward objective and substitutes inverse-distillation regularization for the generally intractable endpoint sequence log-density.The endpoint law marginalizes over all denoising trajectories ending at the completed sample.Source: §3.1 · p. 3, §3 · p. 3
• With an optimal auxiliary, the population inverse-distillation loss upper-bounds sequence-level KL for any sequence length and student sampler, under the proposition’s stated masking, schedule, and finite-reference-NELBO assumptions.The paper’s practical loss is evaluated on the student’s soft predictions at rollout states; its expected term upper-bounds KL for the mixture of one-step completion laws.Source: §3.1 · p. 4

3.2 THE FEW-STEP GENERATOR AS AN MDP

• The K-step student starts from a fully masked sequence, applies its denoiser on a fixed time grid, reveals masked positions with scheduled probabilities, and copies revealed positions.The framework minimizes the sum of the reward policy-gradient loss and weighted inverse-distillation loss, with the latter using student rollout states and soft predictions.Source: §3.2 · p. 4, Figure 2 · p. 4, §3.1 · p. 4
• The generator is a finite-horizon MDP whose reward arrives only at the final step, and reward fine-tuning maximizes expected return over the student’s K-step trajectories.The clipped leave-one-out surrogate uses revealed tokens; for ancestral sampling at the old policy, its negative gradient is an unbiased estimate of the expected-return gradient.Source: §3.2 · p. 5
• On-policy rollout states work best in the reported DNA ablation: forward-masking completed student samples lowers ATAC-Acc from 95.3% to 52.5%.The inverse-distillation regularizer costs one student evaluation and one forward pass through each of the reference and auxiliary models; gradients flow only through student soft-target probabilities.Source: §3.2 · p. 6

Distillation of masked diffusion models

• Masked-diffusion distillation methods reduce sampling steps by matching teacher steps or distributions, whereas IDRF uses inverse distillation to stay close to the teacher while optimizing reward.The cited approaches include SDTT, Duo, DiDi-Instruct, and IDLM.Source: §4 · p. 6

Reward fine-tuning of masked diffusion models

• For masked-diffusion reward fine-tuning, prior approaches include denoising-likelihood approximations, evidence bounds, tractable stepwise MDP policies, recorded-trajectory updates, reward-tilted distributions, and differentiable relaxations.The passage notes DRAKES uses differentiable relaxations for biological sequence design when reward gradients are available.Source: §4 · p. 6

Joint distillation and fine-tuning

• Related joint distillation and fine-tuning methods use integral KL, distribution matching, reward-tilted matching, or soft embeddings; IDRF uses inverse distillation with trajectory-based policy gradients.The passage contrasts Soft-Di[M]O’s use of real images and differentiable rewards with IDRF’s lack of requirements for pretraining data or reward gradients.Source: §4 · p. 6

5.1 REGULATORY-DNA DESIGN

• At K = 8, IDRF reaches 95.3% ATAC-Acc, while its max-confidence variant attains activity 7.31; every baseline uses 128 steps.SEPO + GF has higher activity (7.66) but only 47.4% ATAC-Acc.Source: §5.1 · p. 7, Table 1 · p. 7
• Without inverse distillation, DNA samples collapse into CA/CG-rich sequences, with 37.8% ATAC-Acc and 0.38 3-mer correlation versus 95.3% and 0.68.With the full 128-step sampler, ATAC-Acc falls to 0.1%.Source: §5.1 · p. 7
• As λID increases, DNA activity decreases while ATAC-Acc rises from 37.8% to 95.3% and 3-mer correlation from 0.38 to 0.68 by λID = 0.2.Larger regularization weights trade activity without further gains.Source: Figure 3 · p. 7, §5.1 · p. 8

5.2 CLIP-GUIDED IMAGE GENERATION

• At four steps without guidance, IDRF exceeds the 128-step reference on every held-out image metric, raising ImageReward from −0.73 to 0.88.Without inverse distillation, the student falls below the reference on PickScore and HPSv2.1.Source: §5.2 · p. 8, Table 2 · p. 8

5.3 SENTIMENT STEERING OF REVIEW TEXT

• In review steering, reward-only RLOO leaves 906 of 1,000 samples unterminated, despite lower GenPPL than the reference (76.9 versus 91.7).Its normalized entropy is 0.72, compared with 0.95 for real reviews.Source: §5.3 · p. 9, Figure 5 · p. 9
• At 16 and 32 steps, all evaluated IDRF samples emit EOS; at 32 steps, its GM distance is 2.6× lower than RLOO's (10.9 versus 28.8).IDRF's reward is slightly lower than RLOO's (−0.044 versus −0.008).Source: §5.3 · p. 9, Table 3 · p. 9

6 DISCUSSION AND LIMITATIONS

• Across DNA, image, and text generation, IDRF achieves high reward with fewer sampling steps while mitigating reward hacking.The discussion identifies scaling to larger models and extending to additional domains as future work.Source: §6 · p. 9
• IDRF’s student–auxiliary adversarial min–max game can cause optimization instability and sensitivity to update balance.The auxiliary also adds training cost.Source: §6 · p. 9
• The practical surrogate has a related KL guarantee for a mixture of one-step completion distributions, but control of the final sampling distribution remains open.Source: §6 · p. 9

• AI tools assisted with language editing, code implementation, computational checks, and mathematical analysis; the authors are responsible for the claims, results, and final manuscript.Source: AI USE STATEMENT · p. 10

A SEQUENCE-LEVEL INVERSE-DISTILLATION BOUND

• The student enters the sequence-level proof only through its completed-sequence distribution, so its generation procedure and trajectories are marginalized.Source: A SEQUENCE-LEVEL INVERSE-DISTILLATION… · p. 14
• Random reveal clocks yield a uniform reveal order, and the expected denoising log-loss gives an any-order autoregressive form of the NELBO.The schedule drops out because time serves only as an integration variable.Source: A SEQUENCE-LEVEL INVERSE-DISTILLATION… · p. 14
• A finite reference NELBO ensures a finite gap and pθ ≪ pref; when the reference is the Bayes denoiser of pref, equality holds.Source: A SEQUENCE-LEVEL INVERSE-DISTILLATION… · p. 14

A.1 THE PRACTICAL OBJECTIVE

• The practical objective samples a rollout depth uniformly, uses the student’s state there, and defines soft-target NELBO corruption by masking positions and sampling visible tokens from student predictions.Source: §A.1 · p. 15
• When the auxiliary minimizes its objective, the rollout-state loss equals the population loss and satisfies the practical inverse-distillation bound.The bound assumes the auxiliary objective and regularizer use the same rollout-state law and corruption.Source: §A.1 · p. 15

B SAMPLING AND POLICY-GRADIENT DETAILS

• At a positive noise floor, remaining masks require a final categorical draw whose choices and log-probabilities enter the trajectory and reward loss, adding sampling cost.Source: B SAMPLING AND POLICY-GRADIENT DETAILS · p. 15

C.1 EVALUATION

• The DNA 16× reverse-step ratio compares 128 steps for every baseline with 8 for IDRF.Source: §C.1 · p. 16

C.2 TRAINING CONFIGURATION

• Training uses four auxiliary updates on on-policy states followed by two clipped student updates on a rollout batch, and reports the final checkpoint.The ancestral IDRF student is trained on one NVIDIA RTX 6000 GPU.Source: §C.2 · p. 16

C.3 BASELINE PROVENANCE AND EVALUATION DISCREPANCIES

• Table 1 uses common-protocol measurements for unmarked rows, whereas Table 5 retains published ATAC-Acc values; these are not measurements from the same run.The evaluated baselines include inference-time guidance and sampling methods as well as reward-fine-tuned models.Source: §C.3 · pp. 16–17, Table 5 · p. 17
• SEPO’s released plotting notebook selects GM12878 output 0 rather than HepG2 output 1, and the authors identify this as an evaluation mismatch.They do not establish the execution history of every published result or explain differences in Activity, 3-mer correlation, or Log-Lik.Source: §C.3 · p. 17

C.4 ON-POLICY STATES

• Under both samplers, forward-masking completed samples instead of using rollout states sharply lowers ATAC-Acc and 3-mer correlation.The table compares states used to evaluate the inverse-distillation loss at K = 8.Source: §C.4 · p. 17, Table 6 · p. 18

C.5 COMPOSITIONAL DEGENERATION UNDER REWARD-ONLY FINE-TUNING

• 13 and 5 T bases remained under reward-only fine-tuning at K = 8 and 128, respectively, versus 56,344 for IDRF across 1,920 sequences.CA/CG-related pairs rose from 44.5% under IDRF to 72.7% and 79.6%; 99.3% and 99.7% of reward-only sequences were T-free.Source: §C.5 · p. 18, Table 7 · p. 18

D CLIP-GUIDED MASKGIT

• Image reward uses 100× CLIPScore, while ImageReward, PickScore v1, and HPSv2.1 are not optimized.Evaluation generates 100 images for each of ten classes and averages scores within and across classes.Source: D CLIP-GUIDED MASKGIT · p. 18
• The image student samples with four ancestral steps plus one cleanup call from an ImageNet-pretrained MaskGIT model, with the reference and VQ decoder frozen.The model receives the ImageNet class; the attribute prompt enters only through the reward.Source: D CLIP-GUIDED MASKGIT · p. 18
• Figure 6 presents the first two unfiltered DNA sequences from sampling seed 0 for each method.Line breaks are for display only.Source: Figure 6 · p. 19

E SENTIMENT STEERING OF REVIEW TEXT

• Review experiments use deduplicated Amazon Fashion reviews capped at 126 GPT-2 tokens, with the frozen classifier’s one-star log-probability as reward.The models start from an EOS-aware SFT checkpoint trained on the review data.Source: E SENTIMENT STEERING OF REVIEW TEXT · p. 20
• Repeated words account for a median of 66% of word occurrences in RLOO reviews, 14% in IDRF reviews, and 9% in real one-star reviews.The evaluation measures repetition as the fraction of word occurrences beyond their first occurrence, summarized by the median across reviews.Source: E SENTIMENT STEERING OF REVIEW TEXT · p. 20

E.1 GRADIENT-MOMENT DIAGNOSTIC

• The diagnostic averages ten estimates of the inner product between independent candidate-minus-reference gradient differences, using two disjoint 50-review subsets per set.The estimate is unbiased for the squared distance between mean gradients, so identical distributions score zero in expectation.Source: §E.1 · p. 20
• 2.4–2.7× lower gradient moment than reward-only fine-tuning is reported for IDRF at 32 steps against both held-out references.The table defines lower gradient moment as closer; real-versus-real estimates of 0.03 and 0.06 are consistent with zero.Source: §E.1 · p. 20, Table 10 · p. 21
• The two held-out references contain 1,000 eligible real one-star reviews and 1,000 real reviews across ratings, respectively.Each real-versus-real comparison uses a disjoint sample from the same distribution.Source: §E.1 · p. 20