LongLLaDA: Unlocking Long Context Capabilities in Diffusion LLMs

1 · 动机:为什么 dLLM 的长上下文是个"空白"

1.1 历史脉络:dLLM 这一年发生了什么

2024 下半年到 2025 上半年,diffusion LLM (掩码扩散语言模型) 集中爆发: SEDD/MD4 类的理论简化 (Sahoo 2024 / Ou 2024 / Shi 2024)、LLaDA-8B (Nie 2025) 把 mask diffusion 真正 scale 到 8B 参数、Dream-7B (Ye 2025) 从 Qwen2.5 初始化继续 mask diffusion 训练、LLaDA-1.5 (Zhu 2025) 加 VRPO 偏好优化。多模态 (MMaDA / LLaDA-V / LaViDa)、推理 (d1 / DCoLT)、加速 (dKV-Cache / Fast-dLLM) 都在上面跑。

但这一年的研究全是在 ≤ 4k context 内打转。RLHF、推理 trace、agentic 任务很快会把上下文需求顶到 32k+,而 LLaDA-8B 的 pretrained context 只有 4k。研究者会自然地问:

• dLLM 直接喂 16k token 会爆吗?(perplexity 会不会像 AR 一样炸)
• NIAH 这种 retrieval 任务 dLLM 能跑多远?
• AR 时代积累的那一堆 RoPE 缩放 trick (Linear PI, NTK, YaRN, LongRoPE, ReRoPE...) 在 dLLM 上还成立吗?

本论文就是来回答这三个问题的。

1.2 为什么 dLLM 的长上下文行为本来就该不一样

核心区别一句话: dLLM 是 bidirectional attention + mask prediction,AR 是 causal attention + next-token。这两件事联手让 dLLM 在训练时"看到的相对位置"几乎是 AR 的两倍。

这个差异看似平淡,但它直接决定了:

• RoPE 中频维度在 LLaDA 训练时已经覆盖了一个完整的 sin/cos 周期 ([−4095, 4095]),而 LLaMA3 只覆盖了半个周期 ([0, 8191])。
• OOD 区域大小不同。AR 在 [-T, 0] 整片是没见过的;dLLM 把这片"补"上了。所以 dLLM 在外推时,低频维度面对的 OOD 体积更小。

1.3 别的方案为什么"够呛"

1.4 为什么"做 training-free"是合理的野心

外推方案分两档: (a) continued pretraining 在长数据上重训,(b) training-free / inference-time 只调位置编码。

论文坚持走 (b) 的原因有三:

1. dLLM 训练昂贵且早期。LLaDA-8B 已经是 mask diffusion 在 8B 量级的代表作,长上下文继续训练的成本和经验数据都还没有积累。
2. 位置编码是一个干净的 lever。不动权重就不动 mask diffusion 的训练目标,科学上对照清晰。
3. 验证 scaling law 可迁移性。如果 NTK 这套理论是普适的(几何上来自 RoPE 周期性),它应该不依赖训练目标。论文是在做 existence proof。

需要训练才能突破的领域(论文坦言)是 24k 以上 — 后续工作 UltraLLaDA 把 LLaDA 扩到 128k,但代价是要做 long-data fine-tuning。本论文画出的是训练-免费的天花板。

2 · 背景速查

2.1 RoPE 的关键公式 (复习)

低维 n 小 → Tₙ 短 → 在一个 sequence 内反复转完整周期; 高维 n 大 → Tₙ 长 → 整个序列都不一定转完一圈。

对 LLaDA: d=128, β0=500000, Ttrain=4096 → d_extra = 64。即前 64 维在训练时见过完整周期, 后 64 维没见过。

2.2 NTK 缩放公式

对 LLaDA d/d_extra = 128/64 = 2,所以 λ ∝ t² (sup-linear)。代入 t=8k, 16k, 24k, 32k 得 λ ≈ 4, 14, 31, 55。

3 · 现象 ① · dLLM 的 perplexity 在直接外推下不爆

论文 Figure 1(a): 把 (4k, 8k, 12k, 16k) token 的 PG-19 段直接喂进去算 PPL。

3.1 为什么 PPL 不爆?

原因来自 §1.2 那张图: dLLM 的中频维度在训练时就见过完整 sin/cos 周期,所以即便 |t−s| 超过 Ttrain,中频维度的 cos(θₙ(t−s)) 落在已见过的取值区间内 — 没有 distribution shift。

