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

What it claims

IDRF fine-tunes a few-step masked diffusion student (initialised from a pretrained MDLM-style reference) toward a terminal reward. It keeps the student close to the reference using an inverse-distillation gap instead of a sequence KL. Two properties matter for adoption:

• No reward gradients needed. The reward enters only through policy gradients.
• No pretraining data or reference rollouts needed. The reference is only queried as a frozen denoiser on the student's own states.

The theory contribution is Proposition 1. With an optimal auxiliary denoiser, the population loss L_NELBO(f_ref, p_θ) − min_ψ L_NELBO(f_ψ, p_θ) upper-bounds KL(p_θ ‖ p_ref) for any sequence length and any student sampler. This closes a gap that the IDLM paper left to future work. The bound for the practical objective is weaker: it bounds KL for a mixture of the student's one-step completion laws along rollouts, not for the student's endpoint law.

How it works

flowchart TD
    S[Student G_θ K-step rollout] --> X0[Final sample x0]
    S --> XT[Intermediate state x_tj]
    X0 --> R[Reward → leave-one-out advantage]
    R --> LR[Clipped PG loss L_R]
    XT --> SOFT[Soft prediction G_θ x_tj]
    SOFT --> REF[Frozen reference NELBO]
    SOFT --> AUX[Auxiliary f_ψ NELBO]
    REF --> GAP[Gap L_ID]
    AUX --> GAP
    LR --> TOT[L_R + λ·L_ID → update θ]
    GAP --> TOT
    SOFT -. fit on student states .-> AUX
    TOT --> S

• Initialisation. The student and the auxiliary both start as copies of the reference, so the gap is zero at initialisation. The auxiliary is discarded after training.
• Reward term. The K-step sampler is treated as a finite-horizon MDP with reward only at the last step. The loss is a PPO-style clipped surrogate. Ratios are taken only over tokens revealed at each step. The advantage is the reward minus the mean reward of the other M−1 rollouts in the group. Nothing differentiates through discrete draws or through R. For ancestral sampling the gradient is unbiased at θ = θ_old. For max-confidence decoding the recorded reveal sets are held fixed.
• Regulariser. At a random step j of the student's own rollout, the student's soft prediction is freshly corrupted. Both the reference and the auxiliary score it with a soft-target NELBO.
• Gradient flow. The update is semi-gradient: gradients flow only through the student's soft probabilities.
• Training schedule. Training alternates between auxiliary updates (student frozen) and student updates (auxiliary frozen).
• Extra cost per step. The regulariser costs one student forward pass plus one forward pass each through the reference and the auxiliary.

Evidence

• Step counts. In Figure 1 the students sample in 8 steps (DNA), 4 steps (images) and 32 steps (text). Reward-only fine-tuning at the same step budgets hacks the reward in each domain:
• DNA loses predicted accessibility.
• Images blur.
• Reviews degenerate into repetition and rarely emit EOS.
• DNA setup. The task follows DRAKES: 200-bp enhancers optimised for HepG2 activity. IDRF uses λ_ID = 0.2 and K = 8. Each model is evaluated on 640 sequences × 3 seeds against a held-out oracle.
• DNA baselines. All baselines use 128 steps:

The fetched text cuts off before the IDRF row, so IDRF's own DNA numbers are not available here.

• Ablation: which states to regularise. Regularising on forward-masked completed student samples, instead of on-policy rollout states, drops ATAC-Acc from 95.3% to 52.5% (Appendix C.4). This is the design choice the method turns on.

How to use it

The page gives the training loop in Algorithm 1. It lists no code release, weights, licence, or hardware/compute figures.

Inputs you need:

• a frozen MDLM/SUBS-style reference denoiser without time conditioning
• a scalar terminal reward
• λ
• K
• group size M

The baselines were run through the released DRAKES implementation, and SEPO was evaluated from its released checkpoints.

Limits

• Optimality assumption. The KL guarantee assumes an optimal auxiliary. The practical objective only bounds a mixture-of-completions KL.
• Equality condition. Equality in the bound requires the reference to be the Bayes denoiser of its own endpoint law.
• Truncated text. The fetched text stops partway through Table 1. Image and text results, all IDRF numbers, and any stated limitations section are not visible.
• Mixed-provenance DNA rows. D2-DPO, VRPO, DDPP-IS, diffu-GRPO and SDPO figures are copied from Han et al. (2026), whose evaluation samples in 128 steps. They are not re-measured.
• SEPO metric difference. SEPO's ATAC-Acc was scored with the HepG2 ATAC output and is compared against published values in Appendix C.3.
• Preprint status. This is v1 on arXiv, dated 2 Oct 2026.