高效Transformer中的稀疏令牌路由

# 高效Transformer中的稀疏Token路由
来源：https://arxiv.org/html/2608.20632
Sai Krishna Arthanari, JaeHyeong Chang, Chengzhe Sun, Siwei Lyu
隶属机构：人工智能与数据科学研究所（IAD），纽约州立大学布法罗分校，美国 纽约州布法罗市
邮箱：[email protected], [email protected], [email protected], [email protected]

###### 摘要

高效Transformer研究常以“并非所有token都需要同等计算资源”为由，来推动token剪枝和自适应计算的研究。我们通过SEWN这一双流Transformer架构，利用学习到的门控信号将token路由至轻量级或全容量处理路径，从而端到端地验证了这一主张。实验表明，与参数匹配的基线模型相比，路由机制引入的精度变化可忽略不计，而门控的token重要性信号则高度依赖于其学习方式。在BoolQ数据集上，一个基于静态词表先验的门控未能通过反事实忠实性检验，而一个完全基于上下文的门控则在不改变任务精度的情况下，在两个评估任务上实现了高度显著的分离（p<10⁻¹⁰）。

随后，我们将路由信号转化为显式的计算机制。SEWN-sparse在昂贵的流处理之前执行硬性的top-k选择，在评估的数据集上实现了5.2–8.7倍于BERT-base和2.6–4.4倍于DistilBERT的吞吐量提升，并且在规模更大、上下文更长的27k样本PubMedQA基准测试中，其精度折衷依然保持适度。在相同的反事实掩蔽协议下，我们还将门控与原始注意力、注意力积分（attention rollout）和积分梯度进行了对比。基于注意力的方法能以低成本提供有意义的信号，而积分梯度则产生46倍的推理开销；在所评估的方法中，SEWN-sparse在忠实性与效率之间提供了最佳的权衡。

使用相同架构和瓶颈的随机选择消融实验仅使精度下降2.6个百分点，但使其显著的忠实性消失（p=0.45），这表明效果源于学习到的token排序，而非硬稀疏性本身。对Mixture-of-Depths的直接重新实现也产生了有意义的路由信号，但未能通过我们严格的基于排序的忠实性标准。我们进一步发现，静态先验与上下文门控的结果在RoBERTa上无法被清晰地复现，系统性的先验大小、覆盖范围和内容扫描排除了已测试的干扰因素，表明其依赖于主干网络。最后，跨五个任务类别，效率-精度的权衡在二元问答和情感分类任务中稳健，但在长段落的多项选择阅读理解任务中则失效，在该任务中SEWN和DistilBERT都表现出显著的容量限制。所有实验均采用固定的三随机种子协议，包括负面结果。

###### 索引词：

高效Transformer，token剪枝，稀疏注意力，可解释性，忠实性，知识蒸馏，大数据

## I 引言

Transformer自适应计算和token剪枝方法的前提是：并非每个输入token都需要相同量的处理：功能词、样板文本和冗余上下文通常可以廉价处理，而承载内容的token则需要模型的全部能力。这一前提催生了大量关于token剪枝（PoWER-BERT[5]、Learned Token Pruning[7]、TR-BERT[8]）、token合并（ToMe[9]）和深度自适应计算（Mixture-of-Depths[10]）的工作。

在这一前提下，通常包含两个不同的主张：

1. 效率主张：跳过或廉价处理“不重要”的token可以节省计算资源，而精度损失不大。
2. 可解释性主张：决定哪些token是“重要”的信号，也可作为模型行为的人类可读解释。

本文旨在使用一个具体架构（SEWN——Stream-Efficient Word Network），将路由决策显式化而非隐含在注意力权重中，来严谨、分离且诚实地检验这两个主张。我们并不声称SEWN是一种新的最先进高效架构；直接对比中，DistilBERT在相当或更小的参数预算下，其原始精度优于我们测试的每个SEWN变体（第V-E节）。相反，本文的贡献是方法论性和机制性的：

