🤖 machine learning

Adaptive Computation Depth via Learned Token Routing in Transformers

本文介绍了 Token 选择性注意力（TSA），这是一种轻量级、端到端可微的门控机制，它使 Transformer 能够根据上下文难度动态跳过各个 Token 的冗余层计算，在实现显著效率提升的同时仅造成极小的质量损失，且无需显式的深度正则化。

原作者： Ahmed Abdelmuniem Abdalla Mohammed

发布于 2026-05-08

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Ahmed Abdelmuniem Abdalla Mohammed

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象你经营着一家工厂，传送带上流下的每一件产品都接受完全相同的处理。无论是简单完美的工件，还是复杂破损的工件，每一件物品在每个工序站停留的时间都完全相同。这正是当前标准人工智能模型（Transformer）的工作方式：无论某个词是容易理解还是难以理解，它们都会通过完全相同数量的“思考层”进行处理。

本文介绍了一种名为**令牌选择注意力（Token-Selective Attention, TSA）**的新系统。可以将 TSA 想象成在工厂的每一个工序站都安装了一位智能的自动守门员。

问题：“一刀切”的工厂

在标准人工智能模型中，如果你让它写下“生存还是毁灭（To be, or not to be）”，模型会以相同的脑力资源处理"To"这个词和"question"这个词。

简单的词（如"the"或"is"）就像简单的工件。它们不需要太多处理。
困难的词（如复杂概念或罕见名称）就像破损的工件。它们需要额外的时间和关注来修复。

目前，工厂以与处理困难工件相同的强度来处理简单工件，从而浪费了能量。

解决方案：智能守门员（TSA）

作者在人工智能的每一层之间添加了一个微小、轻量级的“守门员”（一个小型神经网络）。这位守门员逐个查看每个词（令牌），并问道：“这个词需要进入下一个处理站，还是可以直接跳过？”

决策：守门员不仅仅回答“是”或“否”。它给出一个概率分数（就像调光开关）。如果一个词很简单，守门员可能会说：“你可以跳过下一步”，或者“你只需要完成一半的工作”。如果一个词很难，它则会说：“请完整走完所有流程。”
神奇之处：最棒的是，这位守门员能够自我学习。它不需要人类告诉它哪些词很难。它纯粹通过尝试最小化错误来学习。如果跳过某个特定词的某一步导致了错误，守门员就会学会下次保持该词处于激活状态。如果跳过它效果良好，它就会学会节省能量。

实际运作方式

本文在两个主要方面测试了这种方法：

简单逻辑谜题：当人工智能需要复制数字列表时，守门员意识到“这很简单”，并跳过了约 65% 的工作量。当它需要排序数字（这更难）时，它只跳过了约 27% 的工作量。它自然地推断出更困难的任务需要更多步骤。
撰写故事：当被要求模仿莎士比亚写作或总结维基百科文章时，人工智能节省了**14% 到 23%**的计算能力。
- 它学会了标点符号、空格和常见词汇很容易，可以快速处理。
- 它学会了独特的内容词汇需要完整的“工厂”级处理。

结果：更快且更智能

无质量损失：人工智能的写作表现与标准模型一样好，但它使用的能量（计算能力）显著减少。
优于“提前退出”：其他方法试图在人工智能感到“自信”时提前停止整个句子的处理。TSA 更聪明；它阻止单个词经历不必要的步骤，同时让困难的词继续处理。研究发现，在设定为节省相同能量的情况下，TSA 产生的结果优于这些旧方法。
真实速度：当研究人员实际在计算机上运行代码时，他们观察到了真实的速度提升（约快 2.3%），因为计算机确实无需为被跳过的词进行数学运算。

核心结论

本文提出了一种以最聪明的方式让人工智能模型“变懒”的方法。TSA 不再强迫每个词完成相同数量的工作，而是让人工智能实时决定哪些词很简单，哪些词需要额外帮助。这就像拥有一家工厂，其中简单产品飞速通过流水线，而困难产品则获得全套 VIP 待遇，从而在不牺牲质量的情况下节省时间和能量。

技术摘要：基于 Transformer 中学习到的令牌路由的自适应计算深度

问题陈述

标准 Transformer 架构对序列中的每个令牌应用固定数量的层，而不管该令牌的上下文难度或信息密度如何。这种设计优先考虑架构的简洁性，而非每个令牌的适应性，从而导致显著的计算低效。在实践中，可预测上下文中的常见令牌所需的处理远少于新结构中的罕见令牌，但两者在每次前向传播中均获得相同的计算资源。这种低效在推理规模下尤为昂贵，因为主导成本是对所有令牌在所有层进行完整的前向传播。现有的解决方案，如针对 RNN 的自适应计算时间（ACT）、通用 Transformer、深度混合（MoD）和早退（Early Exit）机制，试图解决这一问题，但通常依赖于硬选择、累积预算或事后置信度阈值，这些方法可能无法完全可微或并非最优。

方法论：令牌选择性注意力（TSA）

本文提出了令牌选择性注意力（TSA），这是一种在连续 Transformer 块之间的残差更新上引入每个令牌的学习门控机制。

架构与路由

