Atenção Esparsa com Máscara Dinâmica Treinável

Resumo

Em modelos de linguagem de grande escala, a demanda por modelagem de contextos longos está constantemente aumentando, mas a complexidade quadrática do mecanismo padrão de autoatenção frequentemente se torna um gargalo. Embora os mecanismos existentes de atenção esparsa tenham melhorado a eficiência, eles ainda podem enfrentar problemas como padrões estáticos ou perda de informação. Introduzimos um mecanismo de atenção esparsa com máscara dinâmica treinável, chamado Dynamic Mask Attention (DMA), que utiliza efetivamente a esparsidade consciente do conteúdo e da posição. O DMA alcança isso por meio de duas inovações principais: Primeiro, ele gera dinamicamente máscaras esparsas conscientes do conteúdo a partir das representações de valor, permitindo que o modelo identifique e foque em informações críticas de forma adaptativa. Segundo, ele implementa um cálculo de atenção esparsa consciente da posição que efetivamente ignora regiões de cálculo desnecessárias. Esse design de dupla esparsidade permite que o modelo reduza significativamente a complexidade computacional de informações importantes enquanto retém a informação completa, alcançando um excelente equilíbrio entre fidelidade da informação e eficiência computacional. Verificamos o desempenho do DMA por meio de experimentos abrangentes. Estudos comparativos mostram que o DMA supera a atenção multi-head, a atenção de janela deslizante, a atenção latente multi-head e a atenção esparsa nativa em termos de perplexidade sob as configurações da Lei de Escalonamento Chinchilla. Além disso, em tarefas desafiadoras de recall associativo multi-query, o DMA também demonstra desempenho e eficiência superiores em comparação com esses métodos. Crucialmente, na avaliação de um modelo com 1,7 bilhão de parâmetros, o DMA supera significativamente a atenção multi-head tanto no desempenho de benchmarks padrão quanto na tarefa desafiadora de "agulha no palheiro". Esses resultados experimentais destacam sua capacidade de equilibrar eficiência do modelo e habilidade de modelagem de contextos longos de forma eficaz.

English

In large language models, the demand for modeling long contexts is constantly increasing, but the quadratic complexity of the standard self-attention mechanism often becomes a bottleneck. Although existing sparse attention mechanisms have improved efficiency, they may still encounter issues such as static patterns or information loss. We introduce a trainable dynamic mask sparse attention mechanism, Dynamic Mask Attention, which effectively utilizes content-aware and position-aware sparsity. DMA achieves this through two key innovations: First, it dynamically generates content-aware sparse masks from value representations, enabling the model to identify and focus on critical information adaptively. Second, it implements position-aware sparse attention computation that effectively skips unnecessary calculation regions. This dual-sparsity design allows the model to significantly reduce the computational complexity of important information while retaining complete information, achieving an excellent balance between information fidelity and computational efficiency. We have verified the performance of DMA through comprehensive experiments. Comparative studies show that DMA outperforms multi-head attention, sliding window attention, multi-head latent attention, and native sparse attention in terms of perplexity under Chinchilla Scaling Law settings. Moreover, in challenging multi-query associative recall tasks, DMA also demonstrates superior performance and efficiency compared to these methods. Crucially, in the evaluation of a 1.7B parameter model, DMA significantly outperforms multi-head attention in both standard benchmark performance and the challenging needle-in-a-haystack task. These experimental results highlight its capability to balance model efficiency and long-context modeling ability effectively.

Atenção Esparsa com Máscara Dinâmica Treinável

Trainable Dynamic Mask Sparse Attention

Resumo

Support