Mixture-of-Modules: Reinventing Transformers as Dynamic Assemblies of Modules

在Transformers中，是否总是需要从浅层到深层计算tokens呢？普通Transformers及其变体的持续成功表明了毫无疑问的“是”。然而，在本文中，我们尝试打破深度顺序约定，提出了一种新颖的架构，称为“模块混合”（MoM），它的动机是任何层都可以用于计算token，只要它具有所需的处理能力。MoM的构建始于由多头注意力和前馈网络定义的有限模块集，每个模块都有其独特的参数化。然后，两个路由器从集合中迭代地选择注意力模块和前馈模块来处理token。该选择在token的前向传递中动态地扩展计算图，最终组装成模块。我们展示了MoM不仅提供了一个统一的框架，用于Transformers及其众多变体，而且还提供了一种灵活且可学习的方法，用于减少Transformer参数化中的冗余。我们使用OpenWebText预训练了各种MoM。实证结果表明，不同参数计数的MoMs在GLUE和XSUM基准测试中始终优于普通Transformers。更有趣的是，在固定参数预算的情况下，MoM-large使计算图的深度比GPT-2-large增加了38％以上，从而在GLUE和XSUM上获得了1.4和1的绝对增益。另一方面，MoM-large还使深度减少了60％以上，同时每层涉及更多的模块，与GPT-2-large相比，TFLOPs减少了16％，内存使用减少了43％，同时保持了可比较的性能。

本文试图挑战Transformer中的深度顺序传递假设，探索一种新的模块化深度学习架构Mixture-of-Modules（MoM）

MoM架构通过动态选择不同的attention模块和前馈网络模块进行计算，从而消除了深度顺序传递假设，并提供了一种灵活可学习的方法来减少Transformer参数化中的冗余

实验结果表明，MoM架构在GLUE和XSUM基准测试中都优于vanilla Transformer，MoM-large相比GPT-2-large在相同参数预算下增加了38%的计算图深度，同时保持可比性能，MoM-large还使计算图深度降低了60％，同时每层涉及更多的模块，相比GPT-2-large减少了16％的TFLOPs和43％的内存使用

与该论文相关的研究包括Transformer及其变体以及其他模块化深度学习架构，如Sparse Transformers和Performer