🤖 machine learning

Retrofitting Linear Attention into Diffusion Language Models

This paper introduces block-hybrid attention, a method that linearizes attention over previous blocks while retaining exact softmax within the active block, enabling the efficient retrofitting of pretrained diffusion language models to achieve up to 1.7× faster inference with minimal performance degradation.

Original authors: Jinha Kim, Younghun Roh, Jaeyeon Kim

Published 2026-08-10

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Jinha Kim, Younghun Roh, Jaeyeon Kim

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 the world of artificial intelligence as a massive, bustling library where a super-smart robot librarian is trying to write a story. For years, the standard way this librarian worked was "autoregressive": it would write one word, stop, think about everything it just wrote, write the next word, stop, and think again. It was like building a tower brick by brick, waiting for the glue to dry before adding the next one. While reliable, this method was slow, especially for long stories.

Recently, a new style called "Diffusion" arrived. Instead of building brick by brick, the Diffusion librarian starts with a blank page full of question marks and tries to guess the whole story at once, refining the guesses over and over until the words make sense. This is like having a team of artists painting a mural simultaneously; they can work much faster because they aren't waiting for the previous brick to dry. However, even this fast method has a snag. As the mural gets bigger, the artists have to constantly look back at every single word they've already decided on to make sure the new words fit. It's like trying to remember a 10,000-page book while writing page 101, which eventually slows everyone down and fills up the library's memory.

This is the puzzle a team of researchers tackled. They asked: "Can we make this fast Diffusion method even faster by changing how the librarian remembers the past?" Their answer is a clever trick called "Block-Hybrid Attention." They realized that while the librarian needs to remember the current paragraph perfectly (to keep the sentences flowing), they don't need to re-read every single word of the previous chapters in high definition. Instead, they can summarize the old chapters into a tiny, fixed-size "summary sheet." This allows the librarian to write the new parts of the story at lightning speed without getting bogged down by the weight of the entire history.

The Big Idea: A Two-Speed Memory System

The paper introduces a method to retrofit (or upgrade) an existing, powerful Diffusion Language Model (specifically a 16-billion-parameter model called LLaDA 2.1) so it can use this "summary sheet" strategy without needing to be retrained from scratch.

Think of the model's brain as having 20 different layers of thinking. In the original model, every single layer acts like a super-attentive librarian who reads every single previous word to understand the current one. This is accurate but heavy. The researchers' innovation, Block-Hybrid Attention, splits the work into two modes:

The "Now" Mode (Exact Attention): When the model is working on the current block of words it is actively generating, it keeps its "super-attentive" mode. It uses the standard, heavy-duty memory to look at every word in the current paragraph to ensure the sentences make perfect sense.
The "Past" Mode (Linear Attention): When the model needs to remember what happened in the previous blocks (the finished chapters), it switches to "Linear Attention." Instead of re-reading the whole book, it consults a fixed-size summary state. It's like looking at a book's index or a one-page summary rather than re-reading the entire text. This summary stays the same size no matter how long the book gets.

How They Did It: The Two-Stage Upgrade

The researchers didn't just swap the parts and hope for the best; they used a careful two-stage process to ensure the model didn't lose its intelligence during the upgrade.

Stage 1: The Shadow Training. They took the original, smart model (the "Teacher") and froze it. Then, they replaced 6 of the 20 attention layers with their new "Block-Hybrid" version (the "Student"). The Student was trained to mimic the Teacher's output exactly, using a "corrupted" version of the text as input. It was like having a student shadow a master chef, trying to copy the chef's exact movements and taste the exact same flavors, but only for the specific dishes the student was assigned to change. This stage took about 24 hours.
Stage 2: The Fine-Tuning. Once the Student learned to mimic the Teacher, the researchers let the whole model work together again. They used a technique called LoRA (Low-Rank Adaptation) to make tiny, precise adjustments to the model's connections. This was like giving the upgraded model a quick "polish" to fix any small glitches that happened because the new parts were now talking to the old parts. This stage took about 6 hours.

The Results: Faster, Lighter, and Still Smart

The team tested their upgraded model, which they named LLaDA-HYBRID, to see if it was still good at writing and solving problems, and if it was actually faster.

Does it still know how to think?
Surprisingly, yes. Even though they replaced 6 out of 20 layers with this "summary sheet" method, the model's performance remained very close to the original.

On a coding test called HumanEval, the original model scored 75.6%, and the hybrid model scored 72.0%.
On a math test called CMATH, the original got 88.3%, and the hybrid got 86.7%.
Interestingly, on another coding test called MBPP+, the hybrid model actually improved, going from 57.7% to 63.0%.
The paper suggests that while the model is slightly less perfect on the hardest reasoning tasks (like a difficult science quiz called GPQA-Diamond, where it dropped from 38.9% to 30.8%), it largely preserved its ability to write code and solve math problems.

Is it actually faster?
This is where the magic happens. Because the model no longer has to re-read the entire history of the conversation for every new word, it can handle more requests at once.

When testing how many words the model could generate per second, the hybrid model was 1.7 times faster than the original at its peak performance (handling 128 concurrent requests).
The original model's memory usage grew as the story got longer, eventually hitting a wall where it couldn't handle more users. The hybrid model, thanks to its fixed-size "summary sheet," could handle more concurrent requests before running out of memory.

The Takeaway

This paper suggests that we don't always need to build a new, faster AI from scratch. Instead, we can take a powerful, existing AI and give it a "hybrid memory" system. By keeping its sharp focus on the present moment while summarizing the past into a compact form, we can make these models significantly faster and more efficient. The researchers found that this upgrade could be done in about 60 hours of training on public data, offering a promising path to making AI services faster and cheaper without sacrificing too much of their smarts. However, the authors note that this is just the beginning; they only upgraded 6 layers, and future work might find ways to upgrade even more layers or handle even longer stories.

