Transformer模型在现代机器学习中无处不在，它们被广泛应用于语言模型、图像处理、强化学习等多个领域。然而，尽管Transformer模型取得了巨大的成功，但我们仍然需要思考是否存在更优的架构。
在这篇论文中，作者提出了一种名为Mixture-of-Experts Universal Transformer（MoEUT）的新型架构，该架构结合了Mixture-of-Experts（MoE）方法和Universal Transformer（UT）的优点。MoEUT旨在解决UT在参数-计算效率方面的根本问题，并提供更高效的计算和内存使用。
UT是一种具有共享层参数的Transformer架构，它通过在深度上引入循环来学习组合泛化。UT在处理组合问题方面具有优势，如逻辑推理任务，并在小规模语言建模和翻译任务中表现出色。然而，UT的共享层参数设计在参数-计算效率方面存在根本问题。
具体来说，UT的共享层参数设计会导致参数数量的显著减少，但同时也会增加计算和内存要求。为了补偿参数损失，简单地扩大层大小会使其计算资源要求变得不可行。因此，尽管UT具有潜在的优势，但在实际应用中，特别是在参数主导的任务（如现代语言建模）中，它们并不受欢迎。
为了解决UT的参数-计算效率问题，作者提出了MoEUT架构。MoEUT结合了MoE方法和UT的优点，并引入了两个新的创新：层分组和新颖的层归一化方案。
MoE Feedforward Blocks和MoE Self-Attention Layers: MoEUT使用MoE方法来参数化共享层Transformer的Feedforward块和自注意力层。具体来说，它使用σ-MoE方法来参数化Feedforward块，并使用SwitchHead方法来参数化自注意力层。这些MoE方法允许MoEUT在保持参数效率的同时，提供更高效的计算和内存使用。
Layer Grouping: MoEUT通过将多个具有非共享权重的层堆叠在一起来形成层组，从而解决参数-计算效率问题。这些层组在共享相同参数的意义上被递归地堆叠在一起，从而重新定义了组作为UT中的共享“层”。这种层分组方法可以减少每个σ-MoE中的专家数量，并增加总的注意力头数，从而提供更高效的计算和内存使用。
Peri-LayerNorm Scheme: MoEUT引入了一种新颖的层归一化方案，称为Peri-LayerNorm，以改善信号传播并解决UT中的参数-计算效率问题。Peri-LayerNorm方案不使用主数据路径中的层归一化，而是仅在线性层之前使用层归一化，这些线性层立即跟随sigmoid或softmax激活函数。这种设计可以解决残差增长问题，并提供高效的梯度流动路径。
作者在多个语言建模数据集上进行了实验，包括C4、SlimPajama和peS2o，以及代码生成任务“The Stack”。实验结果表明，MoEUT在所有测试规模上都显著优于标准Transformer，同时使用更少的计算和内存。
具体来说，MoEUT在C4数据集上的性能优于标准Transformer，并且差距随着规模的增加而增大。此外，MoEUT在代码生成任务“The Stack”上也表现出色，并且其零样本性能在多个下游任务（如BLiMP、Children’s Book Test、Lambada、HellaSwag、PIQA和ARC-E）上也得到了验证。
作者将MoEUT与另一个基于MoE的UT模型Sparse Universal Transformer（SUT）进行了比较。SUT也使用MoE层，但与MoEUT存在一些技术差异，如使用不同的专家选择方法和更大的专家大小。实验结果表明，MoEUT在性能和效率方面都显著优于SUT。
此外，作者还对MoEUT的层分组和层归一化方案进行了评估。实验结果表明，MoEUT的层分组方法可以提供更好的性能和效率，而Peri-LayerNorm方案在改善信号传播方面也表现出色。
尽管MoEUT在性能和效率方面都表现出色，但它仍然存在一些局限性。首先，MoEUT的当前实现使用Triton内核，该内核在训练速度方面存在问题。其次，MoEUT在处理大规模数据时可能需要更多的计算资源。
为了解决这些局限性，作者建议在未来工作中使用更优化的CUDA内核来实现MoEUT，并探索在更大规模数据上的训练。此外，作者还建议将MoEUT应用于其他组合设置，以进一步探索其在实际应用中的潜力。
