Two-dimensional early exit optimisation of LLM inference

Resumen Técnico: Optimización de Salida Temprana Bidimensional (2D) en LLM

1. El Problema

Los Grandes Modelos de Lenguaje (LLM) han revolucionado el Procesamiento del Lenguaje Natural (NLP), pero su despliegue en tareas de inferencia (especialmente clasificación) conlleva un costo computacional, de memoria y energético excesivo.

• Limitación actual: Las estrategias de salida temprana (early exit) existentes se centran principalmente en una sola dimensión: detener la inferencia en una capa intermedia del modelo cuando la confianza es suficiente.
• Ineficiencia: Estas estrategias unidimensionales no aprovechan la estructura semántica de las entradas. En muchas tareas de clasificación, la clase correcta puede determinarse con solo una parte del texto (ej. las primeras oraciones de una reseña), pero los modelos actuales procesan todo el texto a través de todas las capas antes de tomar una decisión o detenerse solo por capas.

2. Metodología Propuesta: Estrategia 2D

Los autores proponen una estrategia de salida temprana bidimensional (2D) que coordina simultáneamente la salida por capas (depth) y por oraciones (input trimming).

• Concepto Central: En lugar de procesar todo el documento de una vez, el método divide la entrada en oraciones (unidades semánticas naturales). A medida que se introduce una nueva oración, el modelo no solo procesa esa oración, sino que también activa progresivamente capas más profundas para todo el contexto acumulado hasta ese momento.
• Mecanismo de Progreso:
  • Se define un tamaño de paso $\Delta$ basado en el número total de capas ($L$) y oraciones ($m$).
  • Al procesar la oración $s$, el algoritmo activa las primeras $L_s$ capas, donde $L_s = \min((s + 1) \times \Delta, L)$.
  • Las oraciones anteriores se vuelven a procesar a través de las nuevas capas activadas, creando bloques de procesamiento rectangulares que crecen en 2D (capas vs. oraciones).
• Criterio de Salida:
  • Se utilizan clasificadores ligeros (adaptadores) adjuntos a cada capa.
  • La inferencia se detiene cuando la confianza acumulada (diferencia entre la probabilidad máxima y la segunda máxima) de una clase supera un umbral de aceptación ($\tau_{acc}$).
  • Se utiliza un umbral de ignorancia ($\tau_{ignore}$) para filtrar predicciones de baja confianza antes de acumularlas.
• Entrenamiento:
  • Entrenamiento solo de adaptadores (Adapter-only): Los pesos del LLM se congelan; solo se entrenan los clasificadores de cada capa.
  • Fine-tuning (Ajuste fino): Se ajusta todo el modelo (usando LoRA) con una función de pérdida ponderada que prioriza la precisión en las capas finales pero también fomenta predicciones correctas con entradas parciales.

3. Contribuciones Clave

1. Sinergia Multiplicativa: La integración de la parada por capas con el recorte de entrada por oraciones logra ahorros computacionales multiplicativos (no aditivos). Al reducir tanto la profundidad de la red como la longitud del contexto procesado, se obtienen aceleraciones significativamente mayores que optimizar una sola dimensión.
2. Independencia del Modelo: La estrategia es agnóstica al modelo, requiriendo solo adaptadores de clasificación ligeros. Es compatible con métodos complementarios como cuantización, poda (pruning) y arquitecturas MoE (Mixture of Experts).
3. Evaluación Exhaustiva: Se probó en cuatro LLMs de última generación (Llama 3.1/3.2, Gemma, Qwen) con parámetros entre 3B y 8B, en tres conjuntos de datos de clasificación de sentimientos de dificultad variable.

4. Resultados Experimentales

Los experimentos compararon la estrategia 2D contra la salida temprana óptima por capas y el método LayerSkip.

• Aceleración en Modelos "Vanilla" (sin ajuste fino):
  • En tareas simples (clasificación binaria o ternaria, como reseñas de Steam y MMS), la estrategia 2D logró aceleraciones adicionales de 1.4x a 2.3x en comparación con la salida temprana óptima por capas.
  • En tareas complejas (5 clases, Amazon-5), la ventaja disminuyó (1.1x - 1.5x) o fue negativa en algunos modelos (Llama), indicando que la estrategia funciona mejor cuando la información semántica se acumula de manera predecible.
• Impacto del Fine-Tuning:
  • El ajuste fino del modelo completo redujo la ventaja de la estrategia 2D (aunque mejoró la precisión general).
  • Hallazgo Contraintuitivo: El fine-tuning "aplana" las variaciones de precisión a través de las capas tempranas y las entradas parciales. Esto hace que sea más difícil para el modelo decidir salir temprano, ya que la precisión en las capas iniciales mejora, reduciendo la ganancia de velocidad.
• Comparación con Baselines:
  • La estrategia 2D superó consistentemente a LayerSkip en tareas sencillas, pero LayerSkip fue superior en tareas muy complejas con pérdida de precisión del 2%.

5. Significado e Implicaciones

• Eficiencia en Tareas de Secuencia: El trabajo demuestra que para tareas donde la información semántica se construye incrementalmente (como la clasificación de sentimientos en reseñas), es innecesario procesar todo el texto a máxima profundidad.
• Compromiso Precisión-Velocidad: La estrategia ofrece una degradación elegante de la precisión. Permitir una pequeña pérdida de precisión (ej. 4%) habilita aceleraciones mucho mayores (hasta 6.4x en algunos casos) en comparación con mantener la precisión máxima.
• Dirección Futura: Los autores sugieren que el desafío principal es desarrollar estrategias de fine-tuning que preserven los gradientes de precisión en las capas tempranas para maximizar el beneficio de la salida 2D. Además, proponen extender este enfoque a otras tareas de procesamiento de secuencias más allá de la clasificación de sentimientos.

Conclusión:
La optimización de salida temprana 2D representa un avance significativo en la eficiencia de inferencia de LLMs, logrando ahorros computacionales que superan los límites de las estrategias unidimensionales tradicionales, especialmente en escenarios de alto volumen de datos con tareas de complejidad moderada.