Generalization to Out-of-Distribution transformations

1 INTRODUCTION . Humans have a unique ability to generalize beyond the scope of prior experience ( Chollet , 2019 ; Lake et al. , 2017 ; Marcus , 2001 ) , while artificial agents struggle to apply knowledge to distributions outside the convex hull of their training data ( Santoro et al. , 2018 ; Lake & Baroni , 2018 ) . One way humans seem to achieve such generalization is by learning a set of primitive abstract structures , like the 1D ordinal scale ( Summerfield et al. , 2020 ) and grid-like representations ( Hafting et al. , 2005 ) . These structures can also be thought of as symmetry functions : transformations that are in some way invariant to the specific value of their arguments . As a concrete example , we can imagine moving any object around in space , regardless of its shape , color or size . A fundamental question is how can such symmetry functions be learned ? We hypothesize that during development , humans learn a set of canonical transformations - e.g . the translation , rotation and change in size of objects - that are grounded in the sensorimotor system ( Barsalou , 2008 ) , and learned as a consequence of predicting the sensory results of primitive actions ( Battaglia et al. , 2013 ) . The abilities to translate , rotate or scale arbitrary objects then become our first abstract affordances ( Gibson , 2014 ) . Indeed , infants that spend more time playing with blocks are better at abstract mental rotation tasks ( Schwarzer et al. , 2013 ) . Our aim is to model part of this process by presenting artificial neural networks ( LeCun et al. , 1995 ) with 2-dimensional shapes , and training them to predict the effect of translation , rotation or scaling of that shape in pixel space . This is analogous to predicting tactile or visual signals resulting from a simple movement or saccade ( Rao & Ballard , 1999 ; Wolpert et al. , 1995 ) . We do not explicitly model motor actions , but rather transform the image offline , and feed the result back to the model . We then test the extent to which such predictions generalize OOD along dimensions such as shape , size , location and time . Evidence of OOD generalization would suggest that the network has learned a symmetry function . We assume that , in order to learn symmetry functions , we must introduce principled inductive biases . The first is convolution . Second , to constrain the network to learn a primitive function that can apply to any shape , we assume it requires exposure to a sufficiently diverse set of examples . We therefore operationalize and vary ’ diversity ’ as the number of distinct shapes present in the training set . As the translational invariance built into convolution is well aligned with the task of translation unsurprisingly , a fully convolutional autoencoder with high training set diversity acheives near perfect OOD generalization for translation to unseen shapes in larger grids and new locations ( Figure 1 ) . However the two components weren ’ t sufficient for OOD generalization of rotation ( Figure 2 ) and scaling . To remedy this we introduce two more principle components . First is the effect of ’ iteration ’ during training ( output fed back as input iteratively ) , based on the idea that sequential applications of the same transformation should maintain the identity of an object ( i.e . object permanence ( Piaget , 2006 ) ) . Secondly , motivated by the strong OOD generalization of translation as depicted in Figure 1 , and the fact that translation in log-polar space is equivalent to rotation and scaling in cartesian space ( Figure 1 of ( Esteves et al. , 2017 ; Tootell et al. , 1982 ) ) , we hypothesize that performing convolutions in log polar space would help in the OOD generalization of rotation and scaling . We summarize the following contributions : • We build a flexible data generator that can produce and transform a large variety of simple but structured images - irregular polygons with variable size , location , and complexity . • We show that iteration has the ability to conserve the shape far past the time horizon seen during training and is necessary for better OOD generalization of rotation and scaling transformations along with convolution and high diversity . • We find an interesting tradeoff between diversity and iteration where each could partially make up for less of the other to produce better OOD generalization capabilities . • We novelly propose POLARAE , a fully convolutional autoencoder build using concepts from polar transformer networks ( Esteves et al. , 2017 ) which exploits all the four components demonstrating better OOD generalization performance than a standard autoencoder and a variational autoencoder in cartesian space . 2 RELATED WORK . 2.1 PSYCHOLOGY . Humans can perform mental transformations OOD ( i.e . on unseen shapes ) . For example , humans can mentally translate or rotate shapes at a steady rate ( Shepard & Metzler , 1971 ) or scale abstract distances ( Trope & Liberman , 2010 ) . There are also known mechanisms that might make these transformations general . For example , there is a known impact of diversity on the generalization of properties in psychology , namely , the diversity effect ( Osherson et al. , 1990 ) . Iteration is also a plausible mechanism in humans - psychological data for mental rotation suggest that we transform objects iteratively , since larger angles of rotation elicit longer reaction times ( Shepard & Metzler , 1971 ) . There is also evidence that discrete temporal context updating by recurrent thalamocortical loops serves predictive learning in the brain ( O ’ Reilly et al. , 2014 ) . Basically , iterated operations achieve generality since they are time-invariant . Finally , humans have built-in architectural transforms , such as the transformation of retinal images into a log-polar coordinate system ( Maiello et al. , 2020 ) . 2.2 MACHINE LEARNING . There have been several works that demonstrate that high diversity is necessary for OOD generalization . Xu et al . ( 2020 ) provided theoretical guarantees that ReLU MLP networks extrapolate well with linear functions and sufficient diversity in their training set . Sufficient diversity of input is also required for systematic generalization of neural network based reinforcement learning agents ( Hill et al. , 2019 ) . Madan et al . ( 2021 ) showed that data diversity improves OOD category-viewpoint classification but didn ’ t focus on the generative aspect of it using CNNs . There have also been works collecting large scale controlled synthetic datasets ( Borji et al. , 2016 ; Gross et al. , 2010 ; Qiu et al. , 2017 ) to facilitate learning invariance to different types of transformations in deep neural networks . Following a similar motivation , we want CNNs to be able to learn such transformations and to be able to generate them on OOD examples . The idea of iterative training has also been explored in this domain . Using an iterative training technique similar to ours generative adversarial networks learned 3D rotations ( but not scaling ) ( Galama & Mensink , 2019 ) . However there were limited assessment of how different amounts of iteration during training impacted extrapolation in time . Generally , predicting past the number of iterations seen during training has required building in a conservation law of some sort ( Cranmer et al. , 2020 ; Greydanus et al. , 2019 ) . Kumar et al . ( 2020 ) showed that gradual self training at each iteration helps in domain adaptation . In a different context , Kuchaiev & Ginsburg ( 2017 ) proposed iterative output re-feeding to train deep autoencoders for improved collaborative filtering . Recently , the idea that iteration can achieve OOD generalization has been put forth in the context of problem solving ( Schwarzschild et al. , 2021 ) . Kim et al . ( 2020 ) explored the effectiveness of log-polar space in achieving rotation invariance but was limited to classification and didn ’ t study other types of transformations like scaling . 3 METHOD . 3.1 DATA GENERATOR . All training stimuli were shapes contained within a 64x64 pixel grid space . We constructed irregular N-sided polygonal shapes by first sampling N angular distances between 0 and 2π , and then sampling a radial distance from a centroid ( xcenter , ycenter ) at each of these angles uniformly between 0 and a scale parameter r. This produced a set of vertices ; pixels within the convexity of the vertices were set to 1 , and pixels outside were set to 0 . There was also the option to produce ’ hollow ’ shapes such that an interior cut-out of the shape was set to 0 , in order to produce a test-set of shapes with different distribution from the training set . This data generator was therefore capable of producing a combinatorically large set of possible shapes . Our procedure for shape generation is most similar to Attneave forms ( Attneave & Arnoult , 1956 ) and also bears relation to the method of Fourier descriptors ( Zhang & Lu , 2005 ) , but was selected due to its computational speed and interpretable manipulations of shape parameters . Each shape was used as the input to a neural network , and was transformed in one of the following ways to generate the target for training : for translation , shift 2 pixels to the right ; for rotation , rotate π25 radians clockwise ; for scaling , increase radial length of the vertices by 0.1 ( Figure 3 ) . These transformations were hard-coded but meant to represent a set of innate primitive actions . 3.2 MODEL ARCHITECTURE . We first describe the three baseline models followed by the proposed model : The first model was a fully convolutional autoencoder . The encoder consisted of 3 convolutional layers , the first with 16 3x3 kernels , ( stride=2 , padding=1 ) , the second with 32 3x3 kernels ( stride=2 , padding=1 ) , and the last with 64 7x7 kernels ( stride=1 ) . The decoder had 3 layers ( padding=1 , output padding=1 ) that inverted these operations with transposed convolution layers that were mirror images of the encoder layers . This produced an output with the same dimensions as the input image . All layers were followed by rectified linear ( ReLU ) activations . Being fully convolutional , it could accept any input grid size . All weights were initialized using Xavier uniform initialization . We use AE to denote this model . The second model was a fully convolutional variational autoencoder ( Kingma & Welling , 2013 ) . The encoder and decoder architectures were similar to AE . We use VAE to denote this model . The third model was a β-VAE ( Higgins et al. , 2016 ) . We used β = 4 since that learnt a good disentangled representation of the data generative factors in the original work . The proposed model was a combination of polar transformer networks and the fully convolutional autoencoder described above . More specifically , given an input image the polar origin predictor computed a single channel feature map using a fully convolutional network and the centroid of the heatmap was used as the polar origin . The polar transformer module converted the image to log polar coordinates using the predicted polar origin . The transformed image was then fed to AE . Since the rotation of the input image resulted in vertical shift in log polar space wrapping at the boundary , we used wrap around padding on the vertical dimension and zero padding in the horizontal dimension before applying convolutions for the first two layers of the encoder . An inverse polar transformer module was used to convert the output of the AE from log polar space to the original image space . The inverse polar transformer module used the sample points of the original input image used to compute the log-polar transform and mean approximation since a single point in the original image contributed to multiple points in the log-polar space to convert the image back to cartesian coordinates . Since this is effectively an autoencoder operating in the log-polar space we use POLARAE to denote this model . Figure 4 shows the architecture .

Summary: The authors trained autoencoders and variational autoencoders in cartesian space, as well as autoencoders in log-polar space, with generated data representing a range of interesting transformations. They then tested the ability of the models to extrapolate beyond the learned transformations in pixel space. The authors find that iterative training, data diversity, convolutions, and transformation to log-polar space all improve the generalization performance.