让Transformer按需计算，不同难易的词用不同层数，省算力不丢精度。
Adaptive Computation Depth via Learned Token Routing in Transformers
- 每个词独立决定是否跳过下一层，由轻量门控网络动态控制。
- 在字符级任务中节省14%-23%计算量，质量损失低于0.5%。
- 无需额外训练压力，路由策略可直接用于推理加速。
标准Transformer对所有输入词使用相同层数，无论其上下文难度。本文提出令牌选择性注意力（TSA），在连续Transformer块间引入每令牌的轻量级门控机制，基于残差更新进行学习。该门控为两层MLP，输出连续停止概率，支持端到端可微，仅增加1.7%参数量且不改变基础结构。值得注意的是，TSA在无深度正则化（λ=0）时，仅凭任务损失梯度即可使路由器跳过20%的令牌-层操作。在字符级语言建模任务中，于Tiny-Shakespeare与enwik8数据集上，TSA实现14%-23%的令牌层操作（TLOps）节省，且质量损失低于0.5%。在相同效率下，验证损失比早退出方法低0.7%；所学路由策略可直接应用于推理阶段稀疏执行，实现真实运行时间加速。
原文摘要 · Abstract (English)
Standard transformer architectures apply the same number of layers to every token regardless of contextual difficulty. We present Token-Selective Attention (TSA), a learned per-token gate on residual updates between consecutive transformer blocks. Each gate is a lightweight two-layer multi-layer perceptron (MLP) that produces a continuous halting probability, making the mechanism end-to-end differentiable with 1.7% parameter overhead and no changes to the base architecture. Notably, TSA learns difficulty-proportional routing without any explicit depth pressure: even at $λ=0$ (no depth regularisation), the task-loss gradient alone drives the router to skip 20% of token-layer operations. On character-level language modeling, TSA saved 14-23% of token-layer operations (TLOps) across Tiny-Shakespeare and enwik8 at <0.5% quality loss. At matched efficiency, TSA achieved 0.7% lower validation loss than early exit, and the learned routing transfers directly to inference-time sparse execution for real wall-clock speedup.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。