Exploring Attention Map Reuse for Efficient Transformer Neural Networks
Transformer-based deep neural networks have achieved great success in various
sequence applications due to their powerful ability to model long-range
dependency. The key module of Transformer is self-attention (SA) which extracts
features from the entire sequence regardless of the distance between positions.
Although SA helps Transformer performs particularly well on long-range tasks,
SA requires quadratic computation and memory complexity with the input sequence
length. Recently, attention map reuse, which groups multiple SA layers to share
one attention map, has been proposed and achieved significant speedup for
speech recognition models. In this paper, we provide a comprehensive study on
attention map reuse focusing on its ability to accelerate inference. We compare
the method with other SA compression techniques and conduct a breakdown
analysis of its advantages for a long sequence. We demonstrate the
effectiveness of attention map reuse by measuring the latency on both CPU and
GPU platforms.
Kyuhong Shim, Jungwook Choi, Wonyong Sung
# Exploring Attention Map Reuse for Efficient Transformer Neural Networks

###### Abstract

Transformer-based deep neural networks have achieved great success in various sequence applications due to their powerful ability to model long-range dependency. The key module of Transformer is self-attention (SA) which extracts features from the entire sequence regardless of the distance between positions. Although SA helps Transformer performs particularly well on long-range tasks, SA requires quadratic computation and memory complexity with the input sequence length. Recently, attention map reuse, which groups multiple SA layers to share one attention map, has been proposed and achieved significant speedup for speech recognition models. In this paper, we provide a comprehensive study on attention map reuse focusing on its ability to accelerate inference. We compare the method with other SA compression techniques and conduct a breakdown analysis of its advantages for a long sequence. We demonstrate the effectiveness of attention map reuse by measuring the latency on both CPU and GPU platforms.

Keywords:efficient transformer attention map reuse self attention speech recognition

## 1 Introduction

The ability to learn long-range dependency is essential for various sequence processing tasks such as language modeling, machine translation, text summarizing, question answering, and speech recognition. Deep neural networks (DNNs) have been achieved great success in these complex sequence tasks over traditional handcrafted and rule-based techniques. DNN architectures can be characterized by how the feature extraction mechanism incorporates past or future information. For example, recurrent neural networks (RNNs) such as LSTM [10] encode the entire previous sequence into a single feature vector, which is beneficial for the compact implementation. However, it causes a loss of long-range information since feature representation is restricted to a single vector. In contrast, Transformer [24] models directly access the entire sequence, therefore they are much more advantageous for long-range dependency modeling. Transformer models
have demonstrated excellent performance over RNNs and become the universal choice for most sequence processing applications.

However, Transformer models suffer from quadratic computation and memory complexity to calculate the relationship between every pair of locations. More precisely, the self-attention (SA) module, one of two submodules in Transformer, utilizes an attention mechanism to process a sequence of length \(T\) at once. The attention mechanism computes the correlation between length \(T\) features by the matrix multiplication of the feature matrix (\(T\times d\)) and its transposed form (\(d\times T\)), where \(d\) is the feature vector dimension. This \(O(T^{2})\) quadratic hardware cost is especially problematic when deploying Transformer in practice, especially when the input sequence length is very long. For example, in speech recognition and language modeling, the length of the input sequence is very long and long-range dependency is crucial for accurate prediction. Language modeling models usually take about 256 to 512 previous words as input to predict the next word, and speech recognition models often process much more frames (750 frames for 30 seconds) to transcribe the given utterance. Although these two tasks are core building blocks of many applications, RNN models are still practically favorable because of the heavy computational cost of the Transformer models for resource-limited devices such as mobile and embedded systems.

