Skip to content
Preprint

Retrofitting Linear Attention into Diffusion Language Models

Aug 2026 · 0 citations · 23 references
Computer Science

TL;DR

This work introduces block-hybrid attention, which retains exact softmax attention within the active denoising block while applying linear attention over previous blocks, and shows that pretrained dLLMs can be efficiently linearized for faster inference.

Abstract

Diffusion language models (dLLMs) offer a promising alternative to autoregressive models by accelerating inference through parallel decoding. Recent dLLMs commonly use blockwise semi-autoregressive decoding, generating blocks autoregressively while denoising tokens within each active block in parallel. However, despite KV caching, each denoising step still attends to all previous blocks, repeatedly incurring prefix-attention cost. Motivated by this bottleneck, we ask whether dLLM inference can be further accelerated by linearizing attention over previous blocks. We introduce block-hybrid attention, which retains exact softmax attention within the active denoising block while applying linear attention over previous blocks. We show that this hybrid attention can be retrofitted into a pretrained dLLM with minimal post-training: LLaDA-Hybrid replaces 6 of the 20 attention layers in LLaDA~2.1, a 16B open-source dLLM, largely following LoLCAT (Zhang et al, 2024). The conversion takes only approximately 60 hours while preserving benchmark performance: 72.0% vs. 75.6% on HumanEval, 63.0% vs. 57.7% on MBPP+, and 86.7% vs. 88.3% on CMATH. With a Triton implementation, LLaDA-Hybrid achieves up to $1.7\times$ higher decoding throughput and supports more concurrent requests before exhausting memory, showing that pretrained dLLMs can be efficiently linearized for faster inference. Our code is available at: https://github.com/Diuven/LLaDA-Hybrid.

View source

Similar papers

Preprint Jul 2026

PreDiff-LM: Pretrained Discrete Masked Diffusion Language Modeling with Hybrid Attention

Discrete masked diffusion language models support bidirectional generation and infilling, but adapting pretrained autoregressive (AR) transformers requires reconciling causal pretraining with bidirectional denoising. We study this problem at the level of attention rather than claiming AR-weight reuse itself as novel. PreDiff-LM preserves causal attention within the observed prompt while allowing full bidirectional attention within the masked target. Under a matched GPT-2 Medium, WikiText-103, 90K-step setup, this hybrid mask improves unconditional perplexity from 34.1 to 28.7 and MAUVE from 0.71 to 0.78 over uniform bidirectional attention with the same AR initialization. Attention adaptation also composes with a DiffuGPT-style objective adaptation, reaching 26.9 perplexity. Pretrained initialization reduces the steps required to reach perplexity below 50 from about 350K to 8K, although a compute-matched fine-tuned AR model remains stronger at equal scale (18.9 versus 28.7). Beyond perplexity, PreDiff-LM improves repetition, distributional quality, four zero-shot downstream tasks, and human preference over prior diffusion baselines. The results position hybrid attention as a complementary mechanism for adapting pretrained causal backbones, while making explicit the remaining quality and inference-efficiency gaps to optimized AR models.

Zhengtao Yao, Runhao Li, Xupeng Chen et al. · 0 citations
Preprint Aug 2026

CAI-DLLM: Convergence Aware Inference for Diffusion Language Models

This work proposes CAI-DLLM, a training-free inference method that uses first-step confidence to guide denoising and reduce inference time, and evaluates CAI-DLLM on LLaDA-8B-Instruct and Dream-7B-Instruct across math, code, reasoning, commonsense, and long-context tasks.

Farhana Amin, Sabiha Afroz, D. Nikolopoulos · 0 citations
Preprint Jul 2026

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

This work proposes a simple continual pre-training approach for directly adapting pretrained GPT2 checkpoints to uniform-noise diffusion, and establishes connections among SEDD, MDLM/GIDD, M2S, and Neural CTMC by expressing their conditional losses as a single generalized Kullback--Leibler objective over model reverse rates.

Xiaoyi Jiang, Jingyuan Li, Yixuan Jiang et al. · 0 citations
Preprint Jul 2026

LaCache: Exact Caching and Precision-Adaptive Inference for Diffusion Large Language Models

LaCache is proposed, a training-free acceleration framework that alleviates operator-level redundancy through lossless caching and mixed precision, and inegrates a per-group FP8 quantization strategy for FFN layers, tailored to step-dependent activation distributions across the diffusion process.

Xingru Chen, Zelang Liang, Yongjia Ma et al. · 0 citations
#artificial intelligence Preprint Aug 2026

Trajectory-Level Speculative Decoding for Diffusion Language Models

This work develops a trajectory-level speculative framework that constructs draft denoising trajectories via confidence-stratified tree exploration and verifies them through blockwise parallel evaluation with bidirectional attention masking, and introduces inter-block speculation, exploiting diffusion models'bidirectional structure to perform cross-block lookahead.

Tian-Xiang Pan, Baitao Gong, Mo Guang et al. · 0 citations
Open access Jul 2026

Hierarchical Context Compression for Efficient Long-Sequence Diffusion Language Modeling

HCC directly realizes the hierarchical-compression direction identified as future work by recent segmented-decoding methods, and is validated with controlled and real-text information-retention studies and with direct latency and memory measurements.

Maya Srinivasan, Lucas Bernal · 0 citations