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 work, the authors present a conditional language-specific routing (CLSR) scheme for transformer-based multilingual NMT systems. They introduce a CLSR layer after every transformer encoder and decoder layer; each such layer is made up of hard gating functions conditioned on token representations that will either select a language-specific projection layer or a shared projection layer. Further, a budget is imposed on the language-specific capacity measured by aggregating the number of gates that allow for language-specific computations; this budget constraint forces the network to identify the sub-layers that will benefit most from being language-specific.