注意力即所需：用于深度混合路由的注意力机制
计算机视觉与模式识别
2024-12-31 v1
摘要
深度学习的进展依赖于训练参数规模日益增长的模型，而这又进一步加剧了计算需求。为应对这一问题，深度混合（Mixture-of-Depths, MoD）模型被提出，其通过仅将计算动态分配给输入中最相关的部分，从而能够以高效率部署与训练大参数模型。这些 MoD 模型采用路由机制来决定哪些 token 应由某一层处理，哪些应被跳过。然而，传统 MoD 模型专门使用额外的网络层进行路由，这些层难以训练，并增加了模型的复杂性与部署开销。在本文中，我们提出了一种新颖的基于注意力的路由机制 A-MoD，它利用前一层的现有注意力图来为当前层做出路由决策。与标准路由相比，A-MoD 训练更为高效，因为它不引入任何额外的可训练参数，并且可以轻松地从预训练 Transformer 模型改造而来。此外，它还能提升 MoD 模型的性能。例如，与标准路由以及等 FLOPs 的 ViT 基线相比，我们在 ImageNet 上观察到准确率最高提升 2%。此外，A-MoD 改善了 MoD 训练的收敛性，使迁移学习速度最高提升 2 倍。
引用
@article{arxiv.2412.20875,
title = {Attention Is All You Need For Mixture-of-Depths Routing},
author = {Advait Gadhikar and Souptik Kumar Majumdar and Niclas Popp and Piyapat Saranrittichai and Martin Rapp and Lukas Schott},
journal= {arXiv preprint arXiv:2412.20875},
year = {2024}
}
备注
22 pages, 19 figures