剩下高频维度 (n < d_extra) 跨多个周期,见过完整值;低频维度 (n ≥ d_extra) OOD,但因为 cos 是偶函数(cos(−x) = cos(x)),AR 没见过的负相对位置在 dLLM 里其实可由对称性"猜测"出来 — OOD 体积只有 AR 的一半。

整体效果:整个 q·k·R_{t-s} 的内积分布几乎不漂移 → 输出 logits 分布不漂移 → PPL 稳。

4 · 现象 ② · NIAH 中的 local perception

NIAH 把一个事实("某人最爱的食物是 X")埋在 long context 不同深度,然后问"X 是什么"。score 就是按 (token_limit × depth_percent) 二维网格画热力图。

4.1 解释:为什么是"滑窗"而不是"全军覆没"?

三个观察叠加:

1. 长度外推不爆(§3): query 在 16k 处仍然能产生合理的 attention logits (没有 PPL 爆炸)。所以它能"看见"远方 token 的 representation, 不是 garbage。
2. 但 cos 是偶函数: cos(θ_n · (t−s)) = cos(θ_n · (s−t))。q·k 内积区分不出"前面 4k"和"后面 4k"。
3. monotonic ordering 只到 ±Ttrain: 训练时见过的相对位置只到 ±4095。一旦 |t−s| 超过 4095,某些低频维度的 cos 值进入未见过的 phase, q·k 不可靠。

结合起来 → query 只能可靠地从 |t−s| ≤ 4095 范围内 retrieve = "最近的 4k 窗口"。所以:

• context = 8k, 针埋在 0-50% (远端): 距离 query 4-8k, OOD → 失败
• context = 8k, 针埋在 50-100% (近端): 距离 query 0-4k, 训练见过 → 成功

这就是滑窗的几何解释。

4.2 反向思考:如果是 AR?

AR 在 Ttrain 之外完全崩溃,是因为 PPL 本身爆,attention logits 已经是 garbage。所以即便针就在 query 旁边,AR 也答不出来 (它的 representation 已经坏了)。dLLM 的优势在于:representation 没坏, 只是位置编码区分能力下降。

5 · 机制分析:RoPE 的几何到底发生了什么

5.1 三档维度

关键差异在中频。AR 看到 [0, T] 这半个周期里的 cos/sin → cos 单调递减、sin 单调递增,LLM 学到的是 "线性序号"。换到 [T, 2T] (外推时), cos 进入 OOD,失败。

dLLM 看到 [-T, T] 整个周期 → cos 已经走完一个完整波形 (0→T→2T→3T→4T 在弧度坐标上), 模型对"走过任意 phase 的 cos"有耐受度。这就是 PPL 不爆的原因。

5.2 t-SNE 可视化的佐证

6 · 方法 LongLLaDA = LLaDA + NTK-RoPE

方法本身简单到让人怀疑能不能投 AAAI:把 NTK-aware RoPE 直接套上 LLaDA, β0 → λ·β0, 不动权重, 不动训练目标。

6.1 NTK 在 RoPE 频谱上做什么

6.2 怎么选 λ

论文直接套 Liu et al. 2023b 的scaling law:

对 LLaDA (Ttrain=4k, β0=500000, d_extra=64, d=128):

论文核心 claim: AR 时代的 NTK scaling law 数值上原封不动地适用于 dLLM。Equation 6 里没有任何 dLLM-specific 的项, λ 完全由 RoPE 几何决定。

6.3 反向论证:为什么不是 YaRN 或 LongRoPE?

7 · Worked example: 一个 8k 处的 query 究竟在做什么

设场景: LLaDA-8B (Ttrain=4k), 输入 16k token 的 NIAH context, "针" 埋在 token 1500 (深度 ~10%), query 在 token 16000 末尾。看 query token (位置 t=16000) 怎么 attend 到针 (位置 s=1500)。相对距离 |t−s| = 14500。

7.1 没有 LongLLaDA (直接外推)

1. Query q_{16000} 计算每个 dimension n 的 RoPE-rotated component q_t·R_t。
2. Key k_{1500} 同样旋转得到 k_s·R_s。
3. 内积 q·k 涉及 cos(θ_n · 14500), n=0..63。
4. 对 n < d_extra=64: θ_n · 14500 落在已见过的 [-4095, 4095]·θ_n 区间(因为周期短,反复绕回来) → cos 值合法。
5. 对 n ≥ 64 (低频): θ_n · 14500 这个相位训练时从没见过(训练时 |t−s| ≤ 4095, 而 14500 远超 4095, 低频维度甚至没走过半个周期) → cos 值 OOD → q·k 内积里这部分维度是 garbage。
6. 结果:attention 给针的 logit 是 (合法部分) + (garbage 部分),低频维度噪声主导, 模型分辨不出"针 in 1500 vs distractor in 1700"。
7. NIAH 上回答失败 → 在那个 (16k, 10%-depth) 网格上是红色。

