Title: LARGE SCALE MULTI-DOMAIN MULTI-TASK LEARNING WITH MULTIMODEL

Abstract: Deep learning yields great results across many fields, from speech recognition, image classification, to translation. But for each problem, getting a deep model to work well involves research into the architecture and a long period of tuning.
We present a single model that yields good results on a number of problems spanning multiple domains. In particular, this single model is trained concurrently on ImageNet, multiple translation tasks, image captioning (COCO dataset), a speech recognition corpus, and an English parsing task. Our model architecture incorporates building blocks from multiple domains. It contains convolutional layers, an attention mechanism, and sparsely-gated layers. Each of these computational blocks is crucial for a subset of the tasks we train on. Interestingly, even if a block is not crucial for a task, we observe that adding it never hurts performance and in most cases improves it on all tasks. We also show that tasks with less data benefit largely from joint training with other tasks, while performance on large tasks degrades only slightly if at all.

1. INTRODUCTION:

Recent successes of deep neural networks have spanned many domains, from computer vision (Krizhevsky et al., 2012) to speech recognition (Dahl et al., 2012) and many other tasks. Convolutional networks excel at tasks related to vision, while recurrent neural networks have proven successful at natural language processing tasks, e.g., at machine translation (Sutskever et al., 2014; Bahdanau et al., 2014; Cho et al., 2014) . But in each case, the network was designed and tuned specifically for the problem at hand. This limits the impact of deep learning, as this effort needs to be repeated for each new task. It is also very different from the general nature of the human brain, which is able to learn many different tasks and benefit from transfer learning. The natural question arises:

Can we create a unified deep learning model to solve tasks across multiple domains?

The question about multi-task models has been studied in many papers in the deep learning literature. Natural language processing models have been shown to benefit from a multi-task approach a long time ago (Collobert & Weston, 2008) , and recently multi-task machine translation models (Minh-Thang Luong, 2015) have even been shown to exhibit zero-shot learning when trained on multiple languages (Melvin Johnson, 2016) . Speech recognition has also been shown to benefit from multi-task training (Seltzer & Droppo, 2013) , as have some vision problems, such as facial landmark detection (Zhang Z., 2014) . But all these models are trained on other tasks from the same domain: translation tasks are trained with other translation tasks, vision tasks with other vision tasks, speech tasks with other speech tasks. Multi-modal learning has been shown to improve learned representations in the unsupervised setting (Ngiam et al., 2011) and when used as a-priori known unrelated tasks (Romera-Paredes et al., 2012) . But no competitive multi-task multi-modal model has been proposed, so the above question remains unanswered.

In this work, we take a step toward positively answering the above question by introducing the MultiModel architecture, a single deep-learning model that can simultaneously learn multiple tasks from various domains. Concretely, we train the MultiModel simultaneously on the following 8 corpora:

Code available at redacted. (1) WSJ speech corpus (Consortium et al., 1994) , used for sentence-level speech recognition.

(2) ImageNet dataset (Russakovsky et al., 2015) , used for image classification.

(3) COCO image captioning dataset (Lin et al., 2014) , used for image captioning. (4) WSJ parsing dataset (Marcus et al., 1999) , used for constituency parsing.

(5) WMT English-German translation corpus, used for translation. (6) The reverse of the above: German-English translation. ( 7) WMT English-French translation corpus. (8) The reverse of the above: French-English translation.

These corpora were chosen as they are commonly used for machine learning the respective tasks: speech-to-text, image classification, captioning, parsing and translation. The model learns all of these tasks and achieves good performance: not state-of-the-art at present, but above many task-specific models studied in recent past (see the Section 3 for details). Figure 1 illustrates some decodes taken directly from the model: it is clear that it can caption images, categorize them, translate to French and German and construct parse trees. While the MultiModel is only a first step and will be improved in the future, two key insights are crucial to making it work at all and are our main contributions.

Small modality-specific sub-networks convert into a unified representation and back from it.

To allow training on input data of widely different sizes and dimensions, such as images, sound waves and text, we need sub-networks to convert inputs into a joint representation space. We call these sub-networks modality nets as they are specific to each modality (images, speech, text) and define transformations between these external domains and a unified representation. We design modality nets to be computationally minimal, promoting heavy feature extraction and ensuring that the majority of computation is performed within the domain-agnostic body of the model. Since our model is auto-regressive, modality nets need to both convert the inputs into the unified representation and later convert from this representation into the output space. Two design decisions were important:

• The unified representation is variable-size. While a fixed-size representation is tempting and easier to implement, it creates a bottleneck and limits the performance of the model. • Different tasks from the same domain share modality nets. We avoid creating a sub-network for every task, and prefer only to create one for every input modality. For example, all translation tasks share the same modality-net (and vocabulary), no matter for which language pair. This encourages generalization across tasks and allows to add new tasks on the fly.

Computational blocks of different kinds are crucial for good results on various problems.

The body of the MultiModel incorporates building blocks from mutiple domains. We use depthwiseseparable convolutions, an attention mechanism, and sparsely-gated mixture-of-experts layers. These blocks were introduced in papers that belonged to different domains and were not studied before on tasks from other domains. For example, separable convolutions were introduced in the Xception architecture (Chollet, 2016) and were not applied to text or speech processing before. On the other hand, the sparsely-gated mixture-of-experts (Shazeer et al., 2017) had been introduced for language processing tasks and has not been studied on image problems. We find that each of these mechanisms is indeed crucial for the domain it was introduced, e.g., attention is far more important for languagerelated tasks than for image-related ones. But, interestingly, adding these computational blocks never hurts performance, even on tasks they were not designed for. In fact we find that both attention and mixture-of-experts layers slightly improve performance of MultiModel on ImageNet, the task that needs them the least.

2. MULTIMODEL ARCHITECTURE:

The MultiModel consists of a few small modality-nets, an encoder, I/O mixer, and an autoregressive decoder, as depicted in Figure 2 . As already said above, the encoder and decoder are constructed using 3 key computational blocks to get good performance across different problems:

(1) Convolutions allow the model to detect local patterns and generalize across space.

(2) Attention layers allow to focus on specific elements to improve performance of the model.

(3) Sparsely-gated mixture-of-experts gives the model capacity without excessive computation cost.

We start by describing the architecture of each of these 3 blocks and then introduce the encoder, decoder and the architecture of our modality-nets.

2.1. CONVOLUTIONAL BLOCKS:

To perform local computation, we use blocks of convolutions with ReLU non-linearities and normalization. A block of convolutions gets as input a tensor of shape [batch size, sequence length, feature channels] and returns a tensor of the same shape, processed as follows.

For convolution operations, we use depthwise separable convolutions, studied for images in (Chollet, 2016) , in a way similar to (Kaiser et al., 2017) . Depthwise separable convolutions are a parameterand computationally-efficient variant of the traditional convolution. They are defined by a convolution on each feature channel separately, followed by a pointwise convolution to project to the desired feature depth. We refer the reader to (Chollet, 2016 ) for a complete definition; here we will denote a depthwise separable convolution with weights W h×w corresponding to f kernels of size h × w applied to an input tensor x with stride s and dilated by a factor d (see (Yu & Koltun, 2015) ) as SepConv d,s,f (W, x). Note that subscripts for stride, dilation and output size are omitted when dilation d or stride s are equal to 1, or output size f is equal to the input's feature depth.

We use convolutions in blocks that consist of three components: a ReLU activation of the inputs, followed by a SepConv, followed by layer normalization. Layer normalization (Ba et al., 2016) acts over the h hidden units of the layer below, computing layer-wise statistics for each batch example and normalizing accordingly. These normalized units are then scaled and shifted by scalar learned parameters G and B respectively, producing the final units to be activated by a non-linearity. The complete convolution step is therefore defined as:

ConvStep d,s,f (W, x) = LN (SepConv d,s,f (W, ReLU (x))).

The convolutional steps are composed into blocks by stacking them and adding residual connections (He et al., 2016) as depicted in Figure 3 . We use stacks of four convolutional blocks with two skip-connections between the stack input and the outputs of the second and fourth convolutional steps, and with the first two having 3 × 1 kernels and the next two having 15 × 1 kernels, with the final one dilated by 8 to provide a wide receptive field. We also add 40% dropout at the end of each block, so the complete block is defined as follows:

hidden1(x) = ConvStep(W 3×1 h1 , x) hidden2(x) = x + ConvStep(W 3×1 h2 , hidden1(x)) hidden3(x) = ConvStep(W 15×1 h3 , hidden2(x)) hidden4(x) = x + ConvStep d=8 (W 15×1 h4 , hidden3(x)) ConvBlock(x) = Dropout(hidden4(x), 0.4) during training hidden4(x) otherwise 2.2 ATTENTION BLOCKS

For attention, we use a multi-head dot-product attention mechanism inspired by (Bahdanau et al., 2014) and similar to (Ashish Vaswani, 2017), as depicted in Figure 3 . The inputs to the attention layer are two tensors: a source tensor and a target tensor both with the shape [batch size, sequence length, feature channels] The target tensor is additively composed with a timing signal and mixed using two convolutional blocks. This mixed tensor is then self-attended using a multi-head dot-product attention, which is a dot-product attention with inputs split into g = 8 separate tensors representing each attention head, as shown in Figure 3 . The timing signals are the main difference between this attention mechanism and the ones used previously. They allow this content-based attention to focus based on their position. They are constructed by concatenating sine and cosine curves:

∆(2d) = 1e4 -2d depth timing(t, [2d, 2d + 1]) = [sin(t∆(2d)) 2 cos(t∆(2d))]

where [a|| d b] represent concatenation of a and b along the d th dimension. The source tensor is finally passed through two different pointwise convolutions to generate the memory keys K and values V and the query keys, memory keys and memory values are used to apply the attention mechanism between the self-attended target and the source (see Figure 3 ).

2.3. MIXTURE-OF-EXPERTS BLOCKS:

We use sparsely-gated mixture-of-experts layers of the same kind as introduced in (Shazeer et al., 2017) : A mixture-of-experts layer consists of a number of simple feed-forward neural networks (experts) and a trainable gating network which selects a sparse combination of the experts to process each input. We refer the reader to (Shazeer et al., 2017) for details as we use exactly the architecture described there. In particular, during training we select k = 4 experts out of the whole expert pool and add the additional load-balancing cost as in (Shazeer et al., 2017) . In each of the two mixture-of-experts layers in our model, we use a pool of 240 experts when training on 8 problems jointly, and 60 experts when training on each problem separately.

2.4. ENCODER AND MIXER AND DECODER:

The body of the MultiModel consists of 3 parts: the encoder that only processes the inputs, the mixer that mixes the encoded inputs with previous outputs (autoregressive part), and a decoder that processes the inputs and the mixture to generate new outputs.

The encoder, mixer and decoder are structured similarly to previous fully convolutional sequence to sequence models such as ByteNet (Kalchbrenner et al., 2016) or WaveNet (van den Oord et al., 2016), but differ in the computational blocks that are used. We depict their architecture in Figure 3 .

As can be seen there, the encoder consists of 6 repeated convolutional blocks (described before) with a mixture-of-experts layer in the middle. The mixer consists of an attention block and 2 convolutional blocks. The decoder consists of 4 blocks of convolutions and attention, with a mixture-of-experts layer in the middle. Crucially, the convolutions in the mixer and decoder are padded on the left, so they can never access any information in the future. This allows the model to be autoregressive, and this convolutional autoregressive generation scheme offers large receptive fields over the inputs and past outputs, which are capable of establishing long term dependencies.

To allow the decoder to produce outputs for different tasks even with the same modality, we always start decoding with a command-token, such as To-English or To-Parse-Tree. We learn an embedding vector corresponding to each of the tokens during training.

2.5. MODALITY NETS:

We have 4 modality nets, for language (text data), images, audio, and categorical data. For all predictions, we use the cross-entropy loss, per subword-unit on text, per category on classification.

2.5.1. LANGUAGE MODALITY NET:

Our language-based data is all tokenized using the same vocabulary with 8k subword-units, following the method from (Sennrich et al., 2015) . The language input modality takes a sequence of tokens ending in a termination token. This sequence of tokens is mapped to the correct dimensionality for the body using a learned embedding. On the output side, the language modality takes the decoded output of the body and performs a learned linear mapping, followed by a Sof tmax, resulting in a probability distribution over the token vocabulary.

LanguageM odality in (x, W E ) = W E • x LanguageM odality out (x, W S ) = Sof tmax(W S • x) 2.5.2 IMAGE MODALITY NET

The image input modality is analogous to the Xception entry flow (Chollet, 2016) . The input image's feature depth is gradually deepened using residual convolution blocks which we call ConvRes and define as follows: (Szegedy et al., 2016; Vaswani et al., 2017) .

c1(x, F ) = ConvStep f =F (W 3×3 , x) c2(x, F ) = ConvStep f =F (W 3×3 , c1(x, F )) p1(x, F ) = M axP ool 2 ([3 × 3], c2(x, F )) ConvRes(x, F ) = p1(x, F ) + ConvStep s=2 (W 1×1 , x),

memory cells (Hochreiter & Schmidhuber, 1997) ). Convolutional architectures yielded good results on word-level neural machine translation starting from (Kalchbrenner & Blunsom, 2013) and later in (Meng et al., 2015) . These early models used a standard RNN on top of the convolution to generate the output and had a bottleneck there that hurt performance, especially on longer sentences, similarly to the limitations of RNN sequence-to-sequence models without attention (Sutskever et al., 2014; Cho et al., 2014) . Fully convolutional neural machine translation without this bottleneck was presented in (Kaiser & Bengio, 2016; Kalchbrenner et al., 2016) . The model in (Kaiser & Bengio, 2016 ) (Extended Neural GPU) used a recurrent stack of gated convolutional layers, while the model in (Kalchbrenner et al., 2016) (ByteNet) did away with recursion and used left-padded convolutions in the decoder. This idea, introduced in WaveNet (van den Oord et al., 2016) and also used in MultiModel (see above) significantly improves efficiency. Depthwise separable convolutions were first studied by Sifre (Sifre & Mallat, 2013) and later they were used to get good results on large-scale image classification with Xception (Chollet, 2016) .

3. EXPERIMENTS:

We implemented the MultiModel architecture described above using TensorFlow and trained it in a number of configurations. We focused our experiments so as to answer the following questions:

(1) How far is the MultiModel trained on 8 tasks simultaneously from state-of-the-art results?

(2) How does training on 8 tasks simultaneously compare to training on each task separately?

(3) How do the different computational blocks discussed above influence different tasks?

In answering the above questions, we don't always consider all 8 problems. Especially the 4 translation problems behave very similarly, so we decided to not include them all in each comparison but we focused on the more varied problems instead.

To answer question (1), we compare the performance of the 8-problem MultiModel with state-of-theart results in Table 1 . We use the standard top-5 accuracy metric for ImageNet and the standard BLEU metric for translation (scored with MOSES on newstest2014 while newstest2013 was used as the development set). We did not invest much time yet in tuning hyper-parameters of the MultiModel, so we believe that the difference seen there will become much smaller with more tuning. The results we achieve are similar to the ones task-specific models get without heavy tuning, e.g., on English-French translation we improve on the recent Extended Neural GPU results (Kaiser & Bengio, 2016 ).

To answer question (2), we compare the MultiModel trained jointly with MultiModel trained separately just on a single task. Since we are comparing different instantiations of the same model, we report two internal metrics: the negative log-perplexity and per-token accuracy (measured on the development set). As can be seen from the results in Table 2 , the joint 8-problem model performs similarly to single-model on large tasks, and better, sometimes significantly, on tasks where less data is available, such as parsing.

