훈련 가능한 동적 마스크 희소 어텐션

초록

대규모 언어 모델에서 장기 문맥을 모델링하려는 요구는 지속적으로 증가하고 있지만, 표준 자기 주의(self-attention) 메커니즘의 2차 복잡도는 종종 병목 현상으로 작용합니다. 기존의 희소 주의(sparse attention) 메커니즘은 효율성을 개선했지만, 여전히 정적 패턴이나 정보 손실과 같은 문제에 직면할 수 있습니다. 우리는 학습 가능한 동적 마스크 희소 주의 메커니즘인 Dynamic Mask Attention(DMA)을 소개하며, 이는 내용 인식(content-aware) 및 위치 인식(position-aware) 희소성을 효과적으로 활용합니다. DMA는 두 가지 주요 혁신을 통해 이를 달성합니다: 첫째, 값 표현(value representations)으로부터 내용 인식 희소 마스크를 동적으로 생성하여 모델이 중요한 정보를 적응적으로 식별하고 집중할 수 있게 합니다. 둘째, 불필요한 계산 영역을 효과적으로 건너뛰는 위치 인식 희소 주의 계산을 구현합니다. 이 이중 희소성 설계는 모델이 중요한 정보의 계산 복잡도를 크게 줄이면서도 완전한 정보를 유지할 수 있게 하여, 정보 충실도와 계산 효율성 사이의 탁월한 균형을 달성합니다. 우리는 DMA의 성능을 포괄적인 실험을 통해 검증했습니다. 비교 연구 결과, DMA는 Chinchilla Scaling Law 설정에서 perplexity 측면에서 다중 헤드 주의(multi-head attention), 슬라이딩 윈도우 주의(sliding window attention), 다중 헤드 잠재 주의(multi-head latent attention), 그리고 기본 희소 주의(native sparse attention)를 능가하는 것으로 나타났습니다. 또한, 도전적인 다중 쿼리 연상 회상(multi-query associative recall) 작업에서도 DMA는 이러한 방법들에 비해 우수한 성능과 효율성을 보여주었습니다. 특히, 1.7B 매개변수 모델 평가에서 DMA는 표준 벤치마크 성능과 어려운 needle-in-a-haystack 작업 모두에서 다중 헤드 주의를 크게 앞질렀습니다. 이러한 실험 결과는 DMA가 모델 효율성과 장기 문맥 모델링 능력을 효과적으로 균형 잡을 수 있는 능력을 강조합니다.

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.

훈련 가능한 동적 마스크 희소 어텐션

Trainable Dynamic Mask Sparse Attention

초록

Support