Disentangling 3D Prototypical Networks for Few-Shot Concept Learning

1 INTRODUCTION . Humans can learn new concepts from just one or a few samples . Consider the example in Figure 1 . Assuming there is a person who has no prior knowledge about blue and carrot , by showing this person an image of a blue carrot and telling him “ this is an carrot with blue color ” , the person can easily generalize from this example to ( 1 ) recognizing carrots of varying colors , 3D poses and viewing conditions and under novel background scenes , ( 2 ) recognizing the color blue on different objects , ( 3 ) combine these two concepts with other concepts to form a novel object coloring he/she has never seen before , e.g. , red carrot or blue tomato and ( 4 ) using the newly learned concepts to answer questions regarding the visual scene . Motivated by this , we explore computational models that can achieve these four types of generalization for visual concept learning . We propose disentangling 3D prototypical networks ( D3DP-Nets ) , a model that learns to disentangle RGB-D images into objects , their 3D locations , sizes , 3D shapes and styles , and the background scene , as shown in Figure 2 . Our model can learn to detect objects from a few 3D object bounding box annotations and can further disentangle objects into different attributes through a self-supervised view prediction task . Specifically , D3DP-Nets uses differentiable unprojection and rendering operations to go back and forth between the input RGB-D ( 2.5D ) image and a 3D scene feature map . From the scene feature map , our model learns to detect objects and disentangles each object into a 3D shape code and an 1D style code through a shape/style disentangling antoencoder . We use adaptive instance normalization layers ( Huang & Belongie , 2017 ) to encourage shape/style disentanglement within each object . Our key intuition is to represent objects and their shapes in terms of 3D feature representations disentangled from style variability so that the model can correspond objects with similar shape by explicitly rotating and scaling their 3D shape representations during matching . Project page : https : //mihirp1998.github.io/project_pages/d3dp/ ∗Equal contribution †Work done while at Carnegie Mellon University With the disentangled representations , D3DP-Nets can recognize new concepts regarding object shapes , styles and spatial arrangements from a few human-supplied labels by training concept classifiers only on the relevant feature subspace . Our model learns object shapes on shape codes , object colors and textures on style codes , and object spatial arrangements on object 3D locations . We show in the supplementary how the features relevant for each linguistic concept can be inferred from a few contrasting examples . Thus the classifiers attend only to the essential property of the concept and ignore irrelevant visual features . This allows them to generalize with far fewer examples and can recognize novel attribute compositions not present in the training data . We test D3DP-Nets in few-shot concept learning , visual question answering ( VQA ) and scene generation . We train concept classifiers for object shapes , object colors/materials , and spatial relationships on our inferred disentangled feature spaces , and show they outperform current stateof-the-art ( Mao et al. , 2019 ; Hu et al. , 2016 ) , which use 2D representations . We show that a VQA modular network that incorporates our concept classifiers shows improved generalization over the state-of-the-art ( Mao et al. , 2019 ) with dramatically fewer examples . Last , we empirically show that D3DP-Nets generalize their view predictions to scenes with novel number , category and styles of objects , and compare against state-of-the-art view predictive architectures of Eslami et al . ( 2018 ) . The main contribution of this paper is to identify the importance of using disentangled 3D feature representations for few-shot concept learning . We show the disentangled 3D feature representations can be learned using self-supervised view prediction , and they are useful for detecting and classifying language concepts by training them over the relevant only feature subsets . The proposed model outperforms the current state-of-the-art in VQA in the low data regime and the proposed 3D disentangled representation outperforms similar 2D or 2.5D ones in few-shot concept classification . 2 RELATION TO PREVIOUS WORKS . Few-shot concept learning Few-shot learning methods attempt to learn a new concept from one or a few annotated examples at test time , yet , at training time , these models still require labelled datasets which annotate a group of images as “ belonging to the same category ” ( Koch et al. , 2015 ; Vinyals et al. , 2016b ) . Metric-based few-shot learning approaches ( Snell et al. , 2017 ; Qi et al. , 2018 ; Schwartz et al. , 2018 ; Vinyals et al. , 2016a ) aim at learning an embedding space in which objects of the same category are closer in the latent space than objects that belong to different categories . These models needs to be trained with several ( annotated ) image collections , where each collection contains images of the same object category . Works of Misra et al . ( 2017 ) ; Purushwalkam et al . ( 2019 ) ; Nikolaus et al . ( 2019 ) ; Tokmakov et al . ( 2019 ) compose attribute and nouns to detect novel attribute-noun combinations , but their feature extractors need to be pretrained on large annotated image collections , such as Imagenet , or require annotated data with various attribute compositions . The proposed model is pretrained by predicting views , without the need for annotations regarding object classes or attributes . Our concept classifiers are related to methods that classify concepts by computing distances to prototypes produced by averaging the ( 2D CNN ) features of few labelled examples ( Snell et al. , 2017 ; Qi et al. , 2018 ; Schwartz et al. , 2018 ) . The work of Prabhudesai et al . ( 2020 ) learns 3D prototypes in a self-supervised manner , but they do not disentangle their representation into style and shape codes . We compare against 2D and 3D few shot learning methods and outperform them by a significant margin . The novel feature of our work is that we learn concept prototypes over disentangled 3D shape and 1D style codes as opposed to entangled 3D or 2D CNN features . Learning neural scene representation Our work builds upon recent view-predictive scene representation learning literature ( Tung et al. , 2019 ; Sitzmann et al. , 2019 ; Eslami et al. , 2016 ) . Our scene encoders and decoders , the view prediction objective , and the 3D neural bottleneck and egostabilization of the 3D feature maps is similar to those proposed in geometry-aware neural networks of Tung et al . ( 2019 ) . Sitzmann et al . ( 2019 ) and Eslami et al . ( 2016 ) both encode multiview images of a scene and camera poses into a scene representation , in the form of 2D scene feature maps or an implicit function . Sitzmann et al . ( 2019 ) only considers single-object scenes and needs to train a separate model for each object class . We compare generalization of our view predictions against Eslami et al . ( 2016 ) and show we have dramatically better generalization across number , type and spatial arrangements of objects . Furthermore , the above approaches do not explicitly disentangle style/shape representations of objects . Zhu et al . ( 2018 ) focuses on synthesizing natural images of objects with a disentangled 3D representation , but it remains unclear how to use the learnt embeddings to detect object concepts . Different from most inverse graphics networks ( Tung et al. , 2017 ; Kanazawa et al. , 2018 ) that aim to reconstruct detailed 3D occupancy of the objects , our model aims to learn feature representations that can detect an object across pose and scale variability , and use them for concept learning . Our shape-style disentanglement uses adaptive instance normalization layers ( Huang et al. , 2018 ; Huang & Belongie , 2017 ) that have been valuable for disentangling shape and style in 2D images . Here , we use them in a 3D latent feature space . 3 DISENTANGLING 3D PROTOTYPICAL NETWORKS ( D3DP-NETS ) . The architecture of D3DP-Nets is illustrated in Figure 2 . D3DP-Nets consists of two main components : ( a ) an image-to-scene encoder-decoder , and ( b ) an object shape/style disentanglement encoder-decoder . Next , we describe these components in detail . 3.1 IMAGE-TO-SCENE ENCODER-DECODER . A 2D-to-3D scene differentiable encoder Esc maps an input RGB-D image to a 3D feature map M ∈ Rw×h×d×c of the scene , where w , h , d , c denote width , height , depth and number of channels , respectively . Every ( x , y , z ) grid location in the 3D feature map M holds a c-channel feature vector that describes the semantic and geometric properties of a corresponding 3D physical location in the 3D world scene . We output a binary 3D occupancy map Mocc ∈ { 0 , 1 } w×h×d from M using an occupancy decoder Docc . A differentiable neural renderer Dsc neurally renders a 3D feature map M to a 2D image and a depth map from a specific viewpoint . When the input to D3DP-Nets is a sequence of images as opposed to a single image , each image It in the sequence is encoded to a corresponding 3D per frame map Mt , the 3D rotation and translation of the camera with respect to the frame map of the initial frame I0 is computed and the scene map M is computed by first rotating and translating Mt to bring it to the same coordinate frame as M0 and then averaging with the map built thus far . We will assume camera motion is known and given for this cross frame fusion operation . D3DP-Nets are self-supervised by view prediction , predicting RGB images and occupancy grids for query viewpoints . We assume there is an agent that can move around in static scenes and observes them from multiple viewpoints . The agent is equipped with a depth sensor and knowledge of its egomotion ( proprioception ) provided by the simulator in simulated environments.We train the scene encoders and decoders jointly for RGB view prediction and occupancy prediction and errors are backpropagated end-to-end to the parameters of the network : Lview−pred =‖Dsc ( rotate ( M , vq ) ) − Iq‖1 + log ( 1 + exp ( −Oq ·Docc ( ( rotate ( M , vq ) ) , vq ) ) ) , ( 1 ) where Iq and Oq are the ground truth RGB image and occupancy map respectively , vq is the query view , and rotate ( M , vq ) is a trilinear resampling operation that rotates the content of a 3D feature map M to viewpoint vq . The RGB output is trained with a regression loss , and the occupancy is trained with a logistic classification loss . Occupancy labels are computed through raycasting , similar to Harley et al . ( 2020 ) . We provide more details on the architecture of our model in the supplementary material . We train a 3D object detector that takes as input the output of the scene feature map M and predicts 3D axis-aligned bounding boxes , similar to Harley et al . ( 2020 ) . This is supervised from ground-truth 3D bounding boxes without class labels . 3.2 OBJECT SHAPE/STYLE DISENTANGLEMENT . As the style of an image can be understood as a property which is shared across its spatial dimensions , previous works ( Huang et al. , 2018 ; Karras et al. , 2019 ) use adaptive instance normalization ( Huang & Belongie , 2017 ) as an inductive bias to do style transfer between a pair of images . D3DP-Nets uses this same inductive bias in its decoder to disentangle the style and 3D shape of an object . We believe that 3D shape is not analogous to 3D occupancy , but it is a blend of 3D occupancy and texture ( spatial arrangement of color intensities ) . Given a set of 3D object boxes { bo|o = 1 · · · |O| } where O is the set of objects in the scene , D3DP-Nets obtain corresponding object feature maps Mo = crop ( M , bo ) by cropping the scene feature map M using the 3D bounding box coordinates bo . We use ground-truth 3D boxes at training time and detected boxes at test time . Each object feature map is resized to a fixed resolution of 16 × 16 × 16 , and fed to an object-centric autoencoder whose encoding modules predict a 4D shape code zoshp = Eshp ( M o ) ∈ Rw×h×d×c and a 1D style code zosty = Esty ( M o ) ∈ Rc . A decoder D composes the two using adaptive instance normalization ( AIN ) layers ( Huang & Belongie , 2017 ) by adjusting the mean and variance of the 4D shape code based on the 1D style code : AIN ( z , γ , β ) = γ ( z−µ ( z ) σ ( z ) ) + β , where z is obtained by a 3D convolution on zshp , µ and σ are the channel-wise mean and standard deviation of z , and β and γ are extracted using single-layer perceptrons from zsty . The object encoders and decoders are trained with an autoencoding objective and a cycle-consistency objective which ensure that the shape and style code remain consistent after composing , decoding and encoding again ( see Figure 2 ( b ) ) : Ldis = 1 |O| |O|∑ o=1 ‖Mo −D ( Eshp ( Mo ) , Esty ( Mo ) ) ‖2︸ ︷︷ ︸ autoencoding loss + ∑ i∈O\o Lc−shp ( Mo , Mi ) + Lc−sty ( Mo , Mi ) ︸ ︷︷ ︸ cycle-consistency loss , ( 2 ) where Lc−shp ( Mo , Mi ) = ‖Eshp ( Mo ) −Eshp ( D ( Eshp ( Mo ) , Esty ( Mi ) ) ) ‖2 is the shape consistency loss and Lc−sty ( Mo , Mi ) = ‖Esty ( Mo ) −Esty ( D ( Eshp ( Mi ) , Esty ( Mo ) ) ) ‖2 is the style consistency loss . We further include a view prediction loss on the synthesized scene feature map M̄ , which is composed by replacing each object feature map Mo with its re-synthesized version D ( zoshp , zosty ) , resized to the original object size , as shown in Figure 2 ( c ) . The view prediction reads : Lview−pred−synth = ‖Dsc ( rotate ( M̄ , vt+1 ) ) − It+1‖1 . The total unsupervised optimization loss for D3DP-Nets reads : Luns = Lview−pred + Lview−pred−synth + Ldis . ( 3 )

Summary: This paper describes an approach that learns a disentangled shape and style representation of objects in a self-supervised way from RGB-D images.  The approach is based on various components, like a 3D feature volume, a bounding box detector, and a disentanglement network. Neural rendering (e.g. recomposing the various disentangled parts into an image) is used as the learning signal for self-supervision. Various applications of this representation  are shown, examples are few-shot shape learning and Visual Question Answering.