DARTree: Speculative Diffusion Decoding with Autoregressive Draft Trees

Resumen Técnico: DARTree

Planteamiento del Problema

Los modelos de lenguaje de gran escala (LLM) autorregresivos (AR) modernos son inherentemente secuenciales, generando solo un token por paso mientras requieren acceso a casi todos los parámetros del modelo. Esto crea un cuello de botella limitado por la E/S que restringe la velocidad de inferencia y desperdicia las capacidades de computación paralela del hardware moderno.

La decodificación especulativa aborda esto utilizando un "borrador" (drafter) ligero para proponer múltiples tokens futuros, que luego son verificados en paralelo por el modelo objetivo. Sin embargo, los enfoques existentes enfrentan un compromiso entre la latencia de propuesta y el condicionamiento causal:

1. Borradores Autorregresivos: Producen propuestas de forma secuencial. Aunque mantienen fuertes dependencias causales, su latencia crece linealmente con la longitud del borrador.
2. Borradores de Bloque Basados en Difusión: Predicen un bloque entero de tokens en una sola pasada hacia adelante (forward pass), reduciendo drásticamente la latencia de propuesta. Sin embargo, predicen distribuciones marginales posición a posición condicionadas solo al prefijo verificado, ignorando los tokens realizados dentro del mismo bloque de borrador. Esto crea un "desajuste causal" con la distribución del modelo objetivo.
3. Corrección Causal: Las cabezas de corrección ligeras pueden propagar los tokens seleccionados a lo largo de una única cadena de borrador para restaurar las dependencias intra-bloque. Sin embargo, al combinarse con la verificación de árbol de borrador (que expone múltiples continuaciones candidatas al modelo objetivo), surge un nuevo cuello de botella de eficiencia. Debido a que la puntuación corregida de los hijos de un nodo depende del camino específico tomado para llegar a dicho nodo, la construcción estándar de árboles de mejor primera búsqueda (usando un max-heap) obliga a la cabeza de corrección a operar secuencialmente: se extrae un nodo, se corrige y se insertan sus hijos en el heap antes de que el siguiente nodo pueda ser expandido. Esta interposición de la inferencia de corrección y las operaciones de heap impide el procesamiento por lotes (batching), convirtiendo la construcción del árbol en una fuente significativa de latencia.

Metodología: DARTree

Los autores introducen DARTree, un método de decodificación especulativa libre de entrenamiento que extiende un borrador de bloque paralelo con corrección causal preentrenado de una cadena única a un árbol especulativo. La innovación central es el desacoplamiento de la corrección causal dependiente de la ruta de las operaciones secuenciales de heap.

1. Expansión AR Paralela por Profundidad

En lugar de expandir el árbol nodo por nodo (lo que fuerza la corrección secuencial), DARTree expande el árbol profundidad por profundidad en lotes paralelos:

• Borrador de Bloque Paralelo: El borrador de difusión primero genera representaciones compartidas para todo el bloque de borrador.
• Selección de Candidatos: En cada profundidad $d$, el sistema selecciona los $K$ candidatos principales de los logits base.
• Corrección por Lotes (Batched): Todas las ramas activas en la profundidad actual se procesan simultáneamente. La cabeza de corrección preentrenada evalúa los candidatos para todas las ramas en un solo lote, utilizando el estado del prefijo específico de cada rama.
• Puntuación y Retención: Cada extensión de candidato se puntúa utilizando su probabilidad acumulada corregida. El sistema retiene las $W$ ramas con mayor puntuación para la siguiente profundidad, transportando sus estados de corrección hacia adelante.
• Construcción de Superárbol: Este proceso se repite para todas las profundidades de borrador, construyendo un "superárbol de candidatos de ancho fijo" donde cada profundidad tiene una carga de trabajo regular y tensorizada.

2. Poda Diferida de Mejor Primera Búsqueda

Una vez que se construye el superárbol completo y todos los nodos han sido puntuados, DARTree aplica una estrategia de poda diferida:

