让Transformer按需计算，快21%且损失仅0.2%。
Router-Tuning: A Simple and Effective Approach for Enabling Dynamic-Depth in Transformers
- 只微调路由器，无需重训整个模型，大幅降成本。
- 动态跳过低重要层，实现21%推理加速，性能仅降0.2%。
- 适合追求高效推理的NLP应用，尤其在资源受限场景。
传统Transformer对每个输入词元分配固定计算量，造成资源浪费。为解决此问题，提出动态深度机制Mixture of Depths（MoD），通过跳过不重要层来调整计算深度。现有方法面临两大挑战：一是训练成本高，需同时训练整个模型和路由模块；二是跳过关键层可能降低性能。本文提出Router-Tuning，仅在小数据集上微调路由器，显著降低训练开销；并引入MindSkip，通过注意力机制实现动态深度，保持模型性能的同时大幅提升计算与内存效率。大量实验表明，该方法在保持竞争力的前提下，实现21%的推理速度提升，性能仅下降0.2%。代码已开源。
原文摘要 · Abstract (English)
Traditional transformer models often allocate a fixed amount of computational resources to every input token, leading to inefficient and unnecessary computation. To address this, the Mixture of Depths (MoD) was introduced to dynamically adjust the computational depth by skipping less important layers. Despite its promise, current MoD approaches remain under-explored and face two main challenges: (1) high training costs due to the need to train the entire model along with the routers that determine which layers to skip, and (2) the risk of performance degradation when important layers are bypassed. In response to the first issue, we propose Router-Tuning, a method that fine-tunes only the router on a small dataset, drastically reducing the computational overhead associated with full model training. For the second challenge, we propose MindSkip, which deploys Attention with Dynamic Depths. This method preserves the model's performance while significantly enhancing computational and memory efficiency. Extensive experiments demonstrate that our approach delivers competitive results while dramatically improving the computation efficiency, e.g., 21\% speedup and only a 0.2\% performance drop. The code is released at https://github.com/CASE-Lab-UMD/Router-Tuning.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。