Resumen Técnico: FlashBlock

1. Planteamiento del Problema

Si bien la difusión por bloques ha mejorado la eficiencia de la inferencia en modelos generativos de contexto largo (por ejemplo, modelos de lenguaje de difusión y generación de video) mediante la combinación de la caché KV con la inferencia causal por bloques, sigue sufriendo una sobrecarga significativa en entornos de contexto largo. En la difusión de bloques estándar, cada paso de difusión requiere recomputar la atención sobre toda la caché KV en crecimiento. A medida que la longitud del contexto aumenta, el costo combinado de la computación de atención y el acceso a la caché KV se convierte en el factor dominante en la latencia de inferencia.

El artículo identifica una propiedad poco explorada de la difusión por bloques: la redundancia entre pasos (cross-step redundancy). Específicamente, durante el proceso de eliminación de ruido (denoising), los tokens fuera del bloque actual que se está actualizando permanecen mayormente estáticos, lo que conduce a contribuciones de atención altamente estables de estos tokens "externos al bloque" a través de pasos de difusión adyacentes. Por el contrario, la atención dentro del bloque actual ("interno al bloque") varía significamente a medida que los tokens se refinan o se desenmascaran activamente. Los métodos actuales no logran explotar esta estabilidad, recomputando repetidamente la atención externa al bloque que es redundante.

2. Metodología: FlashBlock

Los autores proponen FlashBlock, un mecanismo que almacena en caché y reutiliza las salidas de atención externa al bloque para reducir la computación y el I/O de memoria sin modificar el proceso de difusión subyacente.

Mecanismo Central

FlashBlock descompone la salida de atención completa en dos componentes:

1. Atención Interna al Bloque ($A_{in}$): Computada sobre los tokens dentro del bloque actual ($J_{in}$), que están cambiando activamente.
2. Atención Externa al Bloque ($A_{out}$): Computada sobre los tokens fuera del bloque actual ($J_{out}$), que permanecen estables a través de los pasos.

El método almacena en caché la salida de atención externa al bloque ($A_{out}$) y su normalizador logarítmico ($L_{out} = \log Z_{out}$) del paso de difusión anterior. En los pasos subsecuentes, en lugar de acceder nuevamente a la caché KV completa para $J_{out}$, el modelo reutiliza los valores almacenados en caché. La salida de atención completa se reconstruye combinando la $A_{out}$ almacenada en caché con la nueva $A_{in}$ recién computada utilizando una composición en el espacio logarítmico (similar a FlashAttention) para asegurar la estabilidad numérica:
$$A_{full} = \frac{e^{L_{out}-m} A_{out} + e^{L_{in}-m} A_{in}}{e^{L_{out}-m} + e^{L_{in}-m}}$$
donde $m = \max(L_{out}, L_{in})$.

Estrategias de Reutilización Adaptativa

• Reutilización basada en umbrales: En modelos de difusión discreta, las actualizaciones de tokens pueden ser abruptas. El método introduce un umbral $\tau$ sobre el número de tokens actualizados ($M^{s+1}$). Si $M^{s+1} < \tau$, se reutiliza la atención en caché; de lo contrario, se recomputa para asegurar la corrección.
• Selección por cabeza (Video): Para la difusión de video, donde la variabilidad entre pasos es mayor, el método estima la similitud por cabeza de atención. Solo las cabezas con una similitud que supere un umbral $\gamma$ reutilizan la caché.
• Compatibilidad con Atención Dispersa (Sparse Attention): FlashBlock es ortogonal a la atención dispersa. Trata los keys no seleccionados en la atención dispersa como el conjunto "externo al bloque", almacenando sus contribuciones residuales en caché. Esto permite que FlashBlock compense la pérdida de información causada por una sparsificación agresiva.

Adaptación de Entrenamiento (Opcional)

Para abordar posibles desajustes distributivos en los modelos de lenguaje de difusión discreta (dLLMs) causados por el desenmascarado abrupto de tokens, los autores proponen una estrategia de destilación consciente de la reutilización (reuse-aware distillation). Un modelo estudiante que emplea FlashBlock es entrenado para igualar la distribución de salida de un modelo profesor congelado que utiliza atención densa, utilizando adaptadores LoRA para mantener la adaptación ligera.

3. Contribuciones Clave

1. Perspectiva Empírica: El artículo proporciona un estudio empírico que revela una fuerte estabilidad entre pasos en la atención externa al bloque, en contraste con la alta variabilidad de la atención interna al bloque.
2. Mecanismo FlashBlock: Un novedoso mecanismo de almacenamiento en caché que explota esta redundancia para reducir la computación de atención y el acceso a la caché KV, implementado mediante agregación en el espacio logarítmico.
3. Sinergia con la Dispersión (Sparsity): Demostración de que FlashBlock sirve como una estrategia complementaria de reutilización residual para la atención dispersa, mitigando la degradación de la calidad causada por la sparsificación agresiva.
4. Ganancias de Eficiencia: Aceleración significativa en el rendimiento (throughput) de inferencia y reducción de la latencia tanto para modelos de lenguaje de difusión como para modelos de generación de video, con un impacto insignificante en la calidad de la generación.

4. Resultados Experimentales

Los experimentos se realizaron en modelos de lenguaje de difusión (Trado-8B/4B) y modelos de difusión de video (LongLive-1.3B).

• Rendimiento (Throughput): FlashBlock logró hasta un 1.44× mayor rendimiento de tokens en modelos de lenguaje de difusión (por ejemplo, aumentando de 312 a 451 tokens/s para un tamaño de bloque de 4).
• Latencia: El tiempo de atención se redujo hasta un 1.6×.
• Calidad: La calidad de la generación se mantuvo comparable a las bases (baselines) densas. En pruebas de razonamiento matemático (GSM8K, MATH500) y de codificación (HumanEval, MBPP), las diferencias de precisión fueron mínimas (dentro de un 0.2% a 1.6% dependiendo del tamaño del bloque).
• Escalabilidad de Contexto Largo: A medida que la longitud del contexto aumentó de 100k a 800k, la brecha de latencia entre FlashBlock y la base se amplió, mostrando FlashBlock aproximadamente la mitad del crecimiento de latencia de la base bajo ciertos umbrales.
• Combinación con Atención Dispersa: Al combinarse con SparseD (al 20%, 30% y 40% de densidad), FlashBlock mejoró consistentemente la precisión sobre SparseD por sí solo (por ejemplo, mejorando la precisión de GSM8K de 34.72% a 42.68% al 20% de densidad), recuperando efectivamente la información perdida por la sparsificación.

5. Significado

El artículo afirma que FlashBlock aborda un cuello de botella crítico en la difusión por bloques de contexto largo: la recomputación redundante de las contribuciones de atención estables. Al desacoplar el contexto estable "externo al bloque" de las actualizaciones dinámicas "internas al bloque", el método reduce sustancialmente los costos de computación y de I/O de memoria asociados con las cachés KV en crecimiento.

Los autores enfatizan que este enfoque es ortogonal a las técnicas existentes de atención dispersa, ofreciendo una estrategia complementaria que mejora tanto la eficiencia como la precisión. Los resultados sugieren que FlashBlock proporciona una solución escalable para la generación eficiente de contexto largo tanto en el dominio de texto como en el de video, permitiendo una inferencia más rápida sin comprometer la calidad del contenido generado.