Ajuste de Enrutador: Un Enfoque Simple y Efectivo para Habilitar la Profundidad Dinámica en 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

Resumen

Los modelos tradicionales de transformadores a menudo asignan una cantidad fija de recursos computacionales a cada token de entrada, lo que resulta en una computación ineficiente y no necesaria. Para abordar esto, se introdujo la Mezcla de Profundidades (MoD) para ajustar dinámicamente la profundidad computacional al omitir capas menos importantes. A pesar de su promesa, los enfoques actuales de MoD siguen siendo poco explorados y enfrentan dos desafíos principales: (1) altos costos de entrenamiento debido a la necesidad de entrenar el modelo completo junto con los enrutadores que determinan qué capas omitir, y (2) el riesgo de degradación del rendimiento al omitir capas importantes. En respuesta al primer problema, proponemos Ajuste de Enrutador, un método que ajusta finamente solo el enrutador en un pequeño conjunto de datos, reduciendo drásticamente la sobrecarga computacional asociada con el entrenamiento completo del modelo. Para el segundo desafío, proponemos MindSkip, que implementa Atención con Profundidades Dinámicas. Este método conserva el rendimiento del modelo mientras mejora significativamente la eficiencia computacional y de memoria. Experimentos extensos demuestran que nuestro enfoque ofrece resultados competitivos al mejorar drásticamente la eficiencia computacional, por ejemplo, un aumento del 21\% en la velocidad y solo una disminución del 0.2\% en el rendimiento. El código está disponible en 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.