DTRNet让多数令牌跳过昂贵的注意力计算,仍保持高性能。
DTRNet: Dynamic Token Routing Network to Reduce Quadratic Costs in Transformers
- 动态路由机制让90%令牌绕过二次复杂度注意力
- 每层仅10%令牌走完整注意力,性能接近全量Transformer
- 适合长序列任务,效率随序列长度提升而显著增强
Transformers在众多任务中表现卓越,但其在每一层对每个令牌均执行二次复杂度的自注意力计算,导致计算成本高昂。本文提出DTRNet(动态令牌路由网络),一种改进的Transformer架构,使令牌能够动态跳过跨令牌混合的二次计算成本,同时接收轻量级线性更新。通过保留MLP模块并仅将大多数令牌的注意力成本降至线性,DTRNet确保每个令牌都得到显式更新,同时大幅降低整体计算量。训练完成后,DTRNet在每层仅约10%的令牌通过注意力路径,仍能保持与完整Transformer相当的性能。在相同FLOPs下,其准确率和内存表现持续优于基于路由的层跳过方法(如MoD和D-LLM),且更少令牌进入完整注意力路径。该方法的效率增益随序列长度增加而提升,为长上下文输入带来显著的FLOPs减少。通过解耦令牌更新与注意力混合,DTRNet大幅降低了计算中的二次项占比,提供了一种简单、高效且可扩展的Transformer替代方案。
原文摘要 · Abstract (English)
Transformers achieve state-of-the-art results across many tasks, but their uniform application of quadratic self-attention to every token at every layer makes them computationally expensive. We introduce DTRNet (Dynamic Token Routing Network), an improved Transformer architecture that allows tokens to dynamically skip the quadratic cost of cross-token mixing while still receiving lightweight linear updates. By preserving the MLP module and reducing the attention cost for most tokens to linear, DTRNet ensures that every token is explicitly updated while significantly lowering overall computation. This design offers an efficient and effective alternative to standard dense attention. Once trained, DTRNet blocks routes only ~10% of tokens through attention at each layer while maintaining performance comparable to a full Transformer. It consistently outperforms routing-based layer skipping methods such as MoD and D-LLM in both accuracy and memory at matched FLOPs, while routing fewer tokens to full attention. Its efficiency gains, scales with sequence length, offering significant reduction in FLOPs for long-context inputs. By decoupling token updates from attention mixing, DTRNet substantially reduces the quadratic share of computation, providing a simple, efficient, and scalable alternative to Transformers.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。