TSA 在预归一化（pre-norm）仅解码器 Transformer 的每个块 $f_l$ （其中 $l = 0, \dots, L-2$ ）之后插入一个轻量级路由器 $r_l$ 。主干块 $f_0$ 始终无条件执行。

路由器设计：每个路由器是一个具有 Sigmoid 输出的两层多层感知机（MLP），为每个令牌生成连续的中止概率 $p_l \in (0, 1)$ 。隐藏层维度为 $d/4$ （向下取整至 16），为基础模型增加了约 1.7% 的参数开销。
门控更新：后续块 $f_{l+1}$ $f_{l + 1}$ 的残差更新被 $(1 - p_l)$ $(1 - p_{l})$ 缩放。
- 如果 $p_l = 0$ ，更新为标准更新。
- 如果 $p_l = 1$ ，状态保持不变，从而有效跳过该块。
- 这种插值是平滑的，确保该机制完全可微，无需直通估计器、Gumbel 采样或强化学习。

训练与正则化

在训练期间，所有层均完全执行；门控缩放残差更新但不跳过计算。总损失定义为 $L = L_{task} + L_{depth}$ 。

深度正则化：添加一项 $L_{depth} = \lambda \cdot \frac{1}{L-1} \sum (1 - p_l)$ ，以温和地鼓励早期中止。
涌现行为：值得注意的是，本文证明，即使在没有显式深度压力（ $\lambda = 0$ ）的情况下，路由机制也能仅由任务损失梯度驱动，学会跳过操作。

推理

在推理时，该机制切换到“稀疏”模式。 $p_l > 0.5$ 的令牌通过 gather/scatter 操作被排除在前馈网络（FFN）之外，而注意力保持密集以保留精确的键值语义。这带来了真实的 FLOPs 节省和挂钟时间加速。

主要贡献

可微令牌路由：一种简单、连续的软门控机制，作用于残差更新，允许每个令牌、每层的路由，而无需对基础 Transformer 进行架构更改。
梯度驱动的涌现：证据表明，与难度成比例的路由仅从任务损失梯度中自然涌现。即使没有深度正则化（ $\lambda=0$ ），路由器也能学会跳过约 20% 的令牌 - 层操作。
效率 - 质量权衡：在字符级语言建模上的验证显示，令牌 - 层操作（TLOps）节省了 14–23%，而质量损失小于 0.5%。
优于早退：在匹配的效率下，TSA 比早退基线实现了更低的验证损失，表明端到端学习的路由决策优于推理时的事后置信度阈值。
现实世界加速：通过稀疏推理，在商用硬件（Apple M1 Pro）上实现了真实的挂钟时间加速（高达 2.3%）。

实验结果

合成任务

在复制和排序任务（长度为 10 的序列）上：

复制任务：路由器学习到该任务是恒等映射，跳过了 54.9% 的 TLOps（ $\alpha = 0.341$ ），且无精度损失。
排序任务：路由器分配了更多计算资源，跳过了 22.5% 的 TLOps（ $\alpha = 0.730$ ），反映了更高的难度。

字符级语言建模

Tiny-Shakespeare：TSA 节省了 22.8% 的 TLOps（ $\alpha = 0.726$ ），验证损失仅增加了 0.006 nats（+0.4%）。
enwik8：在更多样化的语料库（维基百科）上，TSA 节省了 13.9% 的 TLOps（ $\alpha = 0.833$ ），且无质量下降（验证损失略微更好）。
鲁棒性：该机制对正则化超参数 $\lambda$ 在两个数量级（ $[0, 0.1]$ ）范围内具有鲁棒性。在 $\lambda=0.05$ 时，实现了 50.4% 的 TLOps 节省，且质量成本可忽略不计。

与早退的比较

在匹配的活动分数（ $\alpha \approx 0.726$ ）下，TSA 比早退基线实现了0.71% 更低的验证损失。这表明端到端学习路由决策比在推理时应用置信度阈值能产生更好的决策。

挂钟吞吐量

Soft-TSA（训练）：无论 $\alpha$ 如何，均增加了约 1% 的开销。
Sparse-TSA（推理）：在 Tiny-Shakespeare 上实现了2.3% 的加速（ $\alpha=0.726$ ），在 enwik8 上实现了盈亏平衡性能（ $\alpha=0.833$ ）。加速需要批量大小 $\ge 64$ 以克服 CPU-GPU 同步开销。

意义与主张

本文主张，TSA 提供了一种最小化、架构简单的解决方案，以解决固定深度 Transformer 的低效问题。其主要意义在于证明自适应计算深度可以直接从任务目标中学习，而无需关于难度的显式监督或复杂的强化学习信号。

作者强调，路由器充当“学习到的噪声门”，在残差贡献冗余时衰减更新。这使得模型能够根据令牌难度动态分配计算资源，从而实现显著的效率提升（在合成任务上高达 55%，在语言建模上为 14–23%），且质量权衡极小。这项工作表明，与离散选择或事后阈值相比，端到端可微路由是自适应深度的更有效策略，为在商用硬件上实现更高效的推理铺平了道路。

本文在规模方面保持谦逊，指出所有实验均在参数量约为 500 万至 600 万的模型上进行，而扩展到更大模型（1000 万至 1 亿+）的行为仍是一个未解之谜。此外，当前的加速依赖于批量大小，且注意力稀疏性尚未实现。