Considering that a Transformer model is composed of multiple Transformer layers, and each layer consists of multiple (attention) heads that exploit different all-to-all relationships, the complexity increases proportionally to the number of SA heads (\(H\)) and SA layers (\(L\)). Various architectural modifications have been proposed to reduce the heavy computation of SA, where the studies can be mainly categorized into two groups [23]. The first group (Section 2.2) focuses on the output of the attention mechanism, called attention map \(A\in\mathbb{R}^{T\times T}\), whereas the second group (Section 2.3) focuses on reducing the number of \(L\) and \(H\). The first group includes 1) computing only a few elements (sparsely) based on patterns or importance [2, 5], and 2) approximating \(A\) by low-rank factorization, clustering, or kernelization [27, 19, 6]. However, the aforementioned methods cannot utilize the full capability of the modern parallel processing hardware such as graphics processing unit (GPU), digital signal processor (DSP), and neural processing unit (NPU) [18, 13], due to their unstructured computation savings. For example, selectively computing a few important elements may produce a random-like access pattern that depends on the input. Clustering-based methods also require additional K-means or locality-sensitive hashing to group similar ones for a more accurate approximation. On the other hand, the second group reduces the effective number of attention map computations by pruning out SA heads or SA layers. These approaches are much more hardware-friendly because the removal of a large computation block (e.g., attention head, attention layer) is very structured and predictable.

Recently, **attention map reuse** has been proposed for various applications, such as language modeling [30], machine translation [28], and speech recognition [22]. The key idea of the method is to reuse the attention map of \(\ell\)-th layer \(A^{\ell}\) for multiple consecutive layers, \((\ell+1)\)-th to \((\ell+M)\)-th layer, therefore reduc
ing the effective number of SA computation from \(L\) to \(L/M\). This architectural change is highly structured and easy to implement on modern hardware platforms. Especially, for speech recognition, we showed that attention map reuse can be adopted without much degradation of recognition performance [22]. The paper discovered that the reason behind the success of this method is that SA blocks in successive layers perform a similar role for speech recognition and can be merged. However, the analysis on attention map reuse was mainly focused on their behaviors; not much discussion was provided on how the speedup is achieved in terms of the actual implementation.

In this paper, we provide a deeper understanding of attention map reuse in the case of speech recognition. In Section 2, we briefly introduce Transformer and SA architecture and compare the previous SA compression methods with attention map reuse. In Section 3, we analyze the effect of each component of a Transformer model for speech recognition. In Section 4, we evaluate the latency savings of attention map reuse on various configurations and platforms. The results confirm that attention map reuse is a promising inference speedup technique for Transformer-based long-range sequence processing.

## 2 Background and Related Work

### Transformer and Self Attention

We briefly introduce the components of a Transformer layer. A Transformer layer is composed of two submodules: 1) multi-head self-attention (MHSA) and 2) feed-forward (FF). Figure 1 illustrates the overall architecture, and Figure 2 visualizes the internal structure of MHSA submodule.

Figure 1: Illustration of the Transformer-based model which is a stack of \(L\) Transformer layers. For speech recognition, input is human speech and output is a transcribed sentence. Note that every frame is processed together without recurrence.
be a sequence of \(T\) tokens3. For the input \(X=\{x_{1},x_{2},...x_{T}\}\) where \(X\in\mathbb{R}^{T\times d}\), SA for the \(h\)-th head (total \(H\) heads) starts with three linear projections as:

Footnote 3: We exploit the term ‘token’ as a common concept for both natural language processing and speech recognition. Each token represents a (sub-)word feature and a speech frame feature, respectively.

\[Q_{h},K_{h},V_{h}=XW_{Q_{h},K_{h},V_{h}}+b_{Q_{h},K_{h},V_{h}}\quad(Q_{h},K_{h},V_ {h}\in\mathbb{R}^{T\times d_{h}}) \tag{1}\]

where \(Q,K,V\) indicates the query, key, value, and \(W\in\mathbb{R}^{d\times d_{h}},b\in\mathbb{R}^{1\times d_{h}}\) are weight and bias parameters. \(d_{h}=d/H\) is the feature dimension for each attention head. The attention map \(A_{h}\) for the \(h\)-th head is then computed as a scaled dot-product of query and key

\[A_{h}=\text{Softmax}(\frac{Q_{h}K_{h}^{T}}{\sqrt{d_{h}}}).\quad(A\in\mathbb{R} ^{T\times T}) \tag{2}\]

