LongLLaDA: Unlocking Long Context Capabilities in Diffusion LLMs

TL;DR

This paper systematically investigates the long-context capabilities of diffusion LLMs, revealing their stability and local perception advantages, and introduces LongLLaDA, a novel method for extending their context windows without retraining.

Contribution

It provides the first analysis of long-context performance in diffusion LLMs, explains observed phenomena via RoPE scaling theory, and proposes a training-free extrapolation method called LongLLaDA.

Findings

Diffusion LLMs maintain stable perplexity during context extrapolation.

Diffusion LLMs exhibit local perception, enabling retrieval from recent context segments.

LongLLaDA effectively extends context windows, validated by empirical benchmarks.

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.