OmniNet: A unified architecture for multi-modal multi-task learning

1 INTRODUCTION . 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 ) . 2 RELATED WORK . 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 . 3 PROPOSED MODEL . 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 . 3.1 PERIPHERAL NETWORKS . 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 . 3.2 CENTRAL NEURAL PROCESSOR ( CNP ) . 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 ) }

Summary: 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.