SPACE: Unsupervised Object-Oriented Scene Representation via Spatial Attention and Decomposition

1 INTRODUCTION . One of the unsolved key challenges in machine learning is unsupervised learning of structured representation for a visual scene containing many objects with occlusion , partial observability , and complex background . When properly decomposed into meaningful abstract entities such as objects and spaces , this structured representation brings many advantages of abstract ( symbolic ) representation to areas where contemporary deep learning approaches with a global continuous vector representation of a scene have not been successful . For example , a structured representation may improve sample efficiency for downstream tasks such as a deep reinforcement learning agent ( Mnih et al. , 2013 ) . It may also enable visual variable binding ( Sun , 1992 ) for reasoning and causal inference over the relationships between the objects and agents in a scene . Structured representations also provide composability and transferability for better generalization . Recent approaches to this problem of unsupervised object-oriented scene representation can be categorized into two types of models : scene-mixture models and spatial-attention models . In scenemixture models ( Greff et al. , 2017 ; 2019 ; Burgess et al. , 2019 ; Engelcke et al. , 2019 ) , a visual scene is explained by a mixture of a finite number of component images . This type of representation provides flexible segmentation maps that can handle objects and background segments of complex morphology . However , since each component corresponds to a full-scale image , important physical features of objects like position and scale are only implicitly encoded in the scale of a full image and further disentanglement is required to extract these useful features . Also , since it does not explicitly reflect useful inductive biases like the locality of an object in the Gestalt principles ( Koffka , 2013 ) , ∗Visiting Student at Rutgers University . Authors named inside { } equally contributed . Correspondance to zxlin.cs @ gmail.com and sjn.ahn @ gmail.com . the resulting component representation is not necessarily a representation of a local area . Moreover , to obtain a complete scene , a component needs to refer to other components , and thus inference is inherently performed sequentially , resulting in limitations in scaling to scenes with many objects . In contrast , spatial-attention models ( Eslami et al. , 2016 ; Crawford & Pineau , 2019 ) can explicitly obtain the fully disentangled geometric representation of objects such as position and scale . Such features are grounded on the semantics of physics and should be useful in many ways ( e.g. , sample efficiency , interpretability , geometric reasoning and inference , transferability ) . However , these models can not represent complex objects and background segments that have too flexible morphology to be captured by spatial attention ( i.e . based on rectangular bounding boxes ) . Similar to scene-mixture models , previous models in this class show scalability issues as objects are processed sequentially . In this paper , we propose a method , called Spatially Parallel Attention and Component Extraction ( SPACE ) , that combines the best of both approaches . SPACE learns to process foreground objects , which can be captured efficiently by bounding boxes , by using parallel spatial-attention while decomposing the remaining area that includes both morphologically complex objects and background segments by using component mixtures . Thus , SPACE provides an object-wise disentangled representation of foreground objects along with explicit properties like position and scale per object while also providing decomposed representations of complex background components . Furthermore , by fully parallelizing the foreground object processing , we resolve the scalability issue of existing spatial attention methods . In experiments on 3D-room scenes and Atari game scenes , we quantitatively and qualitatively compare the representation of SPACE to other models and show that SPACE combines the benefits of both approaches in addition to significant speed-ups due to the parallel foreground processing . The contributions of the paper are as follows . First , we introduce a model that unifies the benefits of spatial-attention and scene-mixture approaches in a principled framework of probabilistic latent variable modeling . Second , we introduce a spatially parallel multi-object processing module and demonstrate that it can significantly mitigate the scalability problems of previous methods . Lastly , we provide an extensive comparison with previous models where we illustrate the capabilities and limitations of each method . 2 THE PROPOSED MODEL : SPACE . In this section , we describe our proposed model , Spatially Parallel Attention and Component Extraction ( SPACE ) . The main idea of SPACE , presented in Figure 1 , is to propose a unified probabilistic generative model that combines the benefits of the spatial-attention and scene-mixture models . 2.1 GENERATIVE PROCESS . SPACE assumes that a scene x is decomposed into two independent latents : foreground zfg and background zbg . The foreground is further decomposed into a set of independent foreground objects zfg = { zfgi } and the background is also decomposed further into a sequence of background segments zbg = zbg1 : K . While our choice of modeling the foreground and background independently worked well empirically , for better generation , it may also be possible to condition one on the other . The image distributions of the foreground objects and the background components are combined together with a pixel-wise mixture model to produce the complete image distribution : p ( x|zfg , zbg ) = αp ( x|zfg ) ︸ ︷︷ ︸ Foreground + ( 1− α ) K∑ k=1 πk p ( x|zbgk ) ︸ ︷︷ ︸ Background . ( 1 ) Here , the foreground mixing probability α is computed as α = fα ( zfg ) . This way , the foreground is given precedence in assigning its own mixing weight and the remaining is apportioned to the background . The mixing weight assigned to the background is further sub-divided among the K background components . These weights are computed as πk = fπk ( z bg 1 : k ) and ∑ k πk = 1 . With these notations , the complete generative model can be described as follows . p ( x ) = ∫∫ p ( x|zfg , zbg ) p ( zbg ) p ( zfg ) dzfgdzbg ( 2 ) We now describe the foreground and background models in more detail . Foreground . SPACE implements zfg as a structured latent . In this structure , an image is treated as if it were divided into H ×W cells and each cell is tasked with modeling at most one ( nearby ) object in the scene . This type of structuring has been used in ( Redmon et al. , 2016 ; Santoro et al. , 2017 ; Crawford & Pineau , 2019 ) . Similar to SPAIR , in order to model an object , each cell i is associated with a set of latents ( zpresi , z where i , z depth i , z what i ) . In this notation , z pres is a binary random variable denoting if the cell models any object or not , zwhere denotes the size of the object and its location relative to the cell , zdepth denotes the depth of the object to resolve occlusions and zwhat models the object appearance and its mask . These latents may then be used to compute the foreground image component p ( x|zfg ) which is modeled as a Gaussian distribution N ( µfg , σ2fg ) . In practice , we treat σ2fg as a hyperparameter and decode only the mean image µ fg . In this process , SPACE reconstructs the objects associated to each cell having zpresi = 1 . For each such cell , the model uses the z what i to decode the object glimpse and its mask and the glimpse is then positioned on a full-resolution canvas using zwherei via the Spatial Transformer ( Jaderberg et al. , 2015 ) . Using the object masks and z depth i , all the foreground objects are combined into a single foreground mean-image µfg and the foreground mask α ( See Appendix D for more details ) . SPACE imposes a prior distribution on these latents as follows : p ( zfg ) = H×W∏ i=1 p ( zpresi ) ( p ( zwherei ) p ( z depth i ) p ( z what i ) ) zpresi ( 3 ) Here , only zpresi is modeled using a Bernoulli distribution while the remaining are modeled as Gaussian . Background . To model the background , SPACE implements zbgk , similar to GENESIS , as ( z m k , z c k ) where zmk models the mixing probabilities πk of the components and z c k models the RGB distribution p ( x|zbgk ) of the kth background component as a Gaussian N ( µ bg i , σ 2 bg ) . The following prior is imposed upon these latents . p ( zbg ) = K∏ k=1 p ( zck|zmk ) p ( zmk |zm < k ) ( 4 ) 2.2 INFERENCE AND TRAINING . Since we can not analytically evaluate the integrals in equation 2 due to the continuous latents zfg and zbg1 : K , we train the model using a variational approximation . The true posterior on these variables is approximated as follows . p ( zbg1 : K , z fg|x ) ≈ q ( zfg|x ) K∏ k=1 q ( zbgk |z bg < k , x ) ( 5 ) This is used to derive the following ELBO to train the model using the reparameterization trick and SGD ( Kingma & Welling , 2013 ) . L ( x ) = Eq ( zfg , zbg|x ) [ log p ( x|zfg , zbg ) − K∑ k=1 DKL ( q ( z bg k |z bg < k , x ) ‖ p ( z bg k |z bg < k ) ) − H×W∑ i=1 DKL ( q ( z fg i |x ) ‖ p ( zfgi ) ) ] ( 6 ) See Appendix B for the detailed decomposition of the ELBO and the related details . Parallel Inference of Cell Latents . SPACE uses mean-field approximation when inferring the cell latents , so zfgi = ( z pres i , z where i , z depth i , z what i ) for each cell does not depend on other cells . q ( zfg|x ) = H×W∏ i=1 q ( zpresi |x ) ( q ( zwherei |x ) q ( zdepthi |x ) q ( zwhati |zwherei , x ) ) zpresi ( 7 ) As shown in Figure 1 , this allows each cell to act as an independent object detector , spatially attending to its own local region in parallel . This is in contrast to inference in SPAIR , where each cell ’ s latents auto-regressively depend on some or all of the previously traversed cells in a row-major order i.e. , q ( zfg|x ) = ∏HWi=1 q ( zfgi |zfg < i , x ) . However , this method becomes prohibitively expensive in practice as the number of objects increases . While Crawford & Pineau ( 2019 ) claim that these lateral connections are crucial for performance since they model dependencies between objects and thus prevent duplicate detections , we challenge this assertion by observing that 1 ) due to the bottom-up encoding conditioning on the input image , each cell should have information about its nearby area without explicitly communicating with other cells , and 2 ) in ( physical ) spatial space , two objects can not exist at the same position . Thus , the relation and interference between objects should not be severe and the mean-field approximation is a good choice in our model . In our experiments , we verify empirically that this is indeed the case and observe that SPACE shows comparable detection performance to SPAIR while having significant gains in training speeds and efficiently scaling to scenes with many objects . Preventing Box-Splitting . If the prior for the bounding box size is set to be too small , then the model could split a large object by multiple bounding boxes and when the size prior is too large , the model may not capture small objects in the scene , resulting in a trade-off between the prior values of the bounding box size . To alleviate this problem , we found it sometimes helpful to introduce an auxiliary loss which we call the boundary loss . In the boundary loss , we construct a boundary of thickness b pixels along the borders of each glimpse . Then , we restrict an object to be inside this boundary and penalize the model if an object ’ s mask overlaps with the boundary area . Thus , the model is penalized if it tries to split a large object by multiple smaller bounding boxes . A detailed implementation of the boundary loss is mentioned in Appendix C .

Summary: The paper proposes SPACE: a generative latent variable models for scene decomposition (foreground / background separation and object bounding box prediction). The authors state the following contributions relative to prior work in this space: 1) ability to simultaneously perform foreground/background segmentation and decompose the foreground into distinct object bounding box predictions, 2) a parallel spatial attention mechanism that improves the speed of the architecture relative to the closest prior work (SPAIR), 3) a demonstration through qualitative results that the approach can segment into foreground objects elements that remain static across observations (e.g. the key in Montezuma's Revenge).