Search papers, labs, and topics across Lattice.
We introduce Inverse-Distilled Reward Fine-tuning (IDRF), a framework for efficiently fine-tuning masked discrete diffusion models with rewards. IDRF replaces intractable sequence-level KL penalties with an inverse-distillation regularization term, which is proven to upper-bound the KL divergence using an optimal auxiliary denoiser. This approach enables few-step generation with significantly reduced computational cost and improved reward optimization without reference-model rollouts.
Reward hacking is mitigated and sample quality preserved in few-step discrete diffusion models by replacing intractable KL penalties with a novel inverse-distillation regularization.
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.
