Optimisation du routeur : Une approche simple et efficace pour activer la profondeur dynamique dans les transformeurs

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

October 17, 2024

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

cs.AI

Résumé

Les modèles de transformateur traditionnels allouent souvent une quantité fixe de ressources computationnelles à chaque jeton d'entrée, ce qui entraîne une computation inefficace et inutile. Pour remédier à cela, le Mélange de Profondeurs (MoD) a été introduit pour ajuster dynamiquement la profondeur computationnelle en sautant les couches moins importantes. Malgré ses promesses, les approches MoD actuelles restent peu explorées et rencontrent deux principaux défis : (1) des coûts de formation élevés en raison de la nécessité de former l'ensemble du modèle ainsi que les routeurs qui déterminent quelles couches sauter, et (2) le risque de dégradation des performances lorsque des couches importantes sont contournées. En réponse au premier problème, nous proposons l'Accordage des Routeurs, une méthode qui ajuste finement uniquement le routeur sur un petit ensemble de données, réduisant drastiquement la surcharge computationnelle associée à la formation complète du modèle. Pour le deuxième défi, nous proposons MindSkip, qui déploie l'Attention avec des Profondeurs Dynamiques. Cette méthode préserve les performances du modèle tout en améliorant significativement l'efficacité computationnelle et la mémoire. Des expériences approfondies démontrent que notre approche offre des résultats compétitifs tout en améliorant considérablement l'efficacité de calcul, par exemple, une accélération de 21\% et seulement une baisse de performance de 0,2\%. Le code est disponible sur 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.