summaries: The authors present ALBERT, a modification of the BERT architecture with substantially fewer parameters. They show that despite being much smaller, the performance is very strong and achieves state of the art on a variety of different tasks. There are several ideas proposed here: embedding factorization, sharing layers, and sentence ordering as a training objective.

summaries: This paper investigates improving upon BERT by reducing complexity in terms of free parameters and memory footprint as well as computation steps. They propose 2 strategies for doing this: 1) Splitting the embedding matrix into two smaller matrices (going from V x A to V x B + B x A where B <<<< A); 2) layer-wise parameter sharing. They also utilize sentence order prediction to help with training. These coupled with a bunch of other choices such as using the lamb optimizer, certain hyperparameters etc help show dramatic empirical gains across the board on a wide variety of NLP/NLU tasks.

summaries: This paper proposes a new pre-trained BERT-like model called ALBERT. The contributions are mainly 3-fold: factorized embedding parameterization, cross-layer parameter sharing, and intern-sentence coherence loss. The first two address the issue of model size and memory consumption in BERT; the third corresponds to a new auxiliary task in pre-train, sentence-order prediction (SOP), replacing the next sentence prediction (NSP) task in BERT. These modifications lead to a much leaner model and improved performance. As a result, ALBERT pushes the state of the art on GLUE, RACE, and SQuAD while having fewer parameters than BERT-large.

abstractText: Increasing model size when pretraining natural language representations often results in improved performance on downstream tasks. However, at some point further model increases become harder due to GPU/TPU memory limitations and longer training times. To address these problems, we present two parameterreduction techniques to lower memory consumption and increase the training speed of BERT (Devlin et al., 2019). Comprehensive empirical evidence shows that our proposed methods lead to models that scale much better compared to the original BERT. We also use a self-supervised loss that focuses on modeling inter-sentence coherence, and show it consistently helps downstream tasks with multi-sentence inputs. As a result, our best model establishes new state-of-the-art results on the GLUE, RACE, and SQuAD benchmarks while having fewer parameters compared to BERT-large. The code and the pretrained models are available at https://github.com/google-research/ALBERT.

authors:
Alexei Baevski
Michael Auli

title: Adaptive input representations for neural language modeling

venue: arXiv preprint arXiv:1809.10853,

authors:
Shaojie Bai
J. Zico Kolter
Vladlen Koltun

title: Deep equilibrium models

venue: In Neural Information Processing Systems (NeurIPS),

authors:
Roy Bar-Haim
Ido Dagan
Bill Dolan
Lisa Ferro
Danilo Giampiccolo
Bernardo Magnini
Idan Szpektor

title: The second PASCAL recognising textual entailment challenge

venue: In Proceedings of the second PASCAL challenges workshop on recognising textual entailment,

authors:
Luisa Bentivogli
Peter Clark
Ido Dagan
Danilo Giampiccolo

title: The fifth PASCAL recognizing textual entailment challenge

venue: In TAC,

authors:
Daniel Cer
Mona Diab
Eneko Agirre
Iñigo Lopez-Gazpio
Lucia Specia

title: SemEval-2017 task 1: Semantic textual similarity multilingual and crosslingual focused evaluation

venue: In Proceedings of the 11th International Workshop on Semantic Evaluation

authors:
Tianqi Chen
Bing Xu
Chiyuan Zhang
Carlos Guestrin

title: Training deep nets with sublinear memory cost

venue: arXiv preprint arXiv:1604.06174,

authors:
Rewon Child
Scott Gray
Alec Radford
Ilya Sutskever

title: Generating long sequences with sparse transformers

venue: arXiv preprint arXiv:1904.10509,

authors:
Kevin Clark
Minh-Thang Luong
Urvashi Khandelwal
Christopher D Manning
Quoc V Le

title: Bam! born-again multi-task networks for natural language understanding

authors:
Ido Dagan
Oren Glickman
Bernardo Magnini

title: The PASCAL recognising textual entailment challenge

venue: In Machine Learning Challenges Workshop,

authors:
Andrew M Dai
Quoc V Le

title: Semi-supervised sequence learning

authors:
Zihang Dai
Zhilin Yang
Yiming Yang
William W Cohen
Jaime Carbonell
Quoc V Le
Ruslan Salakhutdinov

title: Transformer-xl: Attentive language models beyond a fixed-length context

authors:
Jacob Devlin
Ming-Wei Chang
Kenton Lee
Kristina Toutanova

title: BERT: Pre-training of deep bidirectional transformers for language understanding

authors:
William B. Dolan
Chris Brockett

title: Automatically constructing a corpus of sentential paraphrases

venue: In Proceedings of the Third International Workshop on Paraphrasing (IWP2005),

authors:
Zhe Gan
Yunchen Pu
Ricardo Henao
Chunyuan Li
Xiaodong He
Lawrence Carin

title: Learning generic sentence representations using convolutional neural networks

venue: In Proceedings of the 2017 Conference on Empirical Methods in Natural Language Processing,

authors:
Danilo Giampiccolo
Bernardo Magnini
Ido Dagan
Bill Dolan

title: The third PASCAL recognizing textual entailment challenge

venue: In Proceedings of the ACL-PASCAL Workshop on Textual Entailment and Paraphrasing,

authors:
Aidan N Gomez
Mengye Ren
Raquel Urtasun
Roger B Grosse

title: The reversible residual network: Backpropagation without storing activations

authors:
Linyuan Gong
Di He
Zhuohan Li
Tao Qin
Liwei Wang
Tieyan Liu

title: Efficient training of bert by progressively stacking

venue: In International Conference on Machine Learning,

authors:
Edouard Grave
Armand Joulin
Moustapha Cissé
Hervé Jégou

title: Efficient softmax approximation for gpus

venue: In Proceedings of the 34th International Conference on Machine Learning-Volume

authors:
Barbara J. Grosz
Aravind K. Joshi
Scott Weinstein

title: Centering: A framework for modeling the local coherence of discourse

venue: Computational Linguistics,

authors:
Jie Hao
Xing Wang
Baosong Yang
Longyue Wang
Jinfeng Zhang
Zhaopeng Tu

title: Modeling recurrence for transformer

venue: Proceedings of the 2019 Conference of the North,

authors:
Dan Hendrycks
Kevin Gimpel

title: Gaussian Error Linear Units (GELUs)

venue: arXiv preprint arXiv:1606.08415,

authors:
Felix Hill
Kyunghyun Cho
Anna Korhonen

title: Learning distributed representations of sentences from unlabelled data. In Proceedings of the 2016 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pp. 1367–1377

venue: Association for Computational Linguistics,

title: Coherence and coreference

venue: Cognitive Science,

authors:
Jeremy Howard
Sebastian Ruder

title: Universal language model fine-tuning for text classification

venue: arXiv preprint arXiv:1801.06146,

authors:
Shankar Iyer
Nikhil Dandekar
Kornl Csernai

title: First quora dataset release: Question pairs, January 2017

authors:
Yacine Jernite
Samuel R Bowman
David Sontag

title: Discourse-based objectives for fast unsupervised sentence representation learning

venue: arXiv preprint arXiv:1705.00557,

authors:
Mandar Joshi
Danqi Chen
Yinhan Liu
Daniel S Weld
Luke Zettlemoyer
Omer Levy

title: SpanBERT: Improving pre-training by representing and predicting spans

authors:
Ryan Kiros
Yukun Zhu
Ruslan Salakhutdinov
Richard S. Zemel
Antonio Torralba
Raquel Urtasun
Sanja Fidler

title: Skip-thought vectors

venue: In Proceedings of the 28th International Conference on Neural Information Processing Systems - Volume 2,

authors:
Taku Kudo
John Richardson

title: SentencePiece: A simple and language independent subword tokenizer and detokenizer for neural text processing

venue: In Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing: System Demonstrations,

authors:
Guokun Lai
Qizhe Xie
Hanxiao Liu
Yiming Yang
Eduard Hovy

title: RACE: Large-scale ReAding comprehension dataset from examinations

authors:
Quoc Le
Tomas Mikolov

title: Distributed representations of sentences and documents

venue: In Proceedings of the 31st ICML, Beijing,

authors:
Hector Levesque
Ernest Davis
Leora Morgenstern

title: The Winograd schema challenge

venue: In Thirteenth International Conference on the Principles of Knowledge Representation and Reasoning,

authors:
Xiang Li
Shuo Chen
Xiaolin Hu
Jian Yang

title: Understanding the disharmony between dropout and batch normalization by variance shift

venue: In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition,

authors:
Yinhan Liu
Myle Ott
Naman Goyal
Jingfei Du
Mandar Joshi
Danqi Chen
Omer Levy
Mike Lewis
Luke Zettlemoyer
Veselin Stoyanov

title: RoBERTa: A robustly optimized BERT pretraining approach

authors:
Bryan McCann
James Bradbury
Caiming Xiong
Richard Socher

title: Learned in translation: Contextualized word vectors

venue: Advances in Neural Information Processing Systems

authors:
Tomas Mikolov
Ilya Sutskever
Kai Chen
Greg S Corrado
Jeff Dean

title: Distributed representations of words and phrases and their compositionality

authors:
Allen Nie
Erin Bennett
Noah Goodman

title: DisSent: Learning sentence representations from explicit discourse relations

venue: In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics,

authors:
Jeffrey Pennington
Richard Socher
Christopher Manning

title: Glove: Global vectors for word representation

venue: In Proceedings of the 2014 Conference on Empirical Methods in Natural Language Processing (EMNLP),

authors:
Matthew Peters
Mark Neumann
Mohit Iyyer
Matt Gardner
Christopher Clark
Kenton Lee
Luke Zettlemoyer

title: Deep contextualized word representations. In Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long Papers)

authors:
Alec Radford
Karthik Narasimhan
Tim Salimans
Ilya Sutskever

title: Improving language understanding by generative pre-training

authors:
Alec Radford
Jeffrey Wu
Rewon Child
David Luan
Dario Amodei
Ilya Sutskever

title: Language models are unsupervised multitask learners

venue: OpenAI Blog,

authors:
Colin Raffel
Noam Shazeer
Adam Roberts
Katherine Lee
Sharan Narang
Michael Matena
Yanqi Zhou
Wei Li
Peter J Liu

title: Exploring the limits of transfer learning with a unified text-to-text transformer

venue: arXiv preprint arXiv:1910.10683,

authors:
Pranav Rajpurkar
Jian Zhang
Konstantin Lopyrev
Percy Liang

title: SQuAD: 100,000+ questions for machine comprehension of text

venue: In Proceedings of the 2016 Conference on Empirical Methods in Natural Language Processing,

authors:
Pranav Rajpurkar
Robin Jia
Percy Liang

title: Know what you don’t know: Unanswerable questions for SQuAD

venue: In Proceedings of the 56th Annual Meeting of the Association for Computational Linguistics

authors:
Noam Shazeer
Youlong Cheng
Niki Parmar
Dustin Tran
Ashish Vaswani
Penporn Koanantakool
Peter Hawkins
HyoukJoong Lee
Mingsheng Hong
Cliff Young

title: Mesh-tensorflow: Deep learning for supercomputers

authors:
Tao Shen
Tianyi Zhou
Guodong Long
Jing Jiang
Chengqi Zhang

title: Bi-directional block selfattention for fast and memory-efficient sequence modeling

venue: arXiv preprint arXiv:1804.00857,

authors:
Mohammad Shoeybi
Mostofa Patwary
Raul Puri
Patrick LeGresley
Jared Casper
Bryan Catanzaro

title: Megatron-LM: Training multi-billion parameter language models using model parallelism, 2019

authors:
Richard Socher
Alex Perelygin
Jean Wu
Jason Chuang
Christopher D. Manning
Andrew Ng
Christopher Potts

title: Recursive deep models for semantic compositionality over a sentiment treebank

venue: In Proceedings of the 2013 Conference on Empirical Methods in Natural Language Processing,

authors:
Siqi Sun
Yu Cheng
Zhe Gan
Jingjing Liu

title: Patient knowledge distillation for BERT model compression

venue: arXiv preprint arXiv:1908.09355,

authors:
Christian Szegedy
Sergey Ioffe
Vincent Vanhoucke
Alexander A Alemi

title: Inception-v4, inception-resnet and the impact of residual connections on learning

venue: In Thirty-First AAAI Conference on Artificial Intelligence,

authors:
Iulia Turc
Ming-Wei Chang
Kenton Lee
Kristina Toutanova

title: Well-read students learn better: The impact of student initialization on knowledge distillation

authors:
Ashish Vaswani
Noam Shazeer
Niki Parmar
Jakob Uszkoreit
Llion Jones
Aidan N Gomez
Łukasz Kaiser
Illia Polosukhin

title: Attention is all you need

authors:
Alex Wang
Amanpreet Singh
Julian Michael
Felix Hill
Omer Levy
Samuel Bowman

title: GLUE: A multi-task benchmark and analysis platform for natural language understanding

venue: In Proceedings of the 2018 EMNLP Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP,

authors:
Wei Wang
Bin Bi
Ming Yan
Chen Wu
Zuyi Bao
Liwei Peng
Luo Si

title: StructBERT: Incorporating language structures into pre-training for deep language understanding

authors:
Alex Warstadt
Amanpreet Singh
Samuel R Bowman

title: Neural network acceptability judgments

venue: arXiv preprint arXiv:1805.12471,

authors:
Adina Williams
Nikita Nangia
Samuel Bowman

title: A broad-coverage challenge corpus for sentence understanding through inference

venue: In Proceedings of the 2018 Conference of the North

heading: 1 INTRODUCTION

text:
Full network pre-training (Dai & Le, 2015; Radford et al., 2018; Devlin et al., 2019; Howard & Ruder, 2018) has led to a series of breakthroughs in language representation learning. Many nontrivial NLP tasks, including those that have limited training data, have greatly benefited from these pre-trained models. One of the most compelling signs of these breakthroughs is the evolution of machine performance on a reading comprehension task designed for middle and high-school English exams in China, the RACE test (Lai et al., 2017): the paper that originally describes the task and formulates the modeling challenge reports then state-of-the-art machine accuracy at 44.1%; the latest published result reports their model performance at 83.2% (Liu et al., 2019); the work we present here pushes it even higher to 89.4%, a stunning 45.3% improvement that is mainly attributable to our current ability to build high-performance pretrained language representations.
Evidence from these improvements reveals that a large network is of crucial importance for achieving state-of-the-art performance (Devlin et al., 2019; Radford et al., 2019). It has become common practice to pre-train large models and distill them down to smaller ones (Sun et al., 2019; Turc et al., 2019) for real applications. Given the importance of model size, we ask: Is having better NLP models as easy as having larger models?
An obstacle to answering this question is the memory limitations of available hardware. Given that current state-of-the-art models often have hundreds of millions or even billions of parameters, it is easy to hit these limitations as we try to scale our models. Training speed can also be significantly hampered in distributed training, as the communication overhead is directly proportional to the number of parameters in the model.
Existing solutions to the aforementioned problems include model parallelization (Shazeer et al., 2018; Shoeybi et al., 2019) and clever memory management (Chen et al., 2016; Gomez et al., 2017).
∗Work done as an intern at Google Research, driving data processing and downstream task evaluations.
These solutions address the memory limitation problem, but not the communication overhead. In this paper, we address all of the aforementioned problems, by designing A Lite BERT (ALBERT) architecture that has significantly fewer parameters than a traditional BERT architecture.
ALBERT incorporates two parameter reduction techniques that lift the major obstacles in scaling pre-trained models. The first one is a factorized embedding parameterization. By decomposing the large vocabulary embedding matrix into two small matrices, we separate the size of the hidden layers from the size of vocabulary embedding. This separation makes it easier to grow the hidden size without significantly increasing the parameter size of the vocabulary embeddings. The second technique is cross-layer parameter sharing. This technique prevents the parameter from growing with the depth of the network. Both techniques significantly reduce the number of parameters for BERT without seriously hurting performance, thus improving parameter-efficiency. An ALBERT configuration similar to BERT-large has 18x fewer parameters and can be trained about 1.7x faster. The parameter reduction techniques also act as a form of regularization that stabilizes the training and helps with generalization.
To further improve the performance of ALBERT, we also introduce a self-supervised loss for sentence-order prediction (SOP). SOP primary focuses on inter-sentence coherence and is designed to address the ineffectiveness (Yang et al., 2019; Liu et al., 2019) of the next sentence prediction (NSP) loss proposed in the original BERT.
As a result of these design decisions, we are able to scale up to much larger ALBERT configurations that still have fewer parameters than BERT-large but achieve significantly better performance. We establish new state-of-the-art results on the well-known GLUE, SQuAD, and RACE benchmarks for natural language understanding. Specifically, we push the RACE accuracy to 89.4%, the GLUE benchmark to 89.4, and the F1 score of SQuAD 2.0 to 92.2.

heading: 2 RELATED WORK

heading: 2.1 SCALING UP REPRESENTATION LEARNING FOR NATURAL LANGUAGE

text:
Learning representations of natural language has been shown to be useful for a wide range of NLP tasks and has been widely adopted (Mikolov et al., 2013; Le & Mikolov, 2014; Dai & Le, 2015; Peters et al., 2018; Devlin et al., 2019; Radford et al., 2018; 2019). One of the most significant changes in the last two years is the shift from pre-training word embeddings, whether standard (Mikolov et al., 2013; Pennington et al., 2014) or contextualized (McCann et al., 2017; Peters et al., 2018), to full-network pre-training followed by task-specific fine-tuning (Dai & Le, 2015; Radford et al., 2018; Devlin et al., 2019). In this line of work, it is often shown that larger model size improves performance. For example, Devlin et al. (2019) show that across three selected natural language understanding tasks, using larger hidden size, more hidden layers, and more attention heads always leads to better performance. However, they stop at a hidden size of 1024, presumably because of the model size and computation cost problems.
It is difficult to experiment with large models due to computational constraints, especially in terms of GPU/TPU memory limitations. Given that current state-of-the-art models often have hundreds of millions or even billions of parameters, we can easily hit memory limits. To address this issue, Chen et al. (2016) propose a method called gradient checkpointing to reduce the memory requirement to be sublinear at the cost of an extra forward pass. Gomez et al. (2017) propose a way to reconstruct each layer’s activations from the next layer so that they do not need to store the intermediate activations. Both methods reduce the memory consumption at the cost of speed. Raffel et al. (2019) proposed to use model parallelization to train a giant model. In contrast, our parameter-reduction techniques reduce memory consumption and increase training speed.

heading: 2.2 CROSS-LAYER PARAMETER SHARING

text:
The idea of sharing parameters across layers has been previously explored with the Transformer architecture (Vaswani et al., 2017), but this prior work has focused on training for standard encoderdecoder tasks rather than the pretraining/finetuning setting. Different from our observations, Dehghani et al. (2018) show that networks with cross-layer parameter sharing (Universal Transformer, UT) get better performance on language modeling and subject-verb agreement than the standard
transformer. Very recently, Bai et al. (2019) propose a Deep Equilibrium Model (DQE) for transformer networks and show that DQE can reach an equilibrium point for which the input embedding and the output embedding of a certain layer stay the same. Our observations show that our embeddings are oscillating rather than converging. Hao et al. (2019) combine a parameter-sharing transformer with the standard one, which further increases the number of parameters of the standard transformer.

heading: 2.3 SENTENCE ORDERING OBJECTIVES

text: ALBERT uses a pretraining loss based on predicting the ordering of two consecutive segments of text. Several researchers have experimented with pretraining objectives that similarly relate to discourse coherence. Coherence and cohesion in discourse have been widely studied and many phenomena have been identified that connect neighboring text segments (Hobbs, 1979; Halliday & Hasan, 1976; Grosz et al., 1995). Most objectives found effective in practice are quite simple. Skipthought (Kiros et al., 2015) and FastSent (Hill et al., 2016) sentence embeddings are learned by using an encoding of a sentence to predict words in neighboring sentences. Other objectives for sentence embedding learning include predicting future sentences rather than only neighbors (Gan et al., 2017) and predicting explicit discourse markers (Jernite et al., 2017; Nie et al., 2019). Our loss is most similar to the sentence ordering objective of Jernite et al. (2017), where sentence embeddings are learned in order to determine the ordering of two consecutive sentences. Unlike most of the above work, however, our loss is defined on textual segments rather than sentences. BERT (Devlin et al., 2019) uses a loss based on predicting whether the second segment in a pair has been swapped with a segment from another document. We compare to this loss in our experiments and find that sentence ordering is a more challenging pretraining task and more useful for certain downstream tasks. Concurrently to our work, Wang et al. (2019) also try to predict the order of two consecutive segments of text, but they combine it with the original next sentence prediction in a three-way classification task rather than empirically comparing the two.

heading: 3 THE ELEMENTS OF ALBERT

text: In this section, we present the design decisions for ALBERT and provide quantified comparisons against corresponding configurations of the original BERT architecture (Devlin et al., 2019).

heading: 3.1 MODEL ARCHITECTURE CHOICES

text:
The backbone of the ALBERT architecture is similar to BERT in that it uses a transformer encoder (Vaswani et al., 2017) with GELU nonlinearities (Hendrycks & Gimpel, 2016). We follow the BERT notation conventions and denote the vocabulary embedding size as E, the number of encoder layers as L, and the hidden size as H . Following Devlin et al. (2019), we set the feed-forward/filter size to be 4H and the number of attention heads to be H/64.
There are three main contributions that ALBERT makes over the design choices of BERT.
Factorized embedding parameterization. In BERT, as well as subsequent modeling improvements such as XLNet (Yang et al., 2019) and RoBERTa (Liu et al., 2019), the WordPiece embedding size E is tied with the hidden layer size H , i.e., E ≡ H . This decision appears suboptimal for both modeling and practical reasons, as follows.
From a modeling perspective, WordPiece embeddings are meant to learn context-independent representations, whereas hidden-layer embeddings are meant to learn context-dependent representations. As experiments with context length indicate (Liu et al., 2019), the power of BERT-like representations comes from the use of context to provide the signal for learning such context-dependent representations. As such, untying the WordPiece embedding size E from the hidden layer size H allows us to make a more efficient usage of the total model parameters as informed by modeling needs, which dictate that H E. From a practical perspective, natural language processing usually require the vocabulary size V to be large.1 If E ≡ H , then increasing H increases the size of the embedding matrix, which has size
1Similar to BERT, all the experiments in this paper use a vocabulary size V of 30,000.
V ×E. This can easily result in a model with billions of parameters, most of which are only updated sparsely during training.
Therefore, for ALBERT we use a factorization of the embedding parameters, decomposing them into two smaller matrices. Instead of projecting the one-hot vectors directly into the hidden space of size H , we first project them into a lower dimensional embedding space of size E, and then project it to the hidden space. By using this decomposition, we reduce the embedding parameters from O(V × H) to O(V × E + E × H). This parameter reduction is significant when H E. We choose to use the same E for all word pieces because they are much more evenly distributed across documents compared to whole-word embedding, where having different embedding size (Grave et al. (2017); Baevski & Auli (2018); Dai et al. (2019) ) for different words is important.
Cross-layer parameter sharing. For ALBERT, we propose cross-layer parameter sharing as another way to improve parameter efficiency. There are multiple ways to share parameters, e.g., only sharing feed-forward network (FFN) parameters across layers, or only sharing attention parameters. The default decision for ALBERT is to share all parameters across layers. All our experiments use this default decision unless otherwise specified. We compare this design decision against other strategies in our experiments in Sec. 4.5.
Similar strategies have been explored by Dehghani et al. (2018) (Universal Transformer, UT) and Bai et al. (2019) (Deep Equilibrium Models, DQE) for Transformer networks. Different from our observations, Dehghani et al. (2018) show that UT outperforms a vanilla Transformer. Bai et al. (2019) show that their DQEs reach an equilibrium point for which the input and output embedding of a certain layer stay the same. Our measurement on the L2 distances and cosine similarity show that our embeddings are oscillating rather than converging.
Figure 1 shows the L2 distances and cosine similarity of the input and output embeddings for each layer, using BERT-large and ALBERT-large configurations (see Table 1). We observe that the transitions from layer to layer are much smoother for ALBERT than for BERT. These results show that weight-sharing has an effect on stabilizing network parameters. Although there is a drop for both metrics compared to BERT, they nevertheless do not converge to 0 even after 24 layers. This shows that the solution space for ALBERT parameters is very different from the one found by DQE.
Inter-sentence coherence loss. In addition to the masked language modeling (MLM) loss (Devlin et al., 2019), BERT uses an additional loss called next-sentence prediction (NSP). NSP is a binary classification loss for predicting whether two segments appear consecutively in the original text, as follows: positive examples are created by taking consecutive segments from the training corpus; negative examples are created by pairing segments from different documents; positive and negative examples are sampled with equal probability. The NSP objective was designed to improve performance on downstream tasks, such as natural language inference, that require reasoning about the relationship between sentence pairs. However, subsequent studies (Yang et al., 2019; Liu et al., 2019) found NSP’s impact unreliable and decided to eliminate it, a decision supported by an improvement in downstream task performance across several tasks.
We conjecture that the main reason behind NSP’s ineffectiveness is its lack of difficulty as a task, as compared to MLM. As formulated, NSP conflates topic prediction and coherence prediction in a
single task2. However, topic prediction is easier to learn compared to coherence prediction, and also overlaps more with what is learned using the MLM loss.
We maintain that inter-sentence modeling is an important aspect of language understanding, but we propose a loss based primarily on coherence. That is, for ALBERT, we use a sentence-order prediction (SOP) loss, which avoids topic prediction and instead focuses on modeling inter-sentence coherence. The SOP loss uses as positive examples the same technique as BERT (two consecutive segments from the same document), and as negative examples the same two consecutive segments but with their order swapped. This forces the model to learn finer-grained distinctions about discourse-level coherence properties. As we show in Sec. 4.6, it turns out that NSP cannot solve the SOP task at all (i.e., it ends up learning the easier topic-prediction signal, and performs at randombaseline level on the SOP task), while SOP can solve the NSP task to a reasonable degree, presumably based on analyzing misaligned coherence cues. As a result, ALBERT models consistently improve downstream task performance for multi-sentence encoding tasks.

heading: 3.2 MODEL SETUP

text:
We present the differences between BERT and ALBERT models with comparable hyperparameter settings in Table 1. Due to the design choices discussed above, ALBERT models have much smaller parameter size compared to corresponding BERT models.
For example, ALBERT-large has about 18x fewer parameters compared to BERT-large, 18M versus 334M. An ALBERT-xlarge configuration with H = 2048 has only 60M parameters and an ALBERT-xxlarge configuration with H = 4096 has 233M parameters, i.e., around 70% of BERTlarge’s parameters. Note that for ALBERT-xxlarge, we mainly report results on a 12-layer network because a 24-layer network (with the same configuration) obtains similar results but is computationally more expensive.
This improvement in parameter efficiency is the most important advantage of ALBERT’s design choices. Before we can quantify this advantage, we need to introduce our experimental setup in more detail.

heading: 4 EXPERIMENTAL RESULTS

heading: 4.1 EXPERIMENTAL SETUP

text:
To keep the comparison as meaningful as possible, we follow the BERT (Devlin et al., 2019) setup in using the BOOKCORPUS (Zhu et al., 2015) and English Wikipedia (Devlin et al., 2019) for pretraining baseline models. These two corpora consist of around 16GB of uncompressed text. We format our inputs as “[CLS] x1 [SEP] x2 [SEP]”, where x1 = x1,1, x1,2 · · · and x2 = x1,1, x1,2 · · · are two segments.3 We always limit the maximum input length to 512, and randomly generate input sequences shorter than 512 with a probability of 10%. Like BERT, we use a vocabulary size of 30,000, tokenized using SentencePiece (Kudo & Richardson, 2018) as in XLNet (Yang et al., 2019).
2Since a negative example is constructed using material from a different document, the negative-example segment is misaligned both from a topic and from a coherence perspective.
3A segment is usually comprised of more than one natural sentence, which has been shown to benefit performance by Liu et al. (2019).
We generate masked inputs for the MLM targets using n-gram masking (Joshi et al., 2019), with the length of each n-gram mask selected randomly. The probability for the length n is given by
p(n) = 1/n∑N k=1 1/k
We set the maximum length of n-gram (i.e., n) to be 3 (i.e., the MLM target can consist of up to a 3-gram of complete words, such as “White House correspondents”).
All the model updates use a batch size of 4096 and a LAMB optimizer with learning rate 0.00176 (You et al., 2019). We train all models for 125,000 steps unless otherwise specified. Training was done on Cloud TPU V3. The number of TPUs used for training ranged from 64 to 512, depending on model size.
The experimental setup described in this section is used for all of our own versions of BERT as well as ALBERT models, unless otherwise specified.

heading: 4.2 EVALUATION BENCHMARKS

heading: 4.2.1 INTRINSIC EVALUATION

text: To monitor the training progress, we create a development set based on the development sets from SQuAD and RACE using the same procedure as in Sec. 4.1. We report accuracies for both MLM and sentence classification tasks. Note that we only use this set to check how the model is converging; it has not been used in a way that would affect the performance of any downstream evaluation, such as via model selection.

heading: 4.2.2 DOWNSTREAM EVALUATION

text: Following Yang et al. (2019) and Liu et al. (2019), we evaluate our models on three popular benchmarks: The General Language Understanding Evaluation (GLUE) benchmark (Wang et al., 2018), two versions of the Stanford Question Answering Dataset (SQuAD; Rajpurkar et al., 2016; 2018), and the ReAding Comprehension from Examinations (RACE) dataset (Lai et al., 2017). For completeness, we provide description of these benchmarks in Appendix A.3. As in (Liu et al., 2019), we perform early stopping on the development sets, on which we report all comparisons except for our final comparisons based on the task leaderboards, for which we also report test set results. For GLUE datasets that have large variances on the dev set, we report median over 5 runs.

heading: 4.3 OVERALL COMPARISON BETWEEN BERT AND ALBERT

text:
We are now ready to quantify the impact of the design choices described in Sec. 3, specifically the ones around parameter efficiency. The improvement in parameter efficiency showcases the most important advantage of ALBERT’s design choices, as shown in Table 2: with only around 70% of BERT-large’s parameters, ALBERT-xxlarge achieves significant improvements over BERT-large, as measured by the difference on development set scores for several representative downstream tasks: SQuAD v1.1 (+1.9%), SQuAD v2.0 (+3.1%), MNLI (+1.4%), SST-2 (+2.2%), and RACE (+8.4%).
Another interesting observation is the speed of data throughput at training time under the same training configuration (same number of TPUs). Because of less communication and fewer computations, ALBERT models have higher data throughput compared to their corresponding BERT models. If we use BERT-large as the baseline, we observe that ALBERT-large is about 1.7 times faster in iterating through the data while ALBERT-xxlarge is about 3 times slower because of the larger structure.
Next, we perform ablation experiments that quantify the individual contribution of each of the design choices for ALBERT.

heading: 4.4 FACTORIZED EMBEDDING PARAMETERIZATION

text:
Table 3 shows the effect of changing the vocabulary embedding size E using an ALBERT-base configuration setting (see Table 1), using the same set of representative downstream tasks. Under the non-shared condition (BERT-style), larger embedding sizes give better performance, but not by
much. Under the all-shared condition (ALBERT-style), an embedding of size 128 appears to be the best. Based on these results, we use an embedding size E = 128 in all future settings, as a necessary step to do further scaling.

heading: 4.5 CROSS-LAYER PARAMETER SHARING

text:
Table 4 presents experiments for various cross-layer parameter-sharing strategies, using an ALBERT-base configuration (Table 1) with two embedding sizes (E = 768 and E = 128). We compare the all-shared strategy (ALBERT-style), the not-shared strategy (BERT-style), and intermediate strategies in which only the attention parameters are shared (but not the FNN ones) or only the FFN parameters are shared (but not the attention ones).
The all-shared strategy hurts performance under both conditions, but it is less severe for E = 128 (- 1.5 on Avg) compared to E = 768 (-2.5 on Avg). In addition, most of the performance drop appears to come from sharing the FFN-layer parameters, while sharing the attention parameters results in no drop when E = 128 (+0.1 on Avg), and a slight drop when E = 768 (-0.7 on Avg).
There are other strategies of sharing the parameters cross layers. For example, We can divide the L layers into N groups of size M , and each size-M group shares parameters. Overall, our experimental results shows that the smaller the group size M is, the better the performance we get. However, decreasing group size M also dramatically increase the number of overall parameters. We choose all-shared strategy as our default choice.

heading: 4.6 SENTENCE ORDER PREDICTION (SOP)

text:
We compare head-to-head three experimental conditions for the additional inter-sentence loss: none (XLNet- and RoBERTa-style), NSP (BERT-style), and SOP (ALBERT-style), using an ALBERTbase configuration. Results are shown in Table 5, both over intrinsic (accuracy for the MLM, NSP, and SOP tasks) and downstream tasks.
The results on the intrinsic tasks reveal that the NSP loss brings no discriminative power to the SOP task (52.0% accuracy, similar to the random-guess performance for the “None” condition). This allows us to conclude that NSP ends up modeling only topic shift. In contrast, the SOP loss does solve the NSP task relatively well (78.9% accuracy), and the SOP task even better (86.5% accuracy). Even more importantly, the SOP loss appears to consistently improve downstream task performance for multi-sentence encoding tasks (around +1% for SQuAD1.1, +2% for SQuAD2.0, +1.7% for RACE), for an Avg score improvement of around +1%.

heading: 4.7 WHAT IF WE TRAIN FOR THE SAME AMOUNT OF TIME?

text:
The speed-up results in Table 2 indicate that data-throughput for BERT-large is about 3.17x higher compared to ALBERT-xxlarge. Since longer training usually leads to better performance, we perform a comparison in which, instead of controlling for data throughput (number of training steps), we control for the actual training time (i.e., let the models train for the same number of hours). In Table 6, we compare the performance of a BERT-large model after 400k training steps (after 34h of training), roughly equivalent with the amount of time needed to train an ALBERT-xxlarge model with 125k training steps (32h of training).
After training for roughly the same amount of time, ALBERT-xxlarge is significantly better than BERT-large: +1.5% better on Avg, with the difference on RACE as high as +5.2%.

heading: 4.8 ADDITIONAL TRAINING DATA AND DROPOUT EFFECTS

text:
The experiments done up to this point use only the Wikipedia and BOOKCORPUS datasets, as in (Devlin et al., 2019). In this section, we report measurements on the impact of the additional data used by both XLNet (Yang et al., 2019) and RoBERTa (Liu et al., 2019).
Fig. 2a plots the dev set MLM accuracy under two conditions, without and with additional data, with the latter condition giving a significant boost. We also observe performance improvements on the downstream tasks in Table 7, except for the SQuAD benchmarks (which are Wikipedia-based, and therefore are negatively affected by out-of-domain training material).
We also note that, even after training for 1M steps, our largest models still do not overfit to their training data. As a result, we decide to remove dropout to further increase our model capacity. The
plot in Fig. 2b shows that removing dropout significantly improves MLM accuracy. Intermediate evaluation on ALBERT-xxlarge at around 1M training steps (Table 8) also confirms that removing dropout helps the downstream tasks. There is empirical (Szegedy et al., 2017) and theoretical (Li et al., 2019) evidence showing that a combination of batch normalization and dropout in Convolutional Neural Networks may have harmful results. To the best of our knowledge, we are the first to show that dropout can hurt performance in large Transformer-based models. However, the underlying network structure of ALBERT is a special case of the transformer and further experimentation is needed to see if this phenomenon appears with other transformer-based architectures or not.

heading: 4.9 CURRENT STATE-OF-THE-ART ON NLU TASKS

text:
The results we report in this section make use of the training data used by Devlin et al. (2019), as well as the additional data used by Liu et al. (2019) and Yang et al. (2019). We report state-of-the-art results under two settings for fine-tuning: single-model and ensembles. In both settings, we only do single-task fine-tuning4. Following Liu et al. (2019), on the development set we report the median result over five runs.
The single-model ALBERT configuration incorporates the best-performing settings discussed: an ALBERT-xxlarge configuration (Table 1) using combined MLM and SOP losses, and no dropout.
4Following Liu et al. (2019), we fine-tune for RTE, STS, and MRPC using an MNLI checkpoint.
The checkpoints that contribute to the final ensemble model are selected based on development set performance; the number of checkpoints considered for this selection range from 6 to 17, depending on the task. For the GLUE (Table 9) and RACE (Table 10) benchmarks, we average the model predictions for the ensemble models, where the candidates are fine-tuned from different training steps using the 12-layer and 24-layer architectures. For SQuAD (Table 10), we average the prediction scores for those spans that have multiple probabilities; we also average the scores of the “unanswerable” decision.
Both single-model and ensemble results indicate that ALBERT improves the state-of-the-art significantly for all three benchmarks, achieving a GLUE score of 89.4, a SQuAD 2.0 test F1 score of 92.2, and a RACE test accuracy of 89.4. The latter appears to be a particularly strong improvement, a jump of +17.4% absolute points over BERT (Devlin et al., 2019; Clark et al., 2019), +7.6% over XLNet (Yang et al., 2019), +6.2% over RoBERTa (Liu et al., 2019), and 5.3% over DCMI+ (Zhang et al., 2019), an ensemble of multiple models specifically designed for reading comprehension tasks. Our single model achieves an accuracy of 86.5%, which is still 2.4% better than the state-of-the-art ensemble model.

heading: 5 DISCUSSION

text: While ALBERT-xxlarge has less parameters than BERT-large and gets significantly better results, it is computationally more expensive due to its larger structure. An important next step is thus to speed up the training and inference speed of ALBERT through methods like sparse attention (Child et al., 2019) and block attention (Shen et al., 2018). An orthogonal line of research, which could provide additional representation power, includes hard example mining (Mikolov et al., 2013) and more efficient language modeling training (Yang et al., 2019). Additionally, although we have convincing evidence that sentence order prediction is a more consistently-useful learning task that leads to better language representations, we hypothesize that there could be more dimensions not yet captured by the current self-supervised training losses that could create additional representation power for the resulting representations.

heading: ACKNOWLEDGEMENT

text: The authors would like to thank Beer Changpinyo, Nan Ding, Noam Shazeer, and Tomer Levinboim for discussion and providing useful feedback on the project; Omer Levy and Naman Goyal for clarifying experimental setup for RoBERTa; Zihang Dai for clarifying XLNet; Brandon Norick, Emma Strubell, Shaojie Bai, Chas Leichner, and Sachin Mehta for providing useful feedback on the paper; Jacob Devlin for providing the English and multilingual version of training data; Liang Xu, Chenjie Cao and the CLUE community for providing the training data and evaluation benechmark of the Chinese version of ALBERT models.

heading: A APPENDIX

text:
A.1 EFFECT OF NETWORK DEPTH AND WIDTH
In this section, we check how depth (number of layers) and width (hidden size) affect the performance of ALBERT. Table 11 shows the performance of an ALBERT-large configuration (see Table 1) using different numbers of layers. Networks with 3 or more layers are trained by fine-tuning using the parameters from the depth before (e.g., the 12-layer network parameters are fine-tuned from the checkpoint of the 6-layer network parameters).5 Similar technique has been used in Gong et al. (2019). If we compare a 3-layer ALBERT model with a 1-layer ALBERT model, although they have the same number of parameters, the performance increases significantly. However, there are diminishing returns when continuing to increase the number of layers: the results of a 12-layer network are relatively close to the results of a 24-layer network, and the performance of a 48-layer network appears to decline.
A similar phenomenon, this time for width, can be seen in Table 12 for a 3-layer ALBERT-large configuration. As we increase the hidden size, we get an increase in performance with diminishing returns. At a hidden size of 6144, the performance appears to decline significantly. We note that none of these models appear to overfit the training data, and they all have higher training and development loss compared to the best-performing ALBERT configurations.
5If we compare the performance of ALBERT-large here to the performance in Table 2, we can see that this warm-start technique does not help to improve the downstream performance. However, it does help the 48-layer network to converge. A similar technique has been applied to our ALBERT-xxlarge, where we warm-start from a 6-layer network.
A.2 DO VERY WIDE ALBERT MODELS NEED TO BE DEEP(ER) TOO?
In Section A.1, we show that for ALBERT-large (H=1024), the difference between a 12-layer and a 24-layer configuration is small. Does this result still hold for much wider ALBERT configurations, such as ALBERT-xxlarge (H=4096)?
The answer is given by the results from Table 13. The difference between 12-layer and 24-layer ALBERT-xxlarge configurations in terms of downstream accuracy is negligible, with the Avg score being the same. We conclude that, when sharing all cross-layer parameters (ALBERT-style), there is no need for models deeper than a 12-layer configuration.
A.3 DOWNSTREAM EVALUATION TASKS
GLUE GLUE is comprised of 9 tasks, namely Corpus of Linguistic Acceptability (CoLA; Warstadt et al., 2018), Stanford Sentiment Treebank (SST; Socher et al., 2013), Microsoft Research Paraphrase Corpus (MRPC; Dolan & Brockett, 2005), Semantic Textual Similarity Benchmark (STS; Cer et al., 2017), Quora Question Pairs (QQP; Iyer et al., 2017), Multi-Genre NLI (MNLI; Williams et al., 2018), Question NLI (QNLI; Rajpurkar et al., 2016), Recognizing Textual Entailment (RTE; Dagan et al., 2005; Bar-Haim et al., 2006; Giampiccolo et al., 2007; Bentivogli et al., 2009) and Winograd NLI (WNLI; Levesque et al., 2012). It focuses on evaluating model capabilities for natural language understanding. When reporting MNLI results, we only report the “match” condition (MNLI-m). We follow the finetuning procedures from prior work (Devlin et al., 2019; Liu et al., 2019; Yang et al., 2019) and report the held-out test set performance obtained from GLUE submissions. For test set submissions, we perform task-specific modifications for WNLI and QNLI as described by Liu et al. (2019) and Yang et al. (2019).
SQuAD SQuAD is an extractive question answering dataset built from Wikipedia. The answers are segments from the context paragraphs and the task is to predict answer spans. We evaluate our models on two versions of SQuAD: v1.1 and v2.0. SQuAD v1.1 has 100,000 human-annotated question/answer pairs. SQuAD v2.0 additionally introduced 50,000 unanswerable questions. For SQuAD v1.1, we use the same training procedure as BERT, whereas for SQuAD v2.0, models are jointly trained with a span extraction loss and an additional classifier for predicting answerability (Yang et al., 2019; Liu et al., 2019). We report both development set and test set performance.
RACE RACE is a large-scale dataset for multi-choice reading comprehension, collected from English examinations in China with nearly 100,000 questions. Each instance in RACE has 4 candidate answers. Following prior work (Yang et al., 2019; Liu et al., 2019), we use the concatenation of the passage, question, and each candidate answer as the input to models. Then, we use the representations from the “[CLS]” token for predicting the probability of each answer. The dataset consists of two domains: middle school and high school. We train our models on both domains and report accuracies on both the development set and test set.
A.4 HYPERPARAMETERS
Hyperparameters for downstream tasks are shown in Table 14. We adapt these hyperparameters from Liu et al. (2019), Devlin et al. (2019), and Yang et al. (2019).

title: ALBERT: A LITE BERT FOR SELF-SUPERVISED LEARNING OF LANGUAGE REPRESENTATIONS


summaries: This paper proposes a unified architecture in the context of multi-task learning where they demonstrate that training four tasks (with a variety of modalities like image, text, and videos) together results in about three times compressed model, while maintaining the performance similar to their respective individually trained models. The major components of this archtiecutre are (1) peripheral networks: used to encode the domain specific input into feature representations. (2) central neural processor: a fully attention based encoder-decoder model similar to the Transformer networks which encodes the spatio-temporal information. Further, this paper suggests that their unified architecture enables to perform decent on unseen tasks during its training. In this paper they test such scenarios on video captioning and video question answering.

summaries: This paper describes the "OmniNet" architecture, which is essentially a transformer to convert any 2-dimensional (time and spatial) input into a sequence of output tokens. The idea is that a single model (the "Central Neural Processor") would learn to perform multiple tasks on multiple inputs at the same time. This seems like a reasonable approach, and would allow a single model to be able to process text, image, as well as video or potentially speech input just the same. Dedicated modality-specific "input peripherals" would "normalize" the data appropriately, and a "start token" seems to provide the model with information on what type of input data is to follow.

summaries: The authors propose an extended and unifying learning architecture – OmniNet- based on transformer, which tackles tasks with various modalities such as images, text and videos. This is attained with a spatio-temporal cache mechanism, that can both capture and store temporal information and spatial information. In addition, this proposed framework supports asynchronous multi-task learning with pre-trained neural networks on different modalities. OmiNet’s generalization capability is illustrated and demonstrated with experiments. The proposed model has multiple peripheral networks each majoring on one unique modality of data. These peripheral networks project input data into the same shared format/space that can be uniformly processed in a central neural processor working like a CPU. The central neural processor uses self-attention and RNN for temporal encoding and spatial encoding. The output of the encoding components are stored in temporal and spatial caches respectively. The two caches are then used as input to the spatial temporal decoder for different tasks.

abstractText: Transformer is a popularly used neural network architecture, especially for language understanding. We introduce an extended and unified architecture that can be used for tasks involving a variety of modalities like image, text, videos, etc. We propose a spatio-temporal cache mechanism that enables learning spatial dimension of the input in addition to the hidden states corresponding to the temporal input sequence. The proposed architecture further enables a single model to support tasks with multiple input modalities as well as asynchronous multi-task learning, thus we refer to it as OmniNet. For example, a single instance of OmniNet can concurrently learn to perform the tasks of part-of-speech tagging, image captioning, visual question answering and video activity recognition. We demonstrate that training these four tasks together results in about three times compressed model while retaining the performance in comparison to training them individually. We also show that using this neural network pre-trained on some modalities assists in learning unseen tasks such as video captioning and video question answering. This illustrates the generalization capacity of the self-attention mechanism on the spatio-temporal cache present in OmniNet.

authors:
P. Anderson
X. He
C. Buehler
D. Teney
M. Johnson
S. Gould
L. Zhang

title: Bottom-up and top-down attention for image captioning and visual question answering

venue: IEEE/CVF Conference on Computer Vision and Pattern Recognition,

authors:
Dzmitry Bahdanau
Kyunghyun Cho
Yoshua Bengio

title: Neural machine translation by jointly learning to align and translate

venue: In 3rd International Conference on Learning Representations,

authors:
Eric Battenberg
Jitong Chen
Rewon Child
Adam Coates
Yashesh Gaur
Yi Li
Hairong Liu
Sanjeev Satheesh
David Seetapun
Anuroop Sriram
Zhenyao Zhu

title: Exploring neural transducers for end-to-end speech recognition

authors:
Yaran Chen
Dongbin Zhao
Le Lv
Qichao Zhang

title: Multi-task learning for dangerous object detection in autonomous driving

venue: Information Sciences,

authors:
Kyunghyun Cho
Bart van Merrienboer
Caglar Gulcehre
Dzmitry Bahdanau
Fethi Bougares
Holger Schwenk
Yoshua Bengio

title: Learning phrase representations using RNN encoder– decoder for statistical machine translation

authors:
Ronan Collobert
Jason Weston

title: A unified architecture for natural language processing: Deep neural networks with multitask learning

venue: In Proceedings of the 25th International Conference on Machine Learning,

authors:
Daxiang Dong
Hua Wu
Wei He
Dianhai Yu
Haifeng Wang

title: Multi-task learning for multiple language translation

venue: In ACL,

authors:
Yash Goyal
Tejas Khot
Douglas Summers-Stay
Dhruv Batra
Devi Parikh

title: Making the V in VQA matter: Elevating the role of image understanding in Visual Question Answering

venue: In Conference on Computer Vision and Pattern Recognition (CVPR),

authors:
K. He
X. Zhang
S. Ren
J. Sun

title: Deep residual learning for image recognition

authors:
Lanqing Hu
Meina Kan
Shiguang Shan
Xilin Chen

title: Duplex generative adversarial network for unsupervised domain adaptation

venue: In The IEEE Conference on Computer Vision and Pattern Recognition (CVPR),

authors:
Melvin Johnson
Mike Schuster
Quoc V. Le
Maxim Krikun
Yonghui Wu
Zhifeng Chen
Nikhil Thorat
Fernanda Viégas
Martin Wattenberg
Greg Corrado
Macduff Hughes
Jeffrey Dean

title: Google’s multilingual neural machine translation system: Enabling zero-shot translation

venue: Transactions of the Association for Computational Linguistics,

authors:
Lukasz Kaiser
Aidan N. Gomez
Noam Shazeer
Ashish Vaswani
Niki Parmar
Llion Jones
Jakob Uszkoreit

title: One model to learn them all

authors:
Jin-Hwa Kim
Sang-Woo Lee
Donghyun Kwak
Min-Oh Heo
Jeonghee Kim
Jung-Woo Ha
Byoung-Tak Zhang

title: Multimodal residual learning for visual qa

authors:
Jin-Hwa Kim
Jaehyun Jun
Byoung-Tak Zhang

title: Bilinear attention networks

authors:
Diederik P. Kingma
Jimmy Ba

title: Adam: A method for stochastic optimization

authors:
Kalpesh Krishna
Shubham Toshniwal
Karen Livescu

title: Hierarchical multitask learning for ctcbased speech recognition

authors:
H. Kuehne
H. Jhuang
E. Garrote
T. Poggio
T. Serre

title: HMDB: a large video database for human motion recognition

venue: In Proceedings of the International Conference on Computer Vision (ICCV),

authors:
Jie Lei
Licheng Yu
Mohit Bansal
Tamara L. Berg

title: TVQA: localized, compositional video question answering

authors:
Tsung-Yi Lin
Michael Maire
Serge Belongie
James Hays
Pietro Perona
Deva Ramanan
Piotr Dollár
C. Lawrence Zitnick

title: Microsoft coco: Common objects in context

venue: Computer Vision – ECCV

authors:
Thang Luong
Quoc V. Le
Ilya Sutskever
Oriol Vinyals
Lukasz Kaiser

title: Multi-task sequence to sequence learning

venue: In International Conference on Learning Representations,

authors:
Mitchell Marcus
Grace Kim
Mary Ann Marcinkiewicz
Robert MacIntyre
Ann Bies
Mark Ferguson
Karen Katz
Britta Schasberger

title: The penn treebank: Annotating predicate argument structure

venue: In Proceedings of the Workshop on Human Language Technology,

authors:
Volodymyr Mnih
Adrià Puigdomènech Badia
Mehdi Mirza
Alex Graves
Tim Harley
Timothy P. Lillicrap
David Silver
Koray Kavukcuoglu

title: Asynchronous methods for deep reinforcement learning

venue: In Proceedings of the 33rd International Conference on International Conference on Machine Learning - Volume 48,

authors:
Jiquan Ngiam
Aditya Khosla
Mingyu Kim
Juhan Nam
Honglak Lee
Andrew Y. Ng

title: Multimodal deep learning

venue: In Proceedings of the 28th International Conference on International Conference on Machine Learning,

authors:
Adam Paszke
Sam Gross
Soumith Chintala
Gregory Chanan
Edward Yang
Zachary DeVito
Zeming Lin
Alban Desmaison
Luca Antiga
Adam Lerer

title: Automatic differentiation in PyTorch

venue: In NIPS Autodiff Workshop,

authors:
Benjamin Recht
Christopher Re
Stephen Wright
Feng Niu

title: Hogwild: A lock-free approach to parallelizing stochastic gradient descent

authors:
M.L. Seltzer
J. Droppo

title: Multi-task learning in deep neural networks for improved phoneme recognition

venue: IEEE International Conference on Acoustics, Speech and Signal Processing,

authors:
Rico Sennrich
Barry Haddow
Alexandra Birch

title: Neural machine translation of rare words with subword units. In Proceedings of the 54th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 1715–1725, Berlin, Germany, August 2016

venue: Association for Computational Linguistics. doi: 10.18653/v1/P16-1162. URL https://www.aclweb. org/anthology/P16-1162

authors:
Noam Shazeer
Mitchell Stern

title: Adafactor: Adaptive learning rates with sublinear memory

venue: cost. CoRR,

authors:
Rui Shu
Hung H. Bui
Hirokazu Narui
Stefano Ermon

title: A dirt-t approach to unsupervised domain adaptation

authors:
Karen Simonyan
Andrew Zisserman

title: Two-stream convolutional networks for action recognition in videos

authors:
Drahomı́ra “johanka” Spoustová
Jan Hajič
Jan Raab
Miroslav Spousta

title: Semi-supervised training for the averaged perceptron POS tagger

venue: In Proceedings of the 12th Conference of the European Chapter of the ACL (EACL

authors:
Ilya Sutskever
Oriol Vinyals
Quoc V Le

title: Sequence to sequence learning with neural networks

authors:
Christian Szegedy
Sergey Ioffe
Vincent Vanhoucke

authors:
Ashish Vaswani
Noam Shazeer
Niki Parmar
Jakob Uszkoreit
Llion Jones
Aidan N Gomez
Ł ukasz Kaiser
Illia Polosukhin

title: Attention is all you need

authors:
Zhanpeng Zhang
Ping Luo
Chen Change Loy
Xiaoou Tang

title: Facial landmark detection by deep multi-task learning

venue: In In ECCV

authors:
Wei Zhao
Benyou Wang
Jianbo Ye
Min Yang
Zhou Zhao
Ruotian Luo
Yu Qiao

title: A multi-task learning approach for image captioning

venue: In Proceedings of the Twenty-Seventh International Joint Conference on Artificial Intelligence,

heading: 1 INTRODUCTION

text:
Transformer (Vaswani et al., 2017) is currently one of the best performing models for any sequence transduction tasks, especially those involving natural language. It is originally designed for a single task at a time. In fact, most of the generic deep learning architectures (Wu et al., 2016; Szegedy et al., 2016; Battenberg et al., 2017) that have been designed and developed are able to learn, albeit very well, a single task and handle one task specific input domain like image, text or audio. Furthermore with these models, we often rely on the generalization capability of the trained network to guarantee performance on unseen examples. Transfer learning (Shu et al., 2018; Hu et al., 2018) is another popular paradigm used to adapt the model to learn a related task with similar input domain. The success of neural networks across these challenges is known to be due to their ability in learning effective representations of the data. For example, the self-attention mechanism in Transformers can capture the global temporal dependence in sequential data very well. Naturally, the question arises whether we can extend these architectures, like the Transformer, to be able to learn shared representations from multiple input domains and to be able attend on these representations to perform a multitude of tasks concurrently.
The research into multi-task models that learn to solve varied tasks across a multitude of input domains is not new. Work done in Ngiam et al. (2011) demonstrates an architecture capable of learning a shared representation across audio and video modalities. Similarly in Collobert & Weston (2008) a convolutional architecture has been designed to support a variety of NLP tasks. However, most of these architectures are designed to learn specific set of tasks with known input domains. To the best of our knowledge, there does not exist a single unified architecture that works out of the box for any combination of multi-modal inputs.
To address this gap, we extend Transformer towards a unified architecture, namely OmniNet, which enables a single model to support tasks with multiple input modalities and asynchronous multi-task learning. We consider that most real-life data like image, text, speech, video, etc. is a direct conjunction of spatial and temporal components. Therefore, we employ a spatio-temporal cache mechanism to learn a shared representation of the input data across the spatial (space) and temporal (time) di-
mension. Using a generalized encode() function, OmniNet can process and store spatio-temporal representation for each of the input domains and then decode() predictions across a multitude of tasks. In our experiments, we train a single instance of the OmniNet to solve a number of tasks spanning multiple multi-domains such as part-of-speech tagging, image captioning, visual question answering and video activity recognition. To make our work reproducible, open to scrutiny and further development, we will open source a demonstration of our system implemented using Pytorch (Paszke et al., 2017).

heading: 2 RELATED WORK

text: Multi-task learning has been extensively studied in the literature, with applications to a wide set of problems ranging from natural language processing (NLP) (Collobert & Weston, 2008; Johnson et al., 2017; Dong et al., 2015) to speech recognition (Seltzer & Droppo, 2013; Krishna et al., 2018) to vision (Zhang et al., 2014; Chen et al., 2018; Anderson et al., 2018). It has also found its use in a combination of diverse tasks like image captioning and text translation and parsing (Luong et al., 2016; Zhao et al., 2018). However, most of these architectures assume the set of tasks to be known in advance. Similarly, multi-modal learning has been essential for solving a broad range of interesting problems such as Visual Question Answering (Kim et al., 2016; 2018) and Video Question Answering (Lei et al., 2018). Again, the state-of-the-art models are highly specific to the objective in hand and not easily adaptable to different tasks or domains. (Kaiser et al., 2017) proposed MultiModel architecture for learning multiple tasks but lacks support for multi-modal tasks with more than one input domains such as visual question answering.

heading: 3 PROPOSED MODEL

text:
We propose a unified architecture, namely OmniNet, to enable learning multi-modal tasks with multiple input domains and support generic multi-tasking for any set of tasks. The OmniNet architecture consists of multiple sub-networks, called peripheral networks, connected to a common central neural network called the Central Neural Processor (CNP) (Figure 1). Each peripheral network is used to encode the domain specific input into feature representations. In this work, we describe image, text and video peripherals (Section 3.1). One can add more, say speech peripheral, depending on the task. The output representation of a peripheral network is always a spatio-temporal tensor x ∈ Rt×s×dmodel , where t & s are the temporal and spatial dimensions of the input respectively, and dmodel is the model dimension input to the CNP.
The spatio-temporal representations generated by the peripheral networks corresponding to each input domain are then processed by the CNP. The CNP uses fully attention based encoder-decoder (Bahdanau et al., 2015; Sutskever et al., 2014; Cho et al., 2014) model for sequence transduction similar to the Transformer architecture (Vaswani et al., 2017), which is the state-of-the-art for multiple language modeling tasks (Section 3.2). During the encoding stage, the CNP implements a
generic encode(x, D) function to first process and store the spatio-temporal representations of the input, where x ∈ Rt×s×dmodel is the spatio-temporal tensor produced by the peripheral networks and D ∈ Z : 0 ≤ D < Dlen is the domain id and Dlen is the max number of domains supported by the CNP. The encode() function is called multiple times, once for each multi-modal input from respective peripheral. During the decoding stage, a decode(yshifted, τ ) function is used to decode predictions as softmax probabilities, where yshifted ∈ ZN−1 are the target outputs shifted one timestep to the right, N is the length of the output sequence; τ ∈ Z : 0 ≤ τ < τlen is task id and τlen is the total number of supported tasks. The decoding step is similar to Vaswani et al. (2017), modified to incorporate a two-step attention mechanism over spatial and temporal cache.

heading: 3.1 PERIPHERAL NETWORKS

text:
First, we elaborate on how we support multiple input domains using peripheral networks. A peripheral network can use a pre-trained model from existing literature to ultimately encode a given domain input to a standardized feature representation x ∈ Rt×s×dmodel , where t & s are the temporal and spatial dimensions of the input respectively, and dmodel is the model dimension input to the Central Neural Processor. Here we detail text and vision peripherals and one can add more peripherals or alter the peripheral design depending on the task.
Vision peripheral: This peripheral uses a convolutional neural network to encode image and video inputs in the tasks. For an image of dimension h × w × nc, this peripheral down-samples in to h′ ×w′ × n′c, where h,w, nc are the height, width and number of input channels respectively. For a video, each frame is input to the peripheral to produce F ×h′×w′×n′c, where F is the total number of frames in the video. The encoding vectors are then projected to dimension dmodel using a fully connected layer. The output is then reshaped into a spatio-temporal tensor of x ∈ Rt×h′w′×dmodel , where t = 1 for an image and t = F for a video. In our experiments, we use the pre-trained ResNet-152 model, a variant of ResNet (He et al., 2016) consisting of 152 convolutional layers. We remove the final fully connected and avg-pooling layers to generate spatial feature representations for a given image/video.
Language peripheral: The Language peripheral uses byte-pair encoding (Sennrich et al., 2016) to generate subwords for a given input sentence. The subwords are passed to an embedding layer to generate subword embeddings of dimension demb and projected to dimension dmodel using a fully connected layer. The output is then reshaped into a spatio-temporal tensor x ∈ Rt×1×dmodel , where t equal to number of subwords in the input sentence. As we do not have any spatial dimension in textual data, the spatial dimension of x from a Language peripheral is always 1. In our experiments, We used pre-trained subword embeddings with demb = 300 and vocab size = 25000 from Heinzerling & Strube (2018), which includes pre-trained subword embeddings of over 275 languages, to initialize the weights of the embedding matrix.

heading: 3.2 CENTRAL NEURAL PROCESSOR (CNP)

text:
To process the spatio-temporal information in the input data, the CNP implements a spatial cache Cs, temporal cache Ct and a link array L. The spatial and temporal cache and the link array are a list of elements, initialized as empty before the encoding process. During the encoding stage, an encode() routine takes as input, the tensor x generated from the peripheral and corresponding domain/peripheral id D. This function processes the spatial and temporal information in the input x and stores them into the spatial cache Cs and the temporal cache Ct , respectively and stores their dimensions t& s in the link array. For a given task, this encode() routine is calledK times, whereK is the number of inputs in the task. Note that these inputs can belong to same or different domains.
Encode (x, D): For a given input x ∈ Rt×s×dmodel and domain identifier D, the encode() routine is described in Algorithm 1. Since inputs can come from multiple peripherals, the algorithm first concatenates the input with the domain embedding to ensure a domain-aware encoding of the input (Steps 2 to 3). Steps 4 to 7 process the spatial information in x by unrolling the time dimension and adding these unrolled vectors into the spatial cache. Steps 8 to 10 process the temporal information in x by averaging the spatial dimension of x and then passing the averaged tensor to a self-attention based TemporalEncoder. This TemporalEncoder is similar to the encoder used in Vaswani et al. (2017) as shown in Figure 2 is used to calculate temporal embeddings of the input sequence. The output from the TemporalEncoder is appended to the temporal cache.
Algorithm 1 encode(): Encodes spatial and temporal representations into spatial and temporal cache Require: x ∈ Rt×s×dmodel , D, Cs, L, Ct
1: L← L ∪ (t→ s) 2: Demb ← EmbedLayer(D) 3: x← FC(Concat(x,Demb), dmodel) 4: if s > 1 then 5: S ← Reshape(x, (ts, dmodel)) {where, output S = [S1, . . . , Sts] s.t. Si ∈ Rdmodel is a spatial feature vector.} 6: Cs ← Cs ∪ [S1, . . . , Sts] {Append spatial representations to spatial cache} 7: end if 8: T ← ( ∑s i=1 x[:, i, :])/s 9: T ← TemporalEncoder(T ) {where, output T = [T1, . . . , Tt] s.t. Tj ∈ Rdmodel is the encoding of temporal dimension in x.}
10: Ct ← Ct ∪ [T1, . . . , Tt] {Append temporal representations to temporal cache}
The above encoding routine keeps appending spatio-temporal information to Ct & Cs for each input xk ∈ Rtk×sk×dmodel . Note the superscript k to denote correspondence to k-th input of the task, where k ∈ 1, . . . ,K. After K calls, we have the temporal cache Ct = [T1, . . . , TR], where R =∑K r=1 tr; the spatial cache Cs = [S1, . . . , SP ], where P = { ∑ p tp ∗ sp : p ∈ 1, . . . ,K ∧ sp > 1} and the link array L = [(t1 → s1), . . . , (tK → sK)]. Note that Cs can also be empty in case the encode() is only called with inputs with sk = 1∀k. Next, we use the decode() routine to generate predictions as softmax probabilities.
Decode (yshifted, τ ): The architecture of the decode() function is shown in Figure 2. The decode() takes as argument the output labels yshifted shifted one time step to the right, a task id τ and gen-
erates predictions by attending from the spatial and temporal cache. The decode() function is structured similar to the decoder used in the Transformer architecture Vaswani et al. (2017) and jointly attends on the vectors stored in the temporal and spatial cache. Similar to Vaswani et al. (2017), the decoding first starts by attending over the output embeddings using masked multi-head scaled dot product attention. The attention layer for the temporal cache uses scaled dot-product attention with multiple heads as specified in Vaswani et al. (2017). Attention layer for the spatial cache, uses gated multi-head attention to attend over the elements of the spatial cache. For inputs with both time and space dimension (e.g. video), we want the spatial attention layer to attend more on frames which have relatively high attention scores in the temporal cache attention layer. Therefore, the attention score output from the temporal cache multi-head attention layerA ∈ Rnh×N×R, is used to calculate the tensor G ∈ Rnh×N×P used for gating the attention score output in the spatial attention layer, where nh is the number of heads in multi-head attention as described in Vaswani et al. (2017). The tensor G is calculated using A & L as detailed in Algorithm 2. Given1 Q: the matrix of queries, K: keys of dimension dk and V : values of dimension dv , the scaled dot-product attention for the spatial layer is modified as:
Attention(Q,K, V,G) = ( Softmax ( QKT√ dv ) G ) V (1)
In order to use the same CNP for multiple tasks with varying output vocabularies, we use multiple output embedding layers OutputEmbedLayer1, . . . , OutputEmbedLayerτlen , to generate the output embeddings for each task. At the final layer, we use multiple (FC + Softmax)1, . . . , (FC + Softmax)τlen classification layers for each task. We also calculate a task embedding vector using τ and always start decoding using the task embedding vector.
Algorithm 2 Calculate G using output scores from temporal attention and link array Require: L, A
1: idx← 0 2: for each t, s in L do 3: G← [] 4: if s > 0 then 5: A′ ← A[:, :, idx : idx+ t] 6: A′ ← Expand(A′, (nh, N, t, s)) {where Expand(tensor, dimension) expands tensor according to a given dimension} 7: A′ ← Reshape(A′, (nh, N, ts)) 8: G← G ∪A′ {Append the respective temporal attention scores to G} 9: end if
10: idx← idx+ t 11: end for 12: G← Stack(G) {Stack the list of tensors to construct tensor G of dimension (nh, N, P )}

heading: 3.3 MULTI-TASK LEARNING

text: In order to train a single a model simultaneously on mutiple tasks we used the HogWild training approach as described in Recht et al. (2011). Similar to the approach described in Mnih et al. (2016), the main process holds a global copy of the model. We create separate worker processes for each task, where each process maintains a local copy of a model. At each training iteration, each process starts by synchronizing its local model with the global copy. This is done through forward and backward propagation on its local copy and then copying the locally computed gradients to the global model asynchronously. Each process then calls the global model optimizer asynchronously to update the weights of the global model. Instead of storing the model in CPU as in Mnih et al. (2016) we always store the local copies across multiple GPUs.

heading: 4 TASKS AND SETUP

text:
To evaluate the effectiveness of our proposed framework for tasks spanning diverse modalities, we choose a set covering all possible spatio-temporal data archetypes: Image Captioning, Part-of-
1For brevity, we reuse the notations of Vaswani et al. (2017) in this description
Speech (POS) tagging, Visual Question Answering (VQA) and Video-activity Recognition. Each of these tasks explores a unique potential spatio-temporal configuration of the input containing different values of t and s, where t and s are the temporal and spatial dimensions of the input respectively. This enables us to perform a comprehensive study of the multi-modal and multi-input capabilities of the system. We further elaborate on the properties of these tasks below.
For training, we always use cross-entropy loss with Adam optimizer (Kingma & Ba, 2015) and schedule the learning rate using Noam scheduler (Shazeer & Stern, 2018) similar to (Vaswani et al., 2017) 2. In the vision peripheral, we freeze the layers of pretrained ResNet model. Remaining parts of the architecture (peripherals and CNP) are kept trainable. In this section, we provide details on the datasets used and the model setup for each of these tasks.
Part-of-speech (POS)Tagging: To illustrate the task with only temporal modality (t > 1 & s = 1, where t and s are the temporal and spatial dimensions of the input respectively), we consider POS tagging problem. Given an input sequence of words, the model should produce a sequence of POS tags corresponding to each word. We use Penn Tree-bank3 (Marcus et al., 1994) which contains gold annotations on English WSJ articles by experts. During the encoding stage, each input sentence is processed by the language peripheral to generate a spatio-temporal tensor x ∈ Rt×1×dmodel , where t is the sequence length of subwords in the input. The CNP encode() function is then used to encode x into the temporal cache. Note that spatial cache is empty for text inputs as s = 1. Therefore, the decoding stage is same as that of Transformers to predict the sequence of POS tags.
Image Captioning: This task represents the ones with inputs containing only spatial modality (t = 1 & s > 1). The captioning model is required to predict text caption for a given image. We use the MSCOCO 2014 dataset (Lin et al., 2014) for training and present results on the COCO validation set. During the encoding stage, the input image is resized to 224× 224 and processed by the vision peripheral containing pre-trained ResNet-152 to produce image embeddings x ∈ R1×49×dmodel . x is then input to the encode() function which populates corresponding spatial and temporal cache. The decoding stage uses decode() function with output vocabulary size 25000 to generate the captions.
Visual Question Answering: For the task with inputs from multiple domain, such that each contains either spatial or temporal modality (either t > 1 & s = 1, or t = 1 & s > 1 for each input), we choose the task of visual question answering. Given a question over an image as inputs, the model is supposed to predict the correct answer label. We use the recently introduced VQA v2.0 dataset (Goyal et al., 2017) for this purpose and perform evaluation on the VQA test-dev set. All the images are resized to dimension 224× 224 before training. The encoding stage of this task utilizes two peripherals: the vision peripheral is used to generate a tensor x1 ∈ R1×49×dmodel for the input image. The language peripheral is used to encode the questions into x2 ∈ Rt×1×dmodel , where t is equal to the length of the subwords in the question. The encode() function is the called two times, first with x1 and second with x2 as input. Finally, the decode() with output vocabulary size 3500 is to generate the answers as softmax probabilities in a single decoding step.
Video Activity Recognition: For tasks which contain both spatial and temporal modality in a single input (t > 1 & s > 1), we consider the action recognition task on videos. For this purpose, we use the HMDB dataset (Kuehne et al., 2011). The dataset consists of over 5000 short length clips of real life actions with 51 classes. We present our results on train-test split 1. We use 16 frames per video and resize each of them to 224× 224. During the encoding stage, each frame of the video is passed through the vision peripheral to cumulatively generate a video encoding x ∈ R16×49×dmodel which is then used as input to the encode() function. Finally, the decode() with output vocabulary size 51 is to predict the action as softmax probabilities in a single decoding step.

heading: 5 RESULTS AND DISCUSSION

text:
We present the evaluation on (a) Tasks of various modalities illustrated in Section 4 (b) Multi-tasking setup for these tasks (Table 1) (c) Reuse of the multi-task model for an unseen task (Figure 3). In addition, we also provide some ablation studies on the architecture (Table 2)
2The hyperparameter values used for nh, dmodel, N Layers, dk, dv are same as that specified in Transformer base model (Vaswani et al., 2017).
3https://catalog.ldc.upenn.edu/LDC99T42; We use splits 0-18 as training, 19-21 as development and 22-24 as test sets
Performance of proposed architecture on individual tasks: We choose a set of four tasks with diverse input modalities and combinations as described in previous Section 4. We train the OmniNet model independently across each of the above tasks. Each of the tasks demonstrates unique capabilities of this generic architecture. More specifically, in Table 1 we compare our results with the following state-of-the-art4:- POS tagging: Spoustová et al. (2009); image captioning & VQA: Anderson et al. (2018) and HMDB: Simonyan & Zisserman (2014). It is important to note that we do not perform any hyper-parameter optimization. We believe that, with more computational power, fine tuning the hyperparameters towards these tasks should result in comparable or even improved performance to the state-of-the-art. These results can indeed be used as a baseline for any future work which aims at using a single architecture across various possible spatio-temporal archetypes. It is interesting to note that the model is extensible to a new domain without any modification to the CNP as long as one can add a specific peripheral to convert domain inputs into spatio-temporal tensors. This aspect of the architecture makes it applicable to several popular multi-modal tasks.
Effect of training a diverse set of tasks together: We trained two multi-task models: (1) MULT-3 (POS+VQA+Captioning) and (2) MULT-4 (POS+VQA+Captioning+HMDB), using hogwild approach. While the MULT-3 model attains similar and sometimes better performance, the final MULT-4 model attains slightly reduced performance, when compared to the independent task scores. We believe this is due to the skewness in the size of the HMDB dataset containing only 5000 training samples. However as a tradeoff, adding the HMDB task shows interesting zero-shot results demonstrated below. Using a multi-task model also results in three times reduction in the total number of parameters. That is, when a separate model is used for each task, we have a total of over 450× 106 parameters. Whereas during multi-tasking since a single model is shared, we have a total of over 149 × 106 parameters, while achieving similar performance. Interestingly, the model is able to attend on spatio-temporal components of the inputs from different tasks and concurrently generate predictions across them, thus demonstrating the generalization capability of our architecture.
Towards zero-shot learning: reuse of pre-trained network for unseen task: Sharing representations across multiple tasks provides the benefit to transfer of useful knowledge across multiple domains. Since, image and video are processed by the same vision peripheral, we conducted an experiment to see whether our model pre-trained on all the four tasks (MULT-4) can perform video
4Since most of these tasks are popular challenges, we compare with state-of-the-art which are generically applicable for the respective task instead of the challenge dataset.
captioning and video question-answering without any explicit training on these tasks i.e. zero-shot learning. The results of the evaluation on randomly picked instances from the HMDB test split 1 are shown in Figure 3. Interestingly, the model performs quite well on related actions that were present in the COCO and VQA training set; such as captions related to horse riding and baseball; or questions related to concepts present in VQA. Without training on any video captioning & video QA instance, the model could use the trained information from image captioning, image QA (VQA) and video action recognition (HMDB) apply them on videos to generate meaningful predictions, hence demonstrating the capability of the model to transfer knowledge across related multi-modal tasks. However, on concepts that are not present in the trained datasets, the model either describes the environment in the video or replaces with alternate known concepts. This case study, although not comprehensive, shows the capability of the model to learn shared representations and ability to transfer knowledge across domains. We believe that adding more tasks and domains will lead to more interesting zero-shot learning results in future across a wide range of problems.
Impact of individual architectural components: In order to support different input domains, our architecture introduces spatial cache and link array components to the original Transformer architecture (which only consists of mechanisms to handle temporal data). We conducted an ablation study on each of these components to verify their importance across various tasks as shown in Table 2. The ablation was conducted on the independent (IND) as well the multi-tasking model (MULT-4). The second row ablates the link array from our architecture i.e. removing the multiplication of G in Equation 1. The link array was designed to assist in tasks with inputs such as video, containing both spatial as well as temporal modality in a single input. The total number of spatial components becomes very large as number of frames in the video increases, thereby making it difficult to attend on various spatial regions throughout the video. Using link array the spatial attention layer can attend more on specific important frames in the video. Therefore, removal of link array leads to a huge reduction in performance in HMDB compared to other tasks as they do not have both spatio-temporal modalities for any single input. Removal of spatial cache, on the other hand, has significant effect on performance across all tasks containing spatial modality. Since, image captioning contains primarily spatial modality and hence the BLEU drops significantly after ablation. As other tasks utilize the temporal cache for prediction, in the multi-task setting the captioning task learns to utilize the spatial average of the image stored in the temporal cache for prediction and hence retains some performance after ablation of the spatial cache. On the other hand when trained independently on captioning, the network learns to utilize the information in the spatial cache only, and hence the score drops to zero after ablation. VQA leverages both, spatial information from image and temporal information from question, retains some performance from the use of temporal cache. Note that, POS tagging task is not affected by ablation of any of the components since it only has temporal modality in the input.

heading: 6 CONCLUSIONS AND FUTURE WORK

text: We present a unified neural network architecture OmniNet capable of learning tasks with multiple inputs of varying modalities. The architecture can be further adopted for multi-task learning across any set of tasks containing spatio-temporal data. Sharing one model across multiple tasks also results in a significant reduction in the total number of parameters. We further demonstrate that this shared model can learn robust representations from various spatio-temporal inputs which are reusable for unseen tasks. We believe that this proposed architecture has wide applicability to any task with spatio-temporal inputs. To extend its usability, we would like to introduce new peripherals supporting more domains such as speech. We are also keen on exploring other aspects to the data beyond temporal and spatial dimensions such as graphs and relational data. Further, it would be interesting to investigate scheduling mechanisms for optimizing the multi-tasking framework.