Pay Better Attention to Attention: Head Selection in Multilingual and Multi-Domain Sequence Modeling

1 Introduction . Recent progress on deep learning models , in particular multi-head attention , has brought significant gains to sequence modeling tasks including speech recognition ( Moritz et al. , 2020 ) , text-to-text translation ( Vaswani et al. , 2017 ) , and speech-to-text translation ( Vila et al. , 2018 ; Gangi et al. , 2019 ) . Attention mechanism allows a model to focus on informative parts of the inputs , and multi-head attention computes attention over inputs by multiple heads independently . With each head attending to different information , multi-head attention potentially captures more complicated data patterns and extracts sophisticated knowledge . Sequence modeling has attracted a lot of research interest in multilingual and multi-domain settings , where a model is trained on data in multiple language directions and data from different domains respectively . Key advantages of these settings are better data efficiency and the support of knowledge transfer among languages or domains . This is critical for resource-limited scenarios . For example , multilingual translation enhances the performance of low-resource languages via knowledge transfer from high-resource languages ( Gu et al. , 2018 ; Inaguma et al. , 2019b ) . Given the data scarcity in individual domains , a common practice is to combine the data from various domains to augment the training set ( Wang et al. , 2020d ) . Another appealing aspect of multilingual or multi-domain models is their low deployment and maintenance costs compared with numerous models trained for individual language pairs or domains . Despite the positive knowledge transfer , negative interference has also been observed in multilingual ( or multi-domain ) training especially when languages ( or domains ) are dissimilar . Recent studies reveal from the optimization perspective that conflicting gradients in shared parameters is one cause of interference between languages ( or domains ) ( Yu et al. , 2020 ) . A promising direction for interference 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . mitigation is to design better strategies of parameter sharing . In some previous works , sharing is based on the similarity between languages ( or domains ) , which require expert knowledge or pre-computed relatedness ( Wu et al. , 2019 ) . Recent studies also propose branches and components specific to languages ( or domains ) in addition to shared modules ( Bapna and Firat , 2019 ; Guo et al. , 2020 ) . In this work , we bring the mitigation of language and domain interference under a common umbrella , and tackle it by improving parameter sharing within multi-head attention . We propose strategies to select attention heads for different languages or domains . Instead of sharing everything across languages or domains , our model automatically learns to share heads among a subset of languages or domains . It encourages positive transfer within the subset and preserves their specificity without interference from outside the subset . The major contributions of this work are summarized below : 1 . We propose attention head selection to mitigate language or domain interference ; 2 . The parameter sharing strategies are lightweight and preserve inference efficiency ; 3 . We extensively evaluate attention sharing strategies on various sequence modeling tasks including speech recognition , text-to-text and speech-to-text translation . Consistent gains are achieved across multiple benchmark datasets . The paper is structured as follows . Section 2 discusses related works on sequence modeling in multilingual and multi-domain setting . In Section 3 , we introduce the proposed strategies of head selection in multi-head attention . Section 4 describes the empirical evaluation , followed by a discussion in Section 5 . We conclude this paper in Section 6 . 2 Related Work . Multilingual learning . Multilingual modeling has the potential to improve low-resource language performance through knowledge transfer from high-resource languages , and it draws great interest from researchers in speech recognition and translation ( Pratap et al. , 2020 ; Heigold et al. , 2013 ; Johnson et al. , 2017 ; Dabre et al. , 2020 ; Liu et al. , 2020 ; Inaguma et al. , 2019a ; Li et al. , 2020 ) . Although impressive progress has been made for low-resource or zero-shot tasks , it is also found the multilingual model has inferior performance on high-resource tasks due to multilingual interference . In order to address this issue , some works focus on multilingual models with task-specific parameters . Different parameter sharing strategies are examined on Transformer ( Sachan and Neubig , 2018 ) . Attention dependent on target languages is proposed to enhance multilingual translation ( Blackwood et al. , 2018 ) . Treating multilingual modeling as an adaptation problem , Bapna and Firat ( 2019 ) first build a universal multilingual model for all languages and then finetune newly added adapters for each language pair . Another thread of work is to increase the model capacity to compensate the performance loss in high-resource languages ( Pratap et al. , 2020 ) . Shazeer et al . ( 2017 ) propose mixture-of-experts and select RNN cells based on input tokens . Lepikhin et al . ( 2020 ) integrate a mixture of FFN experts in the GShard model , and later Fedus et al . ( 2021 ) propose Switch Transformer to route tokens to different FFN sub-layers . Different from previous works , we propose strategies of attention sharing among languages in the level of attention heads for multilingual modeling . Multi-domain learning . Similar to multilingual learning , multi-domain learning ( MDL ) can effectively utilize data from different domains but also suffers from interference due to inter-domain heterogeneity ( Saunders , 2021 ; Pham et al. , 2021 ) . Previous works address this issue from two perspectives : optimization and model architecture . In the optimization aspect , attempts have been made to synchronize the learning speed of different tasks ( Chen et al. , 2018 ) , adjust the gradients of individual tasks to alleviate gradient conflicts ( Yu et al. , 2020 ) and apply regularization to achieve better generalization in different domains ( Dakwale and Monz , 2017 ; Khayrallah et al. , 2018 ; Thompson et al. , 2019 ) . In terms of model architecture , domain-specific labels ( Kobus et al. , 2017 ) , word embedding ( Zeng et al. , 2018a ) , sub-networks ( Wang et al. , 2020d ) are adopted to address the issue of domain divergence . The architecture can be specified during the general training with the mixed data from multiple domains ( Wang et al. , 2020d ) or during the finetuning in individual domains ( Bapna and Firat , 2019 ) . In this work , we deal with domain interference by leveraging domain-specific attention heads in multi-head attention . Attention selection . Selective self-attention networks propose to apply masking to the inputs and pay more attention to content words ( Geng et al. , 2020 ) . Liu et al . ( 2021 ) select text-related image regions with attention in multi-modality translation . Compared to these methods , we conduct automatic attention head selection for different tasks and focus on mitigating task interference . 3 Model . In this section , we start with preliminaries of multi-head attention , and introduce our approach to attention interference mitigation . We put multilingual and multi-domain sequence modeling under the same umbrella in this study . For the simplicity of the following discussions , we refer to the two settings as multi-task modeling , where a task is one language or one domain . Different from the standard multi-head attention , our model provides more attention heads than those used in computation . Different subsets of heads are assigned to each task so that partial attention sharing enables knowledge transfer and meanwhile mitigates interference . We introduce latent variables to modulate head selection , and propose strategies to learn the head assignment to different tasks . 3.1 Preliminary . Multi-head attention . As a core component of Transformer , multi-head attention parameterizes each head with key , query and value transformation matrices ( Vaswani et al. , 2017 ) . The token representation is transformed into key , query and value vectors via these transformations . Each head assigns the attention of this token over the input sequence based on the matching between its query vector and key vectors of other tokens . The value vectors are weighted by the attention as the contextualized token representation . It is passed through linear projection as the output of the attention head . Suppose that head h has output x ( h ) . Multi-head attention with H heads yields an output x for the given token , which is the concatenation of all head outputs . x = x ( 1 ) ⊕ · · · ⊕ x ( h ) ⊕ · · · ⊕ x ( H ) , ( 1 ) where ⊕ is the vector concatenation . Interference . Maximal parameter sharing aims to learn universal knowledge across languages ( Wang et al. , 2020e ) and domains ( Zeng et al. , 2018b ) . To capture the task specificity , different languages or domains compete for model capacity , which is observed as the interference in previous studies . The interference results in degraded performance in jointly trained models . However , few works look into the improvement of parameter sharing within multi-head attention . This study explores head selection strategies to mitigate the interference in multilingual and multi-domain models . 3.2 Latent Variable for Head Selection . First , we outline our approach to learn a more general-purpose multi-head attention in Transformer from the Bayesian neural network perspective . Suppose that the input sequence is x and the output sequence is y . For conditional sequence modeling tasks such as machine translation , the posterior of p ( y | x ) can be computed by marginalizing over the posterior of latent variable z , which modulates parameters Θ in the standard Transformer architecture : p ( y | x , Θ ) = Ep ( z|Θ ) [ p ( y | x , z ) ] = ∫ p ( y | x , z ) p ( z|Θ ) dz ( 2 ) Parameterization of zt . In this work , we define zt as modulating the selection of attention heads by task t. We have zt = { z ( h ) t } h where z ( h ) t is a discrete latent variable from Bernoulli distribution indicating whether task t selects attention head h. This modeling choice allows us to prune attention heads , which preserves computation efficiency as well as regularizes training . Marginalizing over zt is intractable given numerous heads in neural models . Therefore , we use variational inference to derive an approximate solution . Suppose that ( xt , yt ) is from task t. Specifically , we learn an inference network qφ ( zt ) , which is paramterized with φ , to approximate the true distribution p ( zt ) and optimize the evidence lower bound ( ELBO ) of p ( y|x ) : log p ( y | x ) ≥ ∑ t ( Eqφ ( zt ) [ log pθ ( yt | xt , zt ) ] − KL ( qφ ( zt ) ‖ p ( zt ) ) ) , ( 3 ) where KL is the KL-divergence between two distributions . In our work , we assume identical probability of each head being selected . Therefore , we have p ( zt = 1 ) = HH′ , where H and H ′ are numbers of selected attention heads and all head candidates . Training and interference . We use the Gumbel-Softmax reparameterization ( Jang et al. , 2017 ) to draw samples of z ( h ) t from the posterior qφ ( z ( h ) t ) . It makes the model end-to-end differentiable , while learning discrete policies of head selection without resorting to policy gradients . We adopt a lightweight estimator of qφ ( z ( h ) t ) by directly learning the logit parameters { φ ( h ) t } : qφ ( z ( h ) t ) = exp ( ( φ ( h ) t ( 1 ) + ( 1 ) ) /τ ) ∑ j∈ { 0,1 } exp ( ( φ ( h ) t ( j ) + ( j ) ) /τ ) , ∼ G ( 0 , 1 ) ( 4 ) where G ( 0 , 1 ) is the Gumbel distribution , and τ is a temperature hyperparameter which increases the discreteness of samples when τ → 0 . We will discuss different head selection strategies in Section 3.3 , which make binary selection decisions based on real-valued posterior qφ ( z ( h ) t ) . 3.3 Attention Selection Strategies Suppose that the output dimension of multi-head attention is d , and the dimension of each attention head is dH . We provide a large pool of H ′ ( H ′ > H ) attention head candidates in every Transformer layer , and H ′ is a hyperparameter controlling the search space size of attention selection strategies . The model requires attention outputs to have a consistent dimension d , so each task needs to select exact H heads among H ′ candidates . We introduce two strategies for the attention head selection : subset strategy and group strategy . Subset strategy . The subset strategy is straightforward , and we compare the posterior { qφ ( z ( h ) t ) : h ∈ [ 1 , H ′ ] } of all H ′ heads given a task t. A subset of H heads with the highest posterior are selected by the task , and there are CH ′ H subset choices . The subset strategy is described in Fig . 1 ( a ) . The binary mask s ( h ) t indicates whether an attention head h is assigned to task t. s ( h ) t = { 1 , h ∈ TopH ( { qφ ( z ( h ) t ) } ) , 0 , otherwise , ( 5 ) where TopH ( · ) returns the top H heads with the highest values . The outputs of the selected heads are concatenated as the attention output . Note that the subset strategy does not consider the order of the attention heads . For example , when head 2 and 3 are selected , head 2 contributes to the beginning part of attention output . With head 1 and 2 selected , the output of head 2 goes to the last part of the attention output . Group strategy . We further propose group strategy to preserve the order of attention heads during head selection . Different from the subset strategy , the group strategy first divides H ′ heads into H groups . As is shown in Fig . 1 ( b ) , each group contains r = H ′ H candidates . Each task could choose one attention head from each group , and has access to H heads per layer . There are rH possible combinations of heads . The group strategy keeps the head order in that heads from group g only contribute to g ’ s corresponding dimensions in the attention output . The head with the highest posterior in its group would be selected by a given task t. We use the binary mask s ( h ) t to indicate the selection of head h in group g. s ( h ) t = { 1 , h = argmax ( { qφ ( z ( h ′ ) t ) : h ′ ∈ g } ) , 0 , otherwise . ( 6 ) The output of group g is : x ( g ) = ∑ h∈g s ( h ) t · x ( h ) . ( 7 ) The outputs of H groups are concatenated as the output of the attention module for task t. With either subset or group strategy , the sequence model is trained to assign attention heads to different tasks to maximize the lower bound in inequality ( 3 ) . The number of additional parameters { φ ( h ) t } introduced by our attention selection is only O ( T × H ′ × L ) , where T is the number of tasks , H ′ is the number of head candidates per layer , and L is the number of layers . It is small compared with the total parameter size of the model , and head selection is thus lightweight and memory efficient . Moreover , the head selection is inherently a pruning process . Regardless of the size of head candidates , only a fixed number of attention heads are involved in computation for a given task . Hence our approach is also computationally efficient in model inference .

Summary: Multi-head attention is an essential component for popular Transformer models. This paper proposes to learn shared and specialized attention heads for different languages and domains. The authors formulate attention selection as latent variables and adopt Gumbel softmax to select attention heads.  Experiments on text-to-text and speech-to-text translation verify the effectiveness of the proposed approach.