HyperGrid Transformers: Towards A Single Model for Multiple Tasks

1 INTRODUCTION . Learning a single multi-task model that performs well across multiple targeted tasks is an attractive proposition for many reasons ( Kaiser et al. , 2017 ; Ruder , 2017 ; Clark et al. , 2019b ) . Although extremely challenging , this paradigm enables a substantial savings in overall parameter costs , along with eliminating the need for maintaining multiple models in production ( Stickland and Murray , 2019 ) . However , achieving state-of-the-art performance on natural language understanding benchmarks today ( Wang et al. , 2018 ; 2019 ) still relies on fine-tuning a new model for every single task . This methodology is infeasible in many situations . Moreover , certain tasks rely on an extensive ensemble of models and/or task-specific fine-tuning tricks ( Liu et al. , 2019b ; Devlin et al. , 2018 ; Clark et al. , 2020 ) . The single-task fine-tuning paradigm is well-established to be the dominant approach ( Raffel et al. , 2019 ) , as training multiple tasks using a single set of parameters can be problematic in many ways , such as catastrophic forgetting ( French and Chater , 2002 ; McCloskey and Cohen , 1989 ; McClelland et al. , 1995 ; Kirkpatrick et al. , 2017 ) or the inherent difficulty of finding a consistently good model for all tasks ( Clark et al. , 2019b ; Wu et al. , 2020 ) . Inevitable task conflicts and difficulty in fitting all models within a set of hard parameters is also a challenging problem for multi-task co-training . In this paper , we propose a new Transformer architecture , the HyperGrid Transformer for efficient modeling of multiple tasks within a single set of model parameters . HyperGrid Transformers rely on a hypernetwork-based ( Ha et al. , 2016 ) module that performs gridwise decomposable hyper projections . This module is task conditioned and dynamically learns to generate weights of the feedforward layers of the Transformer model . Overall , our eventual goal is to dispense with task specific fine-tuning tricks altogether . While neural networks typically maintain the same consistent set of parameters for all input instances , the proposed HyperGrid Transformers introduces instance-specific parameters by conditioning on the current input . This setup enables our model to learn task-specific reparameterization for each input instance , which mitigates several challenges of multi-task co-training . Our proposed HyperGrid module belongs to a family of hypernetworks ( Ha et al. , 2016 ) , in which a side network is responsible for weight generation for the main network . In our case , task-conditioned hypernetworks provide greater flexibility and expressiveness for capturing the dynamics of multiple tasks within a single set of parameters . Specifically , we introduce two novel algorithmic improvements over the existing methods . First , we introduce the notion of grid-wise projections in which we assume a structural layout in vanilla projection layers . For each input sample , our projections dynamically control the parameters in a grid-wise , region-specific manner . The structural segmentation of feed-forward layers is similar in spirit to mixture-of-experts gating ( Shazeer et al. , 2017 ) , albeit at a lower-level . Conversely , standard hypernetworks only consider row-wise re-weighting of weight matrices . Second , we introduce decomposable hyper-projections . The key idea is to learn rich compositional and pairwise interactions between dual hypernetworks . A dual setup is adopted , where we explore different hypernetwork composition variants . We introduce a novel local-global setup , which composes a local instance-specific and task-specific hyper-projection with a task agnostic global state embedding . This is intuitive since this setup is not only highly expressive and flexible but also serves as a factorization of local and global components . To the best of our knowledge , our work is the first to explore this setup with respect to learning conditional parameters . Finally , we conduct extensive experiments on GLUE/SuperGLUE . Our proposed model is able to match the performance of individually fine-tuned state-of-the-art Text-to-Text Transformers ( T5 ) ( Raffel et al. , 2019 ) models with a single model that is learned to fit all GLUE and SuperGLUE tasks at once . Moreover , our single model also outperforms strong baselines that employ ensembling and other task-specific tricks ( Liu et al. , 2019b ; Clark et al. , 2020 ) . Our Contributions The contributions of this paper can be summarized as follows : • We propose HyperGrid Transformers , a form of hypernetwork-based Transformer that learns task-conditioned dynamic weights for its feed-forward layers . • The key novelty behind HyperGrid Transformers is the factorization of local and global components for weight generation . Our weight generation is grid-wise and imbues the model with a structural layout . • We conduct extensive experiments on natural language understanding benchmarks ( GLUE/SuperGLUE ) . With a single model , we match the state-of-the-art T5 model that is finetuned in a per-task fashion ( multiple models ) , resulting in 16x parameter savings . 2 HYPERGRID TRANSFORMERS . This section outlines the key idea of the proposed algorithm . 2.1 HYPERGRID MODULE . HyperGrid operates on weight matrices ( linear transformations ) , i.e. , Y = WX + b . In a hypernetwork formulation , instead of letting W be free weights , we generate W using a parameterized side network H ( . ) . Y = Wx+ b where W = H ( X ) ( 1 ) where W ∈ Rdm×df . In the case where X is a single vector ∈ Rdm , we may parameterize H ( . ) with a simple feed-forward layer . H ( X ) = σ ( UX ) 1 > W ( 2 ) where 1 is a column vector of ones , σ is the sigmoid activation function and U ∈ Rdm×df . The key idea is that the hypernetwork generates a vector , i.e. , UX ∈ Rdf that is broadcast ( multiplied by 1 ) and multiplied by W , acting as a row-wise scaling of W . We are also able to reduce U ∈ Rdm×n where df mod n = 0 and repeat the vector df n times to form the original dimension of df . These methods only consider scaling one dimension of W ( e.g. , row-wise ) . We now consider methods beyond simple row-wise weight scaling . 2.1.1 DECOMPOSABLE GRIDWISE PROJECTIONS . In our method , we propose grid-wise projections that segments W into a grid , i.e. , blocks of dmdr × df dc . We generate blocks by the outer product of Lr ∈ Rdr and Lc ∈ Rdc . Note that dr and dc are user-specified hyperparameters that control the grid-size for the fan-in and fan-out of the output matrix . For simplicity , we consider divisible blocks where dr < dm , dm mod dr = 0 , and dc < df , df mod dc = 0 . In this case : H ( X ) = ψ ( σ ( ( LrX ) ( LcX ) > ) ) W ( 3 ) where ( LrX ) ( LcX ) > ∈ Rdr×dc , ψ ( . ) is a repeat vector function that repeats its input dmdr times on the row axis and dfdc times on the column axis . We name this approach the L2 variant , short for Local-Local Gridwise Projection . Composition between Local and Global Factors The decomposable grid-wise projections learn Lr and Lc from X , which makes it conditioned on local , instance-wise information . Here , we postulate that it may be beneficial for either Lr or Lc to be a global embedding . By keeping Lc as a global , trainable embedding , this can be formulated as : H ( X ) = ψ ( σ ( ( LrX ) G > c ) ) W ( 4 ) where Gc ∈ Rdf . In this case , Lr is conditioned from X , the specific input sample . On the other hand , Gc remains consistent across all input samples . Hence , the outer product is essentially a rich dyadic composition between local and global factors . Local-Global and Global-Local It is easy to see that there are two ways of composing L and G. The above method considers the Local-Global approach where the fan-in uses a local hypernetwork and the global part uses a trainable embedding . An alternative that flips this around to use a GlobalLocal composition is evaluated in our experiments . Namely , this can be expressed as : H ( X ) = ψ ( σ ( ( Gr ( LcX ) > ) ) W ( 5 ) 2.2 DYNAMIC WEIGHT GENERATION WITH HYPERGRID . This section describes how we use task-conditioned hypernetworks to influence and generate the parameters for HyperGrid Transformers . module . Weight Gating The HyperGrid module is added at the position-wise feed-forward layers of the Transformer models . More specifically , we equip the second positional FFN after the ReLU activations with HyperGrid . There are several reasons for doing so . In most Transformer implementations , the fan out of this layer is typically scaled up to very large values ( Raffel et al. , 2019 ) . Hence , the influence on this layer has the greatest potential to benefit the Transformer model . Second , early experiments on both of the positional feed-forward layers yielded no substantial improvements . Hence , we opt to only modify the second positional FFN of the Transformer model . Initialization In our experiments , we take advantage of existing pretrained model ( Raffel et al. , 2019 ) and add additional HyperGrid parameters that are fine-tuned along with the rest of the network . The overall formulation of the HyperGrid-enhanced Transformer can be written as : Yi = Hi ( Xi−1 , Wi ) +Wi ( Xi−1 ) ( 7 ) where i denotes the layer i . We construct a new HyperGrid ( with non-shared parameters ) for each layer . Since W has been pretrained , we also add a residual connection of the original Wi ( Xi−1 ) computation to the mixture . Parameter Costs We note that the parameter counts added by HyperGrid are relatively negligible since dr and dc are small . In the LG setting , the model adds dmdr+dc parameters at each layer . On the GL setting , the parameter cost added is dr + dfdc . The most expensive option is L2 where the added cost is dmdr + dfdc . These added parameter costs are often negligible for large Transformer models . 3 EXPERIMENTAL RESULTS . We conduct experiments on GLUE ( Wang et al. , 2018 ) and SuperGLUE ( Wang et al. , 2019 ) which are consolidated benchmarks of multiple challenging NLP and NLU tasks . While most of the work in this area has been focused on achieving good task-specific performance , our work focuses on trying to get good performance with a single model on all GLUE and SuperGLUE tasks . Therefore , most of our experiments are conducted on a mixture of all GLUE and SuperGLUE tasks . 3.1 EXPERIMENTAL RESULTS . In this section , we discuss the empirical results of our experiments . Further details about the experimental setup can be found in the appendix . Pertaining to parameter counts ( reported as θ in our experiments ) , this is the total parameter cost for serving the entire suite of tasks . If a single model ( finetuned on a single task ) is X parameters , then the costs of serving N tasks in a multi-model setup would be NX . 3.1.1 SINGLE MODEL VERSUS MULTIPLE MODELS . These experiments investigate different model settings , namely using ( 1 ) multiple models ( MM ) or a single model ( SM ) for all tasks . The MM model trains a fresh model on every single task and retains the best checkpoint for all tasks . Meanwhile , the SM model is trained on all tasks at once . We compare them against HGT , our single model HyperGrid Transformer approach . Details about Baselines For the Single Model Baseline , this is the identical T5 model without our HyperGrid layers which serves as a fair comparison to our model . This enables us to observe the effect of HyperGrid directly . As for the multi-model method , this is done by finetuning T5 directly on each task and reporting results from the best checkpoint ( of each task ) . This is a very strong baseline because the model is also allowed to choose specific best checkpoints for each task which the single model approach is prohibited against . As for the sampling strategy in the singlemodel approach ( baseline and our approaches ) , we use a proportionate mix of tasks according to the number of samples in the each task/dataset . Results Table 2 reports results of our experiments on the GLUE and SuperGLUE benchmark . The first key observation is that the single model ( SM ) approach is outperformed by multi model ( MM ) . This is a well known phenomena and therefore multi-model is generally adopted when the absolute best score is desired on every single task . The interesting result is that we are able to come rather close to the performance of single model with our approach . As a result , the multi-model has 16x more parameters . To fit both GLUE and SuperGLUE , this would require 16x the parameters . Given that our goal is to bridge the performance of a single model versus multiple models for multiple tasks , we find that this result is considerably successful . Moreover , we observe that our singlemodel approach outperforms the multi-model baseline by +0.6 % on average across 8 tasks . We observe similar trends as on the GLUE benchmark . Naturally , the best model is the multi-model model which involves finetuning a specialized model for each task . The gap between multi-model and single-model is at 74.8 versus 73.6 . Our approach bridges this gap , improving the single-model score to 74.5 , competitive with the multi-model approach .

Summary: This paper presents a HyperGrid Transformer approach to fine-tuning, where one takes a pre-trained transformer model and then modifies it by introducing hypernetworks that modify the 2nd FFN in each transformer block by generating additional weights conditioned on input. These hyper-networks are trained on all tasks in GLUE/SuperGLUE datasets simultaneously and are task aware through prefixing of a task specific token to input. This allows one to fine-tune only a small number of parameters and end up with a model that performs quite well on all tasks at the same time, not much worse than fine-tuning the entire transformer model on all of these tasks.