[Paper] 主观深度与时间尺度变换器：学习何时何地进行计算

Source: arXiv - 2511.21408v1

概览

Transformer 已成为现代 AI 的主力军，但它们“一刀切”的计算模式——每个 token 在每一层都要与所有其他 token 进行注意力计算——在处理长序列或超大模型时可能非常浪费。本文提出了两种新的 Transformer 变体，学习何时以及在哪里进行计算，在保持性能的同时削减不必要的工作量。

主要贡献

• 主观深度 Transformer（Subjective Depth Transformers，SDT） – 引入交替的 Decision（决策）层和 Dynamic（动态）层，利用贝叶斯惊讶（Bayesian surprise）来决定哪些 token 需要完整的 Transformer 块，哪些可以使用廉价的“先验”近似进行处理。
• 主观时间尺度 Transformer（Subjective Timescale Transformers，STT） – 将该思想扩展到时间维度，让路由器根据学习到的“变化假设”跳过或执行整个 Transformer 块。
• 贝叶斯惊讶信号（预期变化与意外变化）作为门控标准，提供了一种检测数据中新颖性与可预测性的原则性方法。
• 静态计算图 – 尽管存在动态路由，整体图保持静态，简化了在现有硬件和编译器堆栈上的部署。
• 效率提升 – 实验显示，在计算跳过层上 自注意力 FLOPs 减少最高可达 75 %，KV‑cache 使用量约减少 50 %，且仅带来轻微的精度下降。
• 学习动态的实证证据 – 模型在训练早期倾向于新颖性驱动的门控，后期转向预测驱动的门控，呼应了基于惊讶的处理理论预期。

方法论

1. 决策层（SDT） – 为每个 token 计算两套平行表示：

   • 后验（完整的 Transformer 块），捕获丰富的上下文。
   • 先验（轻量线性投影），作为廉价的回退。
     该层还会为每个 token 估计贝叶斯惊讶分数，衡量后验相对于先验的偏离程度。

2. 动态层（SDT） – 使用固定容量的 Top‑K 路由器，选取惊讶分数最高的 K 个 token 进行昂贵的后验计算，其余 token 使用先验。由于路由器的选择在给定分数时是确定性的，整体图保持静态。

3. 转移网络（STT） – 为每个 token 预测一个 残差更新，形成关于该 token 表征随时间变化的假设。

4. 时间路由器（STT） – 将预测的变化与实际变化（同样通过惊讶度量）进行比较。如果 token 被判定为“稳定”，路由器会在该时间步跳过 Transformer 块，复用缓存的 KV‑值；否则执行该块。

5. 训练 – 两种架构均采用端到端的标准语言建模目标进行训练。惊讶驱动的门控是可微的（通过直通估计器），使模型能够从数据中学习最优的路由策略。

结果与发现

• 惊讶动态：早期 epoch 对新颖 token 的门控活动较高；后期路由器趋于只在真正惊讶的输入上触发昂贵计算。
• 精度权衡：困惑度的轻微上升表明可以在不显著损害语言建模质量的前提下裁剪大量计算。
• 硬件友好性：由于计算图形状始终不变，模型在 GPU/TPU 上运行高效，无需自定义内核。

实际意义

• 成本效益推理：服务长文档（如法律合同、代码库）的部署可以通过跳过可预测片段的注意力来降低延迟和 GPU 内存。
• 可扩展训练：在普通硬件上训练大规模解码器模型变得更可行，因为每个批次在自注意力上的时间消耗更少。
• 边缘与移动 AI：静态图设计意味着模型可以使用现有工具链（TensorRT、ONNX Runtime）编译，并在资源受限设备上运行，同时仍能处理可变长度输入。
• 细粒度控制：开发者可以将惊讶阈值暴露为运行时参数，实时在速度与质量之间进行权衡（例如批处理时激进剪枝，交互式聊天时保守门控）。
• 自适应 API 基础：云服务提供商可以根据每次请求实际完成的工作量计费，使成本与模型努力程度相匹配。

局限性与未来工作

• 精度差距有限：虽然计算节省显著，但当前实现仍会导致几百分点的困惑度下降，在高风险应用中可能不可接受。
• 惊讶估计器开销：计算贝叶斯惊讶会带来一定的常数成本；如何通过量化等手段优化此步骤仍是待解问题。
• 向编码器‑解码器或多模态模型的推广：本文聚焦于解码器‑仅语言模型，如何将门控机制扩展到视觉‑语言或语音模型尚未探索。
• 动态硬件支持：尽管图保持静态，实际收益仍依赖于高效的 Top‑K 选择和缓存管理；与硬件调度器的更紧密集成可能进一步提升加速。
• 长期训练动态：从新颖性驱动向预测驱动的门控转变值得更深入的理论分析，并可能激发课程学习策略的创新。

结论：通过让 Transformer 自问“我真的需要计算这个吗？”并用贝叶斯惊讶信号作答，SDT 与 STT 为 更聪明、更廉价、更可适应 的深度学习模型开辟了有前景的道路——这对构建大规模 AI 系统的任何人都是令人振奋的进展。

作者

• Frederico Wieser
• Martin Benfeghoul
• Haitham Bou Ammar
• Jun Wang
• Zafeirios Fountas

论文信息

• arXiv ID: 2511.21408v1
• 分类: cs.LG, cs.AI, cs.CL, cs.IT
• 发布日期: 2025 年 11 月 26 日
• PDF: Download PDF