Resumen Técnico: LEDE – Salidas Dinámicas Impulsadas por la Experiencia para LLMs con Aprendizaje por Refuerzo

1. Planteamiento del Problema

Los Modelos de Lenguaje de Gran Escala (LLMs) sufren una latencia significativa durante la inferencia autorregresiva porque cada token generado debe atravesar todo el stack de capas del transformador. Aunque se ha propuesto la Decodificación Auto-Especulativa (SSD) para acelerar este proceso utilizando las propias capas tempranas de un modelo como un "borrador" interno para generar tokens candidatos en paralelo, los métodos actuales de SSD están fundamentalmente limitados por configuraciones estáticas.

Los enfoques existentes dependen de profundidades de borrador fijas y longitudes de especulación fijas aplicadas uniformemente a lo todo el proceso de generación. Esta rigidez no tiene en cuenta la dispersión contextual, donde la dificultad de predecir tokens varía significativamente dependiendo del contexto local. En consecuencia, los métodos estáticos no pueden equilibrar de manera óptima la relación entre la calidad del borrador (tasa de aceptación) y la velocidad de especulación, lo que conduce a un rendimiento de inferencia subóptimo.

2. Metodología: El Marco de Trabajo LEDE

Los autores proponen LEDE (Learning-based Dynamic Exit), un marco que redefine el control dinámico de la SSD como un Proceso de Decisión de Markov (MDP). En lugar de utilizar reglas rígidas y predefinidas, LEDE emplea aprendizaje por refuerzo (RL) fuera de línea (offline) para entrenar una política que seleccione dinámicamente tanto la capa de salida óptima como la longitud de especulación basándose en el contexto local de la secuencia generada.

2.1. Formulación del Proceso de Decisión de Markov (MDP)

El marco modela el proceso de toma de decisiones durante un solo paso de generación como un MDP con los siguientes componentes:

• Espacio de Estados ($S$): El agente observa un vector de estado $s_l$ en cada capa candidata $l$. Este vector se deriva de la distribución interna de tokens del modelo e incluye:
  • Confianza del Token ($C_k^l$): El logaritmo negativo promedio de los top-$k$ tokens.
  • Entropía del Token ($H^l$): La incertidumbre de la distribución de predicción del siguiente token.
  • Traza Promedio de la Confianza del Token ($C_{trace}$): El logaritmo negativo promedio del top-1 token sobre una ventana deslizante de pasos previos.
• Espacio de Acciones ($A$): En cada capa candidata, la política (una red Q) decide entre dos acciones:
  • Continuar ($a=0$): Proceder a la siguiente capa candidata.
  • Salir ($a=1$): Designar la capa actual como la capa de salida óptima ($l^*$) para la etapa de borrador.
• Función de Recompensa ($R$): Diseñada para optimizar la eficiencia de la inferencia.
  • Recompensa Positiva: Se otorga si el token del borrador es aceptado ($x_t^{l^*} = x_t^L$), escalada por la profundidad ahorrada ($1 + \frac{L-l^*}{L}$).
  • Recompensa Negativa: Se otorga si el token del borrador es rechazado ($-2 + \frac{L-l^*}{L}$).
  • Penalización: Se aplica una pequeña penalización ($-0.01$) por continuar, fomentando las salidas tempranas cuando es apropiado.

2.2. Optimización de la Política mediante RL Fuera de Línea

LEDE utiliza una Red Q Profunda (DQN) entrenada fuera de línea utilizando un buffer de repetición de experiencias de inferencia pasadas:

1. Recolección de Datos: Una fase inicial de exploración recolecta tuplas de experiencia diversas $(s_t, a_t, r_t, s_{t+1})$ donde el agente interactúa con el entorno.
2. Entrenamiento Fuera de Línea: La red Q se entrena minimizando el error de Diferencia Temporal (TD) de n-pasos. El valor objetivo se calcula utilizando una red objetivo estable, equilibrando la recompensa inmediata de una salida temprana contra el valor a largo plazo del estado.
3. Mecanismo de Control Dual: Durante la inferencia en línea, el agente entrenado realiza dos controles adaptativos:
   • Selección Dinámica de la Capa de Salida: Determina la profundidad óptima $l^*$ para comenzar el borrador.
   • Longitud de Borrador Adaptativa: Tras seleccionar $l^*$, el agente reevalúa el estado después de cada token del borrador. Una acción de "salida" continúa el borrador (indicando alta confianza), mientras que una acción de "continuar" (indicando baja confianza) termina el borrador inmediatamente para su verificación.

3. Contribuciones Clave

El artículo describe tres contribuciones principales:

1. Formulación del MDP: Es el primer trabajo que formula el control dinámico de la profundidad del borrador y la longitud de la especulación en la SSD como un problema de MDP, aplicando aprendizaje por refuerzo para derivar una política de control.
2. Marco de Trabajo LEDE: El diseño e implementación de un marco que aprovecha el aprendizaje por refuerzo fuera de línea y la repetición de experiencias para aprender una política consciente del contexto, superando las limitaciones tanto de las configuraciones estáticas como de los métodos basados en heurísticas simples.
3. Validación Empírica: Experimentos exhaustivos que demuestran aceleraciones significativas sobre la generación autorregresiva y ganancias de eficiencia sustanciales sobre las líneas base estáticas.

4. Resultados Experimentales

El marco fue evaluado en modelos Llama-2 (7B, 13B) y Llama-3 (1B), así como en CodeLLaMA (7B, 34B), a través de diversas tareas que incluyen seguimiento de instrucciones, resumen, modelado de lenguaje y generación de código.

• Rendimiento de Aceleración (Speedup):
  • LEDE logra una aceleración promedio de 2.32× sobre la decodificación autorregresiva estándar.
  • Las aceleraciones máximas alcanzan 2.72× (por ejemplo, en Llama-2-7B para modelado de lenguaje).
  • LEDE proporciona una aceleración adicional del 17% sobre la mejor línea base especulativa estática (LayerSkip).
• Comparación con Líneas Base:
  • LayerSkip (Estático): El rendimiento es altamente sensible a las configuraciones fijas y no generaliza bien entre tareas.
  • Métodos Basados en Reglas (LITE, Draft & Verify): Están limitados por umbrales predefinidos, lo que conduce a compromisos subóptimos entre profundidad y tasas de aceptación.
  • LEDE: Al aprender una política de control refinada, supera consistentemente a todas las líneas base en aceleración manteniendo una calidad de generación comparable (medida por Rouge-L).
• Estudio de Ablación: Eliminar ya sea la selección dinámica de la capa de salida o la longitud de borrador adaptativa degrada significamente el rendimiento (reduciendo la aceleración a ~2.0×), confirmando que ambos mecanismos son sinérgicos y esenciales.

5. Significado y Reivindicaciones

El artículo afirma que LEDE representa un cambio conceptual de la inferencia rígida impulsada por heurísticas hacia una inferencia adaptativa impulsada por políticas. Al aprender a co-optimizar la profundidad computacional y la longitud de la especulación basándose en el contexto en tiempo real, el marco permite una decodificación más robusta y eficiente a través de diversos modelos y tareas.

Los autores posicionan este trabajo como un paso hacia la construcción de sistemas de LLM más inteligentes y computacionalmente eficientes. Señalan que el trabajo futuro se centrará en escalar estos beneficios a modelos más grandes (escala de 70B en adelante), pero la contribución actual establece la viabilidad de utilizar el aprendizaje por refuerzo fuera de línea para gestionar dinámicamente el pipeline de inferencia de los modelos de lenguaje de gran escala.