Dr.LLM: Enrutamiento Dinámico de Capas en Modelos de Lenguaje de Gran Escala

Resumen

Los Modelos de Lenguaje de Gran Escala (LLMs) procesan cada token a través de todas las capas de una pila de transformadores, lo que genera un desperdicio de cómputo en consultas simples y una flexibilidad insuficiente para aquellas más complejas que requieren un razonamiento más profundo. Los métodos de profundidad adaptativa pueden mejorar la eficiencia, pero enfoques previos dependen de búsquedas costosas durante la inferencia, cambios arquitectónicos o reentrenamientos a gran escala, y en la práctica a menudo degradan la precisión a pesar de las ganancias en eficiencia. Presentamos Dr.LLM, Enrutamiento Dinámico de Capas para LLMs, un marco adaptable que equipa modelos preentrenados con enrutadores ligeros por capa que deciden omitir, ejecutar o repetir un bloque. Los enrutadores se entrenan con supervisión explícita: utilizando Búsqueda de Árbol de Monte Carlo (MCTS), derivamos configuraciones de capas de alta calidad que preservan o mejoran la precisión bajo un presupuesto de cómputo. Nuestro diseño, que incluye agrupación en ventanas para un enrutamiento estable, pérdida focal con balanceo de clases y enrutadores MLP de cuello de botella, garantiza robustez ante desequilibrios de clases y secuencias largas. En ARC (lógica) y DART (matemáticas), Dr.LLM mejora la precisión hasta en +3.4%p mientras ahorra 5 capas por ejemplo en promedio. Los enrutadores generalizan a tareas fuera del dominio (MMLU, GSM8k, AIME, TruthfulQA, SQuADv2, GPQA, PIQA, AGIEval) con solo un 0.85% de caída en precisión mientras mantienen la eficiencia, y superan a métodos de enrutamiento previos hasta en +7.7%p. En general, Dr.LLM demuestra que los enrutadores supervisados explícitamente adaptan LLMs congelados para inferencia consciente del presupuesto y orientada a la precisión sin alterar los 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: Enrutamiento Dinámico de Capas en Modelos de Lenguaje de Gran Escala

Dr.LLM: Dynamic Layer Routing in LLMs

Resumen

Support