Source-linked AI summary
One Model To Learn Them All
Lukasz Kaiser, Aidan N. Gomez, Noam Shazeer, Ashish Vaswani, Niki Parmar, Llion Jones, Jakob Uszkoreit
TL;DR
The paper addresses whether one model can learn competitive tasks across multiple domains without designing and tuning a separate architecture for each task. It introduces MultiModel, a jointly trained multimodal architecture combining shared representations with convolutional, attention, and mixture-of-experts blocks. The model learns eight tasks, performs well across them, and transfers more strongly to data-limited tasks while large-task performance remains similar.
Problem
Deep learning systems typically require repeated architecture design and tuning for each domain, while competitive multi-task multi-modal models remained unavailable.
Method
MultiModel jointly trains eight tasks using small modality nets, shared variable-size representations, and computational blocks from multiple domains.
Results
The model achieves good performance across all eight tasks, with joint training similar to single-task training on large tasks and better on data-limited tasks.
Takeaways & Limitations
Shared parameters and cross-domain computational blocks support transfer from data-rich tasks to tasks with limited data.
Takeaways & Limitations
Comparisons do not always include all eight problems because the four translation tasks behave similarly.
Abstract
from arXiv · showhide
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
The paper asks whether one competitive model can learn tasks across multiple domains, replacing repeated task-specific architecture design and tuning. It introduces MultiModel, which jointly trains eight tasks using shared multimodal components and domain-specific computational blocks.
- Task-specific architecture design and tuning must be repeated across domains, limiting the broader impact of deep learning.
- The paper asks whether a unified deep learning model can solve tasks across multiple domains.
- MultiModel is a single architecture trained simultaneously on eight corpora spanning speech, vision, captioning, parsing, and translation.
- MultiModel learns all eight tasks with good performance, exceeding many recent task-specific models although not achieving state of the art.
- Small modality nets map images, speech, and text into a shared variable-size representation and back into task outputs.
- The model combines convolutions, attention, and sparsely-gated mixture-of-experts layers, each crucial for some tasks but jointly non-harmful across tasks.
2 MultiModel Architecture
The MultiModel body uses an encoder, an autoregressive mixer, and a decoder built from three complementary computational blocks. Convolutions, attention, and mixture-of-experts provide local processing, selective focus, and capacity-efficient capacity.
- The architecture contains an encoder, an I/O mixer, and an autoregressive decoder.
- Convolutions detect local patterns and generalize across space.
- Attention layers focus on specific elements to improve model performance.
- Sparsely-gated mixture-of-experts layers add model capacity without excessive computation cost.
2.1 Convolutional Blocks
Convolutional blocks perform local computation while preserving tensor shape through ReLU activation, depthwise separable convolution, normalization, residual structure, and dropout.
- A convolutional block maps a [batch size, sequence length, feature channels] tensor to one of the same shape.
- Depthwise separable convolution processes each feature channel separately, then uses a pointwise convolution to project to the desired feature depth.
- Each convolution step applies ReLU, depthwise separable convolution, and layer normalization in sequence.
- Four convolutional steps are stacked with residual connections, 3 × 1 and 15 × 1 kernels, and a final dilation of 8 for a wide receptive field.
- 40% dropout is applied at the end of each convolutional block during training.
2.2 Attention Blocks
The attention block combines target representations with timing signals before multi-head self-attention and cross-attention to source-derived keys and values.
- The attention mechanism receives source and target tensors, each shaped [batch size, sequence length, feature channels].
- Target representations receive timing signals and convolutional mixing before eight-head dot-product self-attention.
- Timing signals encode position through concatenated sine and cosine curves, allowing content-based attention to focus by position.
- The source is transformed into memory keys and values, which attention uses with queries from the self-attended target.
2.3 Mixture-of-Experts Blocks
The MultiModel uses sparsely-gated mixture-of-experts layers, where a trainable gate selects a sparse subset of feed-forward experts for each input.
- A mixture-of-experts layer combines simple feed-forward experts with a trainable gating network that selects which experts process each input.
- During training, the model selects k = 4 experts from the full expert pool and adds load-balancing mechanisms.
- The MultiModel uses the same sparsely-gated mixture-of-experts architecture introduced in prior work.
2.4 Encoder and Mixer and Decoder
The MultiModel body separates processing into an encoder, an autoregressive mixer, and a decoder that generates task-specific outputs.
- The body has three parts: an encoder for inputs, a mixer combining encoded inputs with previous outputs, and a decoder generating new outputs.
- Figure 3 depicts the architecture of the MultiModel body and its computational blocks.
- The encoder contains six repeated convolutional blocks with a mixture-of-experts layer, while the mixer contains attention and two convolutional blocks.
- The decoder produces different task outputs from the same modality by beginning generation with a task-specific command token.
2.5 Modality Nets
Modality-specific input and output networks translate language, images, audio, and categorical data into and out of a shared representation space.
- The architecture has four modality nets for language, images, audio, and categorical data.
- Language outputs apply a learned linear mapping and Softmax to produce probabilities over the token vocabulary.
- The image input modality uses residual convolution blocks to gradually deepen image feature representations, with network depth d = 1024.
- Image processing uses convolutional steps and max-pooling operations to transform the input through progressively deeper representations.
- GlobalAvgPool computes a mean across all spatial and temporal dimensions.
- Audio inputs may be one-dimensional waveforms or two-dimensional spectrograms processed by stacks of eight convolutional residual blocks.
2.6 Related Models
The MultiModel builds on earlier encoder-decoder and fully convolutional sequence-to-sequence models while changing their computational blocks.
- Earlier translation sequence-to-sequence models used recurrent neural networks with long short-term memory cells.
- Later convolutional translation models achieved good word-level results but still placed a standard recurrent network above the convolution.
- The MultiModel draws from this encoder-decoder lineage while adopting a different architecture.
3 Experiments
Experiments evaluate the MultiModel against task-specific systems, separate-task training, cross-task transfer, and computational-block ablations. Joint training preserves performance on large tasks, benefits data-limited tasks, and shows that mixed blocks generally do not hurt performance.
- Experimental design: The experiments compare the eight-task MultiModel with state-of-the-art systems and separately trained versions using shared evaluation settings.The study also examines computational-block ablations and focuses comparisons on varied problems because the four translation tasks behave similarly.
- Comparison with prior systems: The MultiModel achieves good performance across the evaluated tasks, though it is not state-of-the-art and received limited hyperparameter tuning.Its results are similar to task-specific models without heavy tuning, and it improves on previously reported Extended Neural GPU results for English-French translation.
- Joint versus separate training: Joint training performs similarly to single-task training on large tasks but better, sometimes significantly, on data-limited tasks such as parsing.The comparison uses negative log-perplexity and per-token accuracy on development data.
- Transfer learning: Parsing improves when trained jointly with ImageNet, suggesting transfer through computational primitives shared across seemingly unrelated tasks.The authors state that the difference is significant and conjecture it is not caused by over-fitting because dropout and early stopping were used.
- Computational blocks: Ablating mixture-of-experts or attention either leaves performance unchanged or slightly improves it, including on ImageNet, where these blocks were expected to help least.The authors conclude that combining computational blocks from different domains can improve performance across varied tasks.
4 Conclusions
The paper demonstrates that one deep-learning model can jointly learn large-scale tasks from multiple domains. Its approach shares as many parameters as possible, combines computational blocks from different domains, and transfers from data-rich tasks to data-limited ones.
- 4 Conclusions: A single deep-learning model jointly learns multiple large-scale tasks from multiple domains.The authors present this as a first demonstration and identify it as a step toward more general deep-learning architectures.
- 4 Conclusions: The architecture succeeds by maximizing parameter sharing and combining computational blocks from different domains.The paper identifies these design choices as the key to making joint multi-domain learning work.
- 4 Conclusions: Tasks with limited data can receive transfer from tasks with larger amounts of available data.The conclusion frames this transfer as a direction for future general deep-learning architectures.