CausalGate: 因果重要性蒸馏用于Transformer模块剪枝

# CausalGate：面向Transformer模块剪枝的因果重要性蒸馏
来源：https://arxiv.org/html/2607.22720

###### 摘要

现有的大语言模型自适应推理方法依赖观测启发式（如隐藏状态相似度或激活幅度）来丢弃冗余模块。然而，这些基于相关性的度量往往无法捕捉对语义准确性至关重要的微妙非线性结构计算。我们提出CausalGate，一种用于计算高效Transformer推理的干预引导框架。在校准阶段，CausalGate隔离各个注意力和MLP子层，将其相应输出归零，并通过最终logit分布的Kullback-Leibler散度精确衡量语义损伤。为了消除运行时路由开销，该结构重要性层级通过指数移动平均平滑目标和可微成对排序损失蒸馏为一组全局静态轻量标量门。在语言建模和常识推理基准上对TinyLlama-1.1B、Qwen2.5-3B和Llama-3.1-8B进行的评估表明，CausalGate持续优于主流的动态路由和跳层基线，将理论计算节省转化为具体的硬件延迟降低，且零运行时开销。

## 1 引言

自从"注意力即一切"这一范式转变以来（Vaswani等，2017 (https://arxiv.org/html/2607.22720#bib.bib3)），扩展Transformer架构以极高的计算成本解锁了前所未有的语言能力，使大语言模型推理成为关键的部署瓶颈。为了缓解这一问题，模型压缩方法（如结构剪枝和子层跳过（He等，2025 (https://arxiv.org/html/2607.22720#bib.bib5)））旨在永久停用低效用模块，避免运行时动态路由框架（如CALM（Schuster等，2022 (https://arxiv.org/html/2607.22720#bib.bib6)）或深度混合（MoD）（Raposo等，2024 (https://arxiv.org/html/2607.22720#bib.bib4)））特有的高硬件执行开销和分支预测惩罚。现有的结构剪枝度量通常依赖观测代理（如激活幅度或隐藏状态相似度）来判断模块是否可被绕过。然而，此类观测度量有时可能忽略微妙但关键的计算步骤。一个经过模块的隐藏状态向量可能仅表现出微小的几何位移，但该模块可能仍在执行对于事实准确性或逻辑连贯性至关重要的非线性"因果精炼"。通过主要依赖表征特征而非直接的结构必要性，启发式驱动的静态剪枝存在绕过安静必要的模块的风险，从而导致在较高计算节省时性能急剧下降。

为了克服观测代理的局限性，我们借鉴经典科学探究方法：要真正确定一个组件是否对系统至关重要，必须主动对其进行干预，而不仅仅是观察。我们提出CausalGate，一种离线结构校准框架，通过直接且与输入无关的结构干预建立永久的模块选择层级。不同于从运行时表示对齐中估计效用，CausalGate在离线校准阶段隔离各个注意力和MLP模块，将其输出归零，并测量最终词元分布中的Kullback-Leibler散度 \(D_{\text{KL}}\)。该方法基于一个核心假设：模块真正的全局计算重要性明确由其最终logit分布的因果影响定义。通过在子层级别操作这一原则，CausalGate绕过了启发式猜测，直接量化了跳过任何给定模块所导致的精确语义损伤。

CausalGate的一个关键架构区别在于其细粒度操作在子层模块级别，而非整体Transformer层级别。由于多头注意力和MLP模块执行不同功能（注意力在序列间路由上下文，而MLP作为局部键值记忆），独立剪枝比统一丢层提供显著更高的计算灵活性。然而，在运行时计算因果干预是不可行的。为了实现零开销推理，CausalGate将校准过程中发现的结构重要性层级蒸馏为全局静态学习的标量门。我们使用指数移动平均平滑目标来优化这些门以过滤干预噪声，并辅以可微成对排序损失，强制门值严格保留发现的因果顺序。

该优化策略在因果分析与硬件执行效率之间架起桥梁，将全局结构知识编码为确定性推理掩码。由于这些功能子层执行解耦任务，模型建立了在编译时运行的与序列无关的重要性层级。因此，CausalGate消除了运行时路由开销，确保直线推理执行，与局部内存边界和标准硬件循环完美配合。总之，我们的贡献如下：

1. 我们提出了一种新的结构模型剪枝方法，用直接来自校准过程中结构因果干预的 ground-truth 重要性层级取代传统的观测启发式。
2. 我们证明，独立评估和跳过注意力和MLP模块比统一丢层策略能实现显著更灵活、更精确的计算-性能权衡。
3. 我们提出一个利用EMA目标平滑和成对排序损失的优化框架，成功将结构因果层级编码为全局静态门，实现零开销推理。
4. 我们针对主流基线进行了广泛评估，表明CausalGate在较大计算缩减下能更好地保持语言模型质量。

本文其余部分组织如下：第2节 (https://arxiv.org/html/2607.22720#S2) 回顾自适应推理和机制可解释性的相关工作。第3节 (https://arxiv.org/html/2607.22720#S3) 介绍CausalGate架构的结构组件及其干预和训练框架。第4节 (https://arxiv.org/html/2607.22720#S4) 展示我们的实证评估和基线对比分析。最后，第5节 (https://arxiv.org/html/2607.22720#S5) 讨论更广泛的影响、局限性和未来工作方向，第6节 (https://arxiv.org/html/2607.22720#S6) 提供总结性概述。

## 2 相关工作

#### 自适应推理与跳层。

自适应推理根据输入难度或延迟界限动态调整计算。卷积神经网络中的早期机制利用滤波器剪枝、量化（Zhang等，2024 (https://arxiv.org/html/2607.22720#bib.bib9)）、提前退出（Han等，2021 (https://arxiv.org/html/2607.22720#bib.bib8)；Bolukbasi等，2017 (https://arxiv.org/html/2607.22720#bib.bib10)）和动态通道门控（Bejnordi等，2019 (https://arxiv.org/html/2607.22720#bib.bib15)；Gao等，2018 (https://arxiv.org/html/2607.22720#bib.bib13)；Lin等，2017 (https://arxiv.org/html/2607.22720#bib.bib14)）。在视觉Transformer中，像DynamicViT（Rao等，2021 (https://arxiv.org/html/2607.22720#bib.bib18)）和A-ViT（Yin等，2022 (https://arxiv.org/html/2607.22720#bib.bib16)）等框架专注于动态词元减少，而AdaViT（Meng等，2022 (https://arxiv.org/html/2607.22720#bib.bib17)）则根据样本自适应丢弃注意力头或整个块。对于大语言模型，动态计算路由已扩展到词元跳过策略，如深度混合（Raposo等，2024 (https://arxiv.org/html/2607.22720#bib.bib4)）和条件提前退出方案（Rotem等，2023 (https://arxiv.org/html/2607.22720#bib.bib20)）。与我们的工作更接近的是结构性方法，它们侧重于通过均匀跳层（Liu等，2024 (https://arxiv.org/html/2607.22720#bib.bib37)；Luo等，2025 (https://arxiv.org/html/2607.22720#bib.bib38)）、输入-输出相似度启发式（He等，2025 (https://arxiv.org/html/2607.22720#bib.bib5)）、强化学习路由（Liu等，2025 (https://arxiv.org/html/2607.22720#bib.bib39)）或训练注入的残差路由门（Laitenberger等，2025 (https://arxiv.org/html/2607.22720#bib.bib40)；Nair等，2026 (https://arxiv.org/html/2607.22720#bib.bib2)）来绕过冗余段。然而，这些方法主要引入了动态控制流分支或依赖运行时浅层观测代理。相比之下，我们通过利用离线干预蒸馏生成零开销静态模块掩码，将执行边界完全转移到编译时。

#### 结构剪枝。

静态网络压缩会永久丢弃利用不足的权重或通道。对于大语言模型，训练后剪枝通常基于原始幅度（Frantar等，2022 (https://arxiv.org/html/2607.22720#bib.bib30)）、权重-激活乘积（如Wanda（Sun等，2024 (https://arxiv.org/html/2607.22720#bib.bib41)））或结构激活方差来优先考虑通道组、层或单个权重。为了从结构上压缩架构，像LLM-Pruner（Ma等，2023 (https://arxiv.org/html/2607.22720#bib.bib42)）和Sheared LLaMA（Xia等，2024 (https://arxiv.org/html/2607.22720#bib.bib43)）等框架利用一阶梯度近似来隔离冗余耦合组件以便移除。类似地，静态敏感性分析利用泰勒展开或零样本相似性度量（如ShortGPT中的块影响（Men等，2025 (https://arxiv.org/html/2607.22720#bib.bib44)））来映射模块重要性。虽然相关，但这些替代方案严格依赖于观测相关性或局部梯度。CausalGate的根本区别在于，它以解耦的注意力和MLP子层粒度操作化主动干预演算（do-操作），以隔离真实的下游语义偏移。

#### 互补加速范式。

除了层剪枝之外，文献中还包含缓解大语言模型推理瓶颈的正交策略。内存占用通常通过压缩或金字塔KV缓存（Ge等，2024 (https://arxiv.org/html/2607.22720#bib.bib23)；Xiao等，2025 (https://arxiv.org/html/2607.22720#bib.bib24)；Cai等，2024 (https://arxiv.org/html/2607.22720#bib.bib26)）或统一混合精度权重和激活量化框架（Xiao等，2023 (https://arxiv.org/html/2607.22720#bib.bib29)；Lin等，2024 (https://arxiv.org/html/2607.22720#bib.bib33)；Frantar等，2022 (https://arxiv.org/html/2607.22720#bib.bib30)；Park等，2024 (https://arxiv.org/html/2607.22720#bib.bib32)）来管理。同时，推测性和激进解码范式（Leviathan等，2023 (https://arxiv.org/html/2607.22720#bib.bib35)；Stern等，2018 (https://arxiv.org/html/2607.22720#bib.bib36)；Sun等，2021 (https://arxiv.org/html/2607.22720#bib.bib34)）通过不同的草稿验证分支并行预测多个词元来加速自回归生成。这些方法与CausalGate完全互补，可以联合应用以最大化硬件部署效率。

## 3 方法

我们的核心直觉是，移除后显著改变输出分布的模块应予以保留，而影响可忽略的模块可以跳过，且对模型质量影响极小。

请参见图注：图1：TinyLlama骨干网络。
请参见图注：图2：基于干预的因果重要性估计和门学习。
请参见图注：图3：通过跳过低门模块实现计算高效推理。

图4：所提出的干预引导模块选择框架。每个注意力和MLP模块通过将其输出归零独立评估；所得最终logit KL散度定义其因果重要性分数。EMA平滑和成对排序损失将这些信号蒸馏为学习的标量门，允许在目标计算预算下在推理时跳过低效用模块。设一个Transformer模型由 \(L\) 层组成，每层包含一个注意力子模块和一个MLP子模块。我们将每个注意力块和MLP块视为一个独立计算单元，得到一组 \(M\) 个模块 \(\mathcal{M}=\{m_1,m_2,\ldots,m_M\}\)。对于TinyLlama-1.1B，这对应 \(M=44\) 个模块（22个注意力和22个MLP模块）。我们的框架包含三个阶段，如图4 (https://arxiv.org/html/2607.22720#S3.F4) 所示：(1) 基于干预的因果重要性估计，(2) 将因果重要性蒸馏为学习到的模块门，(3) 通过模块选择实现计算高效推理。

### 3.1 因果模块重要性估计

我们通过干预分析估计每个模块的重要性。给定输入序列 \(x\)，我们首先获得原始模型的输出分布 \(p(y \mid x)\)。接下来，我们对模块 \(m_i\) 进行干预，将其输出激活归零，同时保持所有其他模块不变。所得模型产生干预后的输出分布 \(p(y \mid x, m_i \rightarrow 0)\)。模块 \(m_i\) 的因果重要性量化为干预后分布与原始输出分布之间的KL散度：

\[
\Delta_i^{(t)} = D_{\mathrm{KL}}\Big(p(y \mid x, m_i \rightarrow 0) \;\Big\|\; p(y \mid x)\Big),
\]

其中 \(t\) 表示当前训练迭代。\(\Delta_i^{(t)}\) 的值较大表明对模块 \(m_i\) 的干预显著改变了模型的预测分布，表明该模块在生成中起着重要作用。相反，值较小意味着该模块对最终预测的影响相对较小。我们独立计算网络中每个注意力和MLP模块的干预分数。所得分数集形成模块级因果重要性图，表征每个Transformer模块对模型预测行为的相对贡献。

### 3.2 因果重要性蒸馏

尽管干预分数提供了模块重要性的原则性估计，但它们在不同训练迭代之间可能表现出显著变异性。因此，直接使用原始干预分数进行模块选择可能导致不稳定的排序和次优的计算分配决策。为解决此问题，我们将干预衍生的因果重要性蒸馏为一组学习到的模块门，这些门捕获模块相关性的稳定全局排序。对于每个模块 \(m_i\)，我们关联一个可学习标量门

\[
g_i = \sigma(\theta_i),
\]

其中 \(\theta_i\) 是一个可训练logit，\(\sigma(\cdot)\) 表示sigmoid函数。每个门是与一个Transformer模块相关联的单个学习标量，并且跨所有输入和词元位置共享。较高的门值表明该模块应在推理期间优先保留。

#### EMA目标聚合。

干预衍生的重要性估计在不同训练迭代之间可能变化，导致模块排序不稳定。