Understanding the Role of Self Attention for Efficient Speech Recognition

1 INTRODUCTION . Recent advances in end-to-end automatic speech recognition ( ASR ) have been driven by Transformer models ( Vaswani et al. , 2017 ) . Transformer was first introduced for natural language processing ( NLP ) tasks such as neural machine translation ( Vaswani et al. , 2017 ; Ott et al. , 2018 ) , language modeling ( Dai et al. , 2019 ; Rae et al. , 2019 ) , and text generation ( Raffel et al. , 2020 ) . Thanks to its superior performance in processing sequence input , Transformer has been widely adopted in various state-of-art ASR models ( Zhang et al. , 2020b ; Ng et al. , 2021 ; Guo et al. , 2021 ) . Self-attention ( SA ) is a core component of Transformer-based ASR , which dynamically collects information from multiple frames of an audio sequence . However , the computation and memory costs of SA increase quadratically with the length of a sequence , which is particularly problematic for ASR . For example , just a 30-second utterance corresponds to about 750 frames with a widely used window stride of 40ms . Understanding the role of SA may provide essential insights for the efficient design of Transformerbased ASR models . Extensive studies have examined the behavior of SA in the field of NLP ( Kovaleva et al. , 2019 ; Park et al. , 2019 ; Gong et al. , 2019 ; Rogers et al. , 2020 ) . Recently , several studies further attempted to discover the characteristics of SA in the speech domain . Yang et al . ( 2020 ) revealed that self-attention features in the self-supervised audio Transformer are categorized into global , vertical , and diagonal patterns . Zhang et al . ( 2021b ) focused on the diagonality of upper SA layers in ASR models for improving efficiency . However , these prior works revealed limited insights on the patterns discovered in SA , constraining its use for improving model efficiency . Thus , providing a holistic view of the role of SA is desirable for efficient ASR model design . In this work , we reveal that SA plays two distinct roles in the success of Transformer-based ASR models : phonetic and linguistic localization , as illustrated in Figure 1 . First , phonetic localization of lower SA layers attends to the phonologically meaningful global context . Second , linguistic localization of upper SA layers mainly attends to the local context of a near-diagonal attention map . We hypothesize that the phonetic variance in utterances such as variations in pronunciation is standardized in the lower SA layers so that the upper SA layers can identify local linguistic features for accurate transcription . To investigate the behavior of SA layers , we propose phoneme attention relationship ( PAR ) to explain how phoneme localization works quantitatively . Interestingly , we discover that the phonetic localization represents the traditional knowledge on phonetics ; for example , labial , velar , or nasal phonemes tend to attend to each other . Based on this understanding , we propose a practical method for efficient ASR model design . We reuse the attention maps of SA layers while preserving the phonetic localization capability of the lower SA layers , resulting in up to 1.96x speedup in inference and 33 % savings of training time with considerably improved performance in challenging ASR tasks ( LibriSpeech dev/test-other dataset ) . Our contributions can be summarized as follows : • We reveal that SA layers contribute to ASR with two distinct roles : phonetic localization in the lower layers and linguistic localization in the upper layers . This unique distinction leads to an in-depth analysis of phonetic SA for the first time . We further propose phoneme attention relationship ( PAR ) to quantitatively identify the role of phonetic localization . • We propose layer-wise attention map reuse for efficient Transformer-based ASR models . In particular , we discover that attention map reuse is possible in lower SA layers as long as the phonetic localization property quantified by PAR is preserved . We demonstrate with the popular ASR model and dataset that ASR performance can be maintained or slightly improved even if the attention map is reused . • We implement the attention map reuse on real GPU platforms and achieve up to 1.96x inference speedup and 33 % savings in training time , demonstrating that the proposed method is practical . 2 BACKGROUND . 2.1 ASR ENCODER AND SELF-ATTENTION . SA is usually utilized as a module inside the ASR model consisting of stacked Transformer encoder layers . ASR encoder takes a sequence of short-time Fourier-transformed ( STFT ) audio features , known as a ‘ frame ’ , as input and extracts a high-level feature of each frame through multiple layers . As illustrated in Figure 1 , the extracted high-level feature changes for each layer ; stacked SA layers first extract phonetic features from audio features and utilize these features to build linguistic features for the output transcription . We briefly review the SA1 computation procedure . Consider a sequence of d-dimensional column vectors X = { x1 , x2 , ... xT } as input . Each vector corresponds to each frame of speech where the total number of T frames are included . The input feature vector X is projected to query ( Q ) , key ( K ) , and value ( V ) of h-th attention head as follows : qh , i =W Q h xi , kh , i =W K h xi , vh , i =W V h xi ( WQh , W K h , W V h ∈ Rdh×d ) ( 1 ) WQ , K , V indicates projection matrices for Q , K , V , respectively . dh = d/H is the dimension of each attention head where H is the number of attention heads . The attention map ( Ah ) , which represents how much frames attend to each other , is computed by scaled dot-product operation followed by softmax . The resulting attention map takes a form of a 2D matrix where each row is a probability vector . A single element of the attention map ( Ah [ i , j ] ) represents how much i-th frame attends to j-th frame . The attention head ( dh ) is a weighted sum of V using the attention map as weight . Note that each attention head corresponds to a different attention map Ah ; this multi-head design enables focusing on various perspectives within a single SA layer . The output O = { o1 , o2 , ... oT } is computed by the projection ( WO ∈ Rd×d ) on the concatenated attention heads2 . Ah [ i , : ] = Softmax j ( qh , ik T h , j√ dh ) , dh , i = T∑ j=1 Ah [ i , j ] vh , j , oi =W OConcat h ( dh , i ) ( 2 ) SA layer contains O ( d2 ) parameters.As shown in equations , SA requires quadratic computation and memory complexityO ( T 2 ) , in exchange for the ability to access any location in the sequence . When N layers of SA are stacked , the burden proportionally increases . 2.2 PREVIOUS WORK ON SELF-ATTENTION ANALYSIS . ASR considers both phonetic and linguistic aspects to transform audio input to text output . However , the studies on NLP mostly analyze the linguistic characteristics of SA , and the studies on selfsupervised audio representation learning ( SSAL ) mainly focus on the phonetic behaviors of SA . The valuable findings from both domains can not be directly applied to ASR . NLP The behavior of SA has been widely studied in the NLP domain ( Rogers et al. , 2020 ) , mostly focused on BERT ( Devlin et al. , 2019 ) , a self-supervised language representation learning model . Kovaleva et al . ( 2019 ) and Guan et al . ( 2020 ) suggested that attention patterns can be clustered into several groups and the pattern may change depending on the fine-tuning task . Clark et al . ( 2019 ) and Tenney et al . ( 2019 ) observed attention maps that correspond to linguistic concepts of the language . Voita et al . ( 2019 ) also characterized linguistic attention heads and connected the knowledge to efficient model structure . However , studies on NLP only provide analysis on linguistic attention . SSAL Recently , several studies have been introduced to understand how the audio information is encoded in SSAL models , such as CPC ( Oord et al. , 2018 ) , Wav2Vec 2.0 ( Baevski et al. , 2020 ) , Mockingjay ( Liu et al. , 2020 ) , HuBERT ( Hsu et al. , 2021 ) , and Audio ALBERT ( Chi et al. , 2021 ) . Ma et al . ( 2021 ) and Shah et al . ( 2021 ) demonstrated that a wide spectrum of phonetic information is included in these models . Especially , Yang et al . ( 2020 ) categorized attention maps into three categories : global , vertical , and diagonal , where diagonal heads attend to local frames and vertical heads either focus or neglect specific phonemes . However , Yang et al . ( 2020 ) only discovered attention patterns without an explanation on how phonetic feature extraction is achieved with these patterns . ASR Previous works have investigated the redundancy of attention maps mainly based on diagonality . From the observation that attention maps in upper layers show highly diagonal patterns , Zhang et al . ( 2021b ) proposed replacing upper SA layers to feed-forward layers without performance loss . Zhang et al . ( 2021a ) removed SA heads of high diagonality during training as a regularization 1In this paper , SA is equivalently used as a multi-head SA . 2The two terms ‘ head ’ and ‘ attention head ’ are used interchangeably . Figure 2 : Cumulative attention diagonality ( CAD ) of each attention head . Four points for each layer correspond to the CAD of four attention heads . The black line connects the median across the layers . but keep every head for the test time . Similarly , Chang et al . ( 2020 ) introduced adaptive attention span where each attention head equips a different attention span width to reduce the sequence length for the computation , starting from the intuition that some heads only attend to neighboring frames . These approaches mainly focus on reducing the burden of diagonal and concentrated attention , however , diagonality-based analysis has limitations in optimizing phonetic attention . We distinguish SA into two groups and provide proper analysis for each . 3 UNDERSTANDING THE ROLE OF SELF-ATTENTION IN ASR . 3.1 ANALYSIS SETUP . We train and evaluate the model on the LibriSpeech-960 ( Panayotov et al. , 2015 ) dataset . The dataset include two types of data , clean and other , where other contains more challenging utterances . We extract the 80-dimensional log-Mel filterbank feature from a 25ms window with a stride of 10ms . We use 128 sub-word tokens as vocabulary , built on SentencePiece ( Kudo & Richardson , 2018 ) library using the byte-pair encoding ( Sennrich et al. , 2016 ) . The analyses are performed on LibriSpeech test-clean dataset unless specified . We use Conformer-M ( medium ) ( Gulati et al. , 2020 ) as the baseline ASR encoder , trained with CTC ( Graves et al. , 2006 ) loss . Conformer is a variant of Transformer , which combines an additional convolution module to enhance the ability to collect local neighboring features . We follow recently introduced state-of-the-art ASR studies ( Zhang et al. , 2020b ; Ng et al. , 2021 ; Guo et al. , 2021 ) that have adopted Conformer as their ASR encoder . We train the model with AdamW ( Loshchilov & Hutter , 2018 ) optimizer for 200K iterations . Please see Appendix A.1 and A.2 for the model configuration and training details .

Summary: The paper explores how self attention maps in conformer based ASR models are distributed across layers. The authors note that the lower layers of the model capture phonetic information, and the variance of the attention map distribution for the same class of phones across the utterance goes down. The higher layers do not display this behavior. The authors hypothesize that the higher layers capture linguistic information that likely combines information across phones.  Based on the observation that the variance of the attention map distribution goes down, the authors propose reusing the attention map from one layer across multiple layers. This reduces the amount of computation needed for inference, and reduces the training time. Tying attention maps also results in small improvements in performance.