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