跳过冗余层:扩散语言模型与自回归语言模型的表征结构与推理时跳层分析 Skip to the Good Part: Representation Structure & Inference-Time Layer Skipping in Diffusion vs. Autoregressive LLMs
首次系统对比 AR 与扩散 LLM 的层间表征结构,发现扩散目标函数带来可被利用的早层冗余。
前置知识
自回归语言模型 (Autoregressive LM)
通过下一个 token 预测(NTP)训练的语言模型,按从左到右的因果顺序逐 token 生成。每一步只能看到当前位置之前的 token,损失是交叉熵。Qwen2.5、LLaMA、GPT 系列都是典型代表,也是当前工业部署的主流范式。
本文的核心对比对象就是 AR 模型,必须理解其左到右预测机制才能理解为何 AR 表征会产生'近因偏置'。
扩散语言模型 (Diffusion LM, dLLM)
通过全序列去噪训练的语言模型,从一个被掩码或噪声化的完整序列出发,迭代地恢复出原始序列。每个去噪步都基于完整序列的双向上下文进行预测。LLaDA、Dream-7B、SEDD 是代表工作,能并行生成所有 token。
理解 dLLM 的全序列去噪目标是理解为何其表征会更'全局'、早层更冗余的关键。
余弦相似度 (Cosine Similarity) 作为表征探针
衡量两个向量方向一致性的指标,公式为 $\mathrm{sim}(u,v) = \frac{u \cdot v}{\|u\|\|v\|}$,取值 $[-1,1]$。在本文中用于测量相邻层同一 token 隐藏状态的方向稳定性,对幅度变化不敏感。
本文全部分析都基于余弦相似度,需要理解它只捕捉'方向'而忽略'幅度'这一性质对结论的影响。
推理时层跳过 (Inference-time Layer Skipping)
在推理阶段绕过某些 Transformer block,直接将上一层的 hidden state 传给再下一层,相当于把这个 block 从计算图中完全移除,但因为残差连接仍保留,所以不会破坏梯度流和表征连续性。
本文把跳层当作'诊断探针'而非部署策略,需要理解这种静态、无需 KV-cache 共享的剪枝思路。
近因偏置 (Recency Bias)
在 AR 训练下,每个新 token 的隐藏状态相对前一个 token 都会发生显著变化,因为因果掩码让当前 token 主要关注前文最新内容。表现是相邻 token 表示差异大、token-wise 余弦相似度低。这是本文论证 AR 表征'局部性'的核心证据。
本文最核心的发现之一就是 dLLM 的近因偏置显著低于 AR,这是表征结构差异的直接体现。
研究动机
近年来扩散语言模型(dLLM)如 LLaDA、Dream-7B 在各项下游任务上已经能比肩强 AR 基线,但学界对'扩散目标函数是否从根本上重塑了模型的内部表征'这一基础问题仍然缺乏系统认识。现有工作大多聚焦在效率改进(并行解码、verifier 采样、架构优化),却少有人细致拆解过 dLLM 内部隐藏状态到底长什么样。更有意思的是,像 Dream-7B 这类从 AR 预训练初始化、再用扩散目标继续训练的模型,在行为上看似已经'扩散化',但其表征是否真的从 AR 风格切换成了扩散风格?这种'目标函数 vs 初始化'的混淆使得我们难以判断 dLLM 的真实优势到底来源于哪里。
本文的目标是本文瞄准一个具体且可验证的问题:原生扩散目标函数到底有没有给语言模型的内部表征带来结构性变化?为此设计了一套完整的表征分析框架,在层维度和 token 维度同时测量三个模型的差异——Qwen2.5(纯 AR)、LLaDA(纯扩散)、Dream-7B(AR 初始化后扩散微调)。基于观察到的表征冗余,本文进一步提出一个静态、任务无关的推理时跳层策略,作为诊断式探针来量化'这些冗余到底能被压掉多少',从而把表征结构差异和实际可利用的计算节省直接挂钩。
与已有工作不同的是,本文的切入角度有三个独特之处:第一,同时控制'目标函数'和'初始化'两个变量,把 Dream-7B vs Qwen2.5 的对比作为天然对照,能干净地隔离出'目标函数效应'与'初始化偏置';第二,把层跳过从'部署策略'重新定位为'分析探针',不追求提出一个 SOTA 加速方法,而是用跳层后的崩塌曲线作为表征冗余的因果检验;第三,把表征结构差异(余弦相似度、近因偏置)直接翻译成 FLOPs 节省百分比,让'理论观察'和'实际收益'之间有明确的对应关系。
核心方法
直觉上,如果某个相邻层的隐藏状态在方向上几乎没变化,那这一层就只做了微不足道的变换,可以安全跳过。本文正是基于这个直觉展开的。先用余弦相似度量化 LLaDA、Qwen2.5、Dream-7B 在每一层、每一个 token 上的表征变化率,发现 LLaDA 的前 1/3 层存在 $>0.95$ 的高相似度平台,而 AR 类模型则全程持续变化。然后把这个观察翻译成一个静态跳层算法:对所有相邻层计算平均余弦相似度,按相似度从高到低排序,选出 top-K 个非相邻的层作为跳层集合,推理时直接绕过这些 block。
本文的核心创新不在算法本身(top-K 非连续跳层并不新颖),而在把'目标函数—表征结构—可压缩性'这条因果链第一次系统性地串联起来。和已有方法(例如 YOCO 的 cache-once 设计、各种动态路由方案)的本质区别在于:(1) 静态且任务无关——不需要任何 per-task 调优或动态决策;(2) 架构无关——不依赖 KV-cache 共享或参数共享,纯后处理式;(3) 正交互补——和 KV-cache 类优化是乘法关系而非替代关系。更关键的是,本文用同一个探针(跳层后的精度崩塌曲线)反向验证了前面的表征分析结论:原生 dLLM 真有可被利用的早层冗余,AR 模型几乎一点冗余都没有。
方法步骤详情
完整流程包含四个步骤。第一步是收集隐藏状态:对 LLaDA 在多个去噪步 $t \in \{1,\dots,T\}$ 上、对 Qwen2.5 在普通前向传播上,记录所有层的 hidden states $\{h^{(\ell)}\}$。第二步是计算相邻层余弦相似度:对每个 token $i$、每对相邻层 $\ell$ 和 $\ell+1$,计算 $\mathrm{sim}(h^{(\ell)}_i, h^{(\ell+1)}_i) = \frac{h^{(\ell)}_i \cdot h^{(\ell+1)}_i}{\|h^{(\ell)}_i\|\|h^{(\ell+1)}_i\|}$,然后在 prompt 长度内求平均得到 $s_\ell$。第三步是生成跳层集合 $L_{\text{skip}}$(Algorithm 1/2):按 $s_\ell$ 降序遍历,每选中一层 $\ell$ 都要保证 $\ell-1$ 和 $\ell+1$ 不在集合中(避免连续跳层),直到选满 $K$ 层。第四步是推理时执行跳层:对每个 $\ell \in L_{\text{skip}}$,绕过 transformer block、直接 $h^{(\ell+1)} = h^{(\ell-1)}$,由于残差连接已在前面累积,梯度流和表征连续性都得以保留。整个过程无需重新训练,无需修改架构,无需 KV-cache 共享。
技术新颖性
技术上属于'分析驱动压缩'范式,方法层面不算很 fancy,但研究视角的贡献较大:(1) 第一次给出 dLLM vs AR 的层- token 双维度系统对比,填补了之前文献的空白;(2) 通过 Dream-7B 这一'AR-初始化 dLLM'的天然对照实验,给出'初始化偏置在表征层持续存在'的机制性证据,解释为何某些 dLLM 仍保留 AR 风格的解码行为;(3) 揭示了一个反直觉的现象——原生 dLLM 在 MATH500 这种高难度推理上,跳 6 层后准确率反而略升到 97%(绝对值 0.37),可能因为去除了某些过拟合层;(4) 跳层算法虽简单,但通过'禁止连续跳层'的约束(图 6 显示跳层集中在前 40-60%),让结果与冗余定位一致,避免了均匀/随机/连续跳层的崩塌。
实验结果
核心发现一:原生 dLLM(LLaDA-8B-Instruct)在跳 6 层(18.75% FLOPs 削减)后,在 GSM8K、MATH500、MBPP 上保留 93-97% 的基线性能,HumanEval 上保留 83%;即使跳 8 层(25% FLOPs 削减),保留率仍达 60-94%(Appendix Table 3)。核心发现二:AR 模型(Qwen2.5-7B)在跳 2 层(仅 7.14% FLOPs 削减)后就已经崩塌——GSM8K 保留率仅 48%(绝对值从 0.21 跌到 0.10),MBPP 62%,MATH500 因基线仅 0.07 而噪声较大。核心发现三:AR-初始化的 Dream-7B-base 表现介于两者之间,2 层跳时还能与 LLaDA 相当(GSM8K 101%、HumanEval 95%),但 4 层跳时降到 50-80%,6 层跳时跌到 8-25%,显示出扩散微调只能部分缓解 AR 初始化的表征惯性。综合来看,LLaDA 在 18.75% FLOPs 削减下平均保留 91.8%,是 Qwen2.5 在 7.14% 削减下 62.7% 的 1.46 倍——也即'2.6 倍 FLOPs 削减 + 1.5 倍质量保留'。表 4 还给出了真实 wall-time 节省:HumanEval 在跳 6 层时墙钟时间减少 21%,跳 8 层时减少 27%;MBPP 跳 8 层时减少 29%。表 5 显示和 dual-cache 机制叠加后,跳 4 层时仍能保留 88% 准确率。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| GSM8K (小学数学) | Accuracy Retention @ 6 layers skip | LLaDA-8B-Instruct: 94% (0.79->0.74) | Qwen2.5-7B: 不可用 (2层即崩至48%); Dream-7B-base: 22% | 扩散模型在同等FLOPs削减下保留率高42-72个百分点 |
| MATH500 (高难度数学) | Accuracy Retention @ 6 layers skip | LLaDA-8B-Instruct: 97% (0.38->0.37) | Qwen2.5-7B: 0% (从0.07跌至0); Dream-7B-base: 25% | LLaDA甚至在跳层后绝对精度反而提升,呈现近反直觉的鲁棒性 |
| HumanEval (代码合成) | Pass@1 Retention @ 6 layers skip | LLaDA-8B-Instruct: 83% (0.53->0.44) | Qwen2.5-7B: 7% (0.46->0.03); Dream-7B-base: 8% | LLaDA保留率约10倍于基线模型 |
| MBPP (Python编程) | Pass@1 Retention @ 6 layers skip | LLaDA-8B-Instruct: 93% (0.56->0.52) | Qwen2.5-7B: 19% (0.63->0.12); Dream-7B-base: 14% | LLaDA在最大削减率下仍保持93%,远超两类对比模型 |
| Wall-time节省 (HumanEval) | Inference Time Reduction @ 6 layers skip | LLaDA-8B-Instruct: 21% 减少 + 83% 准确率保留 | 无跳层基线 | 8层跳时墙钟时间减少27%,绝对精度仅从0.53降至0.32 |
| 跳过层分布 (HumanEval, 6层跳) | Skipped Layer Index Distribution | 集中在网络前40-60%的层 | 连续跳层导致GSM8K从94%跌至40%、HumanEval从60%跌至8% | 静态非连续跳层算法相比连续跳层避免精度灾难 |
局限与改进
作者承认的局限:(1) 余弦相似度只是'一阶探针',捕捉方向稳定性但不刻画内在信息内容;(2) 分析只覆盖一个原生 dLLM (LLaDA) 和一个 AR-初始化 dLLM (Dream-7B),是否对 SEDD 这类用 uniform transition kernel 的扩散模型同样成立还未知;(3) 未覆盖 RL 后训练或多模态扩散架构;(4) 阈值 $\theta=0.95$ 和 skip budget $K$ 都是基于 LLaDA 的相似度分布选的,对其他模型未必最优。本人观察到的额外局限:(a) 评测只用了 4 个基准(GSM8K、MATH500、HumanEval、MBPP),缺乏阅读理解、对话、翻译等更广任务;(b) Table 1 中 Qwen2.5 在 MATH500 的 143% 保留率虽然作者标注为'基线 0.07 导致噪声',但这种'噪声基线'本身意味着 Qwen2.5 在该任务上几乎不能用,会让 FLOPs 节省对比显得不公平;(c) 跳层虽然省 FLOPs,但实际墙钟时间节省在跳 2 层时甚至为负(-8%/-3%),说明小规模跳层的 kernel launch overhead 反而超过收益;(d) '原生 dLLM 产生全局表征'的因果链证据仍是相关性的,未做控制变量消融(如屏蔽 AR 训练中的某类信号看是否重现 dLLM 风格)。
独立分析的弱点
独立分析的主要弱点及改进方向:(1) 跳层策略是静态的、不感知输入——不同 prompt 的关键层可能不同;改进方向是引入轻量路由器,按 hidden state 特征决定跳哪几层,类似 MoE 的思路。(2) 算法强制禁止连续跳层是基于'连续跳层会崩'的实证观察,但缺乏理论解释;改进方向可以分析 transformer block 的残差结构对连续跳过次数的容忍度,推导最大可连续跳层数上界。(3) 评测只用 4 个英文基准,对中文或代码以外任务的泛化未验证;改进方向应在 C-Eval、HumanEval-X 中文版等更广基准上重复实验。(4) 余弦相似度只测方向、不测信息量,可能存在'方向相同但内容已变'的层被误判为冗余;改进方向是结合 probing task(如直接用某一层 hidden state 做词性标注)量化每层实际承载的信息。(5) 跳层节省 FLOPs 但不直接节省 latency——尤其 dLLM 的 wall-time 还受去噪步数影响;改进方向是联合优化跳层 + 去噪步数,把 layer-skip 和 step-distillation(Deschenaux & Gulcehre 2025)结合。(6) 只在 7-8B 规模验证,扩展到 70B+ 时 AR vs dLLM 的冗余差异是否仍存在未知;大规模实验是必要补充。
未来方向
作者明确提出的方向:(a) 用更高阶探针(如 probing classifier、causal mediation)刻画每层的信息内容,而不是只看方向稳定性;(b) 把分析扩展到 SEDD 等使用不同 transition kernel 的 dLLM,验证'全局表征'假设的普适性;(c) 研究 RLHF 或 DPO 等后训练方法对 dLLM 表征几何的影响,看是否会进一步拉近或拉远 AR 风格;(d) 扩展到多模态扩散架构。基于本文成果可延伸的方向:(e) 既然 LLaDA 早层冗余,是否可以做'早层蒸馏'——用浅层 student 模型逼近 LLaDA 的深层输出;(f) 既然 AR-初始化会留下'表征惯性',是否可以设计一种'重置化'预训练,从 AR 初始化时随机扰动部分层以打破这种惯性;(g) 把层跳过和 dLLM 的并行解码结合,设计'空间(跳层)+ 时间(少去噪步)'的双轴加速;(h) 用 mechanistic interpretability 工具(如 activation patching)精确定位 dLLM 早层到底在编码什么样的全局信息。
复现评估
复现评估:作者在论文中虽未明确声明开源代码,但所有实验都基于公开模型——LLaDA-8B-Base/Instruct(Nie et al. 2025 发布)、Dream-7B-Instruct(Ye et al. 2025a)、Qwen2.5-7B(Yang et al. 2024)都是 HuggingFace 上可下载的标准开源模型,且评测基准 GSM8K、MATH500、HumanEval、MBPP 均为公开。算法本身(Algorithm 1/2)非常简洁,几十行 PyTorch 就能实现,复现门槛很低。算力要求方面,跑 8B 模型在 4×A100 80G 上完整对比实验大约需要 2-3 天 GPU 时间,跑全 8 层跳层扫描附录 Table 3 可能需要更久。最大的复现难点不在算力,而在 dLLM 的解码设置:去噪调度、remasking 策略、temperature annealing 都需要严格对齐各模型官方仓库的 default sampler,否则不同论文之间的对比可能不一致;本文已经声明遵循 Nie et al. 和 Ye et al. 的官方实现。整体复现难度:中等偏低,给有经验的 LLM 工程师 1-2 周可完整跑出主表结果。
论文图表