Attention over Learned Object Embeddings Enables Complex Visual Reasoning

1 Introduction . Despite the popularity of artificial neural networks , a body of recent work has focused on their limitations as models of cognition and reasoning . Experiments with dynamical reasoning datasets such as CLEVRER [ 41 ] , CATER [ 12 ] , and ACRE [ 44 ] show that neural networks can fail to adequately reason about the spatio-temporal , compositional or causal structure of visual scenes . On CLEVRER , where models must answer questions about the dynamics of colliding objects , previous experiments show that neural networks can adequately describe the video , but fail when asked to predict , explain , or consider counterfactual possibilities . Similarly , on CATER , an object-tracking task , models have trouble tracking the movement of objects when they are hidden in a container . Finally , on ACRE , a dataset testing for causal inference , popular models only learned correlations between visual scenes and not the deeper causal logic . Failures such as these on reasoning ( rather than perception ) problems have motivated the adoption of pipeline-style approaches that combine a general purpose neural network ( such as a convolutional block ) with a task-specific module that builds in the core logic of the task . For example , on CLEVRER the NS-DR method [ 41 ] applies a hand-coded symbolic logic engine ( that has the core logic of CLEVRER built-in ) to the outputs of a “ perceptual ” neural front-end , achieving better results than neural network baselines , particularly on counterfactual and explanatory problems . One Model Code : https : //github.com/deepmind/deepmind-research/tree/master/ object_attention_for_reasoning . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . limitation of these pipeline approaches , however , is that they are typically created with a single problem or problem domain in mind , and may not apply out-of-the-box to other related problems . For example , to apply NS-DR to CATER , the entire symbolic module needs to be rewritten to handle the new interactions and task logic of CATER : the custom logic to handle collisions and object removal must be replaced with new custom logic to handle occlusions and grid-resolution , and these changes require further modifications to the perceptual front-end to output data in a new format . This brittleness is not exclusive to symbolic approaches . While Hungarian-matching between object embeddings may be well-suited for object-tracking tasks [ 45 ] , it is not obvious how it would help for causal inference tasks . Here , we describe a more general neural-network-based approach to visual spatio-temporal reasoning problems , which does not rely on task-specific integration of modular components . In place of these components , our model relies on three key aspects : • Self-attention to effectively integrate information over time • Soft-discretization of the input at the most informative level of abstraction – above pixels and local features , and below entire frames—corresponding approximately to ‘ objects ’ • Self-supervised learning , i.e . requiring the model to infer masked out objects , to extract more information about dynamics from each sample . While many past models have applied each individual ingredient separately ( including on the tasks we study ) , we show that it is the combination of all three ingredients in the right way that allows our model to succeed . The resulting model , which we call Aloe ( Attention over Learned Object Embeddings ) , outperforms both pipeline and neural-network-based approaches on three different task domains designed to test physical and dynamical reasoning from pixel inputs . We highlight our key results here : • CLEVRER ( explanatory , predictive , and counterfactual reasoning ) : Aloe achieves significantly higher accuracy than both more task-specific , modular approaches , and previous neural network methods on all question types . On counterfactual questions , thought to be most challenging for neural-only architectures , we achieve 75 % vs 46 % accuracy for more specialised methods . • CATER ( object-permanence ) : Aloe achieves accuracy exceeding or matching other current models . Notably , the strongest alternative models were expressly designed for object-tracking , whereas our architecture is applicable without modification to other reasoning tasks as well . • ACRE ( causal-inference “ beyond the simple strategy of inducing causal relationships by covariation ” [ 44 ] ) : Overall , Aloe achieves 94 % vs the 67 % accuracy achieved by the top neuro-symbolic model . On the most challenging tasks , we achieve , for “ backward-blocking ” inference , 94.48 % ( vs 16.06 % by the best modular , neuro-symbolic systems ) , and , for “ screenoff ” inference , 98.97 % ( vs 0.00 % by a CNN-BERT baseline ) . As we have emphasized , the previous best performing models for each task all contain task-specific design elements , whereas Aloe can be applied to all the tasks without modification . On CLEVRER , we also show that Aloe matches the performance of the previous best models with 40 % less training data , which demonstrates that our approach is data-efficient as well as performant . 2 Methods . A guiding motivation for the design of Aloe is the converging evidence for the value of self-attention mechanisms operating on a finite sequences of discrete entities . Written language is inherently discrete and hence is well-suited to self-attention-based approaches . In other domains , such as raw audio or vision , it is less clear how to leverage self-attention . We hypothesize that the application of self-attention-based models to visual tasks could benefit from an approximate ‘ discretization ’ process , and determining the right level of discretization is an important choice that can significantly affect model performance . At the finest level , data could simply be discretized into pixels ( as is already the case for most machine-processed visual data ) . Pixels are too fine-grained for many applications , however—for one , the memory required to support self-attention across all pixels is prohibitive . Partly for this reason , coarser representations , such as the downsampled “ hyper-pixel ” outputs of a convolutional network , are often used instead ( e.g . [ 27 , 43 ] ) . In the case of videos , previous work considered even coarser discretization schemes , such as frame or subclip level representations [ 35 ] . The neuroscience literature , however , suggests that biological visual systems infer and exploit the existence of objects , rather than spatial or temporal blocks with artificial boundaries [ 5 , 30 , 32 ] . Because objects are the atomic units of physical interactions , it makes sense to discretize on the level of objects . Numerous object segmentation algorithms have been proposed [ 15 , 19 , 29 ] . We chose to use MONet , an unsupervised object segmentation algorithm [ 2 ] . Because MONet is unsupervised , we can train it directly in our domain of interest without the need for object segmentation labels . We emphasize that our choice of MONet is an implementation detail , and in Appendix B , we show that our framework of attention over learned object embeddings also works with other object-segmentation schemes . We also do not need to place strong demands on the object segmentation algorithm , e.g . for it to produce aligned output or to have a built-in dynamics model . To segment each frame into object representations , MONet uses a recurrent attention network to obtain a set of No “ object attention masks ” ( No is a fixed parameter ) . Each attention mask represents the probability that any given pixel belongs to that mask ’ s object . The pixels assigned to the mask are encoded into latent variables with means µti ∈ Rd , where i indexes the object slot and t the frame . These means are used as the object embeddings in Aloe . More details are provided in Appendix A.1 . The self-attention component is a transformer model [ 37 ] operating on a sequence of vectors in Rd : the object representations µti for all t and i , a trainable vector CLS ∈ Rd used to generate classification results ( analogous to the CLS token in BERT [ 9 ] ) , and ( for CLEVRER ) the embedded words wi from the question ( and choice for multiple choice questions ) . For the object representations µti and word embeddings wi , we append a two-dimensional one-hot vector to µti and wi to indicate whether the input is a word or an object . Because the transformer is shared between the modalities , information can flow between objects and words to solve the task , as we show in Section 3.1 . We pass this sequence of vectors through a transformer with NT layers . All inputs are first projected ( via a linear layer and ReLU activation ) to RNH×d , where NH is the number of self-attention heads . We add a relative sinusoidal positional encoding at each layer of the transformer to give the model knowledge of the word and frame order [ 7 ] . The transformed value of CLS is passed through an MLP ( with one hidden layer of size NH ) to generate the final answer . A schema of our architecture is shown in Figure 1 . Note that in the model presented above ( which we call global attention ) , the transformer sees no distinction between objects of different frames ( other than through the position encoding ) . Another intuitive choice , which we call hierarchical attention , is to have one transformer acting on the objects of each frame independently , and another transformer acting on the concatenated outputs of the first transformer ( this temporal division of input data is commonly used , e.g . in [ 35 ] ) . In pseudo-code , global attention can be expressed as out = transformer ( reshape ( objects , [ B , F * N , D ] ) and hiearchical attention as out = transformer1 ( reshape ( objects , [ B * F , N , D ] ) ) out = transformer2 ( reshape ( out , [ B , F , N * D ] ) ) . We study the importance of global attention ( objects as the atomic entities ) vs hierarchical attention ( objects , and subsequently frames as the atomic entities ) . The comparison is shown in Table 1 . 2.1 Self-supervised learning . We explored whether self-supervised learning could improve the performance of Aloe beyond the benefits conveyed by object-level representation , i.e . in ways that support the model ’ s interpretation of scene dynamics rather than just via improved perception of static observations . Our approach is inspired by the loss used in BERT [ 9 ] , where a transformer model is trained to predict certain words that are masked from the input . In our case , we mask object embeddings , and train the model to infer the content of the masked object representations using its knowledge of unmasked objects . Concretely , during training , we multiply each MONet latent µti by a masking indicator , mti ∈ { 0 , 1 } . Let µ′ti be the transformed value of mtiµti after passing through the transformer . We expect the transformer to understand the underlying dynamics of the video , so that the masked out slot µti could be predicted from µ′ti . To guide the transformer in learning effective representations capable of this type of dynamics prediction , we add an auxiliary loss : auxiliary loss = ∑ t , i τtil ( f ( µ ′ ti ) , µ ) , where f is a learned linear mapping to Rd , l a loss function , and τti ∈ { 0 , 1 } are one-hot indicator variables identifying the prediction targets ( not necessarily just the masked out entries , since the prediction targets could be a subset of the masked out entries ) . We propagate gradients only to the parameters of f and the transformer and not to the learned word and CLS embeddings . This auxiliary loss is added to the main classification loss with weighting λ , and both losses are minimized simultaneously by the optimizer . We do not pretrain the model with only the auxiliary loss . We tested two different loss functions for l , an L2 loss and a contrastive loss ( formulas given in Appendix A.2 ) , and six different masking schemes ( settings of mti and τti ) , as illustrated in Figure 2 . This exploration was motivated by the observation that video inputs at adjacent timesteps are highly correlated in a way that adjacent words are not . We thus hypothesized that BERT-style prediction of adjacent words might not be optimal . A different masking strategy , in which prediction targets are separated from the context by more than a single timestep , may stimulate capacity in the network to acquire knowledge that permits context-based unrolls and better long-horizon predictions . The simplest approach would be to set mti = 1 uniformly at random across t and i , fixing the expected proportion of the mti set to 1 ( schema b in Figure 2 ) . The targets would simply be the unmasked slots , τti = 1 −mti . One potential problem with this approach is that multiple objects could be masked out in a single frame . MONet can unpredictably switch object-to-slot assignments multiple times in a single video . If multiple slots are masked out , the transformer can not determine with certainty which missing object to assign to each slot . Thus , the auxiliary loss could penalize the model even if it predicted all the objects correctly . To avoid this problem , we also try constraining the mask such that exactly one slot is masked out per frame ( schema a ) . To pose harder prediction challenges , we can add a buffer between the context ( where mti = 1 ) and the infilling targets ( where τti = 1 ) . For t in this buffer zone , both mti = 0 and τti = 0 ( schemas c–f ) . We choose a single cutoff T randomly , and we set mti = 0 for t < T and mti = 1 for t ≥ T . In the presence of this buffer , we compared prediction ( where the context is strictly before the targets ; schema c , d ) versus infilling ( where the context surrounds the targets ; schema e , f ) . We also compared setting the targets as individual objects ( schema c , e ) versus targets as all objects in the scene ( schema d , f ) . We visually inspect the efficacy of this self-supervised loss in encouraging better representations ( beyond improvements of scores on tasks ) in Appendix D .

Summary: A transformer achieves state of the art in three video+text reasoning tasks. A pre-trained unsupervised method obtains object representations per frame; object representations across frames, word representations (question+options) and a special token is fed to a transformer and a prediction is made using a linear layer on top of the transformed special token. Question answers provide supervision and masking objects provide self-supervision. Ablation experiments are performed: they show the usefulness of object-centric representations and, albeit less so, self-supervision. Model attention and errors seem sensible. It shows that a fully end-to-end network can solve causal reasoning tasks; which until now required symbolic modules such as program generation and execution engines or physical engines.