Router-Afstelling: Een Eenvoudige en Effectieve Benadering voor het Mogelijk Maken van Dynamische Diepte in Transformers

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

Samenvatting

Traditionele transformer modellen wijzen vaak een vastgestelde hoeveelheid rekenbronnen toe aan elk invoertoken, wat leidt tot inefficiënte en onnodige berekeningen. Om dit aan te pakken, werd de Mixture of Depths (MoD) geïntroduceerd om dynamisch de berekeningsdiepte aan te passen door minder belangrijke lagen over te slaan. Ondanks de belofte blijven huidige MoD-benaderingen onderbelicht en worden geconfronteerd met twee belangrijke uitdagingen: (1) hoge trainingskosten als gevolg van de noodzaak om het volledige model samen met de routers die bepalen welke lagen moeten worden overgeslagen, te trainen, en (2) het risico op prestatievermindering wanneer belangrijke lagen worden omzeild. Als reactie op het eerste probleem stellen we Router-Tuning voor, een methode die alleen de router fijnafstemt op een kleine dataset, waardoor de rekenoverhead die gepaard gaat met het volledig trainen van het model drastisch wordt verminderd. Voor de tweede uitdaging stellen we MindSkip voor, dat Attention met Dynamische Dieptes implementeert. Deze methode behoudt de prestaties van het model terwijl de reken- en geheugenefficiëntie aanzienlijk worden verbeterd. Uitgebreide experimenten tonen aan dat onze aanpak concurrerende resultaten oplevert en tegelijkertijd de berekeningsefficiëntie aanzienlijk verbetert, bijvoorbeeld een versnelling van 21% en slechts een prestatiedaling van 0,2%. De code is beschikbaar op 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.