Share or Not? Learning to Schedule Language-Specific Capacity for Multilingual Translation

1 INTRODUCTION . Model architecture design injects inductive biases to neural network layouts , allowing a learning algorithm to favor certain representations over others , independent of the observed data ( Mitchell , 1980 ) . In multilingual neural machine translation ( MNMT ) , where the learning objective is commonly cast as a multi-task learning problem ( Firat et al. , 2016a ; Ha et al. , 2016 ; Johnson et al. , 2017 ) , the inductive bias researchers usually study is deciding on which components of the neural network to share between tasks ( languages ) , and which components to leave specific to the task or language . These components can be entire layer stacks , individual layers or even some sub-layers ( Sachan & Neubig , 2018 ; Blackwood et al. , 2018 ; Wang et al. , 2019 ; Zhu et al. , 2020 ) . Noticeably , the search space of which parameters to share and at which granularity grows rapidly , as we make neural networks large or increase the number of tasks ( languages ) . This rapid expansion of the search space prevents us from exhaustively exploring the choice of sharing patterns in MNMT . The incapability of full-space exploration motivates methods relying on heuristics ( Sachan & Neubig , 2018 ) , that lack flexibility when more languages are covered , or meta-learning ( Platanios et al. , 2018 ) , that are often hard to scale . These limitations hinder their generalization to large-scale multilingual models , which is the very focus of our study . In large scale multilingual models , also known as massively multilingual models ( Aharoni et al. , 2019 ; Arivazhagan et al. , 2019 ; Zhang et al. , 2020b ) , hundreds of languages with varying amounts of training data , difficulty and linguistic properties are jointly trained together in a multi-task setup . While the joint training enables positive ∗Work done while Biao Zhang was interning at Google Research . transfer across languages , it also introduces task-interference between dissimilar languages ( Arivazhagan et al. , 2019 ; Wang et al. , 2020a ; b ) and a capacity bottleneck emerges due to the increased number of languages and data ( Huang et al. , 2019 ; Zhang et al. , 2020b ) . In this paper we adopt an end-to-end data driven approach ( conditional language-specific routing , or CLSR ) which permits directly probing a large section of the search space . We let the network learn the sharing structure from the data itself , by learning to route between language-specific ( LS ) or shared pathways . These two routes determine the mode of operation for the network : when the LS branch is selected , the model is given access to a set of LS layers ( implemented as simple projections per language ) and when the shared branch is chosen , the computation is routed to a layer that is used by all languages . By guiding the ( gating ) decision process with token level activation information , the network flexibly learns to alternate between the two modes and naturally lends itself to a conditional computation approach for multilingual processing ( Bengio et al. , 2013 ; Davis & Arel , 2013 ; Bapna et al. , 2020 ) . The gate states are optimized towards maximizing translation quality , but regularized with a budget constraint to control the amount of LS capacity1 . Reducing the available budget results in fewer gates routing through the LS paths , enforcing CLSR to identify the most crucial sub-layers which allows us to observe and study the importance of each sub-layer for multilingual processing . Our approach is visually depicted in Figure 1 . We verify our proposal on WMT and the massively multilingual OPUS-100 dataset , with models building on the Transformer architecture ( Vaswani et al. , 2017 ) . We explore target-specific and source-specific modeling for one-to-many2 and many-to-one translation , respectively . To measure the degree of each sub-layer ’ s tendency to be language-specific , we propose LSScore metric . Our results show that CLSR successfully navigates the trade-offs in LS modeling , outperforming several strong baselines . Our main findings are summarized below : • Both the amount and the position of LS layers matter for MNMT . The best performance is achieved by distributing 10 % -30 % LS computation to the top and/or bottom encoder/decoder layers . • Feed-forward sub-layers utilize more LS capacity compared to other sub-layers on one-tomany translation . • One-to-many translation benefits more from CLSR ( with target LS parameters ) compared to many-to-one translation ( with source LS parameters ) , particularly when the training data is imbalanced . • The induced sharing pattern learned by CLSR is highly similar across languages . 1We use the term “ the amount of LS capacity ” to refer to the proportion of open gates where CLSR selects to route information through the LS path instead of its shared counterpart , which is directly regularized and guided by the budget constraint p as in Eq . 6 . 2In a one-to-many machine translation setup , a single source side language ( commonly English ) is tasked to be translated into multiple target languages , one at a time . • We can use the learned patterns to hard-code better parameter sharing strategies for multilingual Transformers . 2 RELATED WORK . Our work closely relates to language-specific ( LS ) modeling for multilingual NMT and conditional computation for sequential data which we will recap both here . Early research on MNMT focused on improving shared capacity for separate bilingual models to enhance cross-lingual transfer . These efforts included sharing encoders for one-to-many translation ( Dong et al. , 2015 ) , sharing decoders for many-to-one translation ( Zoph & Knight , 2016 ; Lee et al. , 2017 ) and sharing sub-layers ( attention ) for many-to-many translation ( Firat et al. , 2016a ; b ) . These studies corroborated the feasibility of accommodating multiple languages with shared NMT sub-components , motivating researchers to explore universal MNMT . Ha et al . ( 2016 ) and Johnson et al . ( 2017 ) proposed such an implementation that performs multilingual translation with a single monolithic NMT model where the entire network is shared across languages , thanks to a target language token informing the model which language to translate into . Although this paradigm shows great scalability ( Aharoni et al. , 2019 ) , the language token alone affords little flexibility in handling language diversity with a rigid share-all layout . Follow-up studies thus resort to LS modeling in an attempt to seek a better trade-off between sharing and not sharing . Methods in this category involve specializing neural attentions ( Blackwood et al. , 2018 ; Sachan & Neubig , 2018 ; Wang et al. , 2019 ) , broadening encoder outputs and normalizations ( Zhang et al. , 2020b ) , decoupling multilingual encoders and/or decoders ( Vázquez et al. , 2019 ; Escolano et al. , 2020 ) , using a fixed mix of LS and shared parameters ( Wang et al. , 2018 ) , inserting lightweight adapters ( Bapna & Firat , 2019 ) and separately modeling languages for different clusters ( Tan et al. , 2019 ) , to name a few . Nevertheless , these methods heavily depend on heuristics , providing little evidence about how to optimally distribute LS capacity across the model . By contrast , our proposed CLSR forces the model to learn LS behaviour . It can be treated as a simplified differentiable neural architecture search ( NAS ) model ( Liu et al. , 2019 ) with a search space defined by the presence/absence of LS projections after every transformer sub-layer . However , in contrast with NAS , we utilize conditional computation ( Bengio et al. , 2013 ) to make the choice of executing the shared or LS path conditional on the input representations . This allows us to compare and contrast the choice of LS vs shared paths on different inputs and languages . Conditional computation has previously been successfully applied to adapt the amount of computation to the input in recurrent models ( Graves , 2016 ) and transformers ( Dehghani et al. , 2019 ; Elbayad et al. , 2020 ) , or to significantly scale up model capacity by utilizing sparsely-gated Mixture-of-Experts layers ( Shazeer et al. , 2017 ; Lepikhin et al. , 2020 ) . Zhang et al . ( 2020a ) applied conditional computation to sparsify encoder outputs in sequence-to-sequence models in order to reduce attention costs , while Bapna et al . ( 2020 ) introduced conditional execution of Transformer sub-layers to control the amount of computation expended by the model at inference . Sukhbaatar et al . ( 2019 ) learn parameters that limit the attention spans , in order to make the attention operation more efficient , while Fan et al . ( 2020 ) utilize structured dropout to prune transformer layers at inference . Ruder et al . ( 2019 ) propose the sluice network that learns the inter-task ( shared ) sub-spaces on top of task-specific models . By contrast , CLSR starts with a totally shared model , and learns how to inject task-specific projections into it , which scales more easily to massively multilingual settings . Different from previous studies , we explore conditional computation as an analysis tool to understand the ideal arrangement of LS capacity for MNMT . Utilizing conditional computation to search for better LS sharing patterns in multilingual translation , to the best of our knowledge , has never been investigated before . 3 BACKGROUND : MNMT . Given a source sentence X ′ = { x1 , x2 , . . . , xI } and its target translation Y = { y1 , y2 , . . . , yJ } , we follow Johnson et al . ( 2017 ) to reuse standard bilingual NMT models for multilingual translation by altering the source input with a language token lang , i.e . changing X ′ to X = { lang , x1 , x2 , . . . , xI } . Note that lang denotes the target language in one-to-many translation but source language in many-to-one translation . We model translation from X to Y with Transformer ( Vaswani et al. , 2017 ) . Transformer relies on the following residual-normalization structure ( He et al. , 2015 ; Ba et al. , 2016 ) to smooth informa- tion flow and avoid gradient vanishing and explosion : zl+1 = LN ( zl + f ( zl ) ) , ( 1 ) where l denotes layer depth and LN ( · ) is layer normalization ( Ba et al. , 2016 ) . Function f ( · ) represents the basic building block in Transformer , such as attention network or feed-forward network . The encoder in Transformer is a stack of L identical layers , with each layer involving a self-attention sub-layer ( SAN ) and a feed-forward sub-layer ( FFN ) . The decoder uses a similar structure except for an extra cross-attention sub-layer ( CAN ) inserted in-between the above two sub-layers . 4 CONDITIONAL LANGUAGE-SPECIFIC ROUTING ( CLSR ) . The success of MNMT comes at the cost of expressivity and model ’ s ability to capture languagespecific characteristics . It has been empirically shown that the language signals from language indicator tokens alone are not sufficient ( Arivazhagan et al. , 2019 ) , making architectures dedicated to LS modeling a necessity ( Blackwood et al. , 2018 ; Sachan & Neubig , 2018 ; Zhang et al. , 2020b ) . Nevertheless , the question when and where LS modeling matters most in MNMT still remains to be answered . To this end , we propose conditional language-specific routing ( CLSR ) which specializes f ( · ) and changes the formulation in Equation 1 as follows : zl+1 = LN ( zl + CLSR ( f ( zl ) ) ) . ( 2 ) CLSR learns a hard binary ( scalar-valued ) gate g ( · ) for each input token based on its hidden representation zl ∈ Rd . These gates endow each sub-layer in Transformer with the capability of routing information selectively through either LS path hlang or shared path hshared : CLSR ( f ( zl ) ) = g ( zl ) hlang + ( 1− g ( zl ) ) hshared , ( 3 ) with hlang = f ( zl ) Wlang , hshared = f ( zl ) Wshared , ( 4 ) where Wshared is a weight matrix shared across languages , while parameter Wlang is only used for modeling language lang which endows NMT with source or target LS modeling capacity.3 Intuitively , a closed gate , corresponding to shared capacity , encourages maximal cross-lingual information transfer ; an open gate , corresponding to LS capacity instead , improves language awareness for translation albeit it blocks knowledge transfer . CLSR balances between the two modes as controlled by the gates . Following Bapna et al . ( 2020 ) , we parameterize the gate g ( · ) with a two-layer feed-forward network G ( · ) , and inject zero-mean Gaussian noise during training to discretize it ( Chiu & Raffel , 2018 ) : g ( zl ) = σ ( G ( zl ) + α ( t ) N ( 0 , 1 ) ) , G ( zl ) = Relu ( zlW1 + b ) w2 , ( 5 ) where σ ( · ) is the logistic-sigmoid function , dg is the gating feed-forward hidden dimension , and W1 ∈ Rd×dg , w2 ∈ Rdg are trainable parameters . α is linearly increased along with training steps t. In this way , the gating parameters can be optimized with accurate gradients when α is small at the beginning so as to measure the degree to which each position in each sub-layer benefits from LS modeling . As training progresses , α grows larger , forcing the gating network to emit hard outputs . At inference time , we discretize the gate based on a simple decision rule : g ( zl ) = δ ( G ( zl ) ≥ 0 ) , where δ ( · ) is a Dirac measure . We train the gates based on the standard maximum likelihood objective , along with an additional budget regularization term that enables control over the amount of LS capacity used for translation . Let the set of all CLSR layers in the encoder beMenc and the decoder beMdec . Then the amount of LS computation utilized by a sentence pair ( X , Y ) is given byG ( X , Y ) = ∑ x∈X ∑ m∈Menc gm ( x ) +∑ y∈Y ∑ m∈Mdec gm ( y ) . Given a budget constraint p ∈ [ 0 , 1 ] and a batch of sentence pairs , B , the training loss function of CLSR is formulated below : L ( B ) = ∑ ( X , Y ) ∈B MLE ( X , Y ) + ∣∣∣∣∣ ∑ ( X , Y ) ∈BG ( X , Y ) ∑ ( X , Y ) ∈B ( |X||Menc|+ |Y ||Mdec| ) − p ∣∣∣∣∣ , ( 6 ) 3We maintain a set of language-specific weight matrices in order to compute hlang for each language . To make the number of parameters manageable , we share the set of LS matrices Wlang across all encoder or decoder sub-layers , but distinguish Wlangenc and W lang dec , LS matrices for the encoder and decoder , respectively . Intuitively , the budget constraint tries to regulate the amount of LS computation available to all tokens in the batch as a fraction p of the total LS computation in the model . Given that we make a binary decision for every input for every CLSR layer , this corresponds to a search space of O ( 2|X||Menc|+|Y ||Mdec| ) . This gating space not only grows with respect to the model depth and sub-layer types , but also is highly input dependent . This dependency makes it difficult to search the entire space using heuristic methods ( Blackwood et al. , 2018 ) . The constraint in Equation 6 is imposed upon the aggregated gates . As a consequence , the model can learn to trade-off LS capacity for certain layers and inputs for others . Decreasing the budget encourages gate closure , such that the LS path is chosen only in the critical sub-layers . Thus , a properly learned gating function reveals the activations of LS paths and helps gain insights into the model behavior .

Summary: In this paper, the authors present a study of different aspects of language-specific model capacity for massively multilingual machine translation. To this end, language-specific behaviour is achieved via a combination of conditional computation to decide whether to use language-specific parameters or not and statically assigning experts for each languages. The language specific sub-layers are incorporated throughout the network. The training objective allow budgetary constraints on the amount of language-specific parameters. The paper does a systematic analysis on the role of language specific parameters using the proposed architecture. Based on the analysis, recommendations on design of multilingual NMT architectures are proposed and their efficacy validated experimentally. The study sheds light on the amount of language specific parameter sharing, their distribution in the network, impact of language, etc.