优化Transformer架构：揭示Mixture-of-Depths的优势
文摘
2024-11-15 23:05
北京
在深度学习领域，Transformer架构凭借其卓越的性能成为自然语言处理的中坚力量。然而，随着模型规模的不断扩大，计算资源的需求也随之增加，这不仅限制了模型的可扩展性，还提高了训练和部署的成本。在这样的背景下，Mixture-of-Depths（MoD）技术应运而生，为Transformer模型带来了全新的计算分配策略，旨在提升效率，同时保持或提升模型性能。这篇论文详细介绍了Mixture-of-Depths： https://arxiv.org/pdf/2404.02258![]()
MoD的核心在于动态地在序列的不同位置分配计算资源。与传统Transformer模型在每个token上均匀分配计算量不同，MoD允许模型根据需求在不同层和不同token上优化计算资源的分配。这种动态分配通过限制每层参与自注意力（self-attention）和多层感知机（MLP）计算的token数量来实现，这一限制是在训练过程中由网络学习得到的。MoD的一个显著特点是，尽管它允许动态计算分配，但仍然使用静态计算图。这意味着，与那些需要动态计算图的技术相比，MoD在硬件架构的兼容性上表现更佳，从而提高了硬件的利用率。通过这种方式，MoD在保持计算总量可预测的同时，实现了token级别的动态和上下文敏感的计算分配。在训练过程中，MoD模型展现出与基线性能相匹配的能力，同时在前向传播中所需的FLOPs（浮点运算次数）更少，使得模型在推理时可以更快地进行采样。实验表明，在等效FLOPs和训练时间下，MoD模型的速度比传统Transformer模型快50%以上。MoD的实现涉及几个关键步骤：首先，定义一个静态的计算预算，通过限制参与计算的token数量来减少每个块的计算量；其次，使用每个块的路由器为每个token分配一个标量权重，以表达该token参与计算的偏好；最后，通过选择权重最高的top-k个token来确定参与计算的token。MoD技术还可以与Mixture of Experts（MoE）模型相结合，形成MoDE模型。这种结合不仅能够进一步提升性能，还能在不同计算类型之间进行动态路由，例如将某些token路由到“记忆查找”或“工具使用”功能。Mixture-of-Depths技术为Transformer模型提供了一种新的计算分配策略，通过在模型深度和宽度之间动态分配计算资源，实现了在保持或提升性能的同时减少计算量。MoD不仅是一种提高模型效率的有效方法，也为未来深度学习模型的设计提供了新的思路。
