MSWA: Refining Local Attention with Multi-Scale Window Attention

Language models Bengio et al. (2000) have become a cornerstone of modern Natural Language Processing (NLP). Their primary purpose is to understand and generate human language, making them crucial for applications ranging from machine translation Zhang et al. (2023a) to communicative agents Li et al. (2023). The advent of large-scale pre-trained models has significantly enhanced the performance of these applications.

Among them, Transformer-based models have revolutionized the field. Introduced by Vaswani et al. (2017), the Transformer architecture uses self-attention mechanism to process input sequences in a more parallelizable way. This innovation has led to the development of increasingly large and powerful models, such as GPT-4 Achiam et al. (2023), Llama-3 AI@Meta (2024), and Claude-3 Anthropic (2024). These models, often termed "large language models", leverage vast amounts of data and computational resources to achieve state-of-the-art results on a wide array of NLP tasks.

The attention mechanism which enables the model to capture intricate dependencies across the entire sequence is at the heart of the Transformer’s success. However, the standard self-attention mechanism has quadratic complexity with respect to the sequence length, which poses scalability challenges for longer sequences. To address this issue, many efficient attention variants have been proposed Qiu et al. (2019); Wang et al. (2020); Peng et al. (2021); Hua et al. (2022). For example, sliding window attention Beltagy et al. (2020); Zaheer et al. (2020); Jiang et al. (2023) limit attention to a fixed-size window around each token, making the computation more manageable for long texts. Linear attention methods Choromanski et al. (2020); Katharopoulos et al. (2020); Yang et al. (2023) approximate the attention calculation to reduce complexity from quadratic to linear. These innovations have made it possible to apply Transformer to lengthy documents without prohibitive computational costs.

For the Transformer models, its main computational and memory costs arise from the multi-head self-attention mechanism, focusing on two points: (1) quadratic time complexity over the input length, and (2) linearly increased size of the KV cache during inference. In the following, we analyze the attention mechanism based on the decoder form, as it is widely used in language models.

Given input vectors $\{{\bf x}_{i}\}_{i=1}^{n}\in\mathbb{R}^{D}$, where each $x_{i}$ represents a single input token, $n$ is the sequence length, and $D$ is the dimension of the input vectors, a head in the self-attention layer first maps the input tokens into query vectors $\{{\bf q}_{i}\}_{i=1}^{n}$, key vectors $\{{\bf k}_{i}\}_{i=1}^{n}$, and value vectors $\{{\bf v}_{i}\}_{i=1}^{n}$:

$${\bf q}_{i}={\bf x}_{i}{\bf W}_{q},~{}~{}~{}{\bf k}_{i}={\bf x}_{i}{\bf W}_{k},~{}~{}~{}{\bf v}_{i}={\bf x}_{i}{\bf W}_{v}.$$

(1)

where ${\bf W}_{q}$, ${\bf W}_{k}$, ${\bf W}_{v}\in\mathbb{R}^{D\times d}$ are the mapping matrices, and $d$ is the dimension of each head. The output of this attention head is calculated by:

$$\alpha_{ij}=\frac{{\rm exp}({\bf q}_{i}{\bf k}_{j}^{T}/\sqrt{d})}{\sum_{t=0}^{i}{\rm{exp}}({\bf q}_{i}{\bf k}_{t}^{T}/\sqrt{d})},$$

(2)

$${\bf o}_{i}=\sum_{j=0}^{i}\alpha_{ij}{{\bf v}_{j}}.$$

(3)

Performing the above computation for the entire sequence of length $n$ requires a time complexity of $O(dn^{2})$ and a space complexity of $O(dn)$.

Sliding Window Attention (SWA) is an efficient variant that restricts each token to attend to tokens within a local window of size $w$, as shown in Fig. 1. Given the query, key and value vectors ${\bf q}_{i}$, ${\bf k}_{i}$, ${\bf v}_{i}$, the output of the SWA head is defined as:

$$\alpha_{ij}=\frac{{\rm exp}({\bf q}_{i}{\bf k}_{j}^{T}/\sqrt{d})}{\sum_{t=max(0,i-w)}^{i}{\rm{exp}}({\bf q}_{i}{\bf k}_{t}^{T}/\sqrt{d})},$$

(4)

$${\bf o}_{i}=\sum_{j=max(0,i-w)}^{i}\alpha_{ij}{\bf v}_{j}.$$

(5)

By using this method, the time and space complexity required for each head is reduced to $O(dnw)$ and $O(dw)$, respectively. Considering a Transformer with $l$ layers, each equipped with $h$ attention heads, the time and space complexity will become $O(dn(whl))$ and $O(d(whl))$ respectively. We can see that the computational and memory cost is proportional to $whl$, which is the summation of the window sizes of all the sliding window attention operations from all heads in all layers.

