PRESTO: Prefix-Aligned Tree Drafting for Diffusion Speculative Decoding

Imagina que estás intentando predecir la siguiente palabra en una historia. Durante mucho tiempo, las computadoras más inteligentes (llamadas Modelos de Lenguaje Extensos) hacían esto palabra por palabra, como una persona leyendo un libro en voz alta, deteniéndose después de cada palabra para pensar qué viene después. Esto es preciso, pero es lento. Recientemente, los científicos descubrieron una nueva forma de escribir estas historias usando modelos de "difusión". Piensa en estos como un escultor que comienza con un bloque de mármol y talla toda la estatua de una vez, en lugar de ir quitando trozos diminutos uno por vez. Esto permite que la computadora adivine muchas palabras simultáneamente, lo cual es increíblemente rápido.

Sin embargo, hay un inconveniente. Cuando adivinas muchas palabras a la vez, podrías equivocarte en algunas. Para solucionar esto, existe un truco ingenioso llamado "decodificación especulativa". Es como tener un asistente junior rápido que adivina las siguientes pocas palabras, y luego un jefe superinteligente y lento que comprueba si esos aciertos son correctos. Si el jefe está de acuerdo, acepta todo el lote de palabras instantáneamente, ahorrando muchísimo tiempo. El problema es que el asistente junior (el modelo de difusión) es excelente adivinando palabras individuales, pero no siempre sabe cómo esas palabras encajan entre sí en un orden específico. Es como si el asistente fuera excelente eligiendo ingredientes individuales para un pastel, pero no siempre supiera qué combinación de ingredientes sabrá bien junta.

Aquí es donde entra un nuevo artículo de investigación. Los investigadores, liderados por Zheng Wang y sus colegas, se dieron cuenta de que la forma actual de usar estos asistentes rápidos está dejando mucha velocidad sin aprovechar. Descubrieron que, mientras el asistente puede generar una enorme variedad de combinaciones de palabras, el método actual solo comprueba un único camino, como caminar por un solo pasillo esperando que la puerta al final esté abierta. Los autores proponen un nuevo sistema llamado PRESTO (Prefix-Aligned Tree Drafting). En lugar de caminar por un solo pasillo, PRESTO construye un árbol de posibilidades, explorando muchos caminos diferentes a la vez. Pero aquí está la magia: corrige un fallo fundamental en cómo se mide la confianza del asistente. La confianza original del asistente es "ciega al prefijo", lo que significa que no le importa qué palabras vinieron antes. PRESTO añade una puntuación "alineada al prefijo", que actúa como una brújula, asegurando que los caminos elegidos para la exploración sean los que tienen más probabilidades de ser aceptados por el jefe.

El resultado es un sistema significativamente más rápido. En sus pruebas, PRESTO ayudó a la computadora a aceptar más palabras en cada ronda de adivinación. En algunas de las mejores configuraciones existentes, hizo que todo el proceso fuera 1.5 veces más rápido. En otras configuraciones, proporcionó una aceleración de 1.12 veces. El artículo sugiere que, al tratar el proceso de adivinación como una aventura de ramificación de árboles en lugar de una línea recta, y al asegurarse de que las ramas se elijan basándose en qué tan bien encajan con la historia hasta el momento, podemos obtener lo mejor de ambos mundos: la velocidad de la difusión y la precisión de una revisión cuidadosa.

El Problema: La trampa del "camino único"

Para entender por qué se necesita PRESTO, imagina que estás jugando a un juego de "Mad Libs" con un amigo que intenta adivinar las palabras que faltan. Tu amigo es un modelo de difusión. Es increíble al mirar un espacio en blanco y decir: "¡Apuesto a que la palabra aquí es 'gato'!", o "¿Tal vez sea 'perro'?", o "¿O quizás 'cohete'?". Pueden gritar todas estas opciones al mismoamente.

