Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

MSWA: Refining Local Attention with Multi-ScaleWindow Attention

Yixing Xu, Shivank Nag, Dong Li, Lu Tian, Emad Barsoum

TL;DR

The paper tackles the inefficiency of quadratic self-attention by introducing Multi-Scale Window Attention (MSWA), which assigns diverse window sizes across heads and progressively across layers to capture multi-scale context. By distributing attention budgets as MSWA-h and MSWA-l, MSWA achieves near-SWA efficiency (≈$\frac{7}{8}$ of the total window budget) while enabling richer local representations and long-range integration. Empirical results on language modeling and common-sense reasoning show MSWA outperforms SWA in PPL/bpc and downstream accuracy, and combining MSWA with linear attention yields Transformer-like performance with improved efficiency. The approach is validated with ablations and efficiency experiments, and demonstrated to be compatible with pre-trained LLMs through fine-tuning, indicating practical applicability for scalable, efficient Transformers in NLP tasks.

Abstract

Transformer-based LLMs have achieved exceptional performance across a wide range of NLP tasks. However, the standard self-attention mechanism suffers from quadratic time complexity and linearly increased cache size. Sliding window attention (SWA) solves this problem by restricting the attention range to a fixed-size local context window. Nevertheless, SWA employs a uniform window size for each head in each layer, making it inefficient in capturing context of varying scales. To mitigate this limitation, we propose Multi-Scale Window Attention (MSWA) which applies diverse window sizes across heads and layers in the Transformer. It not only allows for different window sizes among heads within the same layer but also progressively increases window size allocation from shallow to deep layers, thus enabling the model to capture contextual information with different lengths and distances. Experimental results on language modeling and common-sense reasoning tasks substantiate that MSWA outperforms traditional local attention in both effectiveness and efficiency.

MSWA: Refining Local Attention with Multi-ScaleWindow Attention

TL;DR

The paper tackles the inefficiency of quadratic self-attention by introducing Multi-Scale Window Attention (MSWA), which assigns diverse window sizes across heads and progressively across layers to capture multi-scale context. By distributing attention budgets as MSWA-h and MSWA-l, MSWA achieves near-SWA efficiency (≈ of the total window budget) while enabling richer local representations and long-range integration. Empirical results on language modeling and common-sense reasoning show MSWA outperforms SWA in PPL/bpc and downstream accuracy, and combining MSWA with linear attention yields Transformer-like performance with improved efficiency. The approach is validated with ablations and efficiency experiments, and demonstrated to be compatible with pre-trained LLMs through fine-tuning, indicating practical applicability for scalable, efficient Transformers in NLP tasks.

Abstract

Transformer-based LLMs have achieved exceptional performance across a wide range of NLP tasks. However, the standard self-attention mechanism suffers from quadratic time complexity and linearly increased cache size. Sliding window attention (SWA) solves this problem by restricting the attention range to a fixed-size local context window. Nevertheless, SWA employs a uniform window size for each head in each layer, making it inefficient in capturing context of varying scales. To mitigate this limitation, we propose Multi-Scale Window Attention (MSWA) which applies diverse window sizes across heads and layers in the Transformer. It not only allows for different window sizes among heads within the same layer but also progressively increases window size allocation from shallow to deep layers, thus enabling the model to capture contextual information with different lengths and distances. Experimental results on language modeling and common-sense reasoning tasks substantiate that MSWA outperforms traditional local attention in both effectiveness and efficiency.
Paper Structure (37 sections, 11 equations, 3 figures, 6 tables)

This paper contains 37 sections, 11 equations, 3 figures, 6 tables.

Table of Contents

  1. Introduction
  2. Related Works
  3. Large Language Models (LLMs)
  4. Attention Mechanisms
  5. Preliminaries
  6. Self-Attention Mechanism
  7. Sliding Window Attention
  8. Linear Attention Mechanism
  9. Multi-Scale Window Attention
  10. Diverse Window Across Heads
  11. Diverse Window Across Layers
  12. Integrate Diversity of Heads and Layers
  13. Combination with Linear Attention
  14. Experiments
  15. Overview.
  16. ...and 22 more sections

Figures (3)

  • Figure 1: Illustration of Multi-Scale Window Attention mechanism.
  • Figure 2: Combination of MSWA and linear attention.
  • Figure 3: Computational time required by each attention mechanism to predict the next token.