After the softmax operation, each row of \(A_{h}\) becomes a probability distribution of a length \(T\). Intuitively, the \((i,j)\)-th element of the attention map represents how much \(j\)-th token contributes to \(i\)-th token. Then, the outputs of each attention head are concatenated and followed by another linear projection:

\[\text{SA}_{h}(X) =A_{h}V_{h} \tag{3}\] \[\text{MHSA}(X) =\text{Concat}(\text{SA}_{1},...\text{SA}_{h})W_{O}+b_{O}. \tag{4}\]

By exploiting multiple attention heads, Transformer can extract diverse relationships between tokens in a single layer. For example, one head may focus on the syntactic connections while the other head focuses on specific words.

Figure 2: Illustration of the MHSA submodule. The computation flow is identical for each attention head.
The FF submodule is a stack of two linear projections with an intermediate non-linear activation function:

\[\mathrm{FF}(X)=\big{(}\phi(XW_{1}+b_{1})\big{)}W_{2}+b_{2} \tag{5}\]

where \(W_{1}\in\mathbb{R}^{d\times 4d}\), \(W_{2}\in\mathbb{R}^{4d\times d}\), \(b_{1}\in\mathbb{R}^{4d}\) and \(b_{2}\in\mathbb{R}^{d}\) are parameters, and function \(\phi\) can be ReLU, Swish, GELU, etc. Finally, the output of \(\ell\)-th Transformer layer is formulated as below

\[Z^{\ell} =\mathrm{LN}(\mathrm{MHSA}(X^{\ell})+X^{\ell}) \tag{6}\] \[X^{\ell+1} =\mathrm{LN}(\mathrm{FF}(Z^{\ell})+Z^{\ell}). \tag{7}\]

where \(Z\) is the intermediate term and LN indicates the layer normalization [1]. There exist a residual connection that adds each submodule's input and output. As \(T\) increases, the cost of MHSA increases quadratically following \(O(T^{2})\) while the cost of FF increases linearly. Therefore, reducing the MHSA computation is very important for efficient realization of Transformer models.

### Attention Map Sparsification

Numerous studies have proposed techniques to selectively compute elements of the attention map [23]. Patterned attention computation approaches, which select elements in a fixed manner, have been introduced [5, 3, 2]. For example, Sparse Transformer [5] exploits strided pattern that only attends \(\ell\) local positions and positions of stride \(\ell\), resulting in a \(O(T\sqrt{T})\) complexity. Depending on the pattern, these methods can be supported on modern accelerators with custom kernel implementation. Adaptive element selection approaches, which dynamically decide elements to compute, have also been studied [26, 12, 19]. For example, Reformer [12] exploits locality-sensitive hashing for clustering so that the attention map can be computed only within the grouped elements. However, these clustering-based methods require additional computation steps and the access positions dynamically change depending on the input. The aforementioned studies focus on how to reduce the cost of the attention dot product, while not changing the overall structure of the Transformer model.

### Removing Attention-related Blocks

To build an efficient SA mechanism, many studies have focused on the structured removal of a large chunk of computation blocks. The candidate for the removal (pruning) can be attention heads or attention layers. For attention head pruning, previous studies have reported that pruning out certain attention heads does not affect the final performance [15, 25, 34]. Specifically, redundant or less important attention heads in SA can be pruned without degrading the performance for speech recognition [32] and auto-regressive language modeling [21]. Similarly, layer-level pruning has also been studied [8, 20, 11] for natural language processing tasks. LayerDrop [8] randomly omits the residual connection during
training to make the model more robust to layer pruning. Our goal is to provide a comprehensive understanding on the impact of the specific structured removal approach, attention map reuse.

## 3 Transformer for Speech Recognition

### Conformer Architecture