- 提供一个实例，证明对于一个能成功训练的路由机制，其可解释性主张可能在经验上是错误的——我们的第一个门控设计（一个从手工构建的功能词词表初始化的学习分数）在BoolQ数据集上直接未能通过反事实掩蔽测试，尽管该架构达到了正常的任务精度。
- 证明这种特定失败在BERT上可以修复（移除静态词表先验，第V-B节），精度代价接近于零，但在架构上并非没有代价——一个保留相同门控但更精简的单流变体在长上下文输入上再次破坏了相同的忠实性测试（第V-H节），随后我们通过从双流教师模型进行蒸馏来修复它。我们还测试了该修复是否能跨主干网络泛化（它未能清晰地泛化：第V-B节报告了RoBERTa/BoolQ的一个案例，其中我们尝试的任何静态先验大小或内容变体，包括完全没有先验，都未能通过）——报告为依赖于主干网络的结果，而非将其平滑处理为标题主张。
- 在单一共享评估协议下，对路由门控的重要性分数与标准的事后替代方法（原始注意力、注意力积分、积分梯度）进行了直接、成本感知的比较，据我们所知，当基于路由的方法声称具有可解释性的附带好处时，这种做法并不常见（第V-E节）。
- 对我们忠实性结果最简约的替代解释（即一个硬token选择的瓶颈，独立于选择的好坏，本身足以产生强大的反事实掩蔽分数）进行了证伪检验（第V-G节）。结果表明并非如此：使用相同架构的随机选择在相同的机械瓶颈下，显示出非显著的忠实性结果。第二个对照组使用冻结的、非学习的启发式注意力（与几种已发布的剪枝方法使用的思路相同）对token进行排序，虽然通过了忠实性测试，但代价是需要一次完整的额外前向传播来生成排序——这使其比未剪枝的BERT-base本身更慢，且忠实性仍低于SEWN自身的门控。第三个对照组直接重新实现了Mixture-of-Depths（而非进行近似）作为逐层路由基线，显示出一个具有真实信号的路由器，但仍达不到SEWN-sparse的清晰排序。
- 诚实地核算了效率主张在哪些方面成立、在哪些方面不成立，具体到任务类别层面（第V-F节），包括在长段落多项选择阅读理解任务上出现的容量崩溃，这同样严重地影响了DistilBERT——证据表明这是压缩模型普遍存在的深度/容量限制，而非路由特有的假象。

我们全程报告负面结果，因为在调查过程中我们自己的几个工作假设被证伪了，并且我们认为这些证伪结果对社区比经过筛选的成功故事更有用。

## II 相关工作

**Token剪枝和自适应计算。** PoWER-BERT[5]和Length-Adaptive Transformer[6]基于学习到的重要性分数，在各层中逐步消除token。Learned Token Pruning[7]和TR-BERT[8]使用阈值或强化学习策略来决定哪些token得以保留。Token Merging (ToMe)[9]基于相似性合并token而非丢弃，最初用于视觉Transformer。Mixture-of-Depths[10]通过学习的路由器将token路由通过不同数量的Transformer层，这是与SEWN的双流设计最接近的先前框架；我们在第V-G节直接重新实现它（从因果语言模型适配到双向分类）作为第四个基线，而不仅仅是其精神上的启发式对照。我们没有对PoWER-BERT、TR-BERT、Learned Token Pruning或ToMe的已发布实现进行基准测试（我们在第VII节明确说明了这一局限），尽管第V-G节也对比了一个与PoWER-BERT重要性分数精神相同构建的、非学习的基于注意力的排序基线。我们的贡献并非声称在效率/精度前沿上超越这些方法，而是对可比路由机制的忠实性分析，据我们所知，这些被引论文均未对其重要性/重要性分数进行此类分析。

**知识蒸馏。** DistilBERT[2]和TinyBERT[3]通过标准的logit和隐藏状态蒸馏来压缩BERT。我们在整个过程中将DistilBERT作为主要的高效基线，因为它被广泛部署，并且在相当参数数量下，其精度优于我们测试的每个SEWN变体——这是一个不便但重要的基线，提出路由替代方案的论文不应忽略它。注意力迁移[4]专门蒸馏注意力图，在方法论上接近我们自己在第V-H节进行的门控分数蒸馏。

