Ottimizzazione del Router: Un Approccio Semplice ed Efficace per Abilitare la Profondità Dinamica nei Trasformatori

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

October 17, 2024

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

cs.AI

Abstract

I modelli tradizionali di transformer spesso allocano una quantità fissa di risorse computazionali per ogni token di input, portando a una computazione inefficiente e non necessaria. Per affrontare questo problema, è stato introdotto il Mixture of Depths (MoD) per regolare dinamicamente la profondità computazionale saltando strati meno importanti. Nonostante le sue potenzialità, gli attuali approcci MoD rimangono poco esplorati e affrontano due sfide principali: (1) costi elevati di addestramento a causa della necessità di addestrare l'intero modello insieme ai router che determinano quali strati saltare, e (2) il rischio di degradazione delle prestazioni quando vengono saltati strati importanti. In risposta al primo problema, proponiamo Router-Tuning, un metodo che ottimizza solo il router su un piccolo dataset, riducendo drasticamente il sovraccarico computazionale associato all'addestramento completo del modello. Per la seconda sfida, proponiamo MindSkip, che utilizza l'Attention con Dynamic Depths. Questo metodo preserva le prestazioni del modello migliorando significativamente l'efficienza computazionale e di memoria. Esperimenti estesi dimostrano che il nostro approccio fornisce risultati competitivi migliorando notevolmente l'efficienza computazionale, ad esempio, con un aumento della velocità del 21\% e solo una diminuzione delle prestazioni dello 0,2\%. Il codice è disponibile su 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.