VMoBA: Atención Mixta de Bloques para Modelos de Difusión de Video

Resumen

La complejidad cuadrática de los mecanismos de atención completa representa un cuello de botella significativo para los Modelos de Difusión de Video (VDMs, por sus siglas en inglés) que buscan generar videos de larga duración y alta resolución. Si bien se han propuesto diversos métodos de atención dispersa, muchos están diseñados como aceleradores de inferencia sin entrenamiento o no capturan de manera óptima las características espacio-temporales únicas inherentes a los datos de video cuando se entrenan de forma nativa. Este artículo presenta Video Mixture of Block Attention (VMoBA), un novedoso mecanismo de atención dispersa específicamente adaptado para VDMs. Motivado por un análisis en profundidad de los patrones de atención en transformadores de video preentrenados, que revelaron una fuerte localidad espacio-temporal, una importancia variable de las consultas y niveles de concentración específicos por cabeza, VMoBA mejora el marco original de MoBA con tres modificaciones clave: (1) un esquema de partición de bloques recurrente por capas (1D-2D-3D) para adaptarse dinámicamente a diversos patrones de atención espacio-temporal y mejorar la eficiencia; (2) selección global de bloques para priorizar las interacciones más destacadas entre bloques de consulta-clave en toda una cabeza de atención; y (3) selección de bloques basada en umbrales para determinar dinámicamente el número de bloques atendidos en función de su similitud acumulada. Experimentos exhaustivos demuestran que VMoBA acelera significativamente el entrenamiento de VDMs en secuencias más largas, logrando una aceleración de 2.92x en FLOPs y 1.48x en latencia, mientras alcanza una calidad de generación comparable o incluso superior a la atención completa. Además, VMoBA exhibe un rendimiento competitivo en inferencia sin entrenamiento, ofreciendo una aceleración de 2.40x en FLOPs y 1.35x en latencia para la generación de videos de alta resolución.

English

The quadratic complexity of full attention mechanisms poses a significant bottleneck for Video Diffusion Models (VDMs) aiming to generate long-duration, high-resolution videos. While various sparse attention methods have been proposed, many are designed as training-free inference accelerators or do not optimally capture the unique spatio-temporal characteristics inherent in video data when trained natively. This paper introduces Video Mixture of Block Attention (VMoBA), a novel sparse attention mechanism specifically adapted for VDMs. Motivated by an in-depth analysis of attention patterns within pre-trained video transformers, which revealed strong spatio-temporal locality, varying query importance, and head-specific concentration levels, VMoBA enhances the original MoBA framework with three key modifications: (1) a layer-wise recurrent block partition scheme (1D-2D-3D) to dynamically adapt to diverse spatio-temporal attention patterns and improve efficiency; (2) global block selection to prioritize the most salient query-key block interactions across an entire attention head; and (3) threshold-based block selection to dynamically determine the number of attended blocks based on their cumulative similarity. Extensive experiments demonstrate that VMoBA significantly accelerates the training of VDMs on longer sequences, achieving 2.92x FLOPs and 1.48x latency speedup, while attaining comparable or even superior generation quality to full attention. Furthermore, VMoBA exhibits competitive performance in training-free inference, offering 2.40x FLOPs and 1.35x latency speedup for high-res video generation.

VMoBA: Atención Mixta de Bloques para Modelos de Difusión de Video

VMoBA: Mixture-of-Block Attention for Video Diffusion Models

Resumen

Support