Technical Summary: Retrofitting Linear Attention into Diffusion Language Models

Problem Statement
Diffusion Language Models (dLLMs) offer a compelling alternative to autoregressive models by enabling parallel decoding and flexible inference orders. Recent dLLMs typically employ blockwise semi-autoregressive decoding, where blocks are generated sequentially, but tokens within an active block are denoised in parallel. While this allows for KV caching of completed blocks, the standard softmax attention mechanism still incurs a significant bottleneck: every denoising step within an active block must attend to all previous blocks. Consequently, the attention cost grows with the sequence length, and the KV cache size increases linearly with the number of committed tokens, limiting serving throughput and concurrency under fixed memory budgets.

Methodology: Block-Hybrid Attention
To address this, the authors propose block-hybrid attention, a mechanism designed to linearize attention over previously completed blocks while preserving exact bidirectional attention within the active denoising block.

Architecture Design:
- Within-Block (Active): The model retains exact bidirectional softmax attention for tokens currently being denoised in the active block ( $B_{n+1}$ ). This preserves the necessary bidirectional dependencies for the diffusion process.
- Cross-Block (Committed): For all previously committed blocks ( $B_{\le n}$ ), the model replaces softmax attention with linear attention using a fixed-size recurrent state. Specifically, it employs Hedgehog linear attention, which uses a learned positive feature map $\phi(x) = [\text{softmax}(xW), \text{softmax}(-xW)]$ to approximate the softmax kernel.
- State Compression: The keys and values of all committed blocks are summarized into two fixed-size recurrent states, $S_n$ (value-weighted) and $Z_n$ (normalizer), which are updated only when a block is committed. This reduces the cross-block attention complexity from $O(L)$ to $O(1)$ relative to sequence length for the linearized layers.
- Merging: The outputs of the within-block softmax branch and the cross-block linear branch are merged using a shared normalizer and a learned scalar gate, ensuring a balanced attention distribution.
Retrofitting Strategy (Two-Stage Post-Training):
The authors adapt the LoLCATs recipe to retrofit a pretrained dLLM (LLaDA 2.1) with minimal training cost:
- Stage 1: Attention Transfer: The pretrained backbone is frozen. Only the new block-hybrid parameters (feature map matrices $W$ and gates $\alpha$ ) are trained to match the output of the original softmax layers. This is done using an MSE loss on the attention outputs, with the teacher providing hidden states (teacher forcing) to isolate layer-wise approximation errors.
- Stage 2: End-to-End Adaptation: To correct residual distribution shifts caused by composing multiple hybrid layers, the full model is fine-tuned using LoRA (Low-Rank Adaptation) adapters on the attention and MLP projections. The backbone and hybrid parameters remain frozen during this stage, optimizing only the adapters using the standard masked diffusion objective.

Key Contributions

Novel Mechanism: Introduction of block-hybrid attention, the first approach to apply linear attention distillation to a diffusion language model, effectively decoupling cross-block attention cost from sequence length for the selected linearized layers.
Efficient Retrofitting: A two-stage post-training procedure that converts a pretrained 16B dLLM into a hybrid model without full retraining, requiring only approximately 60 hours of training on public instruction data.
Implementation: A fused Triton kernel implementation within the SGLang framework that manages the fixed-size recurrent state alongside paged KV caches, enabling efficient continuous batching.

Experimental Results
The authors evaluated LLaDA-HYBRID, a version of LLaDA 2.1-mini (16B parameters) where 6 out of 20 attention layers were linearized.

Quality Retention: The hybrid model largely preserved the performance of the original teacher model on coding and reasoning benchmarks:
- HumanEval: 72.0% (vs. 75.6% teacher).
- MBPP+: 63.0% (vs. 57.7% teacher, showing an improvement).
- CMATH: 86.7% (vs. 88.3% teacher).
- Note: The most significant degradation was observed on GPQA-Diamond (-8.1 points), suggesting that highly complex reasoning tasks remain sensitive to approximation errors.
Throughput and Concurrency:
- Speedup: LLADA-HYBRID achieved up to 1.73× higher decoding throughput compared to the softmax teacher under SGLang continuous batching (tested on an NVIDIA H200).
- Concurrency: Due to the fixed-size recurrent state in the linearized layers, the model supports more concurrent requests before exhausting memory. At 128 concurrent requests, throughput increased from 2,310.2 to 3,994.4 tokens/s.

Significance and Claims
The paper claims that pretrained dLLMs can be efficiently linearized as an inexpensive, post-hoc retrofit to significantly improve inference speed and serving capacity without sacrificing substantial generation quality. By maintaining bidirectional softmax only where strictly necessary (within the active block) and compressing historical context into a fixed state for a subset of layers, the approach offers a practical path to scaling dLLM serving efficiency. The authors emphasize that this conversion requires only public instruction data and a modest training budget, making linear-attention efficiency accessible for existing dLLMs.

Limitations
The authors note that their retrofit is conservative (only 6 layers linearized), leaving the remaining layers with sequence-dependent costs. Future work could explore more aggressive layer selection, evaluation on long-form generation, and architectures robust to variable block sizes.

The Big Idea: A Two-Speed Memory System

How They Did It: The Two-Stage Upgrade

The Results: Faster, Lighter, and Still Smart

The Takeaway

Technical Summary: Retrofitting Linear Attention into Diffusion Language Models

More like this