Resumen Técnico: Dual-Flow Transformers

Planteamiento del Problema

El coste de servir Modelos de Lenguaje de Gran Escala (LLMs) está cada vez más dominado por la inferencia en lugar del entrenamiento. Sin embargo, las dos fases principales de la inferencia —el prefill (procesamiento del prompt) y el decode autoregresivo (generación de tokens)— exhiben cuellos de botella de hardware fundamentalmente diferentes:

• Prefill: El procesamiento paralelo de los tokens del prompt es típicamente limitado por cómputo (compute-bound), restringido por el rendimiento aritmético.
• Decode: La generación secuencial es típicamente limitada por el ancho de banda de memoria (memory-bandwidth-bound), restringida por la necesidad de leer repetidamente el caché de Clave-Valor (KV) creciente y los pesos del modelo para cada nuevo token.

Las estrategias de escalado convencionales (aumentar el ancho o la profundidad del modelo) elevan los costes para ambas fases simultáneamente porque cada capa añadida se evalúa durante tanto el prefill como el decode. Esto crea una rigidez donde mejorar la calidad de la generación (que a menudo se beneficia de un mayor cómputo) aumenta inevitablemente el coste de procesar el prompt. El artículo plantea si es posible asignar cómputo aprendido adicional específicamente a la predicción de la continuación (decode) preservando el cómputo de todo el prompt y un único caché KV persistente.

Metodología: El Dual-Flow Transformer

Los autores proponen el Dual-Flow Transformer, una arquitectura que desacopla estructuralmente la trayectoria principal de procesamiento del prompt de una trayectoria auxiliar de predicción de la continuación.

Arquitectura Central

El modelo consta de dos trayectorias de residuo que comparten la mayoría de los parámetros del modelo pero difieren en su gestión de estado y alcance de activación:

1. Flujo Primario (Primary Flow): Una trayectoria de Transformer causal completa y estándar. Procesa todo el prompt, escribe el caché KV persistente y determina el estado primario. Es autónomo y no depende del flujo auxiliar.
2. Flujo Auxiliar (Auxiliary Flow): Una trayectoria de residuo que lee el estado y el caché primarios pero nunca escribe en el caché KV persistente y nunca retroalimenta al flujo primario.
   • Activación: El flujo auxiliar se omite durante el procesamiento inicial del prompt (prefill). Se activa solo en la posición final del prompt (para predecir el primer token de continuación) y en cada posición generada subsiguiente.
   • Compartición de Pesos: Ambos flujos comparten todas las matrices principales de atención, MLP y salida.
   • Separación: Utilizan embeddings de tokens distintos y vectores de acoplamiento ligeros y aprendidos.

Mecanismos Clave

• Atención de KV Compartido Asimétrico: El flujo primario realiza una auto-atención causal estándar. El flujo auxiliar genera una nueva consulta ($Q_2$), la cual se acopla con la consulta primaria ($Q_1$) mediante un vector aprendido ($a_\ell$). Crucialmente, el flujo auxiliar atiende a las Claves y Valores (Keys and Values) del flujo primario ($K_1, V_1$). Esto permite que el flujo auxiliar aproveche el estado almacenado en caché sin escribir el suyo propio, permitiendo una ejecución agrupada donde los pesos y las regiones KV se cargan una sola vez para servir a ambos flujos.
• Acoplamiento Ligero: La información fluye unidireccionalmente desde el flujo primario hacia el auxiliar a través de vectores aprendidos ($a_\ell, b_\ell, c_\ell$) en los niveles de consulta, intermedio de MLP y salida de MLP. Esto enriquece la representación auxiliar sin crear dependencias que alterarían la trayectoria primaria.
• Objetivo de Mezcla (Mixture Objective): La predicción final es una mezcla de las distribuciones primaria y auxiliar: $p_{mix} = \alpha_t p_{primary} + (1-\alpha_t) p_{aux}$. El peso de la mezcla $\alpha_t$ se ajusta dinámicamente basándose en la confianza (probabilidad de colisión) de la distribución primaria, asegurando que el flujo primario permanezca dominante (al menos un 50% de peso) durante la inferencia.
• Replay de Router (para MoE): En configuraciones de Mezcla de Expertos (MoE), el flujo auxiliar reutiliza el mismo conjunto de expertos seleccionados que el flujo primario (determinado por el router primario) pero los aplica con sus propios pesos de mezcla. Esto preserva la lógica de selección de expertos del flujo primario mientras aumenta el cómputo realizado sobre los expertos.