Sin embargo, la forma actual de usar a este amigo es muy rígida. Toma la mejor suposición de su amigo, la escribe y luego le pregunta al "jefe" (el modelo objetivo) si esa suposición es correcta. Si el jefe dice "No", todo se desecha y tienes que empezar de nuevo. Si el jefe dice "Sí", avanzas a la siguiente palabra y repites el proceso. Esto se llama redacción lineal. Es como caminar por un bosque y solo mirar el camino que tienes directamente delante de ti.

Los autores observaron que este enfoque es ineficiente. Debido a que el modelo de difusión genera muchas opciones simultáneamente, existe un enorme "espacio combinatorio" de posibilidades. Es como tener un mapa con mil senderos diferentes, pero solo se te permite caminar por uno. El artículo muestra que, al limitarse a un solo camino, el sistema pierde muchas rutas válias. De hecho, en problemas matemáticos como GSM8K, el método actual aceptaba unas 6.5 palabras en promedio, pero los investigadores calcularon que, si pudieran comprobar todos los mejores caminos, habrían podido aceptar casi 10 palabras. ¡Esa es una brecha enorme!

El Desajuste: La brújula "ciega"

Los investigadores profundizaron más y descubrieron una razón específica por la cual simplemente comprobar más caminos (construir un árbol) no funcionaba perfectamente con los métodos antiguos. Identificaron un "desajuste fundamental".

En el mundo de la IA estándar (modelos autorregresivos), la puntuación de confianza de una palabra depende fuertemente de las palabras que la precedieron. Si la frase es "El gato se sentó en el...", el modelo sabe que "tapete" es una palabra muy probable, pero "pizza" no lo es. Esto es alineado al prefijo.

Pero los modelos de difusión funcionan de manera diferente. Generan una probabilidad "marginal" para cada posición de forma independiente. Es como si el modelo dijera: "En la posición 5, 'gato' tiene un 80% de probabilidad", sin importarle si la posición 4 fue "El" o "El rápido marrón". Esto es ciego al prefijo.

Cuando intentas construir un árbol de suposiciones usando estas puntuaciones ciegas, obtienes un problema de clasificación. Podrías elegir un camino que parece genial para la primera palabra, pero terrible para la segunda palabra porque el modelo no se dio cuenta de que la primera palabra cambió el contexto. Es como un GPS que te da direcciones basadas solo en la calle actual, ignorando el hecho de que acabas de girar a la izquierda y ahora estás en una calle de sentido único. El artículo argumenta que usar estas puntuaciones ciegas para construir un árbol conduce a una "clasificación de rutas poco fiable", lo que significa que el sistema explora las ramas equivocadas y pierde el tiempo.

La Solución: PRESTO

PRESTO (Prefix-Aligned Scoring and priority-based Tree search for diffusion Speculative decOding) resuelve esto añadiendo una "corrección" a las puntuaciones del modelo de difusión.

1. Puntuación Alineada al Prefijo: Los autores se dieron cuenta de que necesitaban combinar la fuerte señal "marginal" del modelo de difusión (qué tan probable es una palabra por sí sola) con una señal "condicionada al prefijo" (qué tan probable es dado los términos anteriores). Crearon una nueva fórmula de puntuación que multiplica la probabilidad de difusión por un factor de corrección derivado de un modelo n-grama simple (una herramienta ligera que observa combinaciones de palabras). Esto crea una puntuación que respeta el flujo de la historia.
2. Búsqueda de Árbol Basada en Prioridades: En lugar de simplemente elegir el mejor camino, PRESTO construye un árbol. Utiliza las nuevas puntuaciones corregidas para decidir qué ramas hacer crecer. Prioriza los caminos que tienen mayor potencial de ser aceptados por el jefe. Es como un excursionista que, en lugar de caminar recto, mira un mapa y elige el sendero que tiene más probabilidades de llevarlo a la cima, incluso si ese sendero no es el más obvio al principio.

Los autores probaron dos formas de hacer crecer este árbol: Búsqueda de Haz (Beam Search) (mantener un número fijo de los mejores caminos en cada paso) y Búsqueda de Mejor Primero (Best-First Search) (expandir siempre el único mejor camino encontrado hasta el momento). Descubrieron que para su configuración específica, la Búsqueda de Haz funcionaba tan bien como la más compleja Búsqueda de Mejor Primero, por lo que se quedaron con la opción más simple y eficiente.