7.2 加上 LongLLaDA (λ=14)

1. 把 β0 替换为 14·500000 = 7,000,000。
2. θ_n = 7000000^(−2n/128)。对低频 n=64: θ_64 = 1/√7000000 ≈ 0.000378 (旧是 0.001414, ≈ 3.74× 缩小)。
3. 对 n=64, 相位 θ_64·14500 ≈ 5.48 弧度 (vs 直接外推时 20.5 弧度)。
4. 训练时 LLaDA 见过的最大相位是 θ_64^old · 4095 ≈ 5.79 弧度。5.48 < 5.79 → 这个相位在已见过的范围内, q·k 内积合法。
5. 对所有 n ≥ d_extra 都做了这个校正 → 全部维度的相位回到训练分布内。
6. 高频维度 n=0,1,...: θ 几乎不变 (因为 λ 的 (−2n/d) 次方在 n 小时趋近 1), 局部辨识力保持。
7. 结果: q·k 是 clean 的内积, attention 找到针, NIAH 答出来。

这就是为什么 NTK 这一个简单的 β 缩放,能把 NIAH 全表激活。

8 · 实验关键结果

8.1 NIAH 扩展 (论文 Figure 6, LLaDA-Base)

* λ=31 时 24k 已经能 retrieve, 但中间深度出现 lost-in-the-middle (类似 AR LLM 也有的现象)。

读出:NTK scaling law 不仅形式上适用,数值上 λ ∈ {4, 14, 31} 在 LLaDA 上和 LLaMA 完全等同。这强烈支持论文的核心 claim。

8.2 LongBench (论文 Table 1)

LongBench 8k 平均分: LLaMA3-Instruct 41.9, LLaDA-Instruct+λ=4 40.6。基本打平。说明 dLLM 在通用长文档任务上不弱于 AR。

8.3 RULER (论文 Table 2) — 最有意思的部分

读出:这是论文的"附赠彩蛋"。dLLM 的 bidirectional attention 在 QA-style "回头读"任务上反而占优, 但在 sequential aggregation 这种本质需要 ordering 的任务上吃亏 — 因为 cos 是偶函数,前后向量分不出。

9 · 与同类工作对比

10 · 局限 / 个人 take / 待验证问题

论文承认的局限

• 主要在 LLaDA Series 验证, Dream-v0 因为是从 Qwen2.5 (AR) initialization, 行为不那么纯。
• 未做 fine-tuning 实验 — UltraLLaDA 后续工作填了这个坑。
• 对 dLLM 特有的 sampling strategy / step count 影响只做了简单 ablation。
• aggregation 弱项没给解药 (只是观察到了)。

我的疑问 / 想验证的

1. Bidirectional 特性可量化吗?论文只对比 LLaMA3 vs LLaDA, 但理论上 GPT 系列、Mistral 也会因为 RoPE 几何不同而表现不同。希望看到 d_extra / Ttrain / β0 三维 ablation。
2. NIAH 滑窗的 4k 究竟是怎么决定的?是 Ttrain 决定还是 cos 偶函数决定?如果在 LLaDA 上把 RoPE 改成 NoPE,滑窗还在吗?
3. Aggregation 弱项是否真的源自 cos 偶函数?能否用一个非对称 PE 来修?(类似 ALiBi 的 directional bias)
4. λ=55 失败是因为外推天花板还是 attention entropy?论文引用 Han 2023 / Wang 2024 说 attention entropy 限制 NTK,但这个机制在 dLLM 上是否同样适用还没系统测过。
5. dLLM 的 PPL 不爆和 mask diffusion 训练目标的关系。如果换成 SEDD 这种 score-based diffusion 而非 mask diffusion, bidirectional 仍在,RoPE 几何相同, 是否同样 PPL 稳定?
6. local perception 在 chat / instruct fine-tuning 后是否会消失?表格显示 LLaDA-Instruct 还是有滑窗, 但程度有变化。SFT 数据的长度分布可能放大或抑制这个特性。

记忆点

精读笔记 v1 · 2026-05-07 · 配套论文 PDF 在 /data/szhang967/papers/paper-notes/models/LongLLaDA_2506.14429.pdf