Настройка маршрутизатора: простой и эффективный подход для обеспечения динамической глубины в трансформерах.

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

October 17, 2024

Авторы: Shwai He, Tao Ge, Guoheng Sun, Bowei Tian, Xiaoyang Wang, Ang Li, Dong Yu

cs.AI

Аннотация

Традиционные модели трансформеров часто выделяют фиксированное количество вычислительных ресурсов для каждого входного токена, что приводит к неэффективным и излишним вычислениям. Для решения этой проблемы был представлен метод Смешанных Глубин (MoD), который динамически корректирует вычислительную глубину, пропуская менее важные слои. Несмотря на свои перспективы, текущие подходы MoD остаются мало исследованными и сталкиваются с двумя основными проблемами: (1) высокие затраты на обучение из-за необходимости обучать весь модельный комплекс вместе с маршрутизаторами, которые определяют, какие слои пропускать, и (2) риск снижения производительности при обходе важных слоев. В ответ на первую проблему мы предлагаем метод Тонирования Маршрутизатора, который донастраивает только маршрутизатор на небольшом наборе данных, резко снижая вычислительные издержки, связанные с полным обучением модели. Для решения второй проблемы мы предлагаем метод MindSkip, который использует Внимание с Динамическими Глубинами. Этот метод сохраняет производительность модели, значительно повышая вычислительную и памятьную эффективность. Обширные эксперименты показывают, что наш подход обеспечивает конкурентоспособные результаты, существенно улучшая эффективность вычислений, например, ускорение на 21\% и лишь незначительное снижение производительности на 0.2\%. Код доступен по ссылке 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.