Resumen Técnico: Domino

Enunciado del Problema

La decodificación especulativa acelera la inferencia de Modelos de Lenguaje Grande (LLM) utilizando un modelo borrador ligero para proponer múltiples tokens futuros, los cuales son luego verificados en paralelo por el modelo objetivo. La efectividad de este enfoque está gobernada por un compromiso fundamental entre la calidad del borrador y el costo de generación:

1. Generación Autoregresiva (ej. serie EAGLE): Estos métodos generan tokens de borrador secuencialmente, modelando explícitamente las dependencias causales dentro del bloque de borrador. Esto produce borradores de alta calidad y longitudes de aceptación largas. Sin embargo, generar $k$ tokens requiere $k$ pasadas forward secuenciales y $k$ proyecciones completas de la cabeza del modelo de lenguaje (LM-head) sobre el vocabulario, incurriendo en una sobrecarga secuencial significativa que limita la aceleración.
2. Generación Paralela (ej. DFlash): Estos métodos generan un bloque completo de borrador en una sola pasada forward, eliminando la sobrecarga secuencial y maximizando la utilización de la GPU. Sin embargo, al eliminar las dependencias secuenciales, debilitan el modelado causal intra-bloque, lo que a menudo resulta en longitudes de aceptación más bajas en comparación con los métodos autoregresivos.

Los enfoques existentes luchan por equilibrar estos factores. Reducir los costos de proyección del vocabulario (ej. FR-Spec) o utilizar verificación basada en árboles no resuelve completamente la tensión entre el costo secuencial del modelado causal y la pérdida de calidad de la generación paralela. La pregunta central abordada es: ¿Podemos retener los beneficios de calidad del borrador del modelado de dependencias causales mientras preservamos la eficiencia de bajo costo de la generación de bloques en paralelo?

Metodología: Domino

Los autores proponen Domino, un marco de decodificación especulativa que desacopla el modelado de dependencias causales de la ejecución costosa de borradores autoregresivos. La arquitectura consta de dos componentes principales:

1. Columna Vertebral de Borrador Paralelo

Domino utiliza una columna vertebral de borrador paralela (instanciada a través de la arquitectura DFlash) para generar distribuciones de borrador preliminares para todo el bloque en una sola pasada forward no autoregresiva.

• Entrada: La columna vertebral toma características de contexto del modelo objetivo desde el prefijo verificado y un bloque de borrador enmascarado (que contiene el último token verificado seguido de tokens [MASK]).
• Salida: Produce estados ocultos para todo el bloque, a partir de los cuales se calculan logits base ($L^{base}_i$) utilizando una cabeza de LM del objetivo congelada.

2. Cabeza de Domino (Corrección Causal Ligera)

Para reintroducir información causal sin ejecución secuencial, Domino aplica una cabeza de corrección ligera sobre los logits base paralelos. Esta cabeza opera secuencialmente pero es computacionalmente barata:

• Codificador Causal: Para cada posición de borrador $i$, un GRU ligero resume los embeddings de los tokens borradores precedentes (o tokens de verdad fundamental durante el entrenamiento) en un estado causal $S_{i-1}$. Este estado captura información dependiente del prefijo sin invocar un modelo autoregresivo completo.
• Cabeza de Corrección de Baja Rango: La cabeza calcula una corrección residual en el espacio de logits utilizando un cuello de botella de baja rango. Concatena el estado oculto base $H_i$ y el estado causal $S_{i-1}$, los proyecta en un espacio de baja rango (dimensión $r=256$) y los mapea de vuelta al espacio del vocabulario para producir logits de corrección $\Delta L_i$.
• Logits Finales: Los logits finales del borrador se calculan como $L_i = L^{base}_i + \Delta L_i$. Los tokens se muestrean a partir de esta distribución refinada.

Crucialmente, la corrección se aplica en el espacio de logits en lugar del espacio oculto. Esto evita la necesidad de volver a ejecutar la cabeza completa del LM después de cada actualización causal, manteniendo el cálculo base en paralelo y restringiendo la rama secuencial a una actualización residual de baja rango.

Estrategia de Entrenamiento

