Unsupervised Object Keypoint Learning using Local Spatial Predictability

1 INTRODUCTION . An intelligent agent situated in the visual world critically depends on a suitable representation of its incoming sensory information . For example , a representation that captures only information about relevant aspects of the world makes it easier to learn downstream tasks efficiently ( Barlow , 1989 ; Bengio et al. , 2013 ) . Similarly , when explicitly distinguishing abstract concepts , such as objects , at a representational level , it is easier to generalize ( systematically ) to novel scenes that are composed of these same abstract building blocks ( Lake et al. , 2017 ; van Steenkiste et al. , 2019 ; Greff et al. , 2020 ) . In recent work , several methods have been proposed to learn unsupervised representations of images that aim to facilitate agents in this way ( Veerapaneni et al. , 2019 ; Janner et al. , 2019 ) . Of particular interest are methods based on learned object keypoints that correspond to highly informative ( salient ) regions in the image as indicated by the presence of object parts ( Zhang et al. , 2018 ; Jakab et al. , 2018 ; Kulkarni et al. , 2019 ; Minderer et al. , 2019 ) . Many real world tasks primarily revolve around ( physical ) interactions between objects and agents . Therefore it is expected that a representation based on a set of task-agnostic object keypoints can be re-purposed to facilitate downstream learning ( and generalization ) on many different tasks ( Lake et al. , 2017 ) . One of the main challenges for learning representations based on object keypoints is to discover salient regions belonging to objects in an image without supervision . Recent methods take an information bottleneck approach , where a neural network is trained to allocate a fixed number of keypoints ( and learn corresponding representations ) in a way that helps making predictions about an image that has undergone some transformation ( Jakab et al. , 2018 ; Minderer et al. , 2019 ; Kulkarni et al. , 2019 ) . However , keypoints that are discovered in this way strongly depend on the specific transformation that is considered and therefore lack generality . For example , as we will confirm in our experiments , the recent Transporter ( Kulkarni et al. , 2019 ) learns to prioritize image regions that change over time , even when they are otherwise uninformative . Indeed , when relying on extrinsic object properties ( i.e . that are not unique to objects ) one becomes highly susceptible to distractors as we will demonstrate . In this work , we propose PermaKey a novel representation learning approach based on object keypoints that does not overly bias keypoints in this way . The key idea underlying our approach is to view objects as local regions in the image that have high internal predictive structure ( self-information ) . We argue that local predictability is an intrinsic property of an object and therefore more reliably captures objectness in images ( Alexe et al. , 2010 ) . This allows us to formulate a local spatial prediction problem to infer which of the image regions contain object parts . We perform this prediction task in the learned feature space of a convolutional neural network ( CNN ) to assess predictability based on a rich collection of learned low-level features . Using PointNet ( Jakab et al. , 2018 ) we can then convert these predictability maps to highly informative object keypoints . We extensively evaluate our approach on a number of Atari environments and compare to Transporter ( Kulkarni et al. , 2019 ) . We demonstrate how our method is able to discover keypoints that focus on image regions that are unpredictable and which often correspond to salient object parts . By leveraging local predictability to learn about objects , our method profits from a simpler yet better generalizable definition of an object . Indeed , we demonstrate how it learns keypoints that do not solely focus on temporal motion ( or any other extrinsic object property ) and is more robust to uninformative ( but predictable ) distractors in the environment , such as moving background . On Atari games , agents equipped with our keypoints outperform those using Transporter keypoints . Our method shows good performance even on challenging environments such as Battlezone involving shifting viewpoints where the Transporters ’ explicit motion bias fails to capture any task-relevant objects . As a final contribution , we investigate the use of graph neural networks ( Battaglia et al. , 2018 ) for processing keypoints , which potentially better accommodates their discrete nature when reasoning about their interactions , and provide an ablation study . 2 METHOD . To learn representations based on task-agnostic object keypoints we require a suitable definition of an object that can be applied in an unsupervised manner . At a high level , we define objects as abstract patterns in the visual input that can serve as modular building blocks ( i.e . they are self-contained and reusable independent of context ) for solving a particular task , in the sense that they can be separately intervened upon or reasoned with ( Greff et al. , 2020 ) . This lets us treat objects as local regions in input space that have high internal predictive structure based on statistical co-occurrence of features such as color , shape , etc . across a large number of samples . Hence , our focus is on their local predictability , which can be viewed as an “ intrinsic ” object property according to this definition . For example , Bruce & Tsotsos ( 2005 ) have previously shown that local regions with high self-information typically correspond to salient objects . More generally , self-information approximated via a set of cues involving center-surround feature differences has been used to quantify objectness ( Alexe et al. , 2010 ) . In this paper we introduce Prediction ERror MAp based KEYpoints ( PermaKey ) , which leverages this definition to learn about object keypoints and corresponding representations . The main component of PermaKey is a local spatial prediction network ( LSPN ) , which is trained to solve a local spatial prediction problem in feature space ( light-blue trapezoid in Figure 1 ) . It involves predicting the value of a feature from its surrounding neighbours , which can only be solved accurately when they belong to the same object . Hence , the error map ( predictability map ) that is obtained by evaluating the LSPN at different locations carves the feature space up into regions that have high internal predictive structure ( see rows 4 & 5 in Figure 2 ( a ) ) . In what follows , we delve into each of the 3 modules that constitute our PermaKey system : 1 ) the VAE to learn spatial features 2 ) the LSPN to solve the local spatial prediction task and 3 ) PointNet to group error maps into keypoints ( see Figure 1 for an overview ) . Spatial Feature Embedding To solve the local spatial prediction task we require an informative set of features at each image location . While prior approaches focus on low-level features of image patches ( eg . RGB values Isola et al . ( 2014 ) , or ICA features Bruce & Tsotsos ( 2005 ) ) , we propose to learn features using a Variational Auto-Encoder ( VAE ; Kingma & Welling ( 2014 ) ) . It consists of an Encoder that parametrizes the approximate posterior qφ ( z|x ) and a Decoder that parametrizes the generative model pθ ( xi|z ) , which are trained to model the observed data via the ELBO objective : Le ( θ , φ ) = Eqφ ( z|x ) log [ pθ ( x|z ) ] −DKL [ qφ ( z|x ) ||p ( z ) ] . ( 1 ) The encoder is based on several layers of convolutions that offer progressively more abstract image features by trading-off spatial resolution with depth . Hence , which layer ( s ) we choose as our feature embedding will have an effect on the outcome of the local spatial prediction problem . While more abstract high-level features are expected to better capture the internal predictive structure of an object , it will be more difficult to attribute the error of the prediction network to the exact image location . On the other hand , while more low-level features can be localized more accurately , they may lack the expressiveness to capture high-level properties of objects . Nonetheless , in practice we find that a spatial feature embedding based on earlier layers of the encoder works well ( see also Section 5.3 for an ablation ) . Local Spatial Prediction Task Using the learned spatial feature embedding we seek out salient regions of the input image that correspond to object parts . Our approach is based on the idea that objects correspond to local regions in feature space that have high internal predictive structure , which allows us to formulate the following local spatial prediction ( LSP ) task . For each location in the learned spatial feature embedding , we seek to predict the value of the features ( across the feature maps ) from its neighbouring feature values . When neighbouring areas correspond to the same object- ( part ) , i.e . they regularly appear together , we expect that this prediction problem is easy ( green arrow in Figure 3 ) . In contrast , this is much harder when insufficient context is available ( red arrow Figure 3 : Two different prediction scenarios encountered when attempting to solve the LSP task . In the first case , too little context information is available to predict the center patch from its surroundings , likely yielding a prediction error ( assuming the character is not always present ) . Alternatively , when sufficient context is available as in the second case , prediction becomes easy . in Figure 3 ) , or when features rarely co-occur . Similarly , when parts of an image remain fixed across a large number observations it is easily predictable regardless of how much context is available . We implement the local spatial prediction task by training a neural network ( LSPN in Figure 1 ) to predict a center patch of features from the 8 first-order neighbouring patches using the learned spatial feature embedding at layer l of the VAE encoder . The weights of this network are trained to solve this prediction problem across images and across feature locations using the following objective : Lp ( ψl ) = 1 NHW ∑ x H , W∑ i , j ( fLSPNψl ( A [ ne ( i , j ) ] ) −A [ i , j ] ) 2 , ( 2 ) where A ∈ RH×W×C is the set of feature maps at layer l of the encoder qφ ( z|x ) , and ne ( i , j ) are the coordinates of the first-order neighbours surrounding i , j and H , W are the height and width of the spatial feature maps respectively . The features of the neighbouring locations are concatenated before being fed to the network . We train a separate LSPN for each of the chosen feature embedding specified by the encoder layer ( s ) and record their error at each location to obtain predictability maps that are expected to indicate the presence of objects ( rows 4 & 5 in Figure 2 ( a ) ) . Extracting Keypoints from Error Maps The predictability maps can be thought of as being a superposition of the local predictability errors due to each of the individual objects in the scene . Hence , what remains is to extract the k most salient regions ( belonging to object parts ) as our object keypoints from these predictability maps . We do so by training a PointNet ( Jakab et al. , 2018 ) to reconstruct the error maps through a bottleneck consisting of k fixed-width 2D Gaussian windows with learnable centers µ that are the keypoints ( Lg loss in Figure 1 ) as shown below : Lg ( ω ) = 1 MHW M∑ m H , W∑ i , j ( fPointNetω ( H [ m , i , j ] ) −H [ m , i , j ] ) 2 . ( 3 ) HereH ∈ RM×H×W is the concatenation of layer-wise LSP error maps for various encoder layer ( s ) l ( M in total ) , and H , W are the height and width of the error maps ( potentially up-sampled ) . Since the bottleneck layer has only limited capacity it is forced to consider locations that are most helpful in minimizing reconstruction error . This also requires the PointNet to consider statistical regularities across the predictability maps that can efficiently be described . In that sense this procedure can be thought of as a way of ( spatially ) clustering the local errors belonging to the same object by placing keypoints .

Summary: The authors tackle the problem of self-supervised representation learning, and validate their approach on downstream Reinforcement Learning tasks. Building on the insight that predictability in local patches is a good inductive bias for salient regions that characterize objects, the authors propose a well-reasoned, well-engineered and thoroughly validated pipeline for learning object keypoints without supervision. The authors present a wide range of ablative studies to validate their design choices, and demonstrate the superiority of their method both illustratively as well as quantitatively on a number of standard Atari benchmarks.