Learning Monotonic Alignments with Source-Aware GMM Attention

1 INTRODUCTION . In recent years , transformer models with soft attention have been widely adopted in various sequence generation tasks ( Raffel et al. , 2019 ; Vaswani et al. , 2017 ; Parmar et al. , 2018 ; Karita et al. , 2019 ) . Soft attention does not explicitly model the order of elements in a sequence and attends all encoder outputs for each decoder step . However , the order of elements is crucial for understanding monotonic sequence-to-sequence ( Seq2Seq ) tasks , such as automatic speech recognition ( ASR ) , video analysis , and lip reading . Learning monotonic alignments enables the model to attend to a subset of the encoder output without performance degradation in these tasks . In comparison , soft attention is not suitable for streaming inference applications because the softmax operation needs to wait until all encoder outputs are produced . Figure 1 ( b ) shows the attention plot for soft-attention . Soft attention learns the alignments between queries and keys based on their similarities ; it requires all encoder tokens prior to the attention score calculation . Furthermore , soft attention can not easily decode long-form sequences that are not considered in the training corpus . The Gaussian Mixture Model ( GMM ) attention Graves ( 2013 ) ; Battenberg et al . ( 2020 ) ; Chiu et al . ( 2019 ) have been proposed for learning the monotonic mapping between encoder and decoder states for long-form sequence generation . The GMM attention is a pure location-aware algorithm in which encoder contents are not considered during attention score calculation . However , each element in the encoder output sequence contains different amounts of information and should be attended considering their contents . In figure 1 ( c ) , the GMM attention fails to learn the detailed alignments and attends to many tokens simultaneously . In this study , we adopted the GMM attention mechanism to the modern transformer structure and proposed the Source-Aware Gaussian Mixture Model ( SAGMM ) attention which considers both contents and orders of source sequences . Each component in the SAGMM is multi-modal and discards non-informative tokens in the attention window . For online inference , we propose a truncated SAGMM ( SAGMM-tr ) that discards the long-tail of the attention score in the SAGMM . To the best of our knowledge , this is the first attempt to adopt a GMM-based attention to online sequence generation tasks . Learning accurate monotonic alignments enables the SAGMM-tr to attend to a relevant subset of sequences for each decoder step and improves the performance of the model in terms of streaming and long-form sequence generation tasks . Figure 1 ( d ) shows the monotonic alignments learned by the SAGMM-tr , enabling online inference . Experiments involving streaming and long-form ASR showed substantial performance improvements compared with conventional algorithms without performance degradation in offline in-distribution ASR . Furthermore , we tested the SAGMM-tr in a machine translation task and demonstrated the performance of the proposed algorithm in non-monotonic tasks . 2 SOURCE-AWARE GMM ATTENTION . 2.1 SOFT ATTENTION . Herein , we abbreviate the head index h during attention score calculation for simplicity . In dot-product multi-head soft attention Vaswani et al . ( 2017 ) without relative positional encoding , the attention score from soft attention αSoft is derived from the query matrix Q ∈ RI×d and key matrix K ∈ RJ×d as follows : αSoft = softmax ( QKᵀ/ √ d ) ( 1 ) where d , I , and J are the feature dimension , decoder and encoder sequence length , respectively . The attention context matrix from the h-th head Hh and the multi-head output M are expressed as Hh = αhSoft V h ( 2 ) M = concat [ H1 ; ... ; Hnh ] WO ( 3 ) where αhSoft denotes the αSoft for the h-th head , V h ∈ RJ×d the value matrix for the h-th head , and nh the number of heads . In this study , we adopted relative positional encoding Shaw et al . ( 2018 ) ; Dai et al . ( 2019 ) which provides a stronger baseline for long-form sequence generation for self-attention layers . 2.2 GMM ATTENTION . The previous studies regarding GMM attention Battenberg et al . ( 2020 ) ; Chiu et al . ( 2019 ) were based on early content-based attention ( Cho et al. , 2014 ) . Li et al . ( 2020 ) adopted the GMM attention to the transformer framework , but did not provide detailed descriptions . Herein , we adopt v2 model in Battenberg et al . ( 2020 ) which improved the performance of the original GMM attention mechanism Graves ( 2013 ) . We define the GMM attention as a variant of multi-head attention by considering a Gaussian distribution component as an attention score of single-head in a multi-head mechanism . In the study by Battenberg et al . ( 2020 ) , the value matrix was shared for all Gaussian components , whereas multi-head value matrices were multiplied with the probability from corresponding components in this study to attend to information from different representation subspaces ( Vaswani et al. , 2017 ) . Hence , the multi-head GMM attention introduced here is a more generalized algorithm compared with the early GMM attention . Let us denote the i-th row of Q as Qi ∈ R1×d . The normal distribution parameters for the i-th step are expressed ∆i , σi , φi = ζ ( QiW∆ ) , ζ ( QiWσ ) , QiWφ ( 4 ) µi = ∆i + µi−1 ( 5 ) where ζ ( x ) is a softplus function of x ; W∆ ∈ Rd×1 , Wσ ∈ Rd×1 , and Wφ ∈ Rd×1 . The softplus function was adopted , similar to the study of Battenberg et al . ( 2020 ) in which softplus activation demonstrated better performances than the exponential operation . A mixture component of the GMM attention , from the i-th decoder step to the j-th encoder token , is derived from the αGMM‖i , j is defined as follows : N ( j , µi , σi ) = 1√ 2πσi exp ( − ( j − µi ) 2 2σi ) ( 6 ) αGMM‖i , j = N ( j , µi , σi ) ( 7 ) Hhi = softmaxh ( φhi ) ∑ j αhGMM‖i , jV h j . ( 8 ) where softmaxh denotes the softmax function over heads . The conventional GMM attention mechanism is analogous to integral of source sequences with a uniform axis spacing as shown in Figure 2 ( a ) . In this figure , each rectangle denotes the attention score with specified Gaussian component parameters . The GMM attention assumes that each encoder output is equally important . However , this assumption is not satisfied for many input modalities , e.g . speech and videos from real environments . Moreover , the number of modes in the GMM attention is limited by the number of mixture components . To learn robust alignments for monotonic Seq2Seq tasks , we propose the SAGMM which considers both contents and locations for the attention mechanism . 2.3 SOURCE-AWARE GMM ATTENTION . Figure 2 ( b ) shows the scheme of the proposed SAGMM attention . Compared with the GMM attention , the SAGMM is analogous to an integral of normal distribution with non-uniform spacing based on encoder contents . In the figure 2 ( b ) , the width of rectangle δj controls the model to selectively attend to the informative tokens during the attention score calculation . This content-aware property of the SAGMM enables the model to learn stable monotonic alignments from the training corpus . Furthermore , SAGMM can easily discard non-informative tokens and aggregate information distributed over several remote tokens . In the SAGMM , normal distribution parameter ∆i , σi , φi , and µi are derived from equation 4 and equation 5 . To encode the contents of the source sequences , the weight for each encoder output δj is provided from the j-th row of the key matrix K as follows : δj = sigmoid ( KjWδ ) . ( 9 ) where Wδ ∈ Rd×1 and the sigmoid function are introduced to smoothly bound the maximum weight δj similar to Dong & Xu ( 2020 ) . Subsequently , the probability of normal distribution N ( νj ; µi , σi ) is calculated from its cumulative sum νj expressed as νj = δj + νj−1 ( 10 ) N ( νj ; µi , σi ) = 1√ 2πσi exp ( − ( νj − µi ) 2 2σi ) ( 11 ) Finally , the SAGMM attention score αSAGMM is defined as follows : αSAGMM‖i , j = δj N ( νj ; µi , σi ) ( 12 ) where δj where δj is multiplied to describe uneven step sizes in summation . In the early stage of training , we introduced a length penalty loss to facilitate alignment learning between µ and ν ; it is expressed as Llength = λlength ∗ ( ( µI −min ( I , J ) ) 2 + ( νJ −min ( I , J ) ) 2 ) . ( 13 ) with λlength = 0.0005 where I and J denote the length of the decoder and encoder sequences , respectively . We turned off the Llength after 200K of training steps in the experiments . Finally , we modified 5 as µi = µi−1 + min ( max ( ∆i , 0 ) , 3 ) to scale µi and νj similar to the token indices . This modification facilitates the interpretability of µ and ν . The attention score for each encoder token in the SAGMM was determined independently because they do not rely on the softmax operation over the encoder outputs . It is noteworthy that ∑ j δj N ( νj ; µi , σi ) approximates the integral of the Gaussian distribution . Hence , the sum of the attention weights is approximately 1 , thereby facilitating numerical stability and learning without using softmax . 2.4 SAGMM-TR FOR ONLINE INFERENCE . Since the attention score in SAGMM is generated without softmax normalization over all encoder tokens , we can simply build the attention for streaming inference by cropping the long-tail of the Gaussian distribution . In SAGMM-tr , the normal distribution is truncated to limit the past and future contexts as follows : Ntr ( νj ; µi , σi ) = { 1√ 2πσi exp ( − ( νj−µi ) 2 2σi ) , if µi − 2 √ σi < νj < µi + 2 √ σi 0 , else . ( 14 ) αSAGMM -tr‖i , j = δj Ntr ( νj ; µi , σi ) ( 15 ) Hhi = softmaxh ( φhi ) ∑ j∈µi−2 √ σi < νj < µi+2 √ σi αhSAGMM−tr‖i , jV h j . ( 16 ) Discarding the tokens with threshold of 2 √ σi removes approximately 5 % of the attention score . In online inference with the SAGMM-tr , Hhi can be calculated after νj exceeds µi + 2 √ σi . It is noteworthy that when νj ≥ µi + 2 √ σi satisfies in for a current token j = βi , then the equation satisfies for all j > βi and Hhi can be emitted without for waiting future contexts . We started from the SAGMM model and fine-tuned SAGMM-tr until the performance converged . We discovered that the models with the SAGMM-tr demonstrated slightly better performances than the SAGMM ; hence , we mainly report the results involving the SAGMM-tr model herein . For online speech recognition using the SAGMM-tr , we randomly concatenated the 1-vector after the end of the source sequence with a probability peos in the training stage and performed training to emit the endof-sequence ( EOS ) token only with the utterances containing the 1-vector . The 1-vector concatenation suppresses the EOS token in the long silence part . Finally , we built a unidirectional encoder whose maximum latency was similar to those of Zhang et al . ( 2020 ) ; Dong et al . ( 2019 ) by adopting a block-wise mask on self-attention layers . A detailed explanation on the block-wise masking is provided in the Appendix . Algorithm 1 in the Appendix shows a pseudo code for the SAGMM-tr attention in inference stage . The number of tokens required to 16 was determined by model parameters . We trained and tested the SAGMM-tr with a fixed attention window width c , to demonstrate the performance in environments with maximum latency constraint . In this version , equation 14 - 16 were modified as follows : γi = arg max γi−1≤j < J N ( νj ; µi , σi ) ( 17 ) Ntr ( νj ; µi , σi ) = { 1√ 2πσi exp ( − ( νj−µi ) 2 2σi ) , if γ − c2 < j < γi + c 2 0 , else . ( 18 ) αSAGMM -tr‖i , j = δj Ntr ( νj ; µi , σi ) ( 19 ) Hhi = ∑ γi− c2 < j < γi+ c 2 αSAGMM -tr‖i , jVj . ( 20 ) In equation 20 , we wait c2 additional tokens before producing the output . We compared the performance of the SAGMM-tr with adaptive and fixed window widths through experiments .

Summary: This paper introduces “source-aware” GMM attention and applies it to offline, online, long-form ASR.  The value of source-aware GMM attention appears to be its ability to “ignore” long segments of silence in the input audio, which could potentially be more difficult to do using other attention mechanisms.  Fairly competitive results are presented for offline ASR.  For online ASR, the results are state-of-the-art amongst sequence-to-sequence-based models.