Dr.LLM: Roteamento Dinâmico de Camadas em LLMs

Resumo

Modelos de Linguagem de Grande Escala (LLMs) processam cada token através de todas as camadas de uma pilha de transformadores, resultando em computação desperdiçada em consultas simples e flexibilidade insuficiente para tarefas mais complexas que exigem raciocínio mais profundo. Métodos de profundidade adaptativa podem melhorar a eficiência, mas abordagens anteriores dependem de buscas custosas durante a inferência, alterações arquiteturais ou retreinamento em larga escala, e na prática frequentemente degradam a precisão apesar dos ganhos de eficiência. Apresentamos o Dr.LLM, Roteamento Dinâmico de Camadas para LLMs, uma estrutura retrocompatível que equipa modelos pré-treinados com roteadores leves por camada que decidem pular, executar ou repetir um bloco. Os roteadores são treinados com supervisão explícita: usando Busca em Árvore de Monte Carlo (MCTS), derivamos configurações de camadas de alta qualidade que preservam ou melhoram a precisão sob um orçamento de computação. Nosso design, que inclui pooling em janelas para roteamento estável, perda focal com balanceamento de classes e roteadores MLP com gargalo, garante robustez em cenários de desequilíbrio de classes e sequências longas. No ARC (lógica) e DART (matemática), o Dr.LLM melhora a precisão em até +3,4%p enquanto economiza 5 camadas por exemplo em média. Os roteadores generalizam para tarefas fora do domínio (MMLU, GSM8k, AIME, TruthfulQA, SQuADv2, GPQA, PIQA, AGIEval) com apenas 0,85% de queda na precisão, mantendo a eficiência, e superam métodos de roteamento anteriores em até +7,7%p. No geral, o Dr.LLM demonstra que roteadores supervisionados explicitamente adaptam LLMs congelados para inferência consciente do orçamento e orientada à precisão, sem alterar os pesos base.

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.

Dr.LLM: Roteamento Dinâmico de Camadas em LLMs

Dr.LLM: Dynamic Layer Routing in LLMs

Resumo

Support