**忠实性和事后解释。** Jain和Wallace[11]表明注意力权重通常不能作为忠实的解释，这引发了一系列通过经验评估解释方法而非假设忠实性的工作。ERASER基准测试[12]通过我们此处使用的掩蔽-测量协议，形式化了“全面性”和“充分性”，该协议在此应用于人类标注的推理依据，而非模型的内部路由决策。积分梯度[13]和注意力积分[14]是我们直接作为基准对比的两种事后方法。

## III 架构

### III-A SEWN核心

SEWN使用标准的BERT风格嵌入[1]（词嵌入 + 位置嵌入 + token类型嵌入）来嵌入token，然后通过一个小型MLP（SoftmaxRoutingGate）计算每个token的门控分数。门控存在两种变体：

- 静态先验门控：MLP的上下文分数与一个从手工构建的封闭类（功能词）词表（子词扩展后有269个token id）初始化的、每个词汇表id的偏置项相加，反映了语言学直觉：限定词、介词和助动词承载的内容少于开放类词。
- 上下文门控：相同的MLP，完全没有基于词表的偏置项。门控必须完全从训练信号中学习哪些token应被视为低优先级或高优先级。

门控的输出是[0,1]范围内的标量，解释为P(路由到廉价流)。两个并行流随后处理序列：一个轻量级的功能词流（2层，256维）接收每个token的嵌入，该嵌入按门控值缩放；一个内容流（4层，768维，通过直接从bert-base-uncased移植权重进行BERT初始化）接收每个token的嵌入，该嵌入按1减去门控值缩放。一个分层交叉注意力融合模块在标准池化器/分类器（序列分类）或每个选项评分器（多项选择）头部之前，合并两个流的表示。图1总结了这一数据流。

图1：SEWN核心架构。门控分数g将每个token的嵌入分配给廉价的功能词流和昂贵的、BERT初始化的内容流，然后再进行融合。
### III-B SEWN-sparse：将路由转化为计算声明

上述基础架构仅对token进行重加权；两个流仍然处理每个token，因此实际上没有节省FLOPs。SEWN-sparse则不同，它根据门控推导的重要性硬选择top-k token（总是包括[CLS] token），并且仅将这k个token输入昂贵的内容流；功能流仍然看到完整序列（它设计上就是廉价的，并且保持全长可以保留硬剪枝否则会丢弃的上下文）。融合被重构，使得k长度的内容表示通过一个交叉注意力块查询全长的功能表示，从而在昂贵的流中无需支付O(n²)注意力成本的情况下恢复必要的上下文。图2说明了这一选择过程。

图2：SEWN-sparse token选择。阴影token被门控选中并到达内容流；所有token（无论是否选中）仍然到达廉价的功能流。
### III-C SEWN-lean：移除第二个流

SEWN-lean完全移除功能流和融合模块，仅留下一个内容流，其输入按门控的(1−gate)权重缩放——即门控现在充当模型中唯一的逐token机制，没有第二个流来保留其抑制内容的完整副本。

### III-D SEWN-random-topk（仅用于消融）

架构上与SEWN-sparse相同，唯一区别是输入到内容流的k个token在每次前向传播时是均匀随机抽取的，而非由门控排序。专门用于隔离SEWN-sparse的忠实性是源于学习到的排序，还是仅仅源于硬选择的瓶颈（第V-G节）。

## IV 实验设置

**基线模型。** bert-base-uncased[1]（1.095亿参数）和distilbert-base-uncased[2]（6700万参数），使用单一平坦学习率进行微调。一个深度匹配的BERT-4L（BERT-base的前四个编码器层，5280万参数）用于隔离SEWN的双流设计是否优于相同深度的简单浅层主干网络。

**数据集。** 表I总结了所使用的六个数据集。

表I：实验所用数据集。
**训练方案。** 3个epoch，批大小16，AdamW优化器，基线模型使用平坦学习率2×10⁻⁵；SEWN变体还使用差异化学习率（对随机初始化的子模块：门控、功能流、融合模块、任务头使用8倍），该学习率通过在保留任务上进行学习率倍数验证扫描后固定，用于所有后续运行。除非另有说明，所有报告的结果平均3个随机种子（42, 123, 456）；每个均值旁都报告了每次运行的标准差。

**硬件。** 单块NVIDIA RTX 4090（24GB显存），P