LongLLaDA: Unlocking Long Context Capabilities in Diffusion LLMs

TL;DR

The paper provides the first systematic examination of long-context behavior in diffusion LLMs, revealing stable perplexity and a local-perception mechanism during context extrapolation explained by RoPE dynamics. It introduces LongLLaDA, a training-free NTK-based RoPE extrapolation method that extends diffusion LLM context up to 24k tokens while preserving scaling laws. Across benchmarks, diffusion LLMs broadly match autoregressive models on retrieval, underperform on aggregation tasks, and excel in synthetic QA, underscoring task-dependent strengths. The work thus establishes foundational theory and practical extrapolation techniques for long-context diffusion LLMs, with code available for replication and further research.

Abstract

Large Language Diffusion Models, or diffusion LLMs, have emerged as a significant focus in NLP research, with substantial effort directed toward understanding their scalability and downstream task performance. However, their long-context capabilities remain unexplored, lacking systematic analysis or methods for context extension. In this work, we present the first systematic investigation comparing the long-context performance of diffusion LLMs and traditional auto-regressive LLMs. We first identify a unique characteristic of diffusion LLMs, unlike auto-regressive LLMs, they maintain remarkably stable perplexity during direct context extrapolation. Moreover, where auto-regressive models fail outright during the Needle-In-A-Haystack task with context exceeding their pretrained length, we discover diffusion LLMs exhibit a distinct local perception phenomenon, enabling successful retrieval from recent context segments. We explain both phenomena through the lens of Rotary Position Embedding (RoPE) scaling theory. Building on these observations, we propose LongLLaDA, a training-free method that integrates LLaDA with the NTK-based RoPE extrapolation. Our results validate that established extrapolation scaling laws remain effective for extending the context windows of diffusion LLMs. Furthermore, we identify long-context tasks where diffusion LLMs outperform auto-regressive LLMs and others where they fall short. Consequently, this study establishes the first length extrapolation method for diffusion LLMs while providing essential theoretical insights and empirical benchmarks critical for advancing future research on long-context diffusion LLMs. The code is available at https://github.com/OpenMOSS/LongLLaDA.

Figures (11)

• Figure 1: Comparison of perplexity and retrieval accuracy between the diffusion LLM, LLaDA-8B, and the auto-regressive LLM, LLaMA3-8B, both within and beyond pre-training context length.
• Figure 2: Results of Needle-In-A-Haystack tests needle_in_a_haystack on LLaDA-8B Series nie2025large and LLaMA3-8B Series meta2024llama under direct extrapolation.
• Figure 3: NIAH Results of LLaDA-8B-Base nie2025large with different sampling steps $s$.
• Figure 4: Comparison of trained position embedding interval between LLaDA-8B and LlaMA3-8B. The area within the dashed line represents trained relative position, while that beyond represents the relative position in length extrapolation, with unlearned position embedding values colored in gray.
• Figure 5: Visualization of the QK states from the final layer of LLaMA3-8B-Base meta2024introducing and LLaDA-8B-Base nie2025large for sample from the GovReport subsets in LongBench bai2023longbench. The visualization uses a 2D t-SNE projection van2008visualizing, with each token represented as a point in the image and the position index shown via color changing.