Transformer-based models have been actively employed for state-of-the-art speech recognition, replacing the previous RNN-based or CNN-based models [16, 33] thanks to their ability to extract informative features from the input utterance. In particular, previous works discovered that SA automatically learns to extract useful phonological features [29, 22] during training. The input of a Transformer model is a sequence of audio features (frames) extracted by short-time Fourier-transform (STFT), and the output is a sequence of transcribed words. Because the input and output domains are different, the model internally performs a transformation that turns audio features into text features while passing through a stack of Transformer layers.

Following the previous work [22], we employ Conformer [9], a variant of Transformer widely used for speech recognition. Conformer consists of 4 submodules as illustrated in Figure 3. The main architectural difference between Conformer and Transformer is that Conformer includes two more submodules: an additional FF submodule at the front and the intermediate convolutional (Conv) submodule. Considering that speech is a continuous signal and nearby frames are highly dependent on each other, the convolutional module is beneficial for enhancing the local relationship between frames that might not be emphasized in SA. By utilizing both SA and Conv, Conformer achieved a state-of-the-art recognition performance with much fewer parameters than the Transformer-based model without the Conv submodule [9].

### Breakdown Analysis

To understand the bottleneck of Transformer-based models, we analyze how much resource each submodule takes. We first show the parameter size of each component in Figure 4. We consider the Conformer-M model, which consists of \(L\)=16 layers of hidden dimension of \(d\)=256 and \(H\)=4 attention heads. We

Figure 3: Illustration of a Conformer layer consists of 4 submodules.
can observe that about 66% of parameters come from the FF submodule and SA takes only 21% of parameters. The reason for this imbalance is that each FF includes \(8d^{2}\) parameters (therefore, \(16d^{2}\) parameters for two FFs) while SA has \(5d^{2}\) parameters. In other words, we need to reduce the FF parameters if the target system does not equip enough memory space. However, the slowest submodule is not FF but SA.

Figure 5 presents the inference speed breakdown for different input sequence lengths (we will discuss ReuseSA in the next section). Note that the \(x\)-axis of the figure indicates the number of feature frames extracted with a 40ms stride. The lengths 256, 512, and 1024 can be reinterpreted to about 10, 20, and 40-second input audio lengths, respectively. The results are evaluated on a single RTX-Titan GPU, but the tendency should be similar for CPU platforms. As input length increases, SA requires a quadratic computation cost while FF and

Figure 4: The ratio of the parameter size occupied by each Conformer submodule. The values are based on the Conformer-M model.

Figure 5: Inference time (us) of each Conformer submodule. The order of vertical bars is FF, Conv, SA, and ReuseSA. Attention map reuse replaces the SA module with the ReuseSA module and significantly reduces the inference time. \(x\)-axis is the number of frames in the input sequence.
Conv only need linearly increasing costs. Therefore, SA takes 45% of total inference time for a 10-second input but 67% for a 30-second input. This is highly problematic because many speech-related tasks, such as conference transcription, listening comprehension, and speech-to-speech translation, often process much longer utterances than 30-second.

## 4 Evaluation of Attention Map Reuse

### Attention Map Reuse

From the observation that the behavior of SA is very similar between neighboring layers, attention map reuse only computes a single attention map through \(M\) consecutive SA layers. Figure 6 illustrates the attention reuse procedure. For example, if the attention map from \(1^{\text{st}}\) layer \(A^{1}\) is shared through layers \(2\sim M\), the SA output of those layers can be easily computed as below:

\[\text{SA}_{h}^{i}(X^{i})=A_{h}^{1}V_{h}^{i},\quad i\in\{2,3,...,M\} \tag{8}\]

where \(i\) is the layer index and \(V_{h}^{i}\) is the value computed from each layer. In other words, attention map reuse groups a set of layers that includes one original SA layer at the front and following \(M-1\) reuse SA layers. The remaining parts of the model, such as FF and Conv submodules, are unchanged. By omitting the attention map computation, the effective number of SA calculations can be decreased by \(M\) times. Instead of optimizing each SA mechanism, attention map reuse exploits the characteristic of SA and proposes a new axis for Transformer compression.

