Dr.LLM: Динамическая маршрутизация слоев в языковых моделях

Dr.LLM: Dynamic Layer Routing in LLMs

October 14, 2025

Авторы: Ahmed Heakl, Martin Gubri, Salman Khan, Sangdoo Yun, Seong Joon Oh

cs.AI

Аннотация

Крупные языковые модели (LLM) обрабатывают каждый токен через все слои стека трансформаторов, что приводит к избыточным вычислениям на простых запросах и недостаточной гибкости для более сложных, требующих глубокого анализа. Методы с адаптивной глубиной могут повысить эффективность, но существующие подходы зависят от дорогостоящего поиска на этапе вывода, изменений архитектуры или масштабного переобучения, и на практике часто снижают точность, несмотря на повышение эффективности. Мы представляем Dr.LLM, динамическую маршрутизацию слоев для LLM, — адаптируемую структуру, которая оснащает предобученные модели легковесными маршрутизаторами для каждого слоя, решающими, пропустить, выполнить или повторить блок. Маршрутизаторы обучаются с явным контролем: с использованием метода поиска по дереву Монте-Карло (MCTS) мы получаем высококачественные конфигурации слоев, сохраняющие или улучшающие точность в рамках вычислительного бюджета. Наш дизайн, включающий оконное пуллирование для стабильной маршрутизации, фокусную потерю с балансировкой классов и маршрутизаторы на основе MLP с узким местом, обеспечивает устойчивость при дисбалансе классов и длинных последовательностях. На задачах ARC (логика) и DART (математика) Dr.LLM повышает точность до +3,4% при экономии в среднем 5 слоев на пример. Маршрутизаторы обобщаются на задачи вне домена (MMLU, GSM8k, AIME, TruthfulQA, SQuADv2, GPQA, PIQA, AGIEval) с падением точности всего на 0,85%, сохраняя эффективность, и превосходят предыдущие методы маршрутизации до +7,7%. В целом, Dr.LLM демонстрирует, что явно контролируемые маршрутизаторы адаптируют замороженные LLM для вывода с учетом бюджета и точности без изменения базовых весов.

English

Large Language Models (LLMs) process every token through all layers of a transformer stack, causing wasted computation on simple queries and insufficient flexibility for harder ones that need deeper reasoning. Adaptive-depth methods can improve efficiency, but prior approaches rely on costly inference-time search, architectural changes, or large-scale retraining, and in practice often degrade accuracy despite efficiency gains. We introduce Dr.LLM, Dynamic routing of Layers for LLMs, a retrofittable framework that equips pretrained models with lightweight per-layer routers deciding to skip, execute, or repeat a block. Routers are trained with explicit supervision: using Monte Carlo Tree Search (MCTS), we derive high-quality layer configurations that preserve or improve accuracy under a compute budget. Our design, windowed pooling for stable routing, focal loss with class balancing, and bottleneck MLP routers, ensures robustness under class imbalance and long sequences. On ARC (logic) and DART (math), Dr.LLM improves accuracy by up to +3.4%p while saving 5 layers per example on average. Routers generalize to out-of-domain tasks (MMLU, GSM8k, AIME, TruthfulQA, SQuADv2, GPQA, PIQA, AGIEval) with only 0.85% accuracy drop while retaining efficiency, and outperform prior routing methods by up to +7.7%p. Overall, Dr.LLM shows that explicitly supervised routers retrofit frozen LLMs for budget-aware, accuracy-driven inference without altering base weights.