🤖 machine learning

UNIFUSION: Adapting Autoregressive Language Models into Discrete Diffusion under a Unified Reverse-Rate Objective

The paper proposes UNIFUSION, a unified framework that establishes connections between various discrete diffusion objectives to enable the direct adaptation of pretrained autoregressive models to uniform-noise diffusion, achieving state-of-the-art performance in both generative quality and downstream task accuracy.

Original authors: Xiaoyi Jiang, Jingyuan Li, Yixuan Jiang, Wei Liu, Yi Zhu, Zuoqiang Shi, Pipi Hu

Published 2026-07-28

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Xiaoyi Jiang, Jingyuan Li, Yixuan Jiang, Wei Liu, Yi Zhu, Zuoqiang Shi, Pipi Hu

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a robot to write stories. For a long time, the best way to do this was to teach the robot to write one word at a time, from left to right, like a human typing a sentence. This method, called "autoregressive" learning, is great because it's fast and the robot learns a lot of language rules. But it has a big flaw: once the robot writes a word, it can never go back and change it. If it makes a mistake in the first sentence, the whole story might get ruined because it can't edit its past.

Recently, scientists tried a different approach called "discrete diffusion." Instead of writing word-by-word, this method starts with a complete sentence that has been scrambled or "corrupted" with noise, and the robot learns to clean it up, fixing one word at a time until the sentence is perfect. This is like taking a messy room and slowly organizing it. The problem is that most of these new "clean-up" robots are trained from scratch, wasting all the language knowledge the old "word-by-word" robots already had. Furthermore, many of these new robots use a specific type of noise where they hide words behind a "mask" (like a black box). This means the robot can only fix the hidden parts, not the words that are already visible. The big question was: Can we take a smart, pre-trained "word-by-word" robot and instantly turn it into a powerful "clean-up" robot that can fix any word in a sentence, even the ones that aren't hidden?

This paper, titled "Unifusion," says yes, and it shows how to do it. The researchers discovered that despite the different math formulas used by various "clean-up" robots, they are all secretly trying to solve the same underlying problem: predicting how fast a word should jump from one state to another. They call this the "reverse rate." By realizing that all these different methods are just speaking different dialects of the same language, they created a universal translator. This translator allows them to take the knowledge from a standard "word-by-word" model (like GPT-2) and directly adapt it to a "uniform-noise" diffusion model.

The key innovation is that this new model, Unifusion, doesn't need to hide words behind masks. Instead, it treats every single word in the sentence as editable, just like you can edit any word in a Google Doc. The researchers tested this by taking two versions of a pre-trained model (one with 124 million parameters and one with 355 million) and training them to be "clean-up" artists. They found that as they gave the model more time to think (increasing the number of steps from 16 to 256), the quality of the text improved steadily.

At the maximum of 256 steps, the smaller model (Unifusion-S) achieved a "generative perplexity" (a measure of how confused the model is by the text it creates) of 97.783, while the larger model (Unifusion-M) reached 71.516. More importantly, these models didn't just write fluent text; they also maintained high "unigram entropy," which is a fancy way of saying the text was diverse and didn't get stuck repeating the same phrases or patterns. In fact, at 256 steps, no other model tested at the same size was able to beat Unifusion on both fluency and diversity at the same time.

The paper also showed that this direct conversion works better than the old way of trying to go from "word-by-word" to "masked diffusion" and then to "uniform diffusion." By skipping the middleman, Unifusion reached the same high quality much more efficiently. When tested on real-world logic puzzles and common sense questions (like WinoGrande and SIQA), these new diffusion models outperformed all other diffusion models of their size, proving that they didn't just learn to clean up noise, but actually retained the smart reasoning skills of their original "word-by-word" training.

In short, Unifusion is like taking a master chef who is used to cooking a meal one ingredient at a time and teaching them a new technique where they can taste and adjust the entire pot of soup at once, without having to throw away their years of culinary experience. The result is a model that can generate text that is both fluent and creatively diverse, simply by learning to edit the whole sentence at once.

Technical Summary: Unifusion

Problem Statement

Existing methods for adapting pretrained autoregressive (AR) language models to discrete diffusion typically rely on absorbing mask corruption. While this choice connects naturally to AR training (as AR decoding is a specific token-reveal order), it imposes a significant limitation: once a token is revealed, it cannot be edited during the sampling process. This often leads to generation issues such as repetition, template collapse, or the generation of nonsensical fragments when the sampling budget is increased.

Conversely, uniform-noise diffusion replaces tokens with vocabulary items rather than a special mask, allowing all positions to remain editable throughout sampling. This offers the potential for higher lexical diversity and better control. However, directly adapting AR checkpoints to uniform diffusion has been challenging because existing discrete diffusion models (SEDD, MDLM/GIDD, M2S, Neural CTMC) utilize different prediction parameterizations (e.g., concrete scores, clean tokens, posterior means, exit rates) and distinct loss functions. These differences obscure whether the methods learn fundamentally different reverse processes or merely describe the same process with different outputs, and they prevent the direct reuse of AR outputs for uniform kernels.

