Resumen Técnico: Dr.LLM – Enrutamiento Dinámico de Capas en LLM

Declaración del Problema

Los Modelos de Lenguaje Grandes (LLM) operan actualmente bajo un régimen de profundidad estática, procesando cada token a través de una pila fija de capas transformadoras independientemente de la complejidad de la entrada. Este enfoque conduce a dos ineficiencias principales:

1. Cálculo Desperdiciado: Las consultas sencillas gastan recursos de cálculo innecesarios al recorrer todas las capas.
2. Flexibilidad Insuficiente: Las tareas de razonamiento complejo a menudo requieren un procesamiento más profundo o un refinamiento iterativo, lo que las arquitecturas estáticas no pueden proporcionar sin reentrenamiento o una revisión arquitectónica.

Los métodos existentes de profundidad adaptativa (por ejemplo, salidas tempranas, poda de capas, mezcla de profundidades) enfrentan limitaciones significativas: a menudo intercambian precisión por velocidad, requieren reentrenamiento extenso o modificaciones arquitectónicas, o dependen de una búsqueda costosa en tiempo de inferencia que obstaculiza la escalabilidad. Además, muchos enfoques anteriores degradan la precisión del modelo a pesar de las ganancias en eficiencia.

Metodología: Dr.LLM

Los autores proponen Dr.LLM (Enrutamiento Dinámico de Capas para LLMs), un marco reconfigurable que equipa a LLMs preentrenados congelados con enrutadores ligeros por capa. Estos enrutadores deciden si omitir, ejecutar o repetir un bloque transformador específico para una entrada dada.

1. Arquitectura del Enrutador

• Estructura: Cada bloque transformador $B_\ell$ se empareja con un Perceptrón Multicapa (MLP) ligero compuesto por capas Lineal-GELU-Lineal.
• Entrada: El enrutador opera sobre un resumen compacto de los estados ocultos de la capa anterior. Para garantizar estabilidad en contextos largos, los autores emplean agrupación media con ventanas, dividiendo la secuencia en ventanas contiguas y promediando los estados ocultos dentro de cada ventana.
• Salida: El enrutador genera logits para tres clases: $\{omitir, ejecutar, repetir\}$.
• Sobrecarga: Los enrutadores se ejecutan una vez por secuencia de entrada. La sobrecarga es insignificante (constante en relación con el recuento de tokens) y totalmente compatible con la caché KV.

2. Entrenamiento Supervisado mediante MCTS

A diferencia de los métodos que dependen del aprendizaje por refuerzo o de la búsqueda en tiempo de inferencia, Dr.LLM entrena enrutadores utilizando supervisión explícita derivada de una Búsqueda en Árbol de Monte Carlo (MCTS) fuera de línea.

• Proceso de Búsqueda: Para cada entrada, MCTS explora rutas de ejecución modificadas (omitir o repetir capas) bajo un presupuesto de cálculo. Conserva únicamente las rutas que preservan o mejoran la precisión de la tarea en comparación con la ejecución completa por defecto.
• Conjunto de Datos: La búsqueda genera un conjunto de datos compacto de aproximadamente 4.000 rutas de ejecución optimizadas (tuplas de pregunta, acciones de capa óptimas y respuesta).
• Objetivo de Entrenamiento: Los enrutadores se entrenan de extremo a extremo en este conjunto de datos utilizando Pérdida Focal con pesos de balanceo de clases para manejar el desequilibrio severo entre la clase dominante "ejecutar" y las clases minoritarias "omitir" y "repetir".
• Congelamiento: Crucialmente, los pesos del LLM base permanecen congelados; solo se actualizan los parámetros del enrutador.

3. Inferencia

En el momento de la inferencia, el enrutador entrenado toma decisiones voraces ($\hat{y}_\ell = \arg\max p_\ell$) para cada capa sin realizar ninguna búsqueda.

• Omitir: Elude la capa ($H^{(\ell)} = H^{(\ell-1)}$).
• Ejecutar: Aplica la capa una vez.
• Repetir: Aplica la capa dos veces consecutivas.
• Agrupación: Para manejar la sincronización en el procesamiento por lotes, las secuencias se agrupan según sus decisiones de enrutamiento en cada capa. Los autores señalan que un alto consenso de enrutamiento (80–89% de acuerdo entre secuencias) minimiza la sobrecarga de sincronización.

Contribuciones Clave

