UltraLLaDA: Scaling the Context Length to 128K for Diffusion Large Language Models

该文章提出了UltraLLaDA模型，通过改进的位置嵌入方法和后训练策略，将扩散语言模型（Diffusion LLM）的上下文窗口扩展至128K tokens，显著优于无训练基线模型，为扩散LLM的长上下文能力提升提供了实用方案。一、文章主要内容总结研究背景扩散LLM在多下游任务中潜力显著，但长上下文能力未被充分探索，现有模型难以处理超训练长度的输入（如128K tokens）。自回归LLM的长上下文扩展方法（如RoPE外推）直接应用于扩散LLM时，会因双向注意力特性导致性能偏差，且无训练方法存在“局部感知”偏见，无法有效利用长距离信息。核心方法Diffusion-aware NTK位置嵌入：针对扩散LLM的双向注意力特性，调整RoPE的关键维度和缩放因子，将有效相对跨度（(T_{cap})）设为自回归模型的2倍，减缓RoPE旋转速度，适配128K tokens的长上下文。后训练掩码策略：对比三种数据处理策略以减少跨文档干扰，包括自适应注意力掩码（仅允许文档内注意力）、文档结束符（EOD）拼接（用特殊token标记边界）、直接拼接（无特殊处理，作为基线），最终选择自适应掩码和EOD拼接作为有效方案。实验结果在NIAH（128K检索任务）中，UltraLLaDA全程保持100%