Los Resultados: Más Rápidos y Más Inteligentes

Los autores pusieron a prueba PRESTO en diversas tareas, incluyendo problemas matemáticos (GSM8K, Math500), desafíos de programación (HumanEval, LiveCodeBench) y conversaciones de chat. Utilizaron dos tipos de sistemas:

• Redactores de Difusión Dedicados: Un modelo de difusión pequeño y rápido que adivina para un modelo autorregresivo más grande (como dFlash).
• LLMs de Difusión Auto-especulativos: Un único modelo de difusión que adivina y se comprueba a sí mismo (como Nemotron-Labs-Diffusion).

Los resultados fueron consistentes en todos los ámbitos. PRESTO aumentó consistentemente la Longitud de Aceptación Promedio, que es el número de palabras que el jefe acepta en un solo intento.

• En el sistema dFlash (usando Qwen3-8B), la longitud de aceptación promedio saltó de unas 6.6 palabras a 9.6 palabras. Esto se tradujo en una aceleración de extremo a extremo de 1.5 veces.
• En el sistema Nemotron-Labs-Diffusion, la longitud de aceptación aumentó de 8.8 a 9.9 palabras, lo que resultó en una aceleración de 1.12 veces.

Quizás lo más impresionante es que el artículo demostró que PRESTO funciona incluso cuando el sistema es "estocástico" (aleatorio), lo cual suele ser más difícil de predecir. En estos casos, la aceleración fue aún más pronunciada, con algunos puntos de referencia viendo casi el doble de rendimiento.

Los autores también comprobaron el "costo" de este nuevo método. Descubrieron que el trabajo adicional requerido para construir el árbol y calcular las nuevas puntuaciones era minúsculo, menos del 4% del tiempo total. La gran mayoría del tiempo (más del 90%) se seguía dedicando a la verificación real por parte del modelo jefe. Esto significa que PRESTO es una actualización altamente eficiente que no ralentiza el sistema con sobrecargas.

Lo que PRESTO NO ES

Es importante señalar lo que el artículo no afirma. Los autores declaran explícitamente que aplicar simplemente una estructura de árbol ingenua (sin su puntuación alineada al prefijo) es subóptimo. Si simplemente tomas las puntuaciones puras del modelo de difusión y construyes un árbol, no obtienes todo el beneficio debido a la naturaleza "ciega" de las puntuaciones. PRESTO se trata específicamente de corregir ese desajuste de puntuación.

Además, el artículo no pretende haber resuelto el problema de los modelos de difusión por completo. Reconocen que su método depende de una "señal alineada al prefijo tratable" (como el modelo n-grama que utilizaron) para realizar el trabajo pesado de la corrección. Sugieren que trabajos futuros podrían explorar señales aún más ricas, pero por ahora, su simple corrección es suficiente para ver ganancias masivas.

Por qué esto es importante

En la carrera por hacer que la IA sea más rápida y eficiente, cada brizna de velocidad cuenta. La decodificación especulativa ha sido un tema candente porque nos permite usar los mejores y más precisos modelos sin pagar el pleno penal de tiempo. Sin embargo, los métodos actuales estaban limitados por el hecho de que trataban a los modelos de difusión como máquinas lineales, ignorando su capacidad única para generar muchas opciones simultáneamente.

PRESTO cambia las reglas del juego al tratar a los modelos de difusión como los exploradores de múltiples caminos que son. Al alinear la puntuación con la forma en que el jefe comprueba el trabajo, desbloquea todo el potencial de la generación paralela del modelo de difusión. El resultado es un sistema que no es solo ligeramente más rápido, sino significativamente más eficiente, permitiéndonos generar texto, resolver problemas matemáticos y escribir código a velocidades que antes se consideraban imposibles para este tipo de modelos. Como dicen los autores, han convertido una caminata de "camino único" en una expedición basada en árboles, asegurando que cada paso dado sea un paso hacia la respuesta correcta.