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: The paper proposes to combine self-attention and object-centric representations to be able to solve complex reasoning tasks. The author(s) also propose a self-supervised loss based on masking of the objects' representations to implicitly learn the dynamics, further enhancing the performance on downstream tasks. The results of the proposed method are very impressive, especially on the challenging counterfactuals reasoning task.