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: This paper addresses an inductive bias for relational reasoning tasks to improve generalization performance on out-of-distribution scenarios (so called extrapolation), which the value ranges (continuous variables) or value itself (discrete variables) of the training/test dataset do not overlap. The proposed idea is to project high-dimensional representations onto low-dimensional space and compare those of the targets of relational reasoning. To validate effectiveness of the idea, the authors provide the improved results on three problems: (1) finding the maximum of a set of real numbers, (2) comparing the attributes of visual objects with dSprites dataset, and (3) visual reasoning with Raven progressive matrices (RPM) based on the PGM dataset. The experimental results of each  show better performance than other baseline methods. This research deals with one of important topics in AI/ML research. Out-of-distribution generalization is potentially extended to world models, common sense reasoning, learning and reasoning with small number of instances, and abstraction capability of human intelligence. As Section of related work in this paper also mentions these points, it is understandable to figure out the position of this work.