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: The paper proposes a transformer architecture for multi-task learning, based on sharing some attention heads in multi-head attention between different tasks. The model is supplied with more heads than necessary and a neural module using gumbel-softmax is used to select a subset of these heads for each task.  A variant of this is described, where heads are distributed between groups and exactly one head is selected from each group.