Reproduce “UltraLLaDA: Scaling the Context Length to 128K for Diffusion Large Language Models” on OpenResearch

About this paper

Guangxin He, Shen Nie, Fengqi Zhu, Yuankang Zhao, Tianyi Bai, Ran Yan, Jie Fu, Chongxuan Li and 1 more

Researchers from HKUST, Renmin University of China, University of Chinese Academy of Sciences, and Shanghai AI Lab developed UltraLLaDA, a method that scales diffusion large language models to 128K context length. This approach achieved 100% retrieval accuracy on NIAH-128K and maintained a stable perplexity of 11-12 across all evaluated context lengths, outperforming training-free methods by adapting positional embeddings and using specific masking strategies during post-training.

Key insights

• Diffusion LLMs' inherent bidirectional attention necessitates a specific 'Diffusion-aware NTK' scaling for RoPE, which effectively adapts positional encoding to their unique learning patterns.
• Post-training methods offer substantially superior long-context performance for diffusion LLMs compared to training-free extrapolation, demonstrating the importance of model adaptation to new context ranges.
• Effective handling of document boundaries through masking strategies like Adaptive Attention Masking is critical for preventing cross-document interference when processing packed long-context data with global attention.

Summary by alphaXiv.