A2D2: Fine-Tuning Any-Length Discrete Diffusion for Adaptive Decoding
This work addresses the problem of fine-tuning discrete diffusion models for reward optimization in sequence generation, offering a principled approach for any-length settings.
A2D2 introduces a unified framework for reward-guided fine-tuning of any-length discrete diffusion models, enabling adaptive decoding via joint optimization of insertion and unmasking policies. It achieves improved reward optimization and generation flexibility over prior methods.
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.