Figure 6: Illustration of each attention head for attention map reuse, in case of the \(4\times 4\) reuse configuration. The computed attention map \(A^{\ell}\) from \(\ell\)-th layer is reused for the next three layers. For simplicity, other submodules are omitted.
After applying attention map reuse, the query and key are not needed for the following SA layers and their associated parameters can be removed. To compensate for the reduced parameter size, previous work suggested increasing the hidden dimension of value [22]. We follow the same strategy, which makes the size of \(W_{V_{h}}\) from \(\mathbb{R}^{T\times d_{h}}\) to \(\mathbb{R}^{T\times 2d_{h}}\), and the size of \(W_{O}\) from \(\mathbb{R}^{d\times d}\) to \(\mathbb{R}^{2d\times d}\) (see Figure 6). The number of attention heads is preserved. Note that attention map reuse is not a fine-tuning approach that starts from the fully converged model; we train the modified model from scratch.

### Reuse Analysis

We analyze the effect of attention map reuse using different configurations. Tables 1 and 2 show the inference time of different reuse configurations. The configuration \(A\times B\) indicates that \(A\) layers are grouped to share one attention map and total \(B\) groups of layers exist. The baseline is a 16-layer Conformer-M model [9], which can be represented as the configuration of \(1\times 16\). The results are measured on an RTX Titan GPU and an Intel Xeon Gold 6130 CPU. We evaluate three different attention map reuse configurations, \(2\times 8\), \(4\times 4\), and \(8\times 2\) where the number of attention map computations is 8, 4, and 2, respectively. Although we can combine different numbers of layers to form heterogeneous groups, we find that the number of unique attention map computations determines the total latency of the model. Note that there is a trade-off between the performance and the inference speed if the number of unique attention map computations is too small (see Section 4.3).

\begin{table}
\begin{tabular}{c|c c c c c c c} \hline \hline Config. & 128 & 256 & 384 & 512 & 640 & 768 & 896 & 1024 \\ \hline \(1\times 16\) & 3.32 & 4.16 & 7.66 & 11.77 & 17.11 & 23.10 & 30.27 & 37.82 \\ \(2\times 8\) & 2.40 & 3.38 & 5.89 & 8.61 & 12.11 & 15.87 & 20.27 & 24.51 \\ \(4\times 4\) & 1.99 & 2.91 & 4.92 & 6.93 & 9.53 & 12.18 & 15.28 & 17.98 \\ \(8\times 2\) & 1.25 & 2.56 & 4.11 & 5.75 & 7.72 & 10.03 & 12.66 & 14.63 \\ \hline \hline \end{tabular}
\end{table}
Table 1: GPU Inference time (ms) of a speech recognition model with different attention map reuse configurations. The numbers from 128 to 1024 indicate the frames in the input utterance. Configuration \(1\times 16\) represents the baseline model. Each frame corresponds to a 40ms stride.

\begin{table}
\begin{tabular}{c|c c c c c c c c} \hline \hline Config. & 128 & 256 & 384 & 512 & 640 & 768 & 896 & 1024 \\ \hline \(1\times 16\) & 62.97 & 82.93 & 114.44 & 164.48 & 226.88 & 364.05 & 398.36 & 1033.59 \\ \(2\times 8\) & 54.08 & 70.84 & 87.05 & 124.01 & 152.35 & 255.58 & 315.91 & 630.30 \\ \(4\times 4\) & 50.24 & 63.72 & 77.21 & 100.87 & 135.63 & 182.08 & 231.14 & 398.57 \\ \(8\times 2\) & 42.66 & 55.20 & 67.39 & 96.28 & 116.06 & 136.54 & 191.17 & 281.11 \\ \hline \hline \end{tabular}
\end{table}
Table 2: CPU Inference time (ms) of a speech recognition model with different attention map reuse configurations.
We observe a clear improvement in the inference speed for both GPU and CPU as more layers are grouped. For \(4\times 4\) configuration, 10-second (length 256), 20-second (length 512) and 40-second (length 1024) utterance saves about 38%, 41%, and 52% of the GPU inference time, respectively. The same configuration saves about 33%, 39%, and 61% for the CPU inference time. CPU inference time gain is not as good as that of GPU when the length is short but provides a higher benefit when the length is long. Figure 5 also demonstrates the effect of reuse in the case of \(4\times 4\). Comparing the original SA and the reuse SA, the reuse SA takes about 35% to 45% of the latency of the original SA.