Linear attention replaces the softmax operation in standard attention with a feature map-based dot product, eliminating the computation of $\rm{exp}(\cdot)$ and exchanging the matrix multiplication order, thereby achieving the goal of acceleration and constant memory cost. Specifically, given a kernel function $\phi(\cdot)$ that maps the $q_{i}$ and $k_{j}$ vectors into features, linear attention approximates ${\rm exp}({{\bf q}_{i}{\bf k}_{j}^{T}}/{\sqrt{d}})$ with the dot-product $\phi({\bf q}_{i})\phi({\bf k}_{j})^{T}$. Therefore, the output of an attention head is calculated by:

$$\alpha_{ij}=\frac{\phi({\bf q}_{i})\phi({\bf k}_{j})^{T}}{\sum_{t=0}^{i}\phi({\bf q}_{i})\phi({\bf k}_{t})^{T}},$$

(6)

$${\bf o}_{i}=\sum_{j=0}^{i}\alpha_{ij}{\bf v}_{j}=\frac{\phi({\bf q}_{i})\sum_{j=0}^{i}\phi({\bf k}_{j})^{T}{\bf v}_{j}}{\phi({\bf q}_{i})\sum_{j=0}^{i}\phi({\bf k}_{j})^{T}}.$$

(7)

Based on the exchange of multiplication order, the computational time complexity of linear attention is reduced to $O(d^{2}n)$. In addition, during the auto-regressive inference process, both $\sum_{j=0}^{i}\phi({\bf k}_{j})^{T}{\bf v}_{j}$ and $\sum_{j=0}^{i}\phi({\bf k}_{j})^{T}$ can be written as a variable that is continuously accumulated, requiring only $O(d^{2})$ space complexity.

This section focuses on dynamically changing the window size for each attention head within a layer. We refer to this mechanism as MSWA-h, as shown in the bottom right part of Fig. 1.

Different from SWA where all heads use the same window size $w_{i}$ in the $i$-th layer, in MSWA-h different heads have different scales of window sizes, and the summation of the total window sizes within a layer is less than that in the SWA, which is $w_{i}h$. Specifically, inspired by many hierarchical architecture designs in the CV field Liu et al. (2021), we divide the attention heads into four groups and adjust the receptive field range with a $2\times$ change between each group, resulting in window sizes of $\frac{w_{i}}{4}$, $\frac{w_{i}}{2}$, $w_{i}$, $2w_{i}$, respectively. Therefore, the summation of the total window sizes is:

$$(\frac{w_{i}}{4}+\frac{w_{i}}{2}+w_{i}+2w_{i})\times\frac{h}{4}=\frac{15}{16}w_{i}h.$$

(8)

Leveraging diverse window sizes among the heads within a layer allows the Transformer model to capture the relevant context of different scales simultaneously. This is because the outputs of different heads within an attention layer are concatenated together and then mapped through a matrix to form the final output of the layer, which allows contextual information at different distances to be integrated together. Additionally, considering the allocation of attention resources: all heads will attend to tokens within a distance of $\frac{w_{i}}{4}$ from the current token, while $\frac{3}{4}$ of the heads will attend to tokens within the $\frac{w_{i}}{4}$ to $\frac{w_{i}}{2}$ range, and so on. This implicitly models a long window with weighted emphasis, where the distribution of attention resources gradually decreases from near to far, aligning with the locality of reference characteristic of text.

This section further introduces changing the allocation ratio of the attention window sizes between layers. We refer to this mechanism as MSWA-l, as illustrated in the upper right part of Fig. 1.

To explain more clearly, we still use SWA as a comparison. In SWA, each attention layer has a total window size allocation of $hw$, where $h$ is the number of heads per layer, and $w$ is the base window size, which means that for any layer index $i$, $w_{i}=w$. In MSWA-l, the window size allocation varies across layers. More specifically, we divide all attention layers into several groups, and from shallow to deep, we continuously increase the total window size allocated to the attention layers in each group. We adopt a similar setup to MSWA-h, with four groups having a $2\times$ change between each group, resulting in window size allocations of $\frac{hw}{4}$, $\frac{hw}{2}$, $hw$, and $2hw$, respectively. The total window size resource allocated to all layers in MSWA-l is:

$$(\frac{hw}{4}+\frac{hw}{2}+hw+2hw)\times\frac{l}{4}=\frac{15}{16}whl.$$

(9)

For an attention layer, a larger window size allocation means that the window sizes of the heads in that layer are generally larger, allowing for the perception of a broader range of context. Therefore, gradually increasing the window size allocation from shallow to deep layers enables the model to focus on building local fine-grained information in the initial stages and progressively enhance the capture of long-distance relationships in the later stages. Additionally, the gradually expanding attention window allocation enables the model to continuously integrate local information from previous stages based on a larger receptive field.

In this section, we aim to integrate the two strategies MSWA-h and MSWA-l introduced earlier to construct the final MSWA mechanism.

