[论文解读] Mixture-of-Depths: Dynamically allocating compute in transformer-based language models

David Raposo, Sam Ritter+4 more|arXiv (Cornell University)|Apr 2, 2024

Natural Language Processing Techniques被引用 4

一句话总结

本文提出Mixture-of-Depths（MoD），一种通过学习的top-k路由机制动态分配Transformer语言模型计算量的方法，允许令牌跳过层，从而限制每层的总FLOPs。MoD模型在前向传播中FLOPs减少高达50%，推理速度提升高达50%，同时保持训练效率和硬件兼容性，得益于静态计算图。

ABSTRACT

Transformer-based language models spread FLOPs uniformly across input sequences. In this work we demonstrate that transformers can instead learn to dynamically allocate FLOPs (or compute) to specific positions in a sequence, optimising the allocation along the sequence for different layers across the model depth. Our method enforces a total compute budget by capping the number of tokens ($k$) that can participate in the self-attention and MLP computations at a given layer. The tokens to be processed are determined by the network using a top-$k$ routing mechanism. Since $k$ is defined a priori, this simple procedure uses a static computation graph with known tensor sizes, unlike other conditional computation techniques. Nevertheless, since the identities of the $k$ tokens are fluid, this method can expend FLOPs non-uniformly across the time and model depth dimensions. Thus, compute expenditure is entirely predictable in sum total, but dynamic and context-sensitive at the token-level. Not only do models trained in this way learn to dynamically allocate compute, they do so efficiently. These models match baseline performance for equivalent FLOPS and wall-clock times to train, but require a fraction of the FLOPs per forward pass, and can be upwards of 50\% faster to step during post-training sampling.

研究动机与目标

• 为解决Transformer中计算分布均匀的低效问题，即所有令牌无论计算需求如何均接收相同FLOPs。
• 在不引入动态计算图的前提下，实现在令牌和层维度上的动态、上下文敏感的计算分配。
• 通过智能地将令牌路由至模型中不同深度，实现每前向传播FLOPs减少，同时保持硬件效率和训练速度。
• 探索学习路由是否能在不牺牲训练时间或准确率的前提下提升模型效率和性能。
• 证明可在推理时节省计算量而不造成性能下降，从而实现更快推理或模型扩展。

提出的方法

• 通过路由器头选择每层中参与自注意力和MLP计算的令牌，采用固定k值的top-k路由机制。
• 未被路由器选中的令牌通过残差连接传递，有效跳过Transformer块计算，从而节省FLOPs。
• 该方法强制使用静态计算图，张量大小已知，确保硬件效率和与现有加速器的兼容性。
• 路由决策在训练期间端到端学习，无需辅助平衡损失，利用top-k路由实现稳定性。
• 在自回归采样过程中，通过简单辅助分类器或损失函数蒸馏路由器，使其在无未来令牌信息的情况下模仿训练时的路由决策。
• 该方法可通过调整每层的计算预算（k）实现速度与性能之间的权衡，支持更快推理或模型扩展。

Figure 1: Mixture-of-Depths Transformer. As in mixture-of-experts (MoE) transformers we use a router to choose among potential computational paths. But unlike in MoE transformers the possible choices are a standard block’s computation (i.e., self-attention and MLP) or a residual connection. Since so

实验结果

研究问题

• RQ1 Transformer能否在保持静态计算图的前提下，学习在令牌和层之间动态分配计算？
• RQ2 通过路由实现的动态计算分配是否能在不造成性能下降的情况下显著节省FLOPs？
• RQ3 在等同的训练FLOPs预算下，MoD模型是否能实现比标准Transformer更好的性能或更快的推理速度？
• RQ4 推理期间缺乏未来令牌信息是否会影响路由性能，能否缓解？
• RQ5 该路由机制能否扩展以支持标准模块和残差连接之外的其他计算类型？

主要发现

• MoD模型在相同FLOPs和实际运行时间下训练，性能与基线Transformer相当或更优，证明了更高的计算效率。
• 与标准Transformer相比，MoD模型的前向传播FLOPs减少高达50%，实现无性能损失的更快推理。
• 即使在训练FLOPs和时间保持不变的情况下，MoD模型在微调后采样过程中推理步骤速度提升高达50%。
• 训练期间的top-k路由机制在自回归推理中泛化良好，即使在无未来令牌身份信息的情况下，性能下降也极小。
• MoD模型学会智能路由，对简单预测跳过不必要的计算，同时通过战略性注意力路由保留长距离依赖。
• 该框架具有可扩展性，可支持多种计算类型（如内存查找、工具调用）的路由，从而实现具有可控FLOP成本的新架构设计。

Figure 2: Routing schemes. Tokens are funnelled to the computational path of their choice when using token-choice routing (left). If a given path exceeds its capacity (e.g., more than two tokens in this example) then surplus tokens must be dropped (purple token). The exact token that is ultimately d