AlignNet: Self-supervised Alignment Module

1 INTRODUCTION . Despite head and body movements and the motion of objects , we perceive the world and the objects in it as being spatially stable ( Pylyshyn , 1989 ) . There are two factors that contribute to this stability : Firstly , we typically know where objects are even if we can ’ t see them anymore . Secondly , we are able to keep track of objects over time , even when they disappear and reappear after some delay ( Leslie et al. , 1998 ; Leslie & Kaldy , 2001 ; Noles et al. , 2005 ; Pylyshyn , 1989 ) ( Figure 1A ) . The way we perceive the world can be attributed to three core cognitive capabilities , these are : ( 1 ) our ability to group visual stimuli into objects ( Spelke & Kinzler , 2007 ) , ( 2 ) our notion of object persistence ; from only 2.5 months old , we know that an object that has disappeared from our view continues to exist ( Baillargeon , 2002 ) and ( 3 ) our ability to identify multiple instances of the same object over time ; this capability is often attributed to “ sticky indices ” ( Leslie et al. , 1998 ; Leslie & Kaldy , 2001 ; Noles et al. , 2005 ; Pylyshyn , 1989 ) . These indices may be thought of as pointers that , once assigned to an object , remain attached to it , thus uniquely identifying objects over time and space . Recently , in the field of deep learning , notable progress has been made towards the first of these three core cognitive abilities ; several models have been proposed for extracting multiple object representations ( Engelcke et al. , 2019 ; Greff et al. , 2016 ) from a static visual scene , these include MONet ( Burgess et al. , 2019 ) and IODINE ( Greff et al. , 2019 ) . These models present compelling unsupervised segmentation results . However , while MONet or IODINE can be used to extract objects from an agent ’ s visual input over time , neither model determines the correspondence between one time step and the next . For example , if MONet or IODINE were used to extract entities from the frames shown in Figure 1A , the model would not know that the purple cone at t = 0 was the same as the purple cone at t = 1 . In this paper , we take steps towards developing an agent with the second two core cognitive abilities mentioned above : object persistence and “ sticky indices ” . To this end we propose a model with three key properties : ( 1 ) a persistent slot-wise object-based memory that stores unique objects in unique slots , corresponding to their index , ( 2 ) the ability to assign a new index ( or slot in memory ) to a new object and ( 3 ) the ability to recognise when an entity is an instance of an object that has been seen before , assigning it the same index ( or slot in memory ) that was assigned to that object when it first appeared ( see Figure 3A ) . Our model , the AlignNet , and accompanying object-based memory exhibit these properties . More concretely , the AlignNet is a model for dealing with alignment in the cases where two sets being aligned may contain varying numbers of objects and each set may contain objects that are not present in the other . Additionally , when traditionally solving the alignment problem1 ( e.g . with the Hungarian algorithm ) , an adjacency ( or similarity ) matrix is provided . Construction of this adjacency matrix requires hand-crafting a similarity measure . The AlignNet bypasses the need for hand-crafting a similarity measure , by implicitly learning a similarity measure in order to solve the alignment problem . Finally , once trained , the AlignNet may be used to write single entities to single slots in memory , inducing a persistent object-based memory . Our contributions are as follows : 1 . We are the first to formulate the alignment problem such that a neural network may learn a similarity measure for aligning entities ( Section 3 ) . We propose the self-supervised AlignNet , a model capable of finding entities in one set that correspond with entities in another set , while also being able to deal with the adding and dropping of entities ( Section 3 ) . 2 . We demonstrate that a trained AlignNet may be used to write newly appearing objects to empty slots ( or indices ) in memory and keep track of objects over time ( Section 3.2 & 4.4 ) . 3 . We show that the AlignNet implicitly learns a similarity measure for aligning entities , tolerant to noise , that outperforms hand-crafted similarity measures ( Section 4.1 Figures 6 & 9 ) . 4 . We also show that , by using multiple message passing steps , our model learns to attend sharply to single object slots ( Section 4.3 and Figure 7c ) , while remaining differentiable . 5 . Our experiments , on a partially observable question-answering dataset , demonstrate the benefits of having a persistent object-based memory ( Section 4.5 ) . We begin by introducing the general alignment problem . 2 THE GENERAL ALIGNMENT PROBLEM . Given two sets of entities , U = { u0 , ... , uN−1 } and S = { s0 , ... , sN−1 } , which contain representations for instances of the same N unique objects , the alignment problem is equivalent to finding the aligning indices , I = [ i0 , ... , iN−1 ] , such that , uj ≈ sij for j = 0 , 1 , ... , N − 12 . The representations are only approximately equal because the object representations may be noisy , meaning that the representation for two similar objects may differ . While in the case above , S and U contain similar entities , they may also be thought of as sets of objects observed in a scene , by an agent , at different points in time . As an agent moves within its environment we expect the number of objects within its field of view to change from one time step to the next ( Figure 1A . ) which means that if we are aligning sets of objects extracted at different points in time , each set is likely to have different numbers of objects . This means that the AlignNet , and accompanying object-based memory , needs to be able to account for : • Objects that disappear from one time step to the next , such as an object leaving the field of view after a head movement . We refer to objects that disappear as dropped objects . • New objects that appear from one time step to the next . We refer to new objects as added objects . • Objects that reappear after being hidden for a number of time steps , which requires memory . 1The Alignment problem may also be referred to as the assignment problem . 2In Numpy , Tensorflow and PyTorch this would be this would be implemented as S [ I ] , tf.gather ( S , I ) and torch.gather ( S , I ) , respectively . In this paper , we propose a model that , given two sets of object representations , learns to predict the aligning indices in the general case , where objects may have been added ( Section 3.1.2 ) or dropped 3.1.1 ) . We then propose to use the AlignNet to align sets of objects over time , using an object-based memory ( Section 3.2 ) . 3 THE SELF-SUPERVISED ALIGNMENT MODULE , ALIGNNET .. 3.1 LEARNING TO ALIGN TWO SETS OF OBJECTS VIA SELF-SUPERVISION . We begin with the simple case , where the two sets U and S contain the same objects . We use the term object to refer to a vector representing an object . A set of objects , U , may be converted to an unshuffled list , U , by concatenating3 the elements in an arbitrary order , U = [ u0 , ... , uN−1 ] . To obtain a list of shuffled objects , S = [ s0 , ... , sN−1 ] , we can shuffle elements in U and add noise . Formally , to obtain S we sample a random permutation matrix , A∗ , and some noise to obtain S = UA∗ > + . The true alignment indices , I∗ , are therefore given by I∗ = argmaxA∗ . The argmax is taken over the columns . For an example , see Figure 1B . It may be tempting to find an adjacency matrix , A , that minimises the difference between SA and U . However , firstly , there is no obvious way to quantify this similarity and secondly , it is often the case that an A that minimises this difference is not a permutation matrix . Finally , depending on how the similarity is defined , many different solutions may be found that minimise the difference , when in fact there is only one correct ( alignment ) solution . Rather , our goal is to learn a model , fθ , such that A = fθ ( U , S ) , where each column of A is a categorical distribution over the alignment indices for each element in S , by maximising the expectation , EA∗ log [ fθ ( U , S ) ] . We choose to model fθ as a graph network applied to a bipartite graph4 ( Battaglia et al. , 2018 ) . Graph networks take three inputs : nodes , edges and a global variable . We represent the nodes as the concatenation of shuffled and unshuffled objects , nodes = [ S , U ] , where every node in U is connected to every node in S by two edges , one going in each direction . There are no connections 3 [ · ] is the concatenation operation . 4For details of graph network computation please refer to Battaglia et al . ( 2018 ) between nodes within U or S. We do not use the global variable . In our case |U| = N and |S| = N and the output edges , Eoutput ∈ < 2N×2N . We take the upper-right quarter of this matrix as the predicted adjacency matrix , A ∈ < N×N and apply a softmax across the columns of A . Recall that the predicted alignment indices are given by , I = argmaxA . 3.1.1 LEARNING TO ACCOUNT FOR DISAPPEARING ( DROPPED ) OBJECTS . Here we extend the model to account for objects that disappear between two time steps ; we think of each position in a list as a slot that can either contain an object or be empty . To learn to account for disappearing ( dropped ) objects during self-supervised training , we can construct an S that is not only a shuffled , corrupted version of U but also has some elements that are dropped . Consider an example U = [ A , B , C , D ] and S = [ B , − , − , A ] , where “ − ” represents an empty slot . In this case the aligning indices , I∗ , are ambiguous . The first two aligning indices are straightforward : A has moved to slot three and B has moved to slot zero , leaving slots one and two empty . We could assign C and D to either ( or both ) of the empty slots . This leads to ambiguity in the solution , creating a problem illustrated in Figure 1C . To avoid ambiguity in the solution when objects are dropped , we propose to append an additional null slot to S ( see the Solution in Figure 1C . ) , so that all dropped objects can be indexed by this one null slot , making the solution deterministic . The index of the null slot is N , the number of objects and therefore the aligning indices for the example above are [ 3 , 0 , 4 , 4 ] , since N = 4 . In this updated formulation , the input nodes to the graph network , fθ , are now nodes = [ S , U , ∅ ] , where ∅ is the additional null slot , and the input and output edges are now Einput , Eoutput ∈ < ( 2N+1 ) × ( 2N+1 ) . We take the upper-right quarter of the matrix as a prediction for A ∈ < ( N+1 ) ×N . The adjacency matrix for the problem above is illustrated in Figure 1C . Note that while the columns still sum to one , the rows do not , because the null slot , ∅ , may be aligned with slots in U . This may be formalised as follows : ( 1 ) ∑ i=0 : N Ai , j = 1 for j = 0 , 1 , ... , N − 1 , this ensures that the columns of A are distributions over the shuffled slots , ( 2 ) 0 ≤ ( ∑ j=0 : N−1Ai , j ) ≤ 1 for i = 0 , 1 , ... , N − 1 , this term allows for the special case where if slot si is empty , the sum across row i should be zero , and ( 3 ) ∑ j=0 : N−1AN , j = ndropped , for ndropped ≤ N , this final term , says that if ndropped objects have been dropped from U , then the sum over the N th row will be ndropped .

Summary: This paper presents a model that is able to associate objects seen in a new time step with the objects seen in previous frames and therefore consider the cases of adding new objects, dropping object no longer visible, but keeping them in memory in case they reappear in the future. The model assumes that the location of the objects is given and the only task to perform is the correct association of instances. Results shows that the proposed approach is better than hand-crafted features on different simulated tasks. Additionally the propose model is shown to help to reduce the computational cost for question answering task.