• Selección Global: En lugar de intercalar operaciones de heap con la corrección, DARTree realiza una única selección Top-$B$ global a través de todos los prefijos materializados en el superárbol.
• Equivalencia Teórica: Los autores demuestran (Lema 1) que si se aplica un bono de profundidad no positivo ($\beta \le 0$) para penalizar los prefijos más profundos, el conjunto de nodos seleccionados por una operación Top-$B$ global es idéntico al conjunto seleccionado por un max-heap secuencial de mejor primera búsqueda. Esto permite al sistema evitar por completo el cuello de botella secuencial.
• Verificación: El árbol resultante compacto se pasa al modelo objetivo para una única pasada hacia adelante utilizando atención de árbol (tree attention).

Contribuciones Clave

1. Identificación de un Cuello de Botella Secuencial: El artículo identifica que el acoplamiento de la corrección causal condicionada por la ruta con la construcción de árboles de mejor primera búsqueda nodo a nodo crea un cuello de botella de latencia que impide el procesamiento por lotes eficiente.
2. Construcción por Lotes por Profundidad Libre de Entrenamiento: DARTree introduce un método que aplica la corrección causal a través de múltiples ramas en paralelo en cada profundidad, logrando una aceleración sin pérdidas sin requerir el entrenamiento adicional de la cabeza de corrección.
3. Estrategia de Poda Diferida: Los autores proponen una estrategia de poda de mejor primera búsqueda simplificada que reemplaza las operaciones secuenciales de heap con una única selección global Top-$B$ después de que el superárbol ha sido completamente puntuado, manteniendo la estructura de árbol óptima bajo condiciones de puntuación específicas.

Resultados Experimentales

Los autores evaluaron DARTree en Qwen3-4B y Qwen3-8B a través de siete benchmarks (Matemáticas: GSM8K, MATH-500, AIME25; Código: HumanEval, MBPP; Chat: MT-Bench, Alpaca) a temperaturas $T=0$ y $T=1$.

• Rendimiento: DARTree logró la mayor longitud de aceptación promedio ($\tau$) y aceleración en las cuatro configuraciones de temperatura del modelo.
  • En la mejor configuración (Qwen3-4B, $T=0$, GSM8K), DARTree aceptó 12.97 tokens por ronda de verificación, un aumento del 98.6% sobre DFlash y un 27.9% sobre Domino.
  • Esto resultó en una aceleración sin pérdidas de 9.73× sobre la decodificación autorregresiva medida localmente.
• Comparación:
  • DARTree (podado) superó al método basado en difusión de estado del arte anterior, Domino, por hasta un 46.8% en longitud de aceptación y un 40.1% en aceleración.
  • También superó a DDTree (que utiliza predicciones marginales sin corrección) y a Domino-chain + DDTree (que aplica la corrección solo a una única cadena antes de la construcción del árbol).
• Ablación:
  • Latencia: Un enfoque de corrección secuencial con un heap incurrió en ~70ms por ronda, mientras que DARTree mantuvo ~28ms, demostrando la eficiencia del enfoque de construcción por lotes por profundidad.
  • Generalizabilidad: DARTree mejoró con éxito el rendimiento cuando se aplicó al borrador DSpark (usando una cabeza de corrección de Markov), confirmando que es un método general para borradores paralelos con corrección causal.
  • Hiperparámetros: Aumentar el ancho de capa ($W$) mejoró la aceptación hasta cierto punto (rendimientos decrecientes después de $W=12$), mientras que el presupuesto de verificación ($B$) mostró una aceleración máxima alrededor de $B=64-128$.

Significado

El artículo afirma que DARTree resuelve la tensión fundamental en la decodificación especulativa entre la fidelidad causal (que requiere corrección específica de la ruta) y la eficiencia computacional (que requiere procesamiento por lotes paralelo). Al reestructurar el proceso de construcción del árbol de forma por profundidad en lugar de por nodo, DARTree permite el uso de borradores de difusión de alta capacidad con corrección causal sin incurrir en las penalizaciones de latencia secuencial de la búsqueda tradicional basada en heap. El método se presenta como una mejora libre de entrenamiento que puede aplicarse a las cabezas de corrección preentrenadas existentes, ofreciendo ganancias sustanciales y consistentes en la velocidad de inferencia a través de diversos dominios y temperaturas de decodificación.