1. Enrutamiento Dinámico Supervisado para LLMs Congelados: Introducción de enrutadores por capa que deciden omitir, ejecutar o repetir, entrenados de extremo a extremo solo en 4k rutas derivadas de MCTS sin modificar los pesos base.
2. Supervisión Efectiva de Rutas mediante MCTS: Un algoritmo MCTS consciente de la longitud que descubre configuraciones de capas que preservan o mejoran la precisión fuera de línea, generando un conjunto de datos de supervisión compacto sin búsqueda en tiempo de inferencia.
3. Entrenamiento Ligero de Enrutadores: Un diseño que incluye agrupación media con ventanas para la estabilidad en contextos largos, pérdida focal para el desequilibrio de clases y un enrutador lineal compacto de dos capas, garantizando una sobrecarga computacional insignificante.
4. Ganancias en Precisión y Eficiencia: Mejoras demostradas en la precisión del razonamiento junto con una reducción en la ejecución de capas, superando a los métodos de enrutamiento anteriores de última generación.
5. Generalización Robusta: Mostró que las políticas de enrutamiento aprendidas en tareas específicas de razonamiento se transfieren eficazmente a puntos de referencia fuera de dominio con una degradación mínima de la precisión.

Resultados Experimentales

Rendimiento en Dominio (ARC y DART)

Evaluado en seis modelos (variantes de LLaMA-3.2 y Qwen-2.5), Dr.LLM mejoró consistentemente la precisión mientras reducía el número promedio de capas ejecutadas:

• Precisión: Mejora de hasta +3.4 puntos porcentuales (p.p.) en DART (matemáticas) y +1.0 p.p. en ARC (lógica). Por ejemplo, LLaMA-3B-Base vio un salto del 11.8% al 15.8% en DART.
• Eficiencia: Ahorro de un promedio de 5.0 capas por ejemplo en todos los modelos, con casos específicos que ahorran hasta 11.0 capas.
• Latencia: Logró una reducción del 15.3% en el tiempo de generación del reloj para ejecuciones de 1.000 tokens, con una sobrecarga del enrutador que contribuye menos del 1% a la latencia total.

Generalización Fuera de Dominio

Al aplicarse a diversos puntos de referencia (MMLU, GSM8k, AIME, TruthfulQA, SQuADv2, GPQA, AGIEval, PIQA) sin reentrenamiento:

• Los enrutadores mantuvieron la eficiencia mientras sufrían solo una caída promedio de precisión de 0.85 p.p. en ocho puntos de referencia.
• En algunos casos (por ejemplo, GPQA Diamond con LLaMA-3B), la precisión mejoró realmente (+2.5 p.p.).
• Esto sugiere que las políticas aprendidas capturan redundancias estructurales generales en el cálculo de transformadores en lugar de sobreajustarse a patrones de tareas específicas.

Comparación con la Estado del Arte

Dr.LLM superó a los métodos de enrutamiento anteriores (por ejemplo, FlexiDepth, MindSkip, ShortGPT) en hasta +7.7 p.p. en precisión en puntos de referencia de razonamiento. A diferencia de FlexiDepth, que requiere entrenamiento en cientos de miles de ejemplos y a menudo sacrifica precisión, Dr.LLM logra una mayor precisión con una sobrecarga de entrenamiento significativamente menor (4k ejemplos, una sola GPU) y sin cambios arquitectónicos.

Análisis de Patrones de Enrutamiento

El análisis de las decisiones de enrutamiento reveló comportamientos estructurados:

• Capas Tempranas: Consistentemente ejecutadas (uso estable).
• Capas Intermedias: Frecuentemente omitidas (economizando cálculo).
• Capas Tardías: A menudo repetidas, particularmente para tareas complejas de razonamiento, lo que indica un papel en el refinamiento iterativo.
• Aprendizaje de Características: Los enrutadores entrenados en estados ocultos de capas anteriores superaron significativamente a aquellos entrenados en incrustaciones de entrada estáticas, confirmando que Dr.LLM aprende a mapear estados internos del modelo a decisiones de enrutamiento dinámicamente.

Significado y Afirmaciones

El artículo afirma que Dr.LLM demuestra que los enrutadores explícitamente supervisados pueden reconfigurar LLMs congelados para lograr inferencia consciente del presupuesto y orientada a la precisión. Las conclusiones clave incluyen:

• No Se Requiere Reentrenamiento: Los pesos del modelo base permanecen intactos, preservando las capacidades originales del modelo.
• Eficiencia Positiva en Precisión: A diferencia de muchos métodos adaptativos que intercambian precisión por velocidad, Dr.LLM mejora ambas simultáneamente en tareas pesadas de razonamiento.
• Despliegue Práctico: Al eliminar la necesidad de búsqueda en tiempo de inferencia y modificaciones arquitectónicas, Dr.LLM ofrece una solución práctica y escalable para la inferencia adaptativa.
• Generalizabilidad: Las políticas de enrutamiento aprendidas se transfieren robustamente entre familias de modelos y dominios de tareas, lo que sugiere que la asignación de cálculo adaptativo es una propiedad fundamental y aprendible de las arquitecturas transformadoras.

Los autores concluyen que Dr.LLM proporciona un paso práctico hacia una inferencia adaptativa escalable, conciliando eficiencia, precisión y robustez sin los costos prohibitivos asociados con el reentrenamiento o el rediseño arquitectónico.