详细解读 Transformer的即插即用模块 | MoE插件让ViT模型更宽、更快、精度更高 - 极术社区

本文提出Transformer更宽而不是更深，以实现更高效的参数部署，并将此框架实现为WideNet。首先通过在Transformer块之间共享参数来压缩可训练参数和深度，并用MoE层替换了FFN层。实验表明，WideNet通过较少的可训练参数实现了最佳性能，优于ViT等网络。
作者单位：新加坡国立大学

1简介

Transformer最近在各种任务上取得了令人瞩目的成果。为了进一步提高Transformer的有效性和效率，现有工作中有2种思路：

通过扩展到更多可训练的参数来扩大范围；
通过参数共享或模型压缩随深度变浅。

然而，当可供训练的Token较少时，较大的模型通常无法很好地扩展，而当模型非常大时，则需要更高的并行性。由于表征能力的损失，与原始Transformer模型相比，较小的模型通常会获得较差的性能。

在本文中，为了用更少的可训练参数获得更好的性能，作者提出了一个通过更宽的模型框架来有效地部署可训练参数。特别地，作者通过用MoE替换前馈网络（FFN）来沿模型宽度进行缩放。然后，使用单个层归一化跨Transformer Block共享MoE层。这种部署起到了转换各种语义表征的作用，这使得模型参数更加高效和有效。

为了评估所提的框架，作者设计了WideNet并在ImageNet-1K上对其进行评估。最好的模型在0.72倍可训练参数下的性能比ViT高1.46%。使用0.46×和0.13×参数的WideNet仍然可以分别超过ViT和ViT-MoE0.83%和2.08%。

本文主要贡献

为了提高参数效率，提出了跨Transformer Block的共享MoE层。MoE层可以在不同的Transformer Block中接收到不同的Token表示，这使得每个MoE都能得到充分的训练；
在Transformer Block之间保持单独的标准化层。单独的层具有少量额外的可训练参数可以将输入隐藏向量转换为其他语义。然后，将不同的输入输入到同一Attention层或更强的MoE层，以建模不同的语义信息。
结合以上2种思路，作者提出了更广、更少参数、更有效的框架。然后将用这个框架构建了WideNet，并在ImageNet上对其进行评估。WideNet在可训练参数少得多的情况下，性能大大优于Baseline。

2Methodology

本文研究了一种新的可训练参数部署框架，并在Transformer上实现了该框架。总体结构如图所示。在这个例子中使用Vision Transformer作为Backbone，这意味着在Attention或FFN层之前进行规范化。

WideNet可以很容易地扩展到其他Transformer模型。在WideNet中用MoE层代替FFN层。采用跨Transformer Block的参数共享，以实现更有效的参数部署。在每个MoE层中有一个路由器来选择K个专家来学习更复杂的表示。请注意，为了更多样化的语义表示，层标准化中的可训练参数不是共享的。

2.1 MoE条件计算

2.2 Routing

2.3 Balanced Loading

在基于MoE的Transformer中将每个Token分派给K个专家。在训练期间，如果MoE模型没有规则，大多数Token可能会被分派给一小部分专家。这种不平衡的分配会降低MoE模型的吞吐量。

此外，更重要的是，大多数附加的可训练参数没有得到充分的训练，使得稀疏条件模型在缩放时无法超越相应的稠密模型。因此，为了平衡加载需要避免2件事:

分配给单个专家的Token太多，
单个专家收到的Token太少。

为了解决第1个问题，需要缓冲容量B。也就是说，对于每个专家，最多只保留B个Token，不管分派给该专家多少Token。如果分配了超过B个Token，那么左边的Token将被丢弃:

2.4 Cross Transformer blocks共享MoE

2.5 Individual Layer Normalization

虽然现有的工作表明不同Transformer blocks的激活是相似的，但余弦距离仍然远远大于零。因此，不同于现有的作品在Transformer blocks之间共享所有权重，为了鼓励不同块更多样化的输入表示，这里只共享multi-head attention layer和FFN(或MoE)层，这意味着层标准化的可训练参数在块之间是不同的。

综上所述，框架中的第个Transformer blocks可以写成: