Mint: Matrix-Interleaving for Multi-Task Learning

1 INTRODUCTION . While deep learning has enabled remarkable levels of generalization through the use of function approximators , this comes at the cost of large amounts of data , which remains a critical challenge in deploying deep learning to a number of domains . When combined with deep networks , multitask learning offers the promise of building more powerful representations using less data per task , leading to greater performance and data efficiency . However , multi-task deep learning has also posed considerable challenges . Numerous works have observed that joint training on multiple tasks can actually decrease task performance due to the negative influence of other tasks ( Parisotto et al. , 2015 ; Rusu et al. , 2016a ) . Indeed , training networks entirely independently on each task has remained a strong approach , to the point that multiple multi-task methods have first trained models independently before using them to train a multi-tasking model ( Parisotto et al. , 2015 ; Rusu et al. , 2016a ; Ghosh et al. , 2017 ; Teh et al. , 2017 ; Czarnecki et al. , 2019 ) . Moreover , our experiments in Section 6 indicate that three recently proposed methods for multi-task learning are all surpassed by training models independently per task . However , training independent models will only work well when provided enough data per task , and precludes potential positive data-efficiency gains from multi-task learning , only providing protection against negative transfer . Further , while a number of works have successfully shared parameters , finding an architecture with the appropriate level of parameter sharing for a given problem domain can require a considerable amount of manual engineering . In this work , we aim to develop a multi-task learning method that can perform well both when tasks share very little and when they share a large amount of structure . To address this problem , we consider how a single neural network model can represent two extremes : independent models , when optimization challenges prevail , or a single model with shared weights , when sharing is beneficial . Further , we would like such a model to be able to represent intermediate levels of model sharing , when appliable . One option for performing independent training within a single model is to put separate networks with independent weights into a single model , using the task ID to select which network prediction to output . However , this prevents any sharing . An alternative approach is to condition the model on the task ID , through various conditioning approaches , including additive and multiplicative approaches such as FiLM ( Perez et al. , 2018 ) . In fact , point-wise multiplicative conditioning , as proposed in FiLM , can indeed represent separate networks by selecting which parts of the network to be used for different tasks , as can a number of other approaches in multi-task learning ( Rosenbaum et al. , 2017 ; 2019 ; Fernando et al. , 2017 ) . Yet , these approaches still require an optimization over shared parameters in order to select which parameters are used for each task . These shared parameters can introduce significant optimization challenges . We instead consider how to allow a model to perform optimization on only shared parameters , only disjoint parameters , or any combination thereof . We can achieve this by simply interleaving learned per-task matrices at each layer of a jointly-trained neural network . When optimization over shared parameters is ineffective , the model can still represent a full neural network per task using only the per-task matrices , resulting in independent training ; while using identical per-task matrices results in standard joint training . Intermediately , a mix of shared and per-task parameters may be used . In effect , by incorporating these matrices into the network , the optimizer itself can automatically and dynamically modulate the degree to which a representation is shared between tasks , depending on the problem domain and the optimization progress , and can do so without having to optimize shared parameters . The primary contribution of this paper is a simple yet effective approach for multi-task learning that can represent and smoothly interpolate between independent training and joint training , via matrix interleaving ( Mint ) . We describe how we can implement Mint in deep multi-task models and show its effectiveness in improving data efficiency and generalization in multi-task settings while providing intuition about the reasons why this architecture performs so well . Further , we show that the model can be extended to goal-conditioned reinforcement learning in a straightforward manner by allowing the model to generate the interleaved matrices conditioned on task information such as the goal . We evaluate Mint on sets of tasks with both high and low levels of shared structure and find that it performs well in both settings , performing comparably to or outperforming both joint training and independent training , effectively combining the best elements of both . Further , in comparison to previous methods that use multiplicative interactions for continual learning ( Cheung et al. , 2019 ) and for general conditioning ( Perez et al. , 2018 ) , Mint is better able to separate tasks by avoiding the need to optimize over shared parameters and can empirically produce substantially better performance on a range of challenging multi-task problems . Finally , Mint also outperforms state-of-the-art approaches for multi-task learning while being significantly simpler to implement . 2 PRELIMINARIES . In multi-task learning , the goal is to find a θ-parameterized model fθ that reaches high performance across all training tasks drawn from a task distribution Tk ∼ p ( T ) , i.e . min θ ETk∼p ( T ) [ Lk ( fθ ) ] , where Lk denotes the loss function for task Tk . In this work , we study the setting in which we have a fixed set of K tasks { Ti } Ki=1 , and we wish to obtain high performance on all of the tasks in this set . In Section 4 , we will study this multi-task problem in the context of a reinforcement learning setting . In our multi-task learning setup , we train a model that is conditioned on a task indicator zk which is used to specify a task Tk that should be performed . The task indicator can be represented in a variety of ways , from simple categorical variables to learned task embeddings ( Hausman et al. , 2018 ) . This formulation can be readily extended to a goal-conditioned reinforcement learning setting , where zk indicates the desired goal . The multi-task learning model learns to optimize objectives of all the tasks Tk that the model is trained on . Joint and fully independent training are the two extremes of multi-task learning . Assuming a set of n training tasks , we characterize multi-task training as independent if it optimizes a set of task-specific , disjoint parameters { θk } ni=1 that parameterize the model fθk . We define joint training as finding a single set of task-independent , shared parameters θ of fθ . Note that joint training utilizes the fact that the parameters are shared throughout learning . While joint training has a number of data-efficiency and generalization benefits , it can be difficult to train effectively . 3 MULTI-TASKING WITH INTERLEAVED MATRICES . Considering fully independent and fully joint training as two ends of the spectrum in multi-task learning problems , we want to design an algorithm to get the best of both worlds – the stability of independent training and the parameter-sharing efficiency of jointly trained models . To this end , we propose to build models that allow the neural network to learn how much information should be shared and between which tasks throughout learning , in an adaptive fashion . We describe the details of our approach below . 3.1 MINT : MATRIX-INTERLEAVED NETWORKS . In the case of neural network models , we view representations as network activations at different layers of the network . We aim to introduce a modification to the neural network that would allow the model to either form those activations in a completely task-specific way , in a completely shared way , or in a way that shares to an intermediate degree . To achieve this , we propose a model architecture that transforms the previous layer ’ s representation both in a task-general way and in a task-specific way , in sequence . When two tasks share very little , the network can optimize task-specific weights , while when the tasks share a considerable degree of structure , the network can leverage the shared weights . Since these transformations are task-specific and can be introduced at various layers of the network , they allow for a different amounts of representation shared at different levels of the neural network model . To understand the practical implementation , we consider the activations at layer l of a fully-connected neural network : y ( l ) = σ ( W ( l ) y ( l−1 ) + b ( l ) ) , where W ( l ) is the weight matrix for layer l , b ( l ) is the bias vector for layer l , and σ is the non-linearity at each layer . The Mint layer augments the traditional fully-connected layer with task-specific weight matrix Mk and bias vector βk , where k indexes the task . The forward pass of a Mint layer for some vector of activations y is presented in Definition 1 . Definition 1 . A Mint layer applies an affine transformation to activations y ∈ Rn as follows , yielding new activations a : a = Mint ( y ) = Mky + βk ( 1 ) where Mk and βk are per-layer task-specific matrix and bias , respectively . A neural network augmented with Mint thus contains parameters that are both shared across all tasks as well as parameters that are only used for a particular task for each layer l , i.e . θ = { W ( l ) , b ( l ) , M ( l ) k , β ( l ) k } . See Figure 1 for a visual depiction of the application of Mint . We show how the regular fully-connected layers and Mint layers can be interleaved in Equation 2 and 3 below : y ( l ) = σ ( W ( l ) a ( l−1 ) + b ( l ) ) ( 2 ) a ( l ) = M ( l ) k y ( l ) + βk ( l ) . ( 3 ) Because this layer is fully-differentiable , we can learn the task-specific matrices M ( l ) k and biases β ( 1 ) k jointly with the shared parameters of the model . When we apply Mint to tasks with very large numbers of tasks , or arbitrary task descriptors ( e.g. , goal-conditioning ) , we can train separate neural networks T lM , T l β → M lk , βlk , to output the Mint matrices and biases at every layer , instead of storing independent matrices for each task . In this case , Mint resembles FiLM Perez et al . ( 2018 ) , another method which performs transformation-based task conditioning . In contrast to FiLM , Mint uses a matrix transforming the activations at each layer instead of a point-wise multiplication by a vector . In the next section , we study a theoretical property of Mint that motivates the chosen definition of the Mint layer . We validate the benefits of Mint in our experimental evaluation in Section 6 . 3.2 THE EXPRESSIVE POWER OF MINT . We next aim to theoretically study the expressive power of multi-task learning architectures when the shared parameters can not be optimized effectively ( e.g . due to optimization challenges arising from learning tasks that share little or no structure ) . While a viable approach in this regime is to simply use separate task-specific networks fφi with no shared parameters , this approach makes the strong assumption that the tasks indeed share no exploitable structure . This is often not known a priori . We show in this section that in the ‘ worst case ’ where shared parameters are not useful , the class of functions defined by a Mint network is still as rich as that defined by a set of purely task-specific networks . Thus , Mint makes no sacrifices in terms of expressiveness with respect to the optimal attainable parameters φ∗i and retains the possibility of exploiting shared structure if it does exist . Further , we show that the same is not true of existing methods that are related to Mint . Specifically , we compare the class of functions representable by Mint with 1 . The class of functions defined by a FiLM network ( Perez et al. , 2018 ) , where a layer ’ s activations are ‘ tilted ’ through element-wise multiplication with a task-specific vector γk and addition with a task-specific bias βk . 2 . The class of functions expressable by a task indicator concatenation approach , where a one-hot task indicator vector zk for task Tk is concatenated to each shared layer ’ s input . We begin by considering the optimal parameters for task-specific networks fφ∗i , where φ ∗ i = arg minφ Li ( φ ) , the parameters that minimize the loss function Li for task Ti . We assume that fφ∗i is an L-layer MLP . We also consider L-layer Mint networks fθi ( consisting of L sequential pairs of Mint transform and shared parameter layer as in Equation 1 ) and L-layer FiLM networks ( consisting of L sequential pairs of FiLM transform and shared parameter layer ) . We state Definition 2 to complete our setup . Definition 2 . Let y ( l−1 ) be the activations of layer l − 1 , and let αi = { W ( l ) , b ( l ) } be the shared parameters ( weight matrix and bias vector ) at the l-th layer . We assume that W ( l ) is an arbitrary , fixed invertible matrix . With this setup , we state Lemmas 1 and 2 , which highlight a key difference in the layer-wise expressiveness of Mint and methods such as FiLM and task indicator concatenation . Lemma 1 . For a given αi , applying Mint to y ( l−1 ) can express an arbitrary affine transformation at layer l for each task . Lemma 2 . ( a ) For a given αi , there exist affine transformations that can not be expressed by applying FiLM to y ( l−1 ) . ( b ) Similarly , for a given αi , there exist affine transformations that can not be expressed by a network which concatenates the input y ( l−1 ) of each layer l with a task indicator zk . The proof can be found in Appendix A . We assume only that the shared parameter matrices W ( l ) are invertible , such that they do not lose information about the input . This assumption critically does not require the optimizer to productively update the shared parameters . We use Lemmas 1 and 2 to extend the above layer-wise expressiveness properties of Mint and FiLM to compare the function classes defined by L-layer Mint networks fθi and FiLM networks fψi in the context of optimal L-layer MLP fφ∗i . We state this comparison in Theorem 1 . We use α to denote the subset of θi and ψi that corresponds to the parameters of the network that are shared across tasks . Theorem 1 . For all α , there exists θ∗i such that fθ∗i ( x ) = fφ∗i ( x ) for all x . Further , there exists α and x̄ such that fψi ( x̄ ) 6= fφ∗i ( x̄ ) for all ψi . Because fφ∗i is an MLP ( a composition of affine transformations and activation functions ) , the proof of the first statement follows from applying Lemma 1 to each layer of fφ∗i and fθi ( the task specific and Mint networks , respectively ) . Similarly , Lemma 2 implies that we can construct an α and Li ( and therefore φ∗i ) such that we can find an x̄ satisfying the second statement . That is , for any set of shared parameters in a Mint network and any set of optimal ‘ target ’ MLP , there exist weights for Mint that are effectively equivalent to the target MLP . On the other hand , there exist settings of the shared parameters that make matching the outputs of the target MLP unattainable in the FiLM and task indicator concatenation regimes . Using Lemma 2 , we reach an equivalent conclusion regarding the limited expressiveness of the task indicator concatenation approach .

Summary: The paper propose a single-network approach to multi-task learning by adding a task-specific linear transformation layer after each fully connected layer. The authors prove that the addition of such a layer keeps the expressive power of the network for each task. They also discuss how the linear transformation (parameterized by a transformation matrix and a bias vector) can be represented in a discrete manner in usual multi-task supervised learning and in a continuous manner (by two other neural networks) in goal-conditioned reinforcement learning. Experiments demonstrate the superiority of the proposed single-network approach.