💬 NLP

Prefilling-dLLM: Predictive Prefilling for Long-Context Inference in Diffusion Language Models

El artículo presenta Prefilling-dLLM, un marco de trabajo libre de entrenamiento que acelera la inferencia de contexto largo en modelos de lenguaje de difusión mediante el almacenamiento en caché y la selección dispersa de fragmentos de prefijo relevantes para reducir la complejidad computacional de cuadrática respecto a la longitud de la secuencia completa a cuadrática respecto a la longitud de decodificación, logrando así aceleraciones de vanguardia y mitigando el fenómeno de pérdida en el medio.

Autores originales: Jing Xiong, Qi Han, Shansan Gong, Yunta Hsieh, Chengyue Wu, Chaofan Tao, Chenyang Zhao, Ngai Wong

Publicado 2026-06-10

📖 5 min de lectura🧠 Análisis profundo

CC BY 4.0

Autores originales: Jing Xiong, Qi Han, Shansan Gong, Yunta Hsieh, Chengyue Wu, Chaofan Tao, Chenyang Zhao, Ngai Wong

Artículo original bajo licencia CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Esta es una explicación generada por IA del artículo a continuación. No ha sido escrita ni avalada por los autores. Para mayor precisión técnica, consulte el artículo original. Leer descargo de responsabilidad completo

Imagina que estás intentando resolver un rompecabezas gigante, pero en lugar de mirar la imagen completa de una vez, tienes que reexaminar la caja entera de piezas cada vez que colocas una sola pieza nueva sobre la mesa.

Eso es esencialmente cómo funcionan actualmente los Modelos de Lenguaje de Difusión (dLLMs) cuando manejan historias o documentos largos. Cada vez que el modelo intenta generar una nueva palabra, vuelve a leer y a procesar todo el historial de la conversación desde el principio. A medida que la historia se vuelve más larga, este "releer" se vuelve tan lento y costoso que es como intentar correr un maratón cargando una mochila pesada que se vuelve más pesada con cada paso.

El artículo presenta un nuevo método llamado Prefilling-dLLM para solucionar esto. Así es como funciona, utilizando algunas analogías de la vida cotidiana:

1. El Problema: La trampa del "Releer"

En los modelos de IA tradicionales, una vez que lees una página de un libro, la recuerdas y solo te concentras en la nueva página que estás leyendo. Pero en estos modelos de difusión, la computadora olvida el libro y vuelve a leer el libro entero desde la página 1 hasta la 1,000 cada vez que escribe una sola palabra nueva.

El Resultado: Si el libro es corto, está bien. Si el libro tiene 32,000 páginas, la computadora pasa el 99% de su tiempo simplemente releyendo las páginas antiguas y solo el 1% escribiendo las nuevas.

2. La Solución: El sistema del "Bibliotecario Inteligente"

Los autores proponen un sistema llamado Prefilling-dLLM que actúa como un bibliotecario súper eficiente. En lugar de releer toda la biblioteca cada vez, el bibliotecario hace dos cosas:

Paso A: El "Prefill" (Organizar la Biblioteca)

Antes de que comience la escritura, el bibliotecario toma el texto de entrada masivo (el "prefijo") y lo divide en fragmentos (como capítulos de un libro).

El Truco: El bibliotecario lee cada capítulo una vez, crea una "tarjeta de resumen" (llamada caché KV) para él y la pone en un estante.
La Innovación: El bibliotecario no lee cada palabra de cada capítulo. Utilizan un truco especial para mantener solo las oraciones más importantes de cada capítulo, descartando el relleno. Esto hace que las tarjetas de resumen sean mucho más pequeñas y rápidas de manejar.

Paso B: El "Decode" (Escribir la Historia)

Ahora, cuando el modelo necesita escribir la siguiente palabra, no vuelve a leer toda la biblioteca.

La Selección: El modelo pregunta: "¿Qué capítulos son realmente relevantes para lo que estoy escribiendo ahora mismo?". Utiliza un sistema de puntuación inteligente para elegir solo los pocos capítulos superiores (los fragmentos más relevantes) del estante.
La Eficiencia: Ignora el otro 90% de la biblioteca que no es necesario para esta oración específica. Solo mira las "tarjetas de resumen" relevantes que hizo anteriormente.

3. Por qué esto es un cambio radical

El artículo afirma que este enfoque es una aceleración masiva porque:

No más Releer: El trabajo pesado de leer el texto largo ocurre una sola vez al principio.
Salto Inteligente: Durante la fase de escritura real, el modelo solo mira una fracción diminuta del texto (los fragmentos más relevantes), en lugar de todo el contenido.
La Velocidad: En contextos largos (de 8,000 a 32,000 palabras), este método es de 9 a 28 veces más rápido que los métodos anteriores, manteniendo la precisión de las respuestas.

4. Resolviendo el misterio de "Perdido en el Medio"

Existe un problema conocido con los modelos de IA de largo alcance llamado "Perdido en el Medio" (Lost in the Middle). Imagina a una persona leyendo una larga lista de hechos; recuerda perfectamente los primeros y los últimos, pero olvida por completo lo que hay en medio.

