Ajuste de Roteador: Uma Abordagem Simples e Eficaz para Habilitar Profundidade Dinâmica em Transformadores

Router-Tuning: A Simple and Effective Approach for Enabling Dynamic-Depth in Transformers

October 17, 2024

Autores: Shwai He, Tao Ge, Guoheng Sun, Bowei Tian, Xiaoyang Wang, Ang Li, Dong Yu

cs.AI

Resumo

Os modelos tradicionais de transformadores frequentemente alocam uma quantidade fixa de recursos computacionais para cada token de entrada, resultando em cálculos ineficientes e desnecessários. Para lidar com isso, o Mixture of Depths (MoD) foi introduzido para ajustar dinamicamente a profundidade computacional ao pular camadas menos importantes. Apesar de sua promessa, as abordagens MoD atuais permanecem pouco exploradas e enfrentam dois desafios principais: (1) altos custos de treinamento devido à necessidade de treinar o modelo inteiro juntamente com os roteadores que determinam quais camadas pular, e (2) o risco de degradação de desempenho quando camadas importantes são ignoradas. Em resposta ao primeiro problema, propomos o Ajuste de Roteador, um método que ajusta apenas o roteador em um pequeno conjunto de dados, reduzindo drasticamente o overhead computacional associado ao treinamento completo do modelo. Para o segundo desafio, propomos o MindSkip, que implementa Atenção com Profundidades Dinâmicas. Este método preserva o desempenho do modelo enquanto melhora significativamente a eficiência computacional e de memória. Experimentos extensivos demonstram que nossa abordagem oferece resultados competitivos enquanto melhora drasticamente a eficiência de cálculo, por exemplo, 21% de aceleração e apenas uma queda de desempenho de 0,2%. O código está disponível em https://github.com/CASE-Lab-UMD/Router-Tuning.

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.