AB-Attention acelera la generación de vídeo hasta 2,23× y omite el 92% de la atención

El equipo de InfRec, junto con investigadores de CUHK Shenzhen y de las universidades de Shandong, Tokio y Fudan, además de Berkeley, HKUST y NUS, ha lanzado AB-Attention. Es un método para acelerar la generación de vídeo en transformadores de difusión (Wan2.1 y HunyuanVideo) sin reentrenar los pesos.

La operación más costosa de estos modelos es la atención: cada fragmento del vídeo se compara con todos los demás en el espacio y el tiempo, por lo que el coste crece de forma cuadrática con la longitud de la secuencia. AB-Attention omite el 89–92% de ese cálculo. Aun así, el vídeo resultante se acerca más al de la atención completa, la «de referencia», que el de los competidores.

El PSNR (una métrica de similitud con el vídeo de referencia, expresada en decibelios: cuanto más alto, más parecido) es 1,6–4,4 dB mejor, y la generación completa en una sola tarjeta NVIDIA H20 es entre 1,84 y 2,23 veces más rápida.

Dónde fallan los esquemas dispersos anteriores

La atención dispersa no calcula todos los pares «consulta — clave», sino solo una parte. Los primeros métodos fijaban un patrón de antemano, por ejemplo ventanas 3D locales. Esos patrones no saben qué es importante para un prompt, una cabeza de atención o un paso de eliminación de ruido concretos, así que se ven obligados a conservar mucho cálculo superfluo. Los métodos dinámicos seleccionan bloques de claves y valores según las activaciones actuales, pero los autores identifican tres puntos débiles en ellos:

• Las estimaciones de importancia son toscas o caras. Promediar un bloque en un solo vector mezcla tokens con comportamientos distintos, y la agrupación global (SVG2) exige k-means sobre toda la secuencia y una permutación de tokens. • La selección top-k corta la cola. Un bloque descartado desaparece tanto del numerador de la atención como de la normalización del softmax. Unas estimaciones mejores no lo corregirán: incluso una selección top-k ideal con el 10,8% de las claves conserva solo alrededor del 69% de la masa de atención. • La unidad de selección es la baldosa de cómputo. El bloque se conserva o se descarta entero, aunque solo hagan falta unas pocas de sus claves.

Los tres recursos de AB-Attention

El primero, IBCS (Intra-Block Clustering Selector), estima la importancia de los bloques de claves. Las consultas se dividen en bloques de 128 tokens. Dentro de cada bloque se agrupan en 8 clústeres, cada uno con un centroide y un peso igual a la fracción de consultas que contiene.

Un bloque de claves recibe su puntuación como la atención ponderada de los centroides sobre sus claves promediadas. Los clústeres se usan solo para la estimación; el cálculo en sí sigue el orden original de los tokens. Así, un bloque al que atiende un único clúster no se pierde, como sí ocurre al promediar todo el bloque.

La brecha de exhaustividad de la selección entre el promedio y el top-k ideal se reduce a la mitad.

El segundo, SIP (Sampling with Inclusion-Probability Weighting), sustituye la selección rígida por un muestreo. Un bloque entra en la muestra con una probabilidad proporcional a su puntuación. Los bloques más importantes se toman siempre, y el resto tiene una probabilidad distinta de cero.

Cada bloque seleccionado se multiplica por la inversa de su probabilidad de inclusión: un bloque con un 25% de probabilidad cuenta como cuatro. Es el estimador clásico de Horvitz–Thompson, que ofrece estimaciones insesgadas tanto del numerador como de la normalización del softmax. El error de atención es un 41% menor que el del top-k ideal con el mismo número de bloques.

No hace falta una pasada adicional para la corrección: el cálculo sigue siendo una atención softmax normal con desplazamiento de logits.

El tercero es la selección fina sobre hardware «grueso». Los bloques de claves y valores tienen solo 4 tokens, pero el kernel reúne los bloques elegidos, mediante índices, en baldosas completas de tensor cores. No se reordenan los tokens ni se crean copias innecesarias. Reducir el bloque de 128 a 4 tokens eleva del 59% al 69% la fracción de masa de atención que retiene la selección ideal, sin que aumente el número de baldosas con un presupuesto fijo. El plan de muestreo se recalcula una vez cada cuatro pasos de eliminación de ruido; la atención en sí se calcula en todos.

Cómo se midió y qué resultó

Se compararon tres modelos a 720p: dos versiones de Wan2.1 (81 fotogramas) y HunyuanVideo (129 fotogramas), con 600 prompts del benchmark Penguin Video. Como referencia se usó el vídeo con atención completa, con el mismo prompt y la misma semilla de generación. Además del PSNR, se calcularon SSIM y LPIPS, y cinco evaluaciones del propio vídeo: estética, coherencia del fondo, calidad de imagen, fluidez del movimiento y consistencia del objeto. Los competidores fueron RadialAttention (una máscara estática con decaimiento de energía), SVG2 y DFSAttn, todos con el código oficial y el backend FlashInfer con FlashAttention-3. El primer bloque del transformador y los primeros 10 de los 50 pasos se ejecutan con atención completa en todos los métodos.

• Dispersión: 89,16% en Wan2.1 y 92,35% en HunyuanVideo. • Los competidores calculan entre 1,8 y 3,6 veces más atención. • AB-Attention es el mejor en PSNR, SSIM y LPIPS en los tres modelos. • Aceleración de extremo a extremo: 1,84–2,23× en una sola H20 en BF16.

Qué aporta en la práctica

El método no requiere reentrenamiento, por lo que se puede conectar a pesos ya disponibles. Para los equipos que ejecutan Wan2.1 o HunyuanVideo, supone casi la mitad de tiempo por vídeo con una calidad cercana a la de referencia. La ganancia aumentará en vídeos largos y de alta resolución, donde la atención pesa más en la latencia. También hay limitaciones: las mediciones se hicieron en una H20 y con tres modelos, y tanto el kernel como el plan de muestreo están diseñados para los tensor cores de NVIDIA. El código está abierto en GitHub, en el repositorio InfRec-Cardinal/AB-Attention, así que los resultados pueden verificarse en hardware propio.