Resumen Técnico: N-vium – Transformador de Mezcla de Salidas para Generación Exacta Acelerada

1. Declaración del Problema

Los Modelos de Lenguaje Grandes (LLM) autorregresivos estándar procesan tokens secuencialmente a través de una pila de capas de decodificador. Este paso hacia adelante cumple dos funciones acopladas: (1) construir representaciones Clave-Valor (KV) en cada capa para la atención futura, y (2) proyectar el estado oculto final para predecir el siguiente token. En consecuencia, cada token requiere un paso hacia adelante completo a través de todas las capas, independientemente de si la distribución de predicción ya ha convergido.

Los enfoques existentes para mejorar la eficiencia de la inferencia, como los métodos de Salida Temprana (Early Exit), intentan desacoplar estas funciones adjuntando cabezas de predicción en capas intermedias. Sin embargo, estos métodos suelen tratar las predicciones intermedias como aproximaciones de la distribución de la capa final. Esto conduce a varias desventajas:

• Errores de Aproximación: Las predicciones intermedias suelen ser más débiles que las de profundidad completa, creando una discrepancia entre entrenamiento y prueba.
• Enrutamiento Heurístico: Las decisiones sobre cuándo salir a menudo dependen de umbrales de confianza diseñados a mano que requieren calibración.
• Cachés KV Incompletas: Para ahorrar cómputo, los métodos de salida temprana a menudo omiten el cálculo de las cachés KV para las capas omitidas, requiriendo copias aproximadas o saltos, lo que degrada la calidad.
• Degradación de la Calidad: La velocidad se logra frecuentemente a costa de la perplejidad y la calidad de la generación.

El artículo argumenta que, aunque las observaciones empíricas (por ejemplo, mediante la "lente de logit") muestran que las distribuciones de predicción a menudo se estabilizan antes de la capa final, los métodos actuales no aprovechan esto para una generación exacta sin pérdida de calidad.

2. Metodología: N-vium

Los autores introducen N-vium, un marco de mezcla de salidas que generaliza estrictamente la arquitectura estándar del transformador. En lugar de tratar las salidas tempranas como aproximaciones, N-vium redefine la salida del modelo como una distribución de mezcla aprendida sobre múltiples salidas.

Arquitectura

• Mezcla de Salidas: El modelo particiona las $L$ capas del núcleo en $N_V$ bloques. Al final de cada bloque (unión $k$), un enrutador determina la probabilidad de muestrear el siguiente token desde la distribución $\pi_k$ de esa capa.
• Enrutamiento Adaptativo al Token: Un MLP enrutador procesa el estado oculto $h^{(\ell_k)}_t$ en la unión $k$ para emitir una probabilidad $w_k$. La probabilidad $p_k$ de salir en la unión $k$ se calcula como:
  $$p_k(x_t) = w_k(h^{(\ell_k)}_t) \prod_{j=1}^{k-1} (1 - w_j(h^{(\ell_j)}_t))$$
  donde $w_{N_V} = 1$ por convención.
• Distribución de Mezcla: La distribución final del siguiente token es una suma ponderada de todas las distribuciones de salida:
  $$\pi_{mix}(x_{t+1} | x_{1:t}) = \sum_{k=1}^{N_V} p_k(x_t) \pi_k(x_{t+1} | x_{1:t})$$
• Cabeza Compartida: Todas las salidas comparten la cabeza de proyección final de Modelado de Lenguaje (LM) $W_{lm}$, con MLPs adaptadores ligeros que transforman los estados ocultos intermedios para coincidir con la geometría esperada de la cabeza.

Entrenamiento

• Objetivo: El modelo se entrena para minimizar la pérdida de entropía cruzada de la distribución de mezcla conjunta $\pi_{mix}$, en lugar de optimizar cabezas individuales de forma aislada para igualar la capa final.
• Penalización de Cómputo: Se añade un término de pérdida auxiliar $\mathcal{L}_{compute}$ para alentar al enrutador a asignar masa de probabilidad a salidas más tempranas, guiando al modelo hacia la eficiencia:
  $$\mathcal{L}_{total} = \mathcal{L}_{mix} + \beta \mathcal{L}_{compute}$$
  donde $\mathcal{L}_{compute}$ es la profundidad de capa normalizada esperada.
• Calentamiento: Se emplea una fase de calentamiento del enrutador para evitar el "colapso del enrutador" (donde todos los tokens se enrutan a la capa final) penalizando las desviaciones de una distribución de salida uniforme inicialmente.

Inferencia: Aprovechamiento (Piggybacking)

Una innovación crítica es la técnica de Aprovechamiento (Piggybacking) para garantizar la recuperación exacta de la caché KV sin omitir cómputo:

1. Cálculo Diferido: Cuando un token sale en la unión $k$, el cálculo para las capas $k+1$ a $L$ se difiere.
2. Agrupación en Paralelo: A medida que se procesan los tokens futuros, sus pasos hacia adelante alcanzan las capas diferidas. A estas profundidades, los estados ocultos de los tokens diferidos se concatenan con el estado del token actual y se procesan en paralelo.
3. Resultado: Esto rellena los estados KV faltantes para todos los tokens diferidos en un solo paso paralelo, asegurando que la caché KV esté completa, exacta y sea consistente con el token muestreado. No se omite ningún cálculo; simplemente se reordena de una ejecución secuencial a una ejecución en lotes paralela.

3. Contribuciones Clave

1. Muestreo Exacto de Mezcla: N-vium es el primer método que permite un enrutamiento de profundidad adaptativo al token con muestreo exacto desde la distribución entrenada, evitando los errores de aproximación inherentes a los métodos estándar de salida temprana.
2. Recuperación Completa de KV: Al diferir y agrupar los cálculos de las capas superiores, N-vium recupera cachés KV completas y exactas, una característica a menudo perdida en otros métodos de cómputo adaptativo.
3. Marco de Entrenamiento: Los autores proporcionan una tubería de entrenamiento práctica que optimiza el objetivo de mezcla directamente, incluyendo estrategias para guiar los modelos hacia la eficiencia mediante una penalización de cómputo aprendible.
4. Escalabilidad: El método se valida en escalas de hasta 1.5 mil millones de parámetros, demostrando que el enfoque no está limitado a modelos pequeños.

4. Resultados Experimentales

Los autores pre-entrenaron modelos N-vium (específicamente "Quadrivium" con 4 salidas) en el conjunto de datos C4 utilizando la arquitectura LLaMA 2, comparándolos con líneas base densas coincidentes en parámetros y datos.

• Aceleración vs. Perplejidad:
  • El modelo más grande (1.5 mil millones de parámetros, 48 capas, ancho oculto 1536) logró una aceleración del tiempo real del 57.9% sobre la línea base densa con ninguna degradación de perplejidad (diferencia del 0%).
  • Los modelos más pequeños mostraron aceleraciones consistentes que oscilaron entre el 23% y el 41% con ligeras mejoras en la perplejidad (por ejemplo, un modelo de 36 capas mostró una aceleración del 36.4% con una mejora de perplejidad del 3.7%).
• Frente de Pareto: En los gráficos de calidad-velocidad, N-vium ocupa el cuadrante "superior-derecho" (más rápido y calidad igual o mejor), mientras que métodos como CALM (reimplementado como CALM*) caen en el "inferior-derecho" (más rápido pero peor calidad) y LayerSkip se sitúa en una línea de perplejidad constante por debajo de la línea base.
• Rendimiento en Tareas Posteriores: Después del ajuste fino supervisado (SFT) en la mezcla Tulu 3, los modelos N-vium igualaron o superaron a las líneas base densas en puntos de referencia posteriores (MMLU, ARC, HellaSwag, etc.), confirmando que los modelos siguen siendo controlables y no se sobreajustan al objetivo de entropía cruzada a expensas del rendimiento en tareas.
• Robustez de Hardware: Las mediciones de aceleración en GPUs A100, RTX A6000 y RTX 3090 arrojaron resultados consistentes (~57%), mientras que las GPUs de arquitectura Turing más antiguas mostraron ganancias menores debido a la falta de soporte nativo para bfloat16.

5. Significado y Afirmaciones

El artículo afirma que N-vium representa un cambio de paradigma en la eficiencia de la inferencia. A diferencia de la mayoría de los métodos que mejoran los tokens por segundo reduciendo los FLOPs por token (a menudo mediante aproximación), N-vium mejora la eficiencia aumentando los FLOPs por segundo. Logra esto reorganizando el cálculo secuencial en pasos en lotes paralelos, aprovechando mejor el paralelismo del hardware sin descartar ningún cálculo.

Los autores enfatizan que N-vium es una generalización estricta del transformador estándar: cuando los pesos de enrutamiento para las salidas intermedias son cero, el modelo se comporta exactamente como un transformador estándar. El método satisface cuatro criterios críticos simultáneamente, algo que los autores afirman que ningún otro método existente hace:

1. Inferencia Exacta: Muestreo desde la distribución entrenada sin aproximación.
2. Recuperación Completa de Caché KV: Ninguna capa omitida ni relleno aproximado de caché.
3. Enrutamiento Aprendido: El cómputo por token es asignado por un enrutador entrenado, no por heurísticas.
4. Aceleración en Tiempo de Decodificación: Reduce la latencia del tiempo real durante la generación autorregresiva.

El artículo concluye que, aunque entrenar N-vium incurre en un mayor costo de FLOPs por token (debido a la evaluación de cabezas en múltiples uniones), esta sobrecarga está justificada para la implementación en producción donde los costos de inferencia dominan, y la sobrecarga relativa disminuye a medida que escala el tamaño del modelo.