Methodology

1. Unifying Discrete Diffusion in Reverse-Rate Space

The authors establish a theoretical connection among four major discrete diffusion frameworks: SEDD, MDLM/GIDD, M2S, and Neural CTMC.

Unified Objective: They demonstrate that the conditional losses of all four methods can be expressed as a single generalized Kullback–Leibler (KL) objective over model reverse rates. By mapping each method's specific output (scores, tokens, means, or rates) to the induced off-diagonal reverse rates, the authors show that the underlying optimization target is identical.
Conversion Theorems: The paper derives explicit mathematical conversions to translate a clean-token prediction ( $x_0$ ) into the specific parameterizations required by SEDD, M2S, and Neural CTMC. This creates a shared interface where a single $x_0$ prediction can drive any of these parameterizations.
Kernel Agnosticism: Crucially, the clean-token distribution $x_\theta(x_0 | x_t, t)$ answers the same question regardless of the corruption kernel: "Which vocabulary token occupied this position before corruption?" Once predicted, the known forward kernel (whether absorbing or uniform) converts this prediction into the necessary reverse rates. This bridges the gap between AR initialization and uniform diffusion.

2. The Unifusion Approach

Building on these theoretical connections, the authors propose Unifusion, a simple continual pre-training strategy to adapt pretrained GPT-2 checkpoints directly to uniform-noise diffusion without an intermediate masked diffusion stage.

Initialization: The method leverages the inherent structure of AR models. An AR checkpoint's next-token distribution for position $i$ given $x_{<i}$ is shifted by one position to serve as the initial prediction for the clean token $x_0$ at position $i$ .
Training Dynamics:
- Attention Annealing: The model transitions from causal attention (used in AR pretraining) to fully bidirectional attention over visible corrupted tokens. This allows the model to gradually open the full context during optimization.
- Time Conditioning: Since AR checkpoints lack time conditioning, the model learns time-dependent parameters during fine-tuning while retaining the AR initialization for other weights.
- Corruption Schedule: The model is trained with a linear uniform corruption schedule ( $\alpha_t = 1-t, \beta_t = t$ ) and a uniform prior.

Key Contributions

Theoretical Unification: The paper unifies the conditional losses of SEDD, MDLM/GIDD, M2S, and Neural CTMC in reverse-rate space, proving they are equivalent under the same forward CTMC and noise schedule. It provides explicit conversion formulas between clean-token predictions and all four parameterizations.
Direct AR-to-Uniform Adaptation: By using clean-token $x_0$ predictions as a bridge, the authors enable the direct adaptation of pretrained AR models to uniform diffusion kernels, bypassing the need for intermediate masked-diffusion checkpoints.
Empirical Validation: The study separates rate-space equivalence from finite-budget optimization, showing that while theoretical equivalence exists, native training trajectories differ. Unifusion demonstrates that direct AR-to-uniform training achieves performance comparable to two-stage masked paths but with superior trade-offs in generation quality.

Results

The authors evaluated Unifusion on 124M (Unifusion-S) and 355M (Unifusion-M) parameter models derived from GPT-2.

Generative Perplexity (GenPPL) vs. Entropy: Unifusion steadily improves the trade-off between GenPPL and unigram entropy as the sampling budget increases from 16 to 256 steps.
- At 256 steps, Unifusion-S achieved a GenPPL of 97.783 with an entropy of 5.2626.
- At 256 steps, Unifusion-M achieved a GenPPL of 71.516 with an entropy of 5.6669.
- No evaluated model at the same scale simultaneously outperformed Unifusion on both metrics. For instance, while some baselines achieved lower GenPPL, they suffered from significantly lower entropy (indicating repetition or collapse), whereas Unifusion maintained high lexical diversity.
Downstream Accuracy: Unifusion models achieved the highest zero-shot accuracy among compared diffusion models on WinoGrande, SIQA, and BBH at both scales, outperforming their GPT-2 initializations and other diffusion baselines like DiffuGPT and SEDD.
Sampling Efficiency: The 64-step Unifusion-M point already Pareto-dominated the 256-step DiffGPT-M point, suggesting high efficiency in the sampling budget.

Significance and Claims

The paper claims that Unifusion successfully addresses the "open question" of how to reach a uniform-noise diffusion model directly from an AR checkpoint. By unifying the theoretical landscape of discrete diffusion and leveraging the clean-token interface, the authors demonstrate that:

Uniform diffusion can inherit AR initialization effectively without an intermediate masked stage.
Direct adaptation yields a superior balance between fluency (low GenPPL) and diversity (high entropy) compared to existing methods.
The theoretical equivalence of different discrete diffusion losses allows for flexible implementation choices while maintaining a consistent optimization goal in reverse-rate space.

The work suggests that the choice of corruption kernel (mask vs. uniform) is a critical design decision that impacts generation quality, and that uniform kernels, when properly adapted from AR models, offer a more robust generation process that avoids the repetition issues common in masked diffusion.