The description of MSWA starts with a base window size $w$. In SWA, $w$ is the window size used for all heads across all layers. In contrast, in MSWA, $w$ serves as the basis for window size variation. We denote the base size value of the $i$-th layer as $w_{i}$ and the actual window size of the $j$-th head in the $i$-th layer as $w_{i,j}$. First, we evenly divide all attention layers into four groups. Depending on the group the layer belongs to, the values of $w_{i}$ from shallow to deep are $\frac{w}{4}$, $\frac{w}{2}$, $w$, $2w$, respectively. Further, within each layer $i$, we divide all heads into four groups, each with different $w_{i,j}$ values, denoted as $\frac{w_{i}}{4}$, $\frac{w_{i}}{2}$, $w_{i}$ and $2w_{i}$. Thus, the window size allocation for the entire Transformer model is:

$$\sum_{i=1}^{l}(\frac{w_{i}}{4}+\frac{w_{i}}{2}+w_{i}+2w_{i})\times\frac{h}{4}=\frac{15}{16}\sum_{i=1}^{l}hw_{i},$$

(10)

which can be further derived as:

$$\frac{15}{16}(\frac{hw}{4}+\frac{hw}{2}+hw+2hw)\times\frac{l}{4}\approx\frac{7}{8}whl.$$

(11)

The aforementioned window variation method is demonstrated in Tab. 1. Note that MSWA can benefit from the advantages of both MSWA-h and MSWA-l. When computing within a single layer, it can capture both long-range and short-range contextual information at the same time, and allocate different attention resources to information at various distances. When transitioning the information from one layer to another, MSWA continuously enhances the overall perception scope, integrating previous local information into a broader synthesis.

This section further proposes combining MSWA with an efficient global attention mechanism, i.e., linear attention, as illustrated in Fig. 2.

As we introduced earlier, many efficient mechanisms focus on capturing global information with limited resources. Therefore, they often fail to allocate high importance to relevant local information. Linear attention is a typical example of this issue. As introduced in Sec. 3.3 it stores global sequence information in fixed-size variables $\sum_{j=0}^{i}\phi({\bf k}_{j})^{T}{\bf v}_{j}$ and $\sum_{j=0}^{i}\phi({\bf k}_{j})^{T}$, which can lead to a loss of attention focus Qin et al. (2022).

Therefore, we propose to combine MSWA with linear attention to compensate for its shortcoming and achieve a balance between efficiency and performance. Specifically, we alternately stack MSWA layers and linear attention layers. For example, the $l$ layers in a combined model are evenly divided into four groups, each contains $\frac{l}{12}$ linear attention layers and $\frac{l}{6}$ MSWA layers stacked together. For all MSWA layers in the entire model, we consider them as a whole and utilize the same window size variation method introduced in Sec. 4.3 and adjust the window sizes across layers and heads.

In this section, we evaluate the language modeling capabilities of MSWA mechanism by training models from scratch on Wikitext-103 and enwik8.

In this section, we evaluate the performance of Llama-7B after fine-tuning with local attention patterns and testing on downstream common-sense reasoning tasks. The purpose of this evaluation is to verify the compatibility of MSWA with the current pre-trained LLM and its effectiveness when scaled to a large number of model parameters.

Tab. 4 presents the accuracy results of the models on various downstream benchmarks using 3-shot and 5-shot settings. The experimental results indicate that: 1) The MSWA mechanism demonstrates better common-sense reasoning ability compared to traditional sliding window attention, with average accuracy differences of $+1.9$ and $+7.23$ in 3-shot and 5-shot scenarios, respectively. 2) The MSWA mechanism shows a stronger ability to adapt to different context lengths, with its performance remaining stable across varying shot numbers, whereas SWA’s average accuracy decreases by 4.46 as the number of shots increases.

This section further evaluates the actual computational efficiency of our MSWA mechanism. Unlike the inferable size of the KV cache, the computational speed of Transformer models needs to be measured in practice to obtain realistic results. Therefore, we measure the time required to predict the next token during forward propagation in inference process for various attention mechanisms. We utilize FlashAttention Dao (2023) for the computation of various attention mechanisms.

The computational efficiency is shown in Fig. 3. From the experimental results, we can observe the following: 1) Both SWA and MSWA has a significant efficiency advantage compared to the standard self-attention mechanism. This advantage becomes more pronounced as the batch size increases. 2) Compared to the traditional SWA, MSWA adapts better to larger batch sizes, often achieving better performance as the batch size increases. 3) For larger base window sizes, the efficiency advantage of MSWA becomes even more apparent, making it well-suited for scaling up context lengths.

We conduct a series of ablation studies in this section, mainly focusing on the impact of the base window size on the MSWA mechanism, as well as the effects of other window variation strategies on the MSWA mechanism. In these experiments, we use the same setup from Sec. 5.1.1, only alter the sizes of windows and the variation method.