Para estabilizar el entrenamiento y evitar que la rama de corrección "ataje" a la columna vertebral, los autores introducen dos estrategias clave:

1. Codificación Causal Forzada por el Profesor: A diferencia de los métodos que utilizan prefijos generados por sí mismos (Prueba en Tiempo de Entrenamiento), Domino alimenta al codificador causal con embeddings de tokens de verdad fundamental. Esto alinea la distribución de entrenamiento con el mecanismo de aceptación de la decodificación especulativa, donde las correcciones solo importan cuando los tokens precedentes son verificados como correctos.
2. Plan de Estudios Anclado a la Base: Para evitar que la columna vertebral paralela colapse (es decir, se vuelva débil porque la rama de corrección la compensa), la pérdida de entrenamiento es una suma ponderada de la pérdida de logits base y la pérdida de logits finales. El peso $\lambda_t$ se recocido linealmente de 1 a 0, forzando inicialmente a la columna vertebral a aprender una distribución base fuerte antes de desplazar gradualmente la optimización hacia la distribución final corregida.

Contribuciones Clave

• Desacoplamiento del Modelado Causal de la Ejecución: Domino separa el modelado de dependencias causales de la ejecución costosa de pasos autoregresivos, permitiendo refinamiento causal sin latencia secuencial.
• Corrección Causal Ligera: Al utilizar un codificador basado en GRU y un cuello de botella de baja rango en el espacio de logits, Domino inyecta información dependiente del prefijo con una sobrecarga de parámetros mínima (añadiendo solo ~5.3% de parámetros a la columna vertebral) y un aumento de latencia insignificante.
• Plan de Estudios de Entrenamiento Estabilizado: La introducción de la fuerza del profesor combinada con una pérdida de plan de estudios anclada a la base previene el colapso de la columna vertebral y asegura un aprendizaje robusto tanto de la distribución base paralela como de las correcciones causales.

Resultados Experimentales

Los experimentos se realizaron en los modelos Qwen3-4B y Qwen3-8B en benchmarks de razonamiento matemático (GSM8K, MATH, AIME25), generación de código (HumanEval, MBPP, LiveCodeBench) y diálogo (MT-Bench, Alpaca).

• Aceleración de Extremo a Extremo: Bajo el backend Transformers, Domino logra una aceleración de hasta 5.49× en Qwen3-8B (decodificación voraz), superando a DFlash (4.66×) y EAGLE-3 (2.26×). Específicamente en GSM8K, alcanza una aceleración de 7.92×.
• Longitud de Aceptación: Domino mejora significativamente la longitud promedio de aceptación en comparación con las líneas base paralelas (por ejemplo, aumentando de 6.59 a 10.03 en Qwen3-8B para GSM8K), reduciendo la brecha con los métodos autoregresivos mientras mantiene la eficiencia paralela.
• Rendimiento de Alto Concurrencia: Utilizando SGLang, Domino demuestra un rendimiento superior a través de niveles variables de concurrencia (de 2 a 32 solicitudes), logrando una aceleración de rendimiento de hasta 5.8× sobre la decodificación autoregresiva.
• Estudios de Ablación:
  • Eliminar la cabeza de Domino reduce la aceleración promedio de 3.31× a 2.84×, confirmando la eficacia de la corrección causal.
  • Utilizar fuerza del profesor en lugar de prueba en tiempo de entrenamiento mejora la longitud de aceptación (de 3.80 a 3.96).
  • El plan de estudios anclado a la base aumenta aún más la longitud de aceptación (de 3.96 a 4.19) y previene el colapso de la pérdida de la columna vertebral.

Significado

El artículo afirma que Domino demuestra que la información causal puede reintroducirse efectivamente en la generación de borradores paralelos para lograr una inferencia de LLM más rápida. Al desacoplar el modelado de dependencias del costo de ejecución, Domino logra un equilibrio superior entre la calidad del borrador y la sobrecarga de generación. Los resultados sugieren que la corrección causal ligera es una vía viable para superar las limitaciones de las estrategias de decodificación especulativa puramente autoregresivas y puramente paralelas, ofreciendo una solución práctica para acelerar el servicio de LLM sin comprometer la fidelidad de la distribución de salida.