The $x_0$-parameterization is a specific design choice in the Diffusion-LM training objective designed to bridge the gap between continuous diffusion processes and discrete text, primarily by reducing rounding errors.

The Problem: Rounding Errors in Standard Parameterization

In standard continuous Diffusion Models, the simplified training objective ($L_{simple}$) trains a neural network to predict the mean ($\mu_\theta$) of the previous denoising step. Under this standard parameterization, the constraint that the final continuous vector ($x_0$) must perfectly commit to a single discrete word embedding is only heavily emphasized during the very last denoising steps (when $t$ is near 0). Because the model is not forced to think about the final discrete word early in the denoising process, it often struggles to generate an $x_0$ that lies exactly on a valid word embedding, leading to rounding errors when converting the vectors back to text.

The Solution: Predicting $x_0$ Directly

To force the model to explicitly model the final text structure at every stage of the diffusion process, the researchers re-parameterized the objective. Instead of predicting the intermediate transition mean $\mu_\theta$ or the noise term $\epsilon$ (which is standard in image diffusion), the neural network $f_\theta(x_t, t)$ is trained to directly predict the final, fully denoised state $x_0$ at every single diffusion step $t$.

The objective function is modified to minimize the difference between the true $x_0$ and the model’s prediction: $||x_0 - f_\theta(x_t, t)||^2$. By making $x_0$ the target in every term of the objective, the model quickly learns that its predictions must be precisely centered on valid word embeddings.

Mathematical Equivalence Predicting $x_0$ is mathematically valid because it is equivalent to predicting the mean of the previous step ($\mu_\theta$) up to a scaling constant. Because the forward diffusion process is composed of Gaussian transitions, the tractable Gaussian posterior $q(x_{t-1} | x_0, x_t)$ can be used to calculate the exact mean of $x_{t-1}$ in closed form as long as you have an estimate for $x_0$ and the current state $x_t$. Therefore, at generation time, the model first estimates $x_0$ via $f_\theta(x_t, t)$, and then uses that estimate to sample the next intermediate state $x_{t-1}$.

Key Benefits of $x_0$-parameterization

• Enables the “Clamping Trick”: Because the model explicitly predicts a continuous estimate of the final word vectors ($x_0$) at every intermediate step, this intermediate prediction can be “clamped” (mapped to its nearest actual discrete word embedding) before being used to sample $x_{t-1}$. This forces the latent vectors to commit to specific words early on, making predictions more precise and further reducing rounding errors.
• Robustness to Higher Dimensions: Ablation studies show that while parameterizing by the noise term $\epsilon$ works for small embedding dimensions, it quickly collapses as the embedding dimension increases. The $x_0$-parameterization consistently maintains good performance across different dimensionalities.
• Reduced Sensitivity to Noise Schedules: While the authors designed a custom “sqrt” noise schedule to help text diffusion, they found that once the $x_0$-parameterization was applied, the model became robust enough that the advantage of the custom noise schedule was no longer highly salient.