nablaNABLA: Atenção Adaptativa em Nível de Bloco com Base na Vizinhança

nablaNABLA: Neighborhood Adaptive Block-Level Attention

July 17, 2025

Autores: Dmitrii Mikhailov, Aleksey Letunovskiy, Maria Kovaleva, Vladimir Arkhipkin, Vladimir Korviakov, Vladimir Polovnikov, Viacheslav Vasilev, Evelina Sidorova, Denis Dimitrov

cs.AI

Resumo

Os recentes avanços em arquiteturas baseadas em transformers têm demonstrado sucesso notável em tarefas de geração de vídeo. No entanto, a complexidade quadrática dos mecanismos de atenção completa continua sendo um gargalo crítico, especialmente para sequências de vídeo de alta resolução e longa duração. Neste artigo, propomos o NABLA, um novo mecanismo de Atenção Adaptativa em Blocos de Vizinhança que se adapta dinamicamente aos padrões de esparsidade em transformers de difusão de vídeo (DiTs). Ao aproveitar a atenção em blocos com um limiar adaptativo baseado em esparsidade, o NABLA reduz a sobrecarga computacional enquanto preserva a qualidade gerativa. Nosso método não requer o design de operadores de baixo nível personalizados e pode ser integrado de forma transparente com o operador Flex Attention do PyTorch. Experimentos demonstram que o NABLA alcança até 2,7x mais rapidez no treinamento e inferência em comparação com a linha de base, quase sem comprometer as métricas quantitativas (pontuação CLIP, pontuação VBench, pontuação de avaliação humana) e a queda na qualidade visual. O código e os pesos do modelo estão disponíveis aqui: https://github.com/gen-ai-team/Wan2.1-NABLA.

English

Recent progress in transformer-based architectures has demonstrated remarkable success in video generation tasks. However, the quadratic complexity of full attention mechanisms remains a critical bottleneck, particularly for high-resolution and long-duration video sequences. In this paper, we propose NABLA, a novel Neighborhood Adaptive Block-Level Attention mechanism that dynamically adapts to sparsity patterns in video diffusion transformers (DiTs). By leveraging block-wise attention with adaptive sparsity-driven threshold, NABLA reduces computational overhead while preserving generative quality. Our method does not require custom low-level operator design and can be seamlessly integrated with PyTorch's Flex Attention operator. Experiments demonstrate that NABLA achieves up to 2.7x faster training and inference compared to baseline almost without compromising quantitative metrics (CLIP score, VBench score, human evaluation score) and visual quality drop. The code and model weights are available here: https://github.com/gen-ai-team/Wan2.1-NABLA