### Discussion

#### 4.3.1 Attention Map Reuse in Natural Language Processing

Attention map reuse was applied for neural machine translation [28] and BERT-based language modeling [30]. For machine translation, the average speedup was about 1.3 times without considerable performance degradation [28]. The inference speedup is less than in speech recognition because machine translation usually considers a relatively shorter sequence length \(T\) (shorter than 50) and larger hidden dimension \(d\). If \(d\) is much larger than \(T\), the \(O(T^{2})\) SA computation does not dominate the total inference cost, so the advantage from attention map reuse is limited. For the language model inference, the speedup was about 1.3 times with a marginal performance gain [30]. However, the work mainly focused on the \(2\times 6\) configuration because the more aggressive reuse configuration did not achieve satisfactory performance for GLUE downstream tasks. We conclude that the efficiency of attention map reuse can be maximized when 1) the expected input sequence length is longer than the hidden dimension, and 2) the target task is not very sensitive to a small attention map variation. Speech recognition well fits these conditions because it handles very long sequences and the frame features change slowly along the time axis.

#### 4.3.2 Effect on Performance

In Table 3, we show the word error rate (WER) of attention map reuse on LibriSpeech [17] speech recognition dataset, which includes four evaluation data subsets4. We borrow the result from our original paper [22] on speech recognition to briefly show that attention map reuse can be

\begin{table}
\begin{tabular}{c|c|c c c c} \hline Config. & \#Param (M) & _dev-clean_ & _dev-other_ & _test-clean_ & _test-other_ \\ \hline \(1\times 16\) & 25.45 & 3.1 & 8.3 & **3.2** & 8.4 \\ \(2\times 8\) & 24.92 & **3.0** & **8.2** & 3.3 & **8.2** \\ \(4\times 4\) & 24.66 & **3.0** & **8.2** & 3.3 & **8.2** \\ \(8\times 2\) & 24.52 & 3.3 & 8.8 & 3.6 & 8.7 \\ \hline \end{tabular}
\end{table}
Table 3: Word error rate (%) for different attention map reuse configurations.
employed without affecting the performance. In short, recognition performance is almost preserved for \(2\times 8\) and \(4\times 4\) configurations but not for the \(8\times 2\) case. The authors suggested that the capacity may become insufficient to internalize the necessary information for every layer in a group when too many layers share the same attention map. Note that the number of parameters is almost the same for configurations because we doubled the dimension of value (see Section 4.1).

#### 4.2.3 Attention Computation Reduction in Speech Recognition

Many works have been proposed techniques to reduce the computational cost of SA, especially for speech recognition. Local windowing is the common approach that only exploits a limited range of frames for attention map computation. For example, each frame may only consider neighboring frames (e.g., only accessing past 64 and future 64 frames) as candidates of the attention mechanism; this approach decreases the complexity of SA from \(O(T^{2})\) to \(O(TR)\), where \(R\) is the number of accessible frames. However, restricting the range often lowers the recognition performance over full sequence-based models [31, 7]. On the other hand, several studies designed more efficient Transformer models for speech recognition [4, 14]. These approaches, including faster query-key dot product [14] and time-strided SA [4], are orthogonal to attention map reuse and can be used together.

## 5 Conclusion

In this paper, we analyzed a recently proposed efficient SA compression method, named attention map reuse, for Transformer-based speech recognition. We first perform a detailed analysis of the inference bottleneck of the Conformer model used for speech processing, evaluated on a wide range of input sequence lengths. Our analysis provides a thorough understanding on the burden of SA when using Transformer in practice. Then, we demonstrate the computational savings from attention map reuse on GPU and CPU platforms. We claim that attention map reuse is a very promising method for utilizing Transformer-based models on modern hardware systems.
