在标准的Transformer架构中，每个Token都会经过相同数量的层处理，无论其上下文难度如何。这种均匀处理方式导致计算资源浪费在简单的Token上。针对这一问题，一篇来自arXiv的新论文提出了Token-Selective Attention（TSA），一种通过习得每Token门控来实现自适应计算深度的方法。
TSA的核心是在连续的两个Transformer块之间引入一个轻量级的门控机制。该门控是一个两层多层感知器（MLP），它为一个残差更新产生连续的停止概率。这使得整个机制端到端可微，且仅增加1.7%的参数，无需改变基础架构。值得注意的是，即使在没有任何深度正则化的情况下（λ=0），仅靠任务损失梯度驱动，TSA也能学习到跳过约20%的Token-层操作，表明其学会了与难度成比例的路由。
在字符级语言建模任务上，TSA在Tiny-Shakespeare和enwik8数据集上节省了14-23%的Token-层操作（TLOps），同时质量损失低于0.5%。在与早期退出方法匹配效率的条件下，TSA的验证损失降低了0.7%。更重要的是，学习到的路由可以直接迁移到推理时的稀疏执行，从而实现实际的壁钟时间加速效果。
TSA的门控机制不仅轻量，而且具有可解释性。研究者发现，对于简单的、信息量低的Token，门控倾向于输出较高的停止概率，从而跳过后续层的计算；而对于复杂的、信息量高的Token，门控则允许其通过所有层。这种自适应行为完全由任务损失驱动，无需人为设定阈值。此外，TSA与现有的Transformer变体兼容，可以轻松集成到预训练模型中。实验还表明，TSA的路由决策在训练后可以固定下来，用于推理时的稀疏计算，从而显著提升吞吐量。未来工作可以探索将TSA扩展到其他序列建模任务，如机器翻译和文本生成，进一步验证其通用性。该方法的提出者Ahmed Abdelmuniem Abdalla Mohammed在论文中详细阐述了TSA的设计和实验。该工作为Transformer的高效推理提供了新思路，尤其适合需要动态计算资源的场景。通过简单的门控机制，TSA在几乎不牺牲质量的前提下大幅降低了计算成本，有望在大型语言模型的部署中发挥重要作用。