The large improvement on parsing seen in Table 2 is not that surprising taking into account the large number of text data in translation tasks. But we were curious if training parsing just with ImageNet, a seemingly unrelated task, would also bring any improvements. This is indeed the case, as can be seen in Table 3 . The difference in performance is significant, and since we use both dropout and early stopping, we conjecture that it is not related to over-fitting. Rather, it seems, there are computational primitives shared between different tasks that allow for some transfer learning even between such seemingly unrelated tasks as ImageNet and parsing. To answer question (3), we check how training without the mixture-of-experts layers or without the attention mechanism influences performance on different problems. Since both these mechanisms were designed with machine translation in mind, we check the English-French translation. But we also include ImageNet, since this is the problem that stands the least to benefit from those blocks. In fact, one could expect that removing these blocks will improve performance on ImageNet alone if they were truly useless for this task. In contrast, we see in Table 4 that these blocks either don't affect or slightly improve performance. This leads us to conclude that mixing different computation blocks is in fact a good way to improve performance on many various tasks.

4. CONCLUSIONS:

We demonstrate, for the first time, that a single deep learning model can jointly learn a number of large-scale tasks from multiple domains. The key to success comes from designing a multi-modal architecture in which as many parameters as possible are shared and from using computational blocks from different domains together. We believe that this treads a path towards interesting future work on more general deep learning architectures, especially since our model shows transfer learning from tasks with a large amount of available data to ones where the data is limited. 

:



Figures and Tables:

Figure 1: Examples decoded from a single MultiModel trained jointly on 8 tasks. Red depicts a language modality while blue depicts a categorical modality.

Figure 2: The MultiModel, with modality-nets, an encoder, and an autoregressive decoder.

Figure 3: Architecture of the MultiModel; see text for details.

Comparing MultiModel to state-of-the-art

Table content: <table><tr><td>Problem</td><td colspan="2">MultiModel (joint 8-problem) State of the art</td></tr><tr><td>ImageNet (top-5 accuracy)</td><td>86%</td><td>95%</td></tr><tr><td>WMT EN → DE (BLEU)</td><td>21.2</td><td>28.4</td></tr><tr><td>WMT EN → FR (BLEU)</td><td>30.5</td><td>41.0</td></tr></table>

Comparison of the MultiModel trained jointly on 8 tasks and separately on each task.

Table content: <table><tr><td/><td>Problem</td><td/><td colspan="3">Joint 8-problem log(perpexity) accuracy</td><td colspan="3">Single problem log(perplexity) accuracy</td></tr><tr><td/><td>ImageNet</td><td/><td>1.7</td><td/><td>66%</td><td>1.6</td><td/><td>67%</td></tr><tr><td/><td colspan="2">WMT EN→DE</td><td>1.4</td><td/><td>72%</td><td>1.4</td><td/><td>71%</td></tr><tr><td/><td colspan="2">WSJ speech</td><td>4.4</td><td/><td>41%</td><td>5.7</td><td/><td>23%</td></tr><tr><td/><td>Parsing</td><td/><td>0.15</td><td/><td>98%</td><td>0.2</td><td/><td>97%</td></tr><tr><td>Problem</td><td>log(ppl)</td><td>Alone acc.</td><td>full</td><td colspan="2">W/ ImageNet log(ppl) acc.</td><td>full</td><td colspan="2">W/ 8 Problems log(ppl) acc.</td><td>full</td></tr><tr><td>Parsing</td><td>0.20</td><td colspan="2">97.1% 11.7%</td><td>0.16</td><td colspan="2">97.5% 12.7%</td><td>0.15</td><td>97.9% 14.5%</td></tr></table>

Results on training parsing alone, with ImageNet, and with 8 other tasks. We report log-perplexity, per-token accuracy, and the percentage of fully correct parse trees.

Table content: <table/>

Ablating mixture-of-experts and attention from MultiModel training.

Table content: <table/>

Peer review (iclr2018):
 The paper presents a multi-task architecture that can perform multiple tasks across multiple different domains. The authors design an architecture that works on image captioning, image classification, machine translation and parsing. The proposed model can maintain performance of single-task models and in some cases show slight improvements. This is the main take-away from this paper. There is a factually incorrect statement - depthwise separable convolutions were not introduced in Chollet 2016. Section 2 of the same paper also notes it (depthwise convolutions can be traced back to at least 2012).