Extrapolatable Relational Reasoning With Comparators in Low-Dimensional Manifolds

1 INTRODUCTION . The goal of Artificial Intelligence research , first proposed in the 1950s and reiterated many times , is to create machine intelligence comparable to that of a human being . While today ’ s deep learning based systems achieve human-comparable performances in specific tasks such as object classification and natural language understanding , they often fail to generalise in out-of-distribution ( o.o.d ) scenarios , where the test data distribution differs from the training data distribution ( Recht et al. , 2019 ; Trask et al. , 2018 ; Barrett et al. , 2018 ; Belinkov & Bisk , 2018 ) . Moreover , it is observed that the generalisation error increases as the tasks become more abstract and require more reasoning than perception . This ranges from small drops ( 3 % to 15 % ) in classification accuracy on ImageNet ( Recht et al. , 2019 ) to accuracy only slightly better than random chance for the Raven Progressive Matrices ( RPM ) test ( a popular Human IQ test ) , when testing data are sampled completely out of the training distribution ( Barrett et al. , 2018 ) . In contrast , human brain is observed to generalise better to unseen inputs ( Geirhos et al. , 2018 ) , and typically requires only a small number of training samples . For example , a human , when trained to recognise that there is a progression relation of circle sizes in Figure 1a , can easily recognise that the same progression relation exists for larger circles in Figure 1b , even though such size comparison has not been done between larger circles . However , today ’ s state-of-the-art neural networks ( Barrett et al. , 2018 ; Wang et al. , 2020 ) are not able to achieve the same . Researchers ( Spelke & Kinzler , 2007 ; Chollet , 2019 ; Battaglia et al. , 2018 ; Xu et al. , 2020 ) argue that the human brain developed special inductive biases that adapt to the form of information processing needed for humans , thereby improving generalisation . Examples include convolution-like cells in the visual cortex ( Hubel & Wiesel , 1959 ; Güçlü & van Gerven , 2015 ) for visual information processing , and grid cells ( Hafting et al. , 2005 ) for spatial information processing and relational comparison between objects ( Battaglia et al. , 2018 ) . In this work , we propose a simple yet effective inductive bias which improves o.o.d generalisation for relational reasoning . We specifically focus on a the type of o.o.d called ‘ extrapolation ’ . For extrapolation tasks , the range of one or more data attributes ( e.g. , object size ) from training and test datasets are ( a ) Small Circles ( b ) Big Circles Figure 1 : Size Progression Relations for circles of different sizes . completely non-overlapping . The proposed inductive bias is inspired by neuroscience and psychology research ( Fitzgerald et al. , 2013 ; Chafee , 2013 ; Summerfield et al. , 2020 ) showing that in primate brain there are neurons in the Parietal Cortex which only responds to different specific attributes of perceived entities . For examples , certain LIP neurons fire at higher rate for larger objects , while the firing rate of other neurons correlates with the horizontal position of objects in the scene ( left vs right ) ( Gong & Liu , 2019 ) . These observations show that these neurons map objects ’ attributes such as size and spatial position to the one-dimensional axis of the neurons ’ firing rate . From a computational perspective , this can be viewed as projecting object representations to different low-dimensional manifolds corresponding to different attributes . Based on these observations ( Summerfield et al. , 2020 ) , we hypothesise that these neurons are developed to learn low-dimensional representations of relational structure that are optimised for abstraction and generalisation , and the same inductive bias can be readily adapted for artificial neural networks to achieve similar optimisation for abstraction and generalisation . We test this hypothesis by designing an inductive bias module which projects high-dimensional object representations into low-dimensional manifolds , and make comparisons between different objects in these manifolds . We show that this module can be readily amalgamated with existing architectures to improve extrapolation performance for different relational reasoning tasks . Specifically , we performed experiments on three different extrapolation tasks , including maximum of a set , visual object comparison on dSprites dataset ( Higgins et al. , 2017 ) and extrapolation on RPM style tasks ( Barrett et al. , 2018 ) . We show that models with the proposed low-dimensional comparators perform considerably better than baseline models on all three tasks . In order to understand the effectiveness of comparing in low-dimensional manifolds , we analyse the projection space and corresponding function space of the comparator to show the importance of projection to low-dimensional manifolds in improving generalisation . Finally , we perform the analysis relating to algorithmic alignment theory ( Xu et al. , 2020 ) , and propose an augmentation to the sample complexity criteria used by this theory to better measure algorithmic alignment with generalisation . 2 METHOD . Here , we describe the inductive bias module we developed to test our hypothesis that the same inductive bias of low-dimensional representation observed in Parietal Cortex can be readily adapted for artificial neural network , thus enabling it to achieve similar optimisation for abstraction and generalisation . The proposed module learns to project object representations into low-dimensional manifolds and make comparisons in these manifolds . In Section 2.1 we describe the module in detail . In Section 2.2 , 2.3 and 2.4 we discuss how this module can be utilised for three different relational reasoning tasks , which are : finding the maximum of a set , visual object comparisons and Raven Progressive Matrics ( RPM ) reasoning . 2.1 COMPARATOR IN LOW-DIMENSIONAL MANIFOLDS . The inductive bias module is comprised of low-dim projection functions p and comparators c. Let { oi ; i ∈ 1 . . . N } be the set of object representations , obtained by extracting features from raw inputs such as applying Convolutional Neural Networks ( CNN ) on images . Pairwise comparison between object pair oi and oj can be achieved with a function f expressed as : f ( oi , oj ) = g ( K || k=1 ck ( pk ( oi ) , pk ( oj ) ) ) . ( 1 ) Here pk is the kth projection function that projects object representation o into the kth low dimensional manifold , ck is the kth comparator function that compares the projected representations , || is the concatenation symbol and g is a function that combines theK comparison results to make a prediction . Having K parallel projection functions pk and comparators ck enables a simultaneous comparison between objects with respect to their different attributes . Figure 2 shows an example of comparing sizes of circles by projection onto a 1-dimensional manifold . Both p and c can be implemented as feed forward neural networks . While the comparator c , implemented as a neural network , can theoretically learn a rich range of comparison metrics , we found that adding to c an additional inductive bias of distance measure for the projection , such as vector distance p ( oi ) − p ( oj ) or absolute distance |p ( oi ) − p ( oj ) | , improves the generalisation performance . Let rt ( oi ) be the ground truth mapping function from ith object ’ s representation oi to its tth attributes ( such as colour and size for a visual object ) . If such ground truth labels of object attributes exist , f ( oi , oj ) can be trained to directly predict the differences in attributes by minimising the loss L ( d ( rt ( oi ) , rt ( oj ) ) , f ( oi , oj ) ) , where d is a distance function ( e.g. , rt ( oi ) − rt ( oj ) for continuous attributes or 1rt ( oi ) =rt ( oj ) for categorical attributes ) . However , in real-world datasets , such ground truth attribute labels seldom exist . Instead , in many relational reasoning tasks , learning signals for attribute comparison are only provided implicitly in the training objective . For example , in Visual Question Answering task , an example question might be ‘ Is the object behind Object A smaller ? ’ . The learning signals for the required size and spatial position comparator are provided only through correctness of the answers to the given questions . Thus , the proposed module is only useful and scalable if it can be integrated into neural architectures for relational reasoning and still learn to compare attributes with the weaker , implicit learning signal . Next , we describe 3 examples of such integrations for different relational reasoning tasks , and show in Section 3 that the proposed module can learn relational reasoning tasks with better generalisation capability . 2.2 ARCHITECTURE : MAXIMUM OF A SET . The first task we consider is finding the maximum of a set of real numbers . Formally , given a set { xi ; i ∈ 1 . . . N } where xi is a real number represented as a scalar value , we want to train a function hmax ( { xi , . . . , xN } ) that gives the maximum value in the set . Many neural architectures have been applied on this task , including Deep Sets ( Zaheer et al. , 2017 ) and Set Transformer ( Lee et al. , 2019 ) , but none of them test ( o.o.d ) generalisation capability . In order to test o.o.d generalisation in the extrapolation scenario , we create the training and test sets such that their ranges do not overlap . We sample from the range ( V trainlow , V train high ) for the training set , and from the range ( V test low , V test high ) for the test set , and restrict that V trainhigh < V test low . We integrate the proposed low-dim comparator module with Set Transformer ( Lee et al. , 2019 ) , a state-of-the-art neural architecture for sets . Set Transformer first encodes each element in the set with respect to all other elements with a Multihead Attention Block ( MAB ) , an attention module modified from self-attention used in language tasks ( Vaswani et al. , 2017 ) : ei = encode ( xi ) = MAB ( xi , xj ) . The Set Transformer then uses Pooling with Multihead Attention ( PMA ) to combine all encoded elements of the set as PMA ( e1 , . . . , eN ) . While MAB uses query and key embeddings to generate attention variables , which are then used as weights in the weight sum of value embeddings of elements , we swap the query-key attention mechanism with our low-dim comparator as : ei =MLP ( N∑ j=1 c ( xi , xj ) ) ( 2 ) Here c is the low-dim comparator and MLP is a standard Multi-Layer Perceptron . Note that we directly use the scalar input xi as object representation oi in Equation ( 1 ) as no feature extraction is needed . We then use attention-based pooling to combine projection of xi as ∑N i=1 a ( ei ) p ( xi ) , where a outputs attention values while p is the 1-dim projection function . For detailed architecture configuration , please refer to Appendix A .

Summary: The paper presents two neural network design ideas: low-dimensional projection and arithmetic comparator. By integrating these two ideas into CNNs, the model can solve a set of tasks that require recognizing 1-dimensional properties of objects (such as the size, the color, etc), and making a comparison of these properties. The authors show that their method has strong out-of-distribution generalization: such as generalization to larger objects.