Mixture-of-Depths: Dynamically allocating compute in transformer-based language models
David Raposo Sam Ritter Google DeepMind Blake Richards Google DeepMind McGill University、Mila Timothy Lillicrap Google DeepMind Peter Conway Humphreys Google DeepMind Adam Santoro
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 r
中文速览
大语言模型（Transformer）在处理每个词元（token）时消耗同等计算量，但实际上并非每个词元都需要同等"思考"——这种浪费既拖慢训练又拖慢推理。本文提出"混合深度"方法（Mixture-of-Depths，MoD），让模型在每一层自主决定哪些词元真正需要经过自注意力和MLP计算，其余词元直接走残差连接跳过该层，从而把有限的计算预算集中花在最关键的位置。关键在于：每层参与计算的词元数量 $k$ 事先固定，因此计算图保持静态、张量尺寸已知，对硬件友好，不会引入动态图的工程麻烦。实验结果表明，MoD模型在同等训练FLOPs下能匹配甚至超过标准Transformer的性能，而推理阶段每次前向传播的FLOPs可减少50%以上、采样速度提升50%以上。这项工作说明Transformer本身存在大量可省去的冗余计算，通过让模型学会"按需分配算力"，可以在不损失质量的前提下显著降低训练和部署成本。
原文 arXiv:2404.02258;中英对照 + 大白话阅读 https://aha.fim.ai/paper/2404.02258v1