LayerRoute: Input-Conditioned Adaptive Layer Skipping via LoRA Fine-Tuning for Agentic Language Models

Resumen Técnico: LayerRoute – Omisión de Capas Adaptativa Condicionada por la Entrada mediante el Ajuste Fino de LoRA

Planteamiento del Problema

Los sistemas de IA agéntica modernos orquestan Modelos de Lenguaje de Gran Escala (LLMs) a través de flujos de trabajo heterogéneos que comprenden dos tipos distintos de pasos:

1. Llamadas a herramientas estructuradas: Secuencias cortas y deterministas con baja perplejidad (por ejemplo, consultas a bases de datos).
2. Planificación/razonamiento abierto: Secuencias largas y complejas con alta perplejidad que requieren un razonamiento de múltiples pasos.

Los sistemas de inferencia actuales aplican recursos computacionales idénticos a ambos tipos de pasos, dirigiendo cada entrada a través de todas las capas del transformador. Esto es ineficiente porque la alta confianza en los estados tempranos del flujo residual para las llamadas a herramientas sugiere que las capas profundas contribuyen marginalmente a la predicción final. El artículo aborda la necesidad de un sistema de inferencia que pueda asignar computación dinámicamente según la complejidad de la entrada sin sacrificar la calidad del modelo.

Metodología

Los autores proponen LayerRoute, un adaptador ligero diseñado para el modelo Qwen2.5-0.5B-Instruct (24 capas, tamaño oculto 896). La arquitectura aumenta el esqueleto (backbone) congelado con dos componentes primarios por bloque de transformador:

1. Routers por Capa: Cada bloque $i$ contiene una capa lineal ligera ($w_i^T \bar{h}_i + b_i$) que produce un valor sigmoide suave $\sigma_i$. Este se convierte en una compuerta binaria dura $g_i \in \{0, 1\}$ mediante un umbral de 0.5.

   • Paso hacia adelante (Forward Pass): Si $g_i=1$, el bloque se ejecuta normalmente (con los adaptadores LoRA activos). Si $g_i=0$, el estado oculto pasa sin cambios, incurriendo en cero computación para ese bloque.
   • Entrenamiento: La compuerta dura no diferenciable se gestiona mediante el Estimador de Paso Directo (Straight-Through Estimator o STE), lo que permite que los gradientes fluyan a través de la sigmoide durante la retropropagación mientras se mantienen las decisiones duras durante la inferencia.

2. Adaptadores LoRA: Se aplican adaptadores de bajo rango (Rango 8) a las proyecciones de Consulta, Clave, Valor y Salida ($W_Q, W_K, W_V, W_O$) de los 24 bloques. Los pesos del esqueleto permanecen congelados.

Estrategia de Entrenamiento e Innovaciones:

• Optimización Conjunta: Los pesos del router y las matrices LoRA se entrenan conjuntamente utilizando una función de pérdida combinada: $L = L_{LM} + \lambda \cdot \frac{1}{L}\sum \sigma(s_i)$.
• Regularización de la Compuerta: Un término de regularización ($\lambda=1.0$) penaliza los valores altos de la compuerta suave, evitando que el router colapse hacia una solución trivial donde todas las compuertas permanezcan abiertas.
• Inicialización Sesgada: Para romper la simetría y evitar el problema del "huevo o la gallina" (donde las compuertas no pueden diferenciar hasta que omiten), los autores inicializan los routers de las capas medias (capas 8–16) con un sesgo de $-1.0$ (comenzando cerradas/omitidas) y las capas externas con $+1.0$ (comenzando abiertas). Esto fuerza una diferenciación inmediata basada en la pérdida de modelado de lenguaje.

Contribuciones Clave

El artículo describe cinco contribuciones específicas:

1. Conexiones de salto con compuerta dura por capa utilizando STE para eliminar el desajuste entre entrenamiento e inferencia.
2. Ajuste fino conjunto de LoRA de las proyecciones de atención junto con el aprendizaje del router, permitiendo que la política de enrutamiento co-evolucione con los pesos adaptados.
3. Regularización de la compuerta para forzar el descubrimiento de bloques omitibles en lugar de mantener la computación completa.
4. Inicialización sesgada de los routers de las capas medias para permitir la diferenciación basada en gradientes desde el primer paso de entrenamiento.
5. Validación empírica en Qwen2.5-0.5B-Instruct demostrando un diferencial de omisión significativo entre las llamadas a herramientas y los pasos de planificación, entrenado en menos de 7 minutos.

Resultados Experimentales

El modelo fue entrenado en un conjunto de datos agéntico mixto (Hermes, Glaive, GSM8K, Turing) durante 3,000 pasos (aprox. 6.4 minutos en un A100 40GB) utilizando solo 1.10M de parámetros entrenables (0.22% de los 494M del esqueleto).

• Diferencial de Omisión: LayerRoute logró un diferencial de omisión del 12.91%. Las llamadas a herramientas omitieron el 15.25% de los FLOPs, mientras que los pasos de planificación omitieron solo un 2.34%.
• Mejora de la Calidad: Contrario a los típicos compromisos de eficiencia, el modelo con compuerta mostró una perplejidad mejorada sobre la línea base de capa completa.
  • Llamadas a herramientas: Delta de perplejidad de −1.29.
  • Planificación: Delta de perplejidad de −1.30.
  • Esta mejora se atribuye a que la adaptación LoRA compensa el costo de calidad de la omisión.
• Estudios de Ablación:
  • LayerRoute-NoReg: Colapsó a una omisión casi nula (0.1%), confirmando que la regularización es necesaria.
  • LayerRoute-UniformInit: Invirtió el diferencial de omisión (la planificación omitió más que las llamadas a herramientas), probando el papel crítico de la inicialización sesgada.
  • LayerRoute-BCE: Entrenado con etiquetas explícitas en lugar de optimización conjunta, lo que resultó en diferenciales de omisión significativamente menores (0.3%).

Estructura de la Compuerta: Las compuertas aprendidas convergieron a una estructura estable de dos clústeres: las capas 0–7 y 17–23 permanecieron confiablemente abiertas, mientras que las capas 8–16 (el clúster medio) fueron omitidas de manera confiable para las llamadas a herramientas pero retenidas para la planificación.

Significancia y Reivindicaciones

El artículo afirma que el tipo de paso agéntico puede aprenderse implícitamente a partir del objetivo de modelado de lenguaje sin requerir etiquetas explícitas del tipo de tarea o un cabezal de clasificación separado.

• Eficiencia: Logra una reducción de FLOPs del 15.2% para las llamadas a herramientas preservando (y mejorando ligeramente) la calidad.
• Adaptabilidad: El sistema distingue con éxito entre secuencias cortas y estructuradas y secuencias largas y pesadas de razonamiento, aplicando computación solo donde es necesario.
• Eficiencia de Parámetros: Todo el proceso de adaptación requiere menos del 0.22% de los parámetros del esqueleto y puede completarse en minutos.

Limitaciones Reconocidas:
Los autores señalan que el clúster de la capa media (8–16) forma un bloque de omisión monolítico sin diferenciación interna, lo que limita el potencial de ahorro de FLOPs. Además, la tasa de omisión para los pasos de planificación sigue siendo baja (2.34%), ofreciendo un beneficio mínimo para las cargas de trabajo pesadas en planificación. La evaluación está actualmente limitada al tamaño del modelo 0.5B.