La Solución del Artículo: Los autores descubrieron que al dividir el texto en fragmentos y añadir un "token de anclaje" especial (como un encabezado de capítulo) al inicio de cada fragmento, el modelo puede encontrar información en cualquier parte del texto, incluso en el medio. Es como tener un índice que funciona perfectamente, de modo que el modelo nunca se pierde, sin importar qué tan larga sea la historia.

5. El equilibrio del "Tamaño del Fragmento"

El artículo también probó qué tan grandes deberían ser estos "capítulos".

Demasiado Grandes: Si los fragmentos son enormes, el modelo podría perder detalles importantes en medio del texto.
Demasiado Pequeños: Si los fragmentos son minúsculos, la computadora pasa demasiado tiempo gestionando la lista de fragmentos.
El Punto Dulce: Descubrieron que para textos muy largos, usar muchos fragmentos más pequeños funciona mejor para mantener la precisión alta.

Resumen

Prefilling-dLLM es como cambiar a un estudiante que vuelve a leer un libro de texto de 500 páginas cada vez que responde una pregunta de un examen, por un estudiante que crea una guía de estudio condensada y de forma inteligente una vez, y luego solo consulta las páginas específicas que necesita al responder.

¿El resultado? La IA puede manejar conversaciones y documentos mucho más largos sin ralentizarse hasta detenerse, y no olvida los detalles importantes ocultos en medio del texto.

Resumen Técnico: Prefilling-dLLM

Planteamiento del Problema

Los Modelos de Lenguaje de Gran Escala de Difusión (dLLMs) ofrecen una alternativa prometedora a los modelos autorregresivos (AR) al permitir la generación paralela de tokens mediante la eliminación iterativa de ruido. Sin embargo, los dLLMs enfrentan una ineficiencia crítica en escenarios de contexto largo: en cada paso de eliminación de ruido, el modelo debe volver a codificar todo el prefijo de entrada para computar la autoatención. Esto resulta en una complejidad computacional que escala cuadráticamente con la longitud total de la secuencia ( $O((L_p + L_d)^2 \cdot T)$ , donde $L_p$ es la longitud del prefijo, $L_d$ es la longitud de decodificación y $T$ es el número de pasos de eliminación de ruido). A diferencia de los modelos AR, donde la fase de prellenado (prefill) está limitada por cómputo y la fase de decodificación está limitada por memoria, la inferencia de los dLLM permanece limitada por cómputo durante todas las iteraciones porque la secuencia completa (prefijo + respuesta) se procesa conjuntamente en cada paso. Esta recomputación repetida del prefijo se vuelve prohibitiva para aplicaciones de contexto largo (por ejemplo, 8K–32K+ tokens).

Metodología: PREFILLING-dLLM

Los autores proponen PREFILLING-dLLM, un marco de desagregación de prellenado-decodificación libre de entrenamiento diseñado para eliminar la recomputación del prefijo e introducir dispersión (sparsity). El método opera en tres etapas:

1. Prellenado fragmentado (Chunked Prefill) y Caché de KV

En lugar de procesar todo el prefijo como un único bloque, el prefijo de entrada se particiona en $N$ fragmentos (chunks) no solapados de tamaño fijo $C$ .

Procesamiento Independiente: Cada fragmento se procesa de forma independiente con atención bidireccional intra-fragmento. Se antepone un token especial de Inicio de Secuencia (BOS) a cada fragmento para actuar como delimitador.
Puntuación Predictiva: Para determinar qué fragmentos son relevantes sin tener objetivos de verdad fundamental (ground-truth), el marco genera etiquetas pseudo (pseudo-labels $m$ $m$ ) mediante una pasada de difusión preliminar. Luego, califica cada fragmento utilizando dos señales:
- Puntuación de Auto-información ( $s_I$ ): La log-verosimilitud negativa de la ventana de consulta condicionada al fragmento.
- Puntuación de Pseudo-etiqueta ( $s_P$ ): La verosimilitud de las pseudo-etiquetas condicionada al fragmento y la consulta.
Selección: Se seleccionan los $K$ fragmentos más relevantes basándose en la puntuación combinada ( $s = s_I + s_P$ ). Las representaciones KV de estos fragmentos seleccionados se almacenan en caché.

2. Dispersión de Dos Niveles

El marco introduce dispersión tanto a nivel de fragmento como de token para reducir la complejidad:

Dispersión Inter-fragmento: Solo los $K$ fragmentos superiores participan en el mecanismo de atención durante la decodificación, en lugar del conjunto completo de $N$ fragmentos.
Dispersión Intra-fragmento: Dentro de cada fragmento seleccionado, los tokens se clasifican por su peso de atención bidireccional acumulado con los tokens de consulta. Solo se retienen en el caché KV los $B$ tokens más relevantes por fragmento, eliminando los tokens irrelevantes.

3. Decodificación con Reutilización de KV Estática

Durante el proceso iterativo de eliminación de ruido:

