Learning to Skip the Middle Layers of Transformers

条件计算是一种提升Transformer效率的常用策略。现有方法通常针对个别模块（例如混合专家层），或彼此独立地跳过某些层。然而，可解释性研究表明，Transformer 的中间层表现出更高的冗余性，而较早的层则负责将信息聚合到各个token位置上。受这些观察的启发，我们提出了一种新颖的架构，能够从中间向外动态地跳过数量可变的若干层。具体来说，一个学习得到的门控机制根据输入决定是否绕过对称的中间若干层块，同时一个带门控的注意力机制阻止后续token关注被跳过的token位置。我们使用“三明治”或“每层归一化”方案控制残差中的归一化，并通过自适应正则化损失来控制门控的稀疏性。我们的目标是减少对“简单”token的计算需求，并可能促进一种多层级表示结构的自然形成。但在我们测试的规模下，与层数更少的密集模型基线相比，我们的方法未能在验证交叉熵和估计FLOPs之间的权衡上取得改善。我们已公开代码，地址为：https://github.com/tim-lawson/skip-middle。

论文试图解决Transformer模型计算效率低的问题，特别是通过动态跳过中间层来减少处理“简单”token时的冗余计算。这是一个当前深度学习领域中备受关注的问题，因为大模型的计算成本越来越高。

提出了一种新的架构，基于输入动态决定是否跳过中间若干对称分布的Transformer块，并引入了门控注意力机制防止后续token关注被跳过的token位置。相比已有方法（如MoE或独立跳层），该方法更注重不同层之间的依赖关系和信息流动，具有结构设计上的创新性。

{采用学习到的门控机制控制层跳跃，并通过'夹心'或'每层归一化'策略管理残差归一化。,使用自适应正则化损失控制门控稀疏性，以优化计算资源分配。,实验目标是验证是否能在保持性能的同时降低FLOPs，但目前在所测试规模下未优于基线模型。,代码已开源（https://github.com/tim-lawson/skip-middle），便于复现与进一步研究。}

{"Shazeer et al., 'Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer' (2017)","Lepikhin et al., 'GShard: Scaling Giant Models with Conditional Computation and Automatic Partitioning' (2020)","Fedus et al., 'Switch Transformers: Scaling to Trillion Parameter Models via Sparse Activation' (2022)","Liu et al., 'Deep Gradient Compression: Reducing the Communication Bandwidth for Distributed Training' (2018)","Xin et al., 'Dynamic Transformer: Learn to Predict Sequence Length' (2020)"}