Costes de Inferencia vs. Entrenamiento

• Inferencia: La fase de prefill no incurre en ningún coste adicional comparado con un Transformer estándar; el flujo auxiliar se omite por completo. La fase de decode incurre en un coste de cómputo adicional (aproximadamente el doble de aritmética por token) pero no aumenta el tamaño del caché KV ni el conjunto de pesos únicos leídos, ya que los dos flujos se agrupan.
• Entrenamiento: Ambos flujos se evalúan a lo largo de toda la secuencia, lo que resulta en aproximadamente 2 veces más FLOPs de entrenamiento que el modelo primario. Sin embargo, el incremento en el número de parámetros es mínimo (solo embeddings separados y vectores de acoplamiento).

Contribuciones Clave

1. Arquitectura: Introducción de un diseño de KV compartido asimétrico que desacopla el cómputo de todo el prompt del cómputo de la continuación. El flujo primario gestiona la persistencia del estado, mientras que el flujo auxiliar proporciona una trayectoria "residual" para una predicción mejorada sin alterar el estado primario.
2. Calidad Empírica: Demostración de reducciones consistentes en la pérdida de validación a través de tres ejes experimentales:
   • Escalado de Datos (NanoGPT): Dual-Flow supera a los Transformers estándar con presupuestos de tokens equivalentes, con un suelo de pérdida ajustado más bajo.
   • Escalado de Tamaño de Modelo (Dense LLaMA): Las ganancias persisten a medida que aumenta el tamaño del modelo.
   • MoE Disperso (Sparse MoE): Validado en configuraciones MoE de estilo LLaMA con replay de router.
3. Asignación de Cómputo Específica por Fase: A diferencia de los diseños de flujo paralelo previos que aplican una asignación de cómputo uniforme, Dual-Flow permite configurar de forma independiente los fan-outs de experto del primario y el auxiliar. Esto permite un espacio de intercambio entre el coste de prefill, el coste de decode y la calidad predictiva, permitiendo optimizar para cargas de trabajo específicas (por ejemplo, fijar el coste de prefill mientras se aumenta el cómputo de decode).

Resultados

• Pérdida de Validación: En las comparaciones de presupuesto de tokens igualados en experimentos de NanoGPT, Dual-Flow logra consistentemente una pérdida de validación menor que los Transformers estándar. El suelo de pérdida asintótica ajustado disminuyó de 2.9416 (Transformer) a 2.9013 (Dual-Flow).
• Eficiencia de Cómputo: En una comparación de cómputo de entrenamiento igualado (donde un modelo Dual-Flow con la mitad de los pasos de datos realiza aproximadamente los mismos FLOPs de backbone compartido que un Transformer con el doble de pasos), Dual-Flow logró una pérdida de validación menor, lo que sugiere que asignar cómputo a través de flujos que interactúan es más efectivo que simplemente aumentar los datos para un solo flujo.
• Ablaciones: El análisis de componentes mostró que la construcción Dual mínima (embeddings separados, KV compartido) ya supera al arte previo más cercano (PHD-2). Añadir el acoplamiento entre flujos y la mezcla de probabilidad mejoró aún más el rendimiento.

Significado y Reivindicaciones

El artículo afirma que Dual-Flow establece una oportunidad estructural para la asignación de expertos específica por fase. Al desacoplar la trayectoria del prompt de la trayectoria de la continuación, la arquitectura permite que los sistemas de servicio traten el cómputo en tiempo de generación como un recurso escalable sin aumentar proporcionalmente el coste del procesamiento del prompt.

Específicamente, los autores destacan que en modelos MoE, los flujos primario y auxiliar pueden tener fan-outs de expertos independientes. Esto crea un control de intercambio de tres vías:

1. Coste de Prefill: Determinado por el uso de expertos del flujo primario.
2. Coste de Decode: Determinado por la suma del uso de expertos primarios y auxiliares.
3. Calidad Predictiva: Mejorada al aumentar el cómputo total asignado a la fase de continuación.

Esto permite una flexibilidad de despliegue específica donde un sistema puede preservar los costes de cómputo del prompt mientras gasta el presupuesto de decode disponible en más expertos auxiliares, o viceversa, sin cambiar el conjunto de parámetros compartidos ni la estructura de un único caché primario. El trabajo respalda la premisa de que el cómputo del siguiente token es un recurso de escalado útil que puede asignarse internamente dentro del modelo en lugar de externamente mediante trazas de razonamiento más largas o múltiples muestreos.