Las representaciones KV almacenadas en caché de los $K$ fragmentos seleccionados permanecen estáticas y se reutilizan en todos los $T$ pasos de eliminación de ruido.
El modelo solo vuelve a computar las representaciones KV para la consulta y los tokens de respuesta en evolución.
Esto reduce la complejidad por paso de $O((L_p + L_d)^2 \cdot T)$ a $O(N \cdot C^2 + (L_d^2 + K \cdot C) \cdot T)$ , desacoplando efectivamente el costo de decodificación de la longitud total del prefijo.

Optimización del Kernel de Atención

Para abordar los cuellos de botella de memoria causados por la carga de cachés KV de fragmentos no contiguos, los autores implementan un kernel Split-S FlexAttention. Este kernel paraleliza la decodificación sobre los cachés KV de fragmentos almacenados de forma no contigua mediante la partición de la dimensión KV y el procesamiento de fragmentos de forma independiente antes de fusionar los resultados mediante una reducción log-sum-exp, evitando la necesidad de reunir los datos en una memoria contigua.

Hallazgos Clave y Resultados

Rendimiento y Eficiencia

Aceleración: En LongBench e InfiniteBench, PREFILLING-dLLM logra una aceleración de 9.1× a 28.0× en longitudes de contexto de 8K a 32K en comparación con las líneas base de Transformers estándar.
Calidad: El método alcanza una calidad de vanguardia entre los métodos de aceleración de dLLM. En LongBench, logra una puntuación promedio de 34.59 (Dream-7B) y 37.02 (UltraLLaMA), superando a líneas base como Fast-dLLM, Sparse-dLLM y dKV-Cache.
Escalabilidad de Contexto Largo: En InfiniteBench (contextos >128K), el método alcanza un 43.62% de precisión promedio en Dream-7B, superando significativamente a la línea base más fuerte (Fast-dLLM v2 con 30.32%). Notablemente, alcanza un 95.42% de precisión en la tarea de recuperación de Passkey y un 70.00% en la recuperación de Números.

Abordaje de "Lost-in-the-Middle"

Fenómeno: El artículo observa que los dLLM exhiben un fenómeno de "pérdida en el medio" (lost-in-the-middle) bajo extrapolación de contexto, donde la precisión disminuye para la información ubicada en el medio del contexto, mostrando un sesgo de decaimiento monotónico hacia la respuesta en lugar de la curva en forma de U vista en los modelos AR.
Mitigación: Al seleccionar fragmentos basados en la relevancia predictiva en lugar de la proximidad posicional, PREFILLING-dLLM elimina este sesgo. El método mantiene una alta precisión de coincidencia exacta en todas las posiciones del contexto, incluso a longitudes de 32K y 128K.

Perspectivas del Mecanismo

Anclas de Atención: Los tokens BOS antepuestos a cada fragmento actúan como anclas de atención periódicas. El análisis muestra que estos tokens distribuyen la masa de atención de manera uniforme a través del contexto (formando crestas estables de baja magnitud) en lugar de degenerar en "sumideros de atención" (attention sinks) que absorben una masa de atención desproporcionada. Este mecanismo permite una recuperación invariante a la posición.
Compromiso del Tamaño del Fragmento: Existe un compromiso entre el tamaño del fragmento y la precisión de la recuperación. Los fragmentos más pequeños (por ejemplo, 256 tokens) mantienen una alta precisión incluso en contextos de 128K, mientras que los fragmentos más grandes (por ejemplo, 4096 tokens) se degradan bruscamente más allá de 32K. La selección de múltiples fragmentos es esencial para la inferencia de contexto largo.

Significancia y Reivindicaciones

El artículo afirma que PREFILLING-dLLM cierra la brecha entre la naturaleza limitada por cómputo de la inferencia de dLLM y la eficiencia de las arquitecturas de desagregación de prellenado-decodificación utilizadas en modelos AR. Sus principales contribuciones son:

Desagregación Libre de Entrenamiento: Introduce un marco que almacena en caché las representaciones KV del prefijo una sola vez y las reutiliza, eliminando el costo de recomputación cuadrática del prefijo en cada paso de eliminación de ruido sin requerir el reentrenamiento del modelo.
Prellenado Disperso: Demuestra que el prellenado disperso (selección de fragmentos y tokens relevantes) puede superar a la atención densa en dLLMs, reduciendo la complejidad mientras mejora o mantiene la calidad de la generación.
Optimización a Nivel de Sistema: Proporciona un kernel de atención especializado (Split-S FlexAttention) que aprovecha la estructura fragmentada del prefijo para lograr reducciones significativas de latencia (hasta 10.2× sobre el vanilla FlexAttention) durante la fase de decodificación.
Robustez ante Contextos Largos: Muestra que este enfoque mitiga el fenómeno de "lost-in-the-middle" y permite que los dLLM escalen efectivamente a 128K+ tokens con alta precisión de recuperación, una capacidad previamente inexplorada en la aceleración de dLLM.

Los autores señalan limitaciones, incluyendo la naturaleza estática de la selección de fragmentos (sin reselección dinámica durante la decodificación), la necesidad de un ajuste de tamaño de fragmento y $K$ específico para cada tarea, y que la evaluación actual se limita a benchmarks en inglés sobre arquitecturas de modelos específicas (Dream-7B y UltraLLaMA).