S3TA: A Soft, Spatial, Sequential, Top-Down Attention Model

1 INTRODUCTION . Traditional RL agents and image classifiers rely on some combination of convolutional and fully connected components to gradually process input information and arrive at a set of policy or class logits . This sort of architecture is very effective , but does not lend itself to easy understanding of how decisions are made , what information is used and why mistakes are made . Previous efforts to visualize deep RL agents ( Greydanus et al . ( 2017 ) ; Zahavy et al . ( 2016 ) ; Wang et al . ( 2015 ) ) focus on generating saliency maps to understand the magnitude of policy changes as a function of a perturbation of the input . This can uncover some of the “ attended ” regions , but may be difficult to interpret . For example , it can ’ t reveal certain types of behavior when the agent makes decisions based on components absent from a frame . Our mechanism provides a more direct interpretation by making the attention a core part of the network . In this work we present a soft , spatial , sequential and top-down attention model ( S3TA , pronounced SETA ) . This model enables us to build agents and classifiers that actively select important , taskrelevant information from visual inputs by sequentially querying and receiving compressed querydependent summaries to generate appropriate outputs . To do this , the model generates attention maps , which can uncover some of underlying decision process used to solve the task . By observing and analyzing the resulting attention maps we can make educated guesses at how the system solves a task and where and why it might be failing . In the RL domain , we observe that the attention focuses on the key components of each level : tracking the region ahead of the player , focusing on enemies and important moving objects . In supervised learning , we observed that the attention sequentially focuses on different portions of the input to build up confidence in a classification or resolve ambiguity between different class labels . We also find that our model maintains competitive performance on both learning paradigms while providing interpretability . 2 MODEL . Our model , outlined in Figure 1 , queries a large input tensor through an attention mechanism and uses the returned compressed answer ( a low dimensional summary of the input ) to produce its output . We refer to this full query-answer system as an attention head . Our system can implement multiple attention heads by producing multiple queries and receiving multiple answers . An observation X ∈ RH×W×C at time t ( here an RGB frame of height H and width W ) is passed through a “ vision core ” . The vision core is a multi-layer convolutional network visθ followed by a recurrent layer with state svis ( t ) such as a ConvLSTM ( Shi et al . ( 2015 ) ) , which produces an output tensor Ovis ∈ Rh×w×c : Ovis , svis ( t ) = visθ ( X ( t ) , svis ( t− 1 ) ) ( 1 ) The vision core output is then split along the channel dimension into two tensors : the “ Keys ” tensor K ∈ Rh×w×Ck and the “ Values ” tensor V ∈ Rh×w×Cv , with c = CV + CK . To the keys and values tensors we concatenate a spatial basis — a fixed tensor S ∈ Rh×w×CS which encodes spatial locations ( see below for details ) . A recurrent neural network ( RNN ) with parameters φ produces N queries , one for each attention head . The RNN sends its state sRNN from the previous time step t− 1 into a “ Query Network ” . The query network Qψ is a multi-layer perceptron ( MLP ) with parameters ψ whose output is reshaped into N query vectors qn of size Ck + CS such that they match the channel dimension of K : q1 ... qN = Qψ ( sRNN ( t− 1 ) ) ( 2 ) Similar to Vaswani et al . ( 2017 ) , we take the inner product between each query vector qn and all spatial locations in the keys tensor K to form the n-th attention logits map Ãn ∈ Rh×w : Ãni , j = ∑ c qnc Ki , j , c ( 3 ) where K ∈ Rh×w×Ck+CS is the concatenation along the channel dimension of K and S. We then take the spatial softmax to form the final normalized attention map An : Ani , j = exp ( Ãni , j ) ∑ i , j exp ( Ã n i , j ) ( 4 ) Each attention map An is broadcast along the channel dimension of the values tensor V , pointwise multiplied with it and then summed across space to produce the n-th answer vector an ∈ R1×1×Cv+Cs : anc = ∑ i , j Ani , jVi , j , c ( 5 ) where V ∈ Rh×w×Cv+CS is the concatenation along the channel dimension of V and S Finally , the N answer vectors an , and the N query form the input to the RNN core to produce the next RNN state sRNN ( t ) and output o ( t ) for this time step : o ( t ) , sRNN ( t ) = RNNφ ( a1 , ... , an , q1 , ... , qn , sRNN ( t− 1 ) ) ( 6 ) The exact details for each of the networks , outputs and states are given in Section 4 and the Appendix . It is important to emphasize several points about the proposed model . First , the model is fully differentiable due to the use of soft-attention and can be trained using back-propagation . Second , the query vectors are a function of the RNN core state alone and not the observation — this allows for a “ top-down ” mechanism where the RNN can actively query the input for task-relevant information rather than having to filter out large amounts of information . Third , the spatial sum ( equation 5 ) is a severe spatial bottleneck , which forces the system to make the attention maps in such a way that information is not “ blurred ” out during summation . The summation of the values tensor of shape h×w×Cv to an answer of shape 1×1×Cv is invariant to permutation of spatial position , which emphasizes the need for the spatial basis . Due to the spatial structure being lost during the spatial summation , the only way the RNN core can know and reason about spatial positions is by using the channels coming from the spatial basis 1 . We postulate that the query and answer structure can have different “ modes ” — the system can ask “ where ” ( “ what ” ) something is by sending out a query with zeros in the spatial channels of the query and non-zeros in the channels corresponding to the keys ( which are input dependent ) . It can then read the answer from the spatial channels , localizing the object of interest . Conversely it can ask “ what is in this 1The vision core might have some ability to produce information regarding absolute spatial positioning , but due to its convolutional structure it is limited . particular location ” by zeroing out the content channels of the query and putting information on the spatial channels , reading the content channels of the answer and ignoring the spatial channels . This is not a dichotomy as the two can be mixed ( e.g . “ find enemies in the top left corner ” ) , but it does point to an interesting “ what ” and “ where ” separation , which we discuss in Section 4.1.5 . 2.1 THE SPATIAL BASIS . The spatial basis S ∈ Rh×w×CS such that the channels at each location i , j encode information about the spatial position . Adding this information into the values of the tensor allows some spatial information to be maintained after the spatial summation ( equation 5 ) removes the structural information . Following Vaswani et al . ( 2017 ) and Parmar et al . ( 2018 ) we use a Fourier basis type of representation . Each channel ( u , v ) of S is an outer product of two Fourier basis vectors . We use both odd and even basis functions with several frequencies . For example , with two even functions one channel of S with spatial frequencies u and v would be : Si , j , ( u , v ) = cos ( πui/h ) cos ( πvj/w ) ( 7 ) where u , v are the spatial frequencies in this channel , i , j are spatial locations in the tensor and h , w are correspondingly the height and width of the tensor . We produce all the outer products such that the number of channels in S is ( U + V ) 2 where U and V are the number of spatial frequencies we use for the even and odd components ( 4 for both throughout this work , so 64 channels in total ) . The spatial basis can also be learned as another parameter of the model — while we tested this in some cases we did not observe that this makes a big difference in performance and for brevity this is not done in this work . 3 RELATED WORK . There is a vast literature in recurrent attention models . They have been applied with some success to question-answering datasets ( Hermann et al. , 2015 ) , text translation ( Vaswani et al. , 2017 ; Bahdanau et al. , 2014 ) , video classification and captioning ( Shan & Atanasov , 2017 ; Li et al. , 2017 ) , image classification and captioning ( Mnih et al. , 2014 ; Chung & Cho , 2018 ; Fu et al. , 2017 ; Ablavatski et al. , 2017 ; Xiao et al. , 2015 ; Zheng et al. , 2017 ; Wang et al. , 2017 ; Xu et al. , 2015 ; Ba et al. , 2014 ) , text classification ( Yang et al. , 2016 ; Shen & Lee , 2016 ) , generative models ( Parmar et al. , 2018 ; Zhang et al. , 2018 ; Kosiorek et al. , 2018 ) , object tracking ( Kosiorek et al. , 2017 ) , and reinforcement learning ( Choi et al. , 2017 ) . These attention mechanisms can be grouped by whether they use hard attention ( e.g . Mnih et al . ( 2014 ) ; Ba et al . ( 2014 ) ; Malinowski et al . ( 2018 ) ) or soft attention ( e.g . Bahdanau et al . ( 2014 ) ) and whether they explicitly parameterize an attention window ( e.g . Jaderberg et al . ( 2015 ) ; Shan & Atanasov ( 2017 ) ) or use a weighting mechanism ( e.g . Vaswani et al . ( 2017 ) ; Hermann et al . ( 2015 ) ) . Our work introduces a novel architecture which builds on existing methods . We use a soft key , query , and value type of attention similar to Vaswani et al . ( 2017 ) and Parmar et al . ( 2018 ) , but instead of doing “ self ” -attention where the queries come from the input ( together with the keys and values ) we have a different , top-down source for them . This enables the system to be both state/context dependent and input dependent . Furthermore the output of the attention model is highly compressed and has no spatial structure ( other than the one preserved using the spatial basis ) , unlike in “ self ” attention where each pixel attends to every other pixel and the structure is preserved . Finally , we apply the attention sequentially in time similar to Xu et al . ( 2015 ) but with a largely different attention mechanism . Of existing models , the MAC model ( Hudson & Manning , 2018 ) is the closest to ours . There are several differences between our model and MAC . First , MAC was built to solve CLEVR ( Johnson et al. , 2017 ) ; major parts of it are geared for that dataset . Specifically the “ control ” unit is built to expect a guiding question for the reasoning process — this may not always exist , such as in the case of RL or classic supervised learning where the systems needs to come up with its own queries to produce the required output . Another difference is the use of a pre-trained ResNet-101 ( Wang et al. , 2017 ) as the visual backend ; we train the visual core to co-adapt with the top-down mechanism such that it learns to produce useful keys and values for different queries . Finally , MAC does not use a spatial basis . It can still reason about space to some extent through the fully connected layers , but there is not a clear separation between space and content as in our model .

Summary: This paper proposed a soft, spatial, sequential, top-down attention model which enable the agent and classifier actively select important, task relative information to generate the appropriate output. Given the observations, the proposed method uses a convLSTM to produce the key and value tensor. Different from multi-head attention, the query vector is generated in a top-down fashion. The authors proposed to augment the spatial feature with Fourier bases which is similar to previous work. The authors verify the proposed method on both reinforcement learning and supervised learning. On reinforcement learning, the proposed method outperformed the feedforward baseline and LSTM baseline. On reinforcement learning task, the proposed method achieves compelling result with more interpretable attention map that shows the model's decision.