R-MONet: Region-Based Unsupervised Scene Decomposition and Representation via Consistency of Object Representations

1 INTRODUCTION . In recent years , supervised object detection and segmentation ( He et al . ( 2017 ) ; Ren et al . ( 2015 ) ; Fan et al . ( 2019 ) ; Liao et al . ( 2001 ) ; Lin et al . ( 2017 ) ; Ronneberger et al . ( 2015 ) ) have made great progress with the extensive human labels . However , these supervised methods are still unable to take the advantage of massive unlabeled vision data . Unsupervised learning of scene representation starts to become a key challenge in computer vision . The breakthrough ( Burgess et al . ( 2019 ) ; Greff et al . ( 2019 ) ; Eslami et al . ( 2016 ) ; Crawford & Pineau ( 2019 ) ; Engelcke et al . ( 2020 ) , Greff et al . ( 2017 ) ; Van Steenkiste et al . ( 2018 ) ; Pathak et al . ( 2016 ) ; Lin et al . ( 2020 ) ) in the unsupervised scene decomposition and representation learning proves that a complex visual scene containing many objects can be properly decomposed without human labels . It proves that there is still much useful information that can be discovered in those unlabeled data . Recent approaches to address the unsupervised scene decomposition and representation learning can be categorized into two groups : models which explicitly acquire disentangled position and scale ( i.e . bounding boxes ) representation of objects ( Eslami et al . ( 2016 ) ; Crawford & Pineau ( 2019 ) ; Lin et al . ( 2020 ) ) and models implicitly encode objects ’ geometric representation into segmentation masks or entangle it with object appearance representations ( Burgess et al . ( 2019 ) ; Greff et al . ( 2019 ) ; Engelcke et al . ( 2020 ) ; Greff et al . ( 2017 ) ; Van Steenkiste et al . ( 2018 ) ) . In the former type of models , the scene is explicitly encoded into the object-oriented spatial encoding and appearance encoding . A decoder will generate the scene with explicitly defined object encoding for representation learning . This type of models can not use rectangular bounding boxes to fully represent complex objects with flexible morphology . In the other type of models , the scene is decomposed into a finite number of object segmentation masks which can better represent complex objects with its pixel-topixel alignment . However , this type of models only use segmentation masks as the pixel-wise object mixture weights . They do not utilize the geometric information in the segmentation masks and still entangle object position and appearance representations in the scene generation step . Also , those models tend to decompose the entire scene in the image which does not use the locality benefit of objects . Inspired by the observation that foreground segmentation masks and bounding boxes both contain object geometric information and should be consistent with each other , a method called R-MONet ( Region-based Multiple Object Net ) is proposed in this paper . R-MONet uses the spirit of MONet ( Burgess et al . ( 2019 ) ) and S4Net ( Fan et al . ( 2019 ) ) by using a single stage , non-iterative network ( spatial attention module ) for generating object geometric representations in both bounding boxes and segmentation masks . Then , a variational autoencoder ( VAE ) ( Kingma & Welling ( 2013 ) ) is used for encoding object appearance representations and regenerating the scene for training . To ensure the consistency between bounding boxes and foreground segmentation masks , the bounding boxes generated from spatial attention module is supervised with the pseudo bounding box generated by Multi-Otsu thresholding method ( Liao et al . ( 2001 ) ) on foreground segmentation masks . More than that , the foreground instance segmentation is only performed in the bounding box area instead of the full image to take advantage of the spatial locality and make scene generation less complex . The contributions of this paper are : - We introduce an effective single stage , non-iterative framework to generate object geometric representations in both bounding boxes and segmentation masks for unsupervised scene decomposition and representation learning . - We propose a self-supervised method that can better utilize object geometric information by ensuring the consistency between bounding boxes and foreground segmentation masks . This approach can improve the scene decomposition performance compared with the stateof-art . - We design a new segmentation head that can preserve global context and prevent coordinate misalignment in small feature maps which improves the foreground segmentation performance . 2 RELATED WORKS . There many influential works ( Burgess et al . ( 2019 ) ; Greff et al . ( 2019 ) ; Eslami et al . ( 2016 ) ; Crawford & Pineau ( 2019 ) ; Engelcke et al . ( 2020 ) ; Greff et al . ( 2017 ) ; Van Steenkiste et al . ( 2018 ) ; Pathak et al . ( 2016 ) ; Lin et al . ( 2020 ) ) in unsupervised scene decomposition in recent years . Some models tend to explicitly factor an object representation into spatial and appearance encodings such as ’ what ’ , ’ where ’ , ’ presence ’ , etc . with the help of VAE ( Kingma & Welling ( 2013 ) ) . Influential related models include AIR ( Eslami et al . ( 2016 ) ) and its successor SPAIR ( Crawford & Pineau ( 2019 ) ) . AIR uses the Recurrent Neural Network as the encoder to decompose a complex scene into objects ’ representation but it suffers from the iteration speed . SPAIR improves its bounding box average precision and running speed by using Convolution Neural Network as the encoder to generate objects ’ representation in parallel . However , these models have not been tested on photorealistic 3D object dataset and bounding boxes can not fully represent flexible morphology like foreground segmentation masks . The other type of models tend to decompose each object into its own representation without explicit positional encoding and use segmentation masks to mix object reconstruction . Influential models such as MONet ( Burgess et al . ( 2019 ) ) which leverages a UNet ( Ronneberger et al . ( 2015 ) ) variant as an iterative attention network for segmentation mask generation and Spatial Broadcast Decoder ( Watters et al . ( 2019 ) ) for representation learning via scene reconstruction . Spatial Broadcast Decoder replaces deconvolutional network with transform by tiling ( broadcast ) the latent vector across space , concatenate fixed X- and Y- “ coordinate ” channels . This decoder provides better disentanglement between positional and non-positional features in the latent distribution . IODINE ( Greff et al . ( 2019 ) ) tackles this problem with its amortized iterative refinement of foreground and background representation . However its iterative refinement process will heavily impact the speed of training and inference . GENESIS ( Engelcke et al . ( 2020 ) ) uses the similar idea as MONet but with different latent encoding in different iterative steps . These models all focus on decomposing the entire scene which does not leverage the spatial locality around each object . SPACE ( Lin et al . ( 2020 ) ) is the closest to our work in spirit . This model leverages the encoder similar to SPAIR to process foreground objects in parallel with explicit positional encoding and adapts the segmentation masks for background modeling . However , different from R-MONet , it does not use the shared information in both bounding boxes and segmentation masks . It is worth mentioning that S4Net ( Fan et al . ( 2019 ) ) provides an end-to-end single shot object detection and instance segmentation framework simplified from Mask R-CNN ( He et al . ( 2017 ) ) for our spatial attention module . It uses the ROI Masking method inspired by the fact that the segmentation can benefit from the background features surrounding the foreground instance . In this method , an ROI mask is generated by enlarging the proposed region on the feature maps . The mask values inside the original proposed region are set to 1 . The mask values outside the original proposed region but inside the enlarged proposed region are set to -1 . The rest values on the mask are set to 0 . The regional features are extracted by multiplying the ROI mask and the feature maps . This method is proved to have better segmentation compared with ROI Align ( He et al . ( 2017 ) ) in both their experiments and our use case ( discussed in Section 4.1 ) . 3 METHOD . In this section , the proposed model , R-MONet is described . The R-MONet will decompose the scene into K different components ( K-1 for the foreground and 1 component for the background ) and each of them is indexed by k. K also limits the max number of objects this model can detect . The network architecture of R-MONet is presented in Figure 1 . 3.1 INFERENCE . Inference of object geometric representations . A spatial attention network parameterized by ψ provides the object geometric representation in both segmentation masks { mk } and foreground object bounding box { tk } with the input image x . It is represented by equation 1 . At the same time , { mk } is also the mixing probability in spatial Gaussian model . Since this mixing probability is learned by the network conditioned on x , we refer to this distribution as qψ ( mk|x ) . { mk } k∈ [ 1 , K ] , { tk } k∈ [ 1 , K−1 ] = fψ ( x ) ( 1 ) This spatial attention network is a convolutional neural network similar to S4Net ( Fan et al . ( 2019 ) ) with small adjustment . In this module , ResNet18 ( He et al . ( 2016 ) ) and Feature Pyramid Network ( FPN ) ( Lin et al . ( 2017 ) ) are used as the backbone . The feature maps are extracted for the following Region Proposal Network ( RPN ) ( Ren et al . ( 2015 ) ) . The bounding boxes proposed by RPN will be filtered with Non-Maximum suppression ( NMS ) ( Neubeck & Van Gool ( 2006 ) ) . Only ROIs with top K-1 prediction scores will be selected for further processing . The ROI Masking ( Fan et al . ( 2019 ) ) method then uses the ROI selected by RPN and transforms feature maps from the backbone to the segmentation head . The segmentation intermediate outputs { αk } ( not segmentation masks ) from the segmentation branch will be transformed into the segmentation masks { mk } via the stick-breaking process ( Sethuraman ( 1994 ) ) to make sure ∑K k=1mk = 1 which represents all pixels are explained in the segmentation masks ( softmax can also be used ) . The scope sk represents the proportion of pixels still unexplained . The details about stick-breaking process is in Appendix A . Inference of object representation latents . In the VAE encoder parameterised by φ , we use variational inference to get an approximate latent posterior qφ ( zk|x , mk ) which is a Gaussian distribution conditioned on both input image x and segmentation masks { mk } . As pointed out in MONet ( Burgess et al . ( 2019 ) ) , conditioning on { mk } can provide geometric heuristic information for probabilistic inference . Then , we sample from this posterior to infer object representation latents { zk } .

Summary: This paper presents a variation of the MONet model where an additional Region Proposal Network generates bounding boxes for various objects in the scene. An additional loss is introduced during training to make the segmentations produced by the MONet segmenter consistent with the proposed bounding boxes. Results are demonstrated on multi d-Sprites and CLEVR with modest performance gains.