UltraLLaDA: Scaling the Context Length to 128K for Diffusion Large Language Models
About
Diffusion LLMs have attracted growing interest, with plenty of recent work emphasizing their great potential in various downstream tasks; yet the long-context behavior of diffusion LLMs remains largely uncharted. We present a case study of post-training techniques for extending the context window of diffusion LLMs (i.e., LLaDA) without retraining from scratch. We show that a simple modification to the standard Rotary Positional Embeddings (RoPE) extension effectively accommodates the probabilistic modeling inherent in the diffusion process, enabling stable scaling to longer context ranges. We further compare masking strategies used during post-training and analyze their impact on optimization stability and long-range recall. Instantiating these insights, we introduce UltraLLaDA, a diffusion LLM with a 128K-token context window that, in our empirical evaluation on long-context tasks, significantly outperforms training-free baselines. Our experimental results highlight the special positional extension as a key lever for scaling diffusion LLMs to extended contexts and offer practical guidance for practitioners seeking 128K-scale context via efficient post-training.
Related benchmarks
| Task | Dataset | Result | Rank | |
|---|---|---|---|---|
| Long-context language modeling | LongBench 16K context length | NrtvQA Score12.1 | 6 | |
| Long-context Reasoning | RULER 8K context length | NIAH Score98.28 | 6 | |
| Long-context Reasoning | RULER 16k context length | NIAH93 | 6 | |
| Long-context Reasoning | RULER 4k context length | NIAH97.31 | 6 | |
| Long-context Reasoning | RULER 32k context length | NIAH92.78 | 3 |