点击下方卡片，关注「集智书童」公众号
想要了解更多：
前沿AI视觉感知全栈知识👉「分类、检测、分割、关键点、车道线检测、3D视觉（分割、检测）、多模态、目标跟踪、NerF」
行业技术方案 👉「AI安防、AI医疗、AI自动驾驶」
AI模型部署落地实战 👉「CUDA、TensorRT、NCNN、OpenVINO、MNN、ONNXRuntime以及地平线框架」
欢迎扫描上方二维码，加入「 集智书童-知识星球 」，日常分享论文、学习笔记、问题解决方案、部署方案以及全栈式答疑，期待交流！
免责声明
凡本公众号注明“来源：XXX（非集智书童）”的作品，均转载自其它媒体，版权归原作者所有，如有侵权请联系我们删除，谢谢。
深度学习的发展驱动着训练模型参数数量的不断增加，这反过来又提高了计算需求。为了解决这个问题，混合深度（MoD）模型被提出，该模型能够动态地将计算仅分配给输入中最相关的部分，从而在推理和训练过程中以高效率部署大参数模型。
这些MoD模型利用路由机制来决定哪些 Token 应由某一层处理或跳过。然而，传统的MoD模型采用额外的网络层专门用于路由，这难以训练，并增加了模型的复杂性和部署开销。
在本文中，作者介绍了一种基于注意力的路由机制A-MoD，该机制利用前一层现有的注意力图来做出当前层的路由决策。
与标准路由相比，A-MoD允许更高效的训练，因为它没有引入额外的可训练参数，并且可以轻松地从预训练的transformer模型中调整过来。
此外，它还能提高MoD模型的表现。例如，作者发现与标准路由和isoFLOP ViT Baseline 相比，在ImageNet上提高了高达2%的准确率。
此外，A-MoD提高了MoD的训练收敛速度，从而实现了高达2倍的迁移学习速度。
- 引言
增加模型规模使得基于transformer的深度学习模型能够在多个领域实现最先进的性能，包括计算机视觉（Dosovitskiy等人，2021年）和自然语言处理 ——甚至解锁了新兴的能力。然而，这些大型模型的计算成本带来了显著挑战。因此，达到帕累托最优模型以最大化效率和性能至关重要（见图1）。
Jacobs等人（1991年）最初通过专家混合的方式引入了条件计算，为在保持浮点运算次数（FLOPs）不变的情况下增加模型大小奠定了基础，通过仅根据输入动态激活模型参数的子集，即所谓的“专家”。这一原则使得网络规模可以扩展到极其庞大，并在当前大语言模型（LLMs）的前沿得到了应用（Jiang等人，2024年）。
与Dosovitskiy等人；Wang等人；He等人（2016年）提出的标准深度学习模型相比，动态模型在研究关注度上较少，且在标准GPU架构的性能和运行时间Pareto前沿上通常还不具备竞争力。在此，作者专注于进一步推进动态计算领域的发展。
近期，Raposo等人（2024年）引入了混合深度（MoD）作为专家混合的一种变体。在MoD模型中，通过仅处理层中的一部分 Token 来动态降低计算成本，其余 Token 则跳过该层（见图2a）。与具有等效FLOPs的 Baseline 相比，MoD在语言任务上表现优异。MoD的一个关键组件是其路由器，它接收 Token 作为输入，并根据用户定义的容量确定哪些 Token 应进入或跳过一层。路由器通常由一个与模型共同训练的线性层组成（见图2b）。
路由机制由于以下几个原因对模型性能产生重大影响。首先，路由过程会在训练过程中引入噪声，因为路由是一个离散的决策，通常在多个层次和每个 Token 上进行操作。其次，路由器依赖于额外的层，因此，在将预训练的 Baseline 模型适应到MoD模型时，需要从头开始重新训练。最后，路由器为稀疏化模型增加了轻微的计算开销。
因此，在本篇论文中，作者提出了并解答了这样一个问题：能否在MoD模型中基于模型内部已有的信息来优化路由机制，而不是在路由器中引入额外的可训练参数？作者发现答案存在于常用Transformer架构（Vaswani等人，2017；Dosovitskiy等人，2021）中的注意力机制中。
作者假设注意力图可以用来估计一个 Token 的重要性，通过平均它与其它 Token 的交互。基于此，作者提出聚合注意力图中的信息，并将其用作MoD中 Token 路由的重要性度量。作者称作者的方法为MoD的注意力路由：A-MoD（见图2c）。作者发现，A-MoD可以在各种模型规模和任务中，持续优于MoD网络中的标准路由（如图1所示）。作者的A-MoD不仅参数免费，还可以几乎无需额外训练，将现成的预训练transformer模型适应到MoD中。作者进一步通过实证验证了作者的方法，并显示由A-MoD计算的路径得分与 Token 重要性估计的相关性比标准路由器计算的路由得分更好。
本文在将MoD应用于视觉领域的应用上取得了重要进展。作者的主要贡献包括：
-
作者发现MoD（多输出深度学习）不仅可行，而且在视觉任务中具有优势，提供了实证证据，表明它在FLOPs（浮点运算次数）和性能方面均优于传统模型。
-
作者引入了A-MoD，这是一种基于注意力图的参数自由路由方法，用于MoD的计算 Token 重要性。
-
作者证明A-MoD在多个数据集上优于标准路由器，在MoD的微调和迁移学习中表现出色。在迁移学习的案例中，A-MoD展示了从密集预训练模型迁移的MoD模型更快收敛。
-
与标准MoD相比，A-MoD始终选择重要的 Token ，并且路由决策与通过移除 Token 估计的留一法 Token 重要性相关。
-
相关研究工作
注意力机制（Bahdanau，2014）使模型能够通过恒定的操作数学习序列中的长距离依赖关系。Transformer（Vaswani等人，2017）在语言领域利用了注意力机制，并已成为事实上的标准模型。Dosovitskiy等人（2021）进一步将Transformer应用于视觉领域，将图像块视为 Token ，引入了视觉Transformer（ViT）。对于图像而言，注意力图已被证明会集中在关键区域，如图像中的物体。在本文中，作者利用这一特性在神经网络内部有效路由 Token 。此外，作者使用数据高效的图像Transformer DeiT-T和DeiT-S（Touvron等人，2021）而非普通的ViT-T和ViT-S模型，因为小型的ViT在较小数据集上训练时泛化能力不佳（Dosovitskiy等人，2021）。
自三十多年前问世以来，专家混合（MoE）已被应用于多种模型类型。Shazeer等人（2016年）将MoE引入到扩展 Transformer 架构中（Ludziejewski等人，2024年）。此后，MoE在视觉和语言任务中取得了广泛的实证成功（Puigcerver等人，2024年；Jain等人，2024年；Fedus等人，2022a；Riquelme等人，2021年）。在训练MoE网络时面临的一个主要挑战是训练不稳定（Zoph等人，2022年；Fedus等人，2022b）。Raposo等人（2024年）最近引入了深度混合（MoD）架构，其中每个 Transformer 块仅处理部分 Token ，与大型 Transformer 模型相比，实现了有利的计算性能权衡。Liang等人（2022年）也通过融合 Token 来诱导稀疏性，而无需完全跳过 Token 。在当前形式下，MoE和MoD都使用专门的路由网络来决定整个网络的哪些组件处理哪些 Token 。MoE和MoD之间的区别在于，MoE模型由多个独立的专家组成，这些专家独立处理 Token 。相比之下，MoD模型在每个层中仅选择两个专家，其中一个是该层本身，另一个是恒等函数，如图2所示。
对于大多数条件计算块（Cai等人，2024年），需要路由机制。在MoE模型的 Transformer 中，路由器的作用是将 Token 与专家匹配，以最大化性能。对于像Switch Transformers（Fedus等人，2022b）或MoD这样的单专家模型，路由器决定一个 Token 是否能够从专家的处理中受益，或者是否会被跳过。已经提出了各种方法（Liu等人，2024年），例如具有 Token 选择或专家选择路由（Zhou等人，2022年）的学习路由器，通过解决线性规划来匹配 Token 与专家（Lewis等人，2021年），通过哈希输入来匹配专家（Roller等人，2021年），或者使用强化学习来做出路由决策（Clark等人，2022年；Bengio等人，2015年；2013年）。显式地学习路由器是目前最先进的技术，在大多数情况下优于其他方法（Dikkala等人，2023年）。然而，这种方法主要在路由参数数量较多的情况下有效，并且容易产生训练不稳定性（Ramachandran和Le，2019年）。因此，训练出能持续产生强大性能的路由器仍然是一个未解决的问题。
作者的研究聚焦于提升MoD架构。作者提出了一种基于注意力图的创新路由机制，从而消除了对标准路由器的需求。 Token 以参数自由的方式进行路由，无需额外的计算开销。
- 方法
在本文中，作者将详细描述所采用的方法论，该方法论旨在通过综合实验和理论分析来探究人工智能在特定领域的应用效果。作者将首先介绍实验设计，随后阐述用于数据处理的算法，并最终讨论如何评估模型性能的指标。
在本节中，作者介绍了混合深度（MoD）架构，并提出了基于注意力的MoD路由算法，即A-MoD，该算法可用于提升其路由性能。
3.1 深度混合
作者的研究聚焦于视觉Transformer（Dosovitskiy等人，2021年；Touvron等人，2021年）。在此，给定一个由 Token X表示的输入，输出预测通过一个由L个Transformer块组成的模型计算得出，这些Transformer块由一组可学习的权重进行参数化。每个Transformer块包含一个具有H个头的多头自注意力（MHSA），其后跟随一个两层全连接网络，并使用GeLU激活函数（MLP）。
在MoD中，Raposo等人（2024）提出了一种基于transformer架构的变体，该变体基于这样一个假设：在模型中，单个 Token 需要不同数量的计算量。具体来说，MoD层只处理所选的重要 Token 子集，而其余 Token 则跳过这一层。从经验上看，这种方法可以在与vanilla ViT相当的计算预算下提升性能。
是否跳过某一层的决定由路由算法估计的token重要性分数决定。传统上，标准路由通过额外的层来计算这些重要性分数（参见第3.2节）。相比之下，作者的A-MoD算法直接从前一层的注意力图中计算分数，无需额外参数（参见第3.3节）。
3.2 标准路由
考虑单个MoD层，计算输入 Token 的重要性分数的标准方法需要额外的路由网络，如图2b所示。通常，路由器是一个将 Token 向量投影到表示其重要性分数的标量的线性层（如Raposo等人（2024年）所引入）。形式上，作者考虑由一组参数参数化的第个transformer层，其输入
代表长度为的 Token 序列。现在，作者可以估计 Token 的重要性分数如下：
在额外的线性路由网络中，是该网络的参数。这些 Token 将根据下述方程中的得分来决定是跳过还是处理：
在此， 表示所有 Token 重要性得分 的第 个百分位数。 可以用容量 来定义，即 ，其中 是MoD层的容量。为了在反向传播过程中学习 Token 重要性得分，将 Transformer 层的输出与重要性得分 相乘，以确保它可以接收非零梯度。
3.3 注意力路由
（请注意，由于“ATTENTIONROUTING”这一术语在AI领域可能有多重含义，我无法确定其确切含义。如果您能提供更多上下文或定义，我可以进行更准确的翻译。）
与标准路由不同，作者提出了，一种基于注意力计算路由分数的方法，无需额外的可训练参数。A-MoD利用前一层注意力图来确定当前MoD层的路由分数，如图2c所示。上一层的第个头的注意力图可按以下方式计算：Vaswani等人（2017年）提出：
和 分别是从前一层计算得到的 Query 矩阵和键矩阵，它们属于 ，其中 表示 Query 和键的嵌入维度。
.
在 中， 表示在计算第 个输出时，考虑了第 个 Token 的信息量。将 在所有行上聚合，可以得到第 个 Token 与其他所有 Token 的相关性度量。因此，在 中，作者提出通过平均所有行和注意力头上的对应注意力值来计算 Token 重要性分数，具体公式如下：
根据上述的分数计算，第层的输出可以计算如下：
作者注意到，在中，作者不将 Token 得分与输出相乘，因为注意力图已经在之前的层中可学习。这保留了原始的 Token 输出，有助于从标准预训练预训练权重进行适应时加快训练速度。作者还尝试了一种将相乘的变体，但这并未导致性能提升，因此为了简化而被移除。对于方程2中的标准路由，这个乘积项是正确计算路由参数梯度的必要条件。相比之下，-MoD移除了路由参数，从而使得MoDs的后置适应性调整更加容易，并消除了路由得分的训练不稳定性。
- 实验
4.1 训练设置与概述
在作者的实验中，作者对A-MoD进行了系统评估，并从实证上展示了其在MoD标准路由上的优势。作者跨越多种模型架构和多个图像分类任务进行了评估。在每次实验中，作者训练了一个MoD，该模型是从一个未经过训练的预训练 Transformer 模型中改编而来。作者在同一数据集上对改编的MoD模型进行了微调，并在不同的数据集上进行了迁移学习实验。
训练设置：作者评估了A-MoD在四种不同规模的视觉Transformer架构上的表现：DeiT-Tiny、DeiT-Small（Touvron等人，2021年）、ViT-Base和ViT-Large（Dosovitskiy等人，2021年）。每个MoD架构均来自ImageNet-1k（Russakovsky等人，2015年）上的预训练预训练权重。从该预训练权重开始，作者按照公式（5）的描述，以50%和12.5%的容量训练MoD模型，即每个MoD层分别处理50%和12.5%的token。遵循Raposo等人（2024年）的方法，作者在MoD架构中交替使用MoD层和密集层，即每两层中有一层是MoD层。作者还在第4.5节分析了仅在模型后期层放置MoD层的影响。
作者针对ImageNet-1k数据集对MoD模型进行微调。对于每个案例，作者将作者的A-MoD与标准路由进行比较。作者还比较了两种MoD变体与isoFLOP原始视觉Transformer。这个isoFLOP模型通过适当地减少原始模型中层的数量，以匹配其MoD对应物的FLOPs数量而获得。仅减少层仍然允许isoFLOP模型受益于预训练预训练权重的权重。每个模型都使用AdamW优化器（Loshchilov & Hutter, 2017）在100个epoch内进行训练，批大小为128，学习率为，采用线性预热后跟随余弦退火。作者通过在附录中的图12(a)、12(b)和13所示的操作中进行搜索后确定了这个学习率方案。
表1：A-MoD在ImageNet上，无论是在50%还是12.5%的容量下，相对于MoD和isoFLOP Baseline ，都表现出更好的性能。
为了进一步研究A-MoD的优势，作者在较小的斯坦福汽车、牛津宠物（Parkhi等，2012年）和Flowers102数据集上进行了迁移学习，以实现图像分类。在此，每个模型使用SGD（随机梯度下降）训练200个周期，批大小为64，学习率为0.01，并采用余弦退火法。
最后，作者对比了标准路由和A-MoD计算的路由分数，以及一个衡量每个token重要性的参考分数。这一分析使作者能进一步区分A-MoD的优势。作者采用留一法（Hastie等，2009年）来估计token的重要性。具体来说，作者测量在MoD层移除某个token时模型损失的变化。这使作者能为每个MoD层中的每个token分配一个参考重要性分数，针对每张输入图像。然后，作者将这些分数与每个MoD层和token的作者的基于注意力的路由和标准路由的权重进行关联。总的来说，A-MoD不仅选择了视觉上相关的token，而且路由分数也与留一法的token重要性有很强的相关性。
4.2A-MoDIM通过改进性能，提升了微调效果
在微调过程中，作者对每个MoD模型在ImageNet上进行训练。在所有作者考虑的视觉Transformer模型（参数量从5000万到3亿不等）中，A-MoD大多优于标准路由。具有50%和12.5%容量的MoD的结果展示在表1中。通过附录中图3（50%容量）和图9（12.5%容量）所示的训练曲线，作者强调A-MoD收敛速度更快。
对于DeiT-Tiny模型50%容量的情况（见图3(a)），A-MoD的表现优于MoD超过2%，在其他更大的模型上则超过1%。同样，对于12.5%的容量，A-MoD在DeiT-Tiny和小型模型上优于标准路由，在大容量变体中表现相当。尽管A-MoD在ViT-Base模型12.5%容量时略逊于标准路由，但如图9(c)（附录中的收敛图）所示，它需要更少的迭代次数才能收敛，并在第20次迭代时已达到峰值。总的来说，表1以及图3中的训练曲线证实，A-MoD可以优于使用标准路由以及isoFLOP Baseline 的MoDs。具体来说，A-MoD在较小的DeiT-Tiny和DeiT-Small上性能提升更大，并且能够使所有模型更快地收敛。
根据公式（5）所述，A-MoD可以仅基于注意力图来计算路由分数，它不会将每个MoD块输出与路由分数相乘，因此基本上保留了 Token 输出。这两个特性使得从预训练预训练权重通过注意力路由微调的A-MoD能够在最少的训练下收敛。图3展示了A-MoD能够实现更快的收敛速度，与标准路由相比，大幅减少了所需的训练时间。在某些情况下，A-MoD无需训练即可达到合理的准确率。
图3(c)中，A-MoD未经过任何训练即达到了78%的准确率。附录中的表3报告了所有从预训练预训练权重调整而来的MoD的准确率，并突显出A-MoD总是从比标准路由更高的准确率开始。这是可能的，因为模型使用已经学习到的注意力图来估计最不重要的 Token ，使得最终准确率受到的影响最小，这一点在第4.4节中得到进一步证实。相比之下，标准路由将层输出乘以路由分数，并且需要从头开始学习路由，因为它基于额外的层。这些因素导致了收敛速度较慢。
在附录的图11中，作者对比了A-MoD与一种修改方案，该方案将MoD模块的输出与注意力路由得分相乘，以验证A-MoD是否能够从额外的学习梯度中获益，类似于标准的路由，即使用公式（2）而非公式（5）。然而，对于A-MoD而言，将路由得分乘以输出反而降低了适应后的MoD模型的准确性，并且在没有任何训练的情况下略微减缓了收敛速度。
学习率稳定性分析为了研究作者针对A-MoD和MoD的训练在学习率方面的稳定性，作者对不同的学习率进行了全面测试，并跟踪了性能表现。作者发现，在所有测试的单个学习率中，A-MoD的表现都优于MoD，如图13所示。
4.3 基于A-MoD的迁移学习加速收敛
作者现在研究了A-MoD在从ImageNet-1k迁移学习任务到三个较小的图像分类数据集（OxfordIIT-Pets、Stanford Cars和Flower102）中的应用。这些任务具有挑战性，因为预训练模型必须在有限的数据上进行训练，同时适应具有降低容量的MoD架构。图4展示了A-MoD与Flower102数据集上的MoD在准确率曲线上的比较。Stanford Cars和OxfordIIT-Pets数据集的结果见附录图10。
在所有数据集和模型架构中，作者发现A-MoD相较于标准路由收敛速度更快，并且在大多数情况下都超越了标准路由。作者通过测量模型达到峰值精度范围内（误差在2%以内）所需的训练轮数来分析收敛性。图4和图10中的黑色虚线有助于可视化A-MoD和标准路由的收敛情况。对于在Flowers数据集上使用的ViT-Large模型（见图4(d)），A-MoD在第35轮训练时达到了94.5%的精度，而标准路由则需要第100轮才能达到相同的精度。同样，在DeiT-S模型在Pets数据集上的应用中（见附录中的图10(f)），A-MoD在第25轮训练时达到了90%的精度，而标准路由则需要70轮才能达到相同的精度，实现了约2倍的速度提升。这些观察结果在架构和数据集上都是一致的，突显了A-MoD在将MoD模型从预训练预训练权重迁移方面的有效性。
图5：与MoD相比，A-MoD展示了更具意义的路由。路由可视化：DeiT-Small在ImageNet上的50%容量示例。每个示例展示了每个MoD层中由标准MoD（顶部）和A-MoD（底部）选择的 Token ，白色区域表示跳过的部分。每列代表一个MoD层，深度从左到右增加。
作者在附录的表4和表5中，将A-MoD与标准路由方法在迁移学习任务中与isoFLOP模型进行了比较，分别针对50%和12.5%的容量。作者发现MoD模型在迁移任务中无法与isoFLOP模型的表现相匹配。作者将这一点视为MoD框架在图像任务迁移学习中的局限性，无论采用何种路由机制。在4.5节中，作者提出了一种可能的解决方案，以超越isoFLOP模型。
4.4 注意力路由识别关键 Token
为了理解为什么A-MoD比标准路由有所改进，作者研究了路由得分及其与留一法（Hastie等人，2009） Token 重要性的相关性。作者的目标是估计一个 Token 的重要性与其由标准或A-MoD路由器分配的路由得分之间的关系。基于作者的实证结果，作者推测A-MoD的权重与 Token 重要性相比标准路由有更好的相关性，从而使得A-MoD能够始终选择最相关的 Token 。
首先，作者通过可视化ImageNet1k中的单个示例的路由来验证这一说法，如图5所示。该图突出了每个MoD层中路由器选择的图像 Patch 。对于A-MoD（底部），路由器从第三层MoD开始选择构成鸟轮廓和面部的 Token 。相比之下，标准路由（顶部）在最后几层中选择了更多构成背景的 Token 。
可视化图6中最后一层的注意力图也证实了A-MoD能够专注于图像中的物体，作者将此作为路由分数。对于DeiT-Small，最后一层每个头的注意力图识别出了A-MoD所关注的鸟的轮廓，但在标准路由方面却显得有些吃力。然而，正如Darcet等人（2024年）所展示的，注意力图并不总是学习到具有语义意义的分数。这一点在较大的模型中尤为明显，注意力分数往往集中在单个块（ Token ）上（参见附录中的图15）。
图6：A-MoD生成的注意力图比MoD更具意义。注意力可视化：DeiT-Small在ImageNet上使用50%容量的示例。每个示例中，顶部行为标准路由的最后一个MoD层的注意力图，底部行为A-MoD的注意力图。每一列代表最后一层的注意力头。
图7：A-MoD在路由得分与leave-one-out Token 重要性之间显示出更高的相关性。ImageNet中路由得分与分层leave-one-out Token 重要性的相关系数和p值。
为了量化作者的定性观察结果，作者计算了路由得分与 Token 重要性估计之间的相关性。对于一个 Token 的重要性，作者计算了在原始的Transformer中省略该 Token 后模型损失的变动，即“留一法” Token 重要性。损失的大幅度变化意味着 Token 的重要性较高，作者预期该 Token 会有一个更高的路由得分。图7展示了标准路由和A-MoD的路由得分与 Token 重要性的相关性，以及相应的p值。
作者观察到，A-MoD计算的路由得分与 Token 重要性始终保持高度相关性，这表明注意力路由会对重要的 Token 分配更高的分数。相比之下，标准路由有时甚至与 Token 重要性呈负相关，这意味着它可能会对不太重要的 Token 分配更高的分数。此外，观察到的所有A-MoD的p值均低于，而在标准路由的情况下，它们（在某些层中甚至大于0.5）具有显著性，这暗示了在标准路由的情况下存在更高的不确定性。
4.5 模拟层差异的影响
在迄今为止的实验中，作者采用了与Raposo等人（2024年）类似的方法，在交替层中使用MoD层。这种模型架构在ImageNet-1k上为作者提供了Pareto最优的结果（见表1）。作者进行了一项消融研究，以探究仅在较深层引入MoD层而保持初始层密集是否具有优势，特别是在视觉任务如分类中，学习低级特征可能至关重要。为了验证MoD层是否能够从早期层中额外学习特征中受益，作者从第4层开始交替引入MoD层，同时保持前四层保持密集。
图8的结果显示，保持前四层密集连接可以提升在斯坦福汽车数据集上训练的DeiT-Small和ViT-Base模型的效果。额外的浮点运算次数（FLOPs）使得在该训练模式下能够获得更好的学习效果，如图8所示。经过这一修改，A-MoD能够匹配相应的isoFLOP Baseline ，即使是对于迁移学习任务也是如此。这突显了一种解决第4.3节中提到的A-MoD局限性的潜在方法，虽然代价是增加了额外的FLOPs。
图8：仅在使用深层网络时，A-MoD提升了性能。仅在最后8层引入MoDs，在斯坦福汽车数据集上与isoFLOP性能相当。
作者提出了一种名为A-MoD的混合深度（MoD）变体，它使用注意力路由而非标准路由。
为了计算MoD层的 Token 重要性，A-MoD利用其前一层的注意力图，从而在不增加额外参数的情况下实现注意力路由。
在从预训练预训练权重进行训练的情况下，利用训练好的注意力信息也能导致比传统MoD更高的训练稳定性和更快的收敛速度。此外，作者通过实证研究表明，A-MoD在不同模型配置和数据集上均优于标准MoD，并且能够做出更好的路由决策。
扫码加入👉「集智书童」交流群
（备注： 方向+学校/公司+昵称 ）
点击下方“ 阅读原文 ”，
了解更多AI学习路上的 「武功秘籍」