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

🃏 Defining Unmasking and Insertion Quality 🔮

A key challenge in any-length discrete diffusion is that model performance is highly sensitive to the chosen insertion and unmasking trajectory. We formally define the quality of a step taken in the decoding process.

Unmasking Quality

We define the unmasking quality as the probability that the unmasked token is sampled from the unmasking posterior given the context from the rest of the sequence:

We train a parameterized model $\mu_\phi : \mathcal{X} \to [0,1]$ to predict the unmasking quality by minimizing the Unmasking Quality Loss (UQL):

During inference, the unmasking quality determines which tokens are inconsistent and should be re-masked. We prove that maximizing unmasking quality upper-bounds the probability of successful parallel unmasking, directly minimizing the compounding parallelization error (CPE) — the gap between sequential single-token unmasking and the factorized approximation made when unmasking many tokens at once.

Insertion Quality

The insertion quality is defined as the probability that an inserted mask token is likely to be decoded into a true token in the corresponding gap of the target sequence:

We train a parameterized model $\nu_\phi$ by minimizing the Insertion Quality Loss (IQL):

Maximizing insertion quality provides an upper bound for the probability of reconstructing a clean sequence, justifying its use for adaptive removal of low-quality insertions.

---

🃏 Adaptive Joint Decoding Loss 🔮

We define the optimal reward-tilted path measure for any-length MDMs as:

To optimize toward $\mathbb{P}^\star$, we derive the Radon–Nikodym derivative for the joint insertion–unmasking CTMC path measures. This yields our Adaptive Joint Decoding (AJD) loss:

where $W^v := \log \frac{\mathrm{d}\mathbb{P}^\star}{\mathrm{d}\mathbb{P}^v}$ is the log Radon–Nikodym derivative used to reweight trajectories sampled from an arbitrary path measure $\mathbb{P}^v$ toward the target, and the normalizing constant is approximated as $Z \approx \mathbb{E}_{\mathbb{P}^v}[e^{W^v}]$.

The AJD loss provably yields the optimal unmasking and insertion generators that minimize trajectory-induced error while generating the reward-tilted distribution. It jointly optimizes:

• The unmasking policy $f_\theta$ and insertion policy $g_\theta$
• The unmasking quality predictor $\mu_\phi$ and insertion quality predictor $\nu_\phi$

---

🃏 Off-Policy RL & Alternating Optimization 🔮

We optimize the AJD loss with an off-policy RL strategy. Each iteration samples a batch of sequences $\boldsymbol{x}_1$ from the detached path measure $\mathbb{P}^{\bar\theta, \bar\phi}$, computes their log RND $W^{\bar\theta,\bar\phi}$, and stores them in a replay buffer; intermediate states are then drawn from the interpolant to evaluate the AJD loss. Because the policy parameters $\theta$ vastly outnumber the quality-head parameters $\phi$, we alternate between updating $\theta$ (with $\phi$ frozen) and $\phi$ (with $\theta$ frozen). This schedule also acts as implicit loss balancing, so the weighting $\lambda_{\text{quality}}$ never needs tuning.

---

🃏 Adaptive Inference 🔮

Starting from an empty sequence, at each discrete time step $[t_k, t_{k+1}]$ A2D2 performs:

1. Adaptive Unmasking: Sample tokens to unmask via $f_\theta$, predict unmasking quality via $\mu_\phi$, and re-mask low-quality tokens that fall below a threshold (or until the expected number of masked tokens is reached).
2. Adaptive Insertion: Insert masks by sampling counts $I^\ell_t \sim \text{Poisson}(g_\theta(\boldsymbol{x}_t,t)[\ell]\cdot\Delta t)$ per gap, predict insertion quality via $\nu_\phi$, and remove low-quality insertions.

Generation stops when no masks remain and the insertion expectation falls below $0.5$, or when the total number of time steps is reached.