Atención Esparsa de Máscara Dinámica Entrenable

Trainable Dynamic Mask Sparse Attention

August 4, 2025

Autores: Jingze Shi, Yifan Wu, Bingheng Wu, Yiran Peng, Liangdong Wang, Guang Liu, Yuyu Luo

cs.AI

Resumen

En los modelos de lenguaje de gran escala, la demanda para modelar contextos largos está en constante aumento, pero la complejidad cuadrática del mecanismo de autoatención estándar suele convertirse en un cuello de botella. Aunque los mecanismos de atención dispersa existentes han mejorado la eficiencia, aún pueden enfrentar problemas como patrones estáticos o pérdida de información. Introducimos un mecanismo de atención dispersa con máscara dinámica entrenable, denominado Atención con Máscara Dinámica (Dynamic Mask Attention, DMA), que utiliza eficazmente la dispersión consciente del contenido y la posición. DMA logra esto mediante dos innovaciones clave: primero, genera dinámicamente máscaras dispersas conscientes del contenido a partir de las representaciones de valor, permitiendo que el modelo identifique y se enfoque en información crítica de manera adaptativa. Segundo, implementa un cálculo de atención dispersa consciente de la posición que omite eficazmente regiones de cálculo innecesarias. Este diseño de doble dispersión permite que el modelo reduzca significativamente la complejidad computacional de la información importante mientras retiene la información completa, logrando un equilibrio excelente entre la fidelidad de la información y la eficiencia computacional. Hemos verificado el rendimiento de DMA mediante experimentos exhaustivos. Estudios comparativos muestran que DMA supera a la atención multi-cabeza, la atención de ventana deslizante, la atención latente multi-cabeza y la atención dispersa nativa en términos de perplejidad bajo los ajustes de la Ley de Escalamiento Chinchilla. Además, en tareas desafiantes de recuperación asociativa multi-consulta, DMA también demuestra un rendimiento y eficiencia superiores en comparación con estos métodos. Crucialmente, en la evaluación de un modelo de 1.7B parámetros, DMA supera significativamente a la atención multi-cabeza tanto en el rendimiento de benchmarks estándar como en la tarea desafiante de "aguja en un pajar". Estos resultados experimentales destacan su capacidad para equilibrar eficazmente la eficiencia del modelo y la capacidad de modelar contextos largos.

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.