A LIGN N ET : S ELF - SUPERVISED ALIGNMENT MODULE
Anonymous authors
Paper under double-blind review
A BSTRACT
The natural world consists of objects that we perceive as persistent in space and time, even though these objects appear, disappear and reappear in our ﬁeld of view as we move.
This can be attributed to our notion of object persistence – our knowledge that objects typically continue to exist, even if we can no longer see them – and our ability to track objects.
Drawing inspiration from the psychology literature on ‘sticky indices’, we propose the AlignNet, a model that learns to assign unique indices to new objects when they ﬁrst appear and reassign the index to subsequent instances of that object.
By introducing a persistent object-based memory, the AlignNet may be used to keep track of objects across time, even if they disappear and reappear later.
We implement the AlignNet as a graph network applied to a bipartite graph, in which the input nodes are objects from two sets that we wish to align.
The network is trained to predict the edges which connect two instances of the same object across sets.
The model is also capable of identifying when there are no matches and dealing with these cases.
We perform experiments to show the model’s ability to deal with the appearance, disappearance and reappearance of objects.
Additionally, we demonstrate how a persistent object-based memory can help solve question-answering problems in a partially observable environment.
I NTRODUCTION
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 ﬁeld of deep learning, notable progress has been made towards the ﬁrst 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 ﬁrst 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 problem 1 (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:
We are the ﬁrst 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 ﬁnding 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).
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).
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).
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.
Our experiments, on a partially observable question-answering dataset, demonstrate the beneﬁts of having a persistent object-based memory (Section 4.5).
We begin by introducing the general alignment problem.
T HE GENERAL ALIGNMENT PROBLEM
Given two sets of entities, U = { u 0 , ..., u N − 1 } and S =
{ s 0 , ..., s N − 1 } , which contain representations for instances of the same N unique objects, the alignment problem is equivalent to ﬁnding the aligning indices , I =
[ i 0 , ..., i N − 1 ] , such that, u j ≈
s i j for j = 0 , 1 , ..., N − 1 2 .
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 ﬁeld 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 ﬁeld 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.
Note that each of the four columns of A ∗ is a one-hot vector corresponding to the four indices in I ∗ .
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).
T HE S ELF - SUPERVISED A LIGNMENT M ODULE , A LIGN N ET . 3. L EARNING 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 unshufﬂed list, U , by concatenating 3 the elements in an arbitrary order, U =
[ u 0 , ..., u N − 1 ] .
To obtain a list of shufﬂed objects, S =
[ s 0 , ..., s N − 1 ] , we can shufﬂe elements in U and add noise.
Formally, to obtain S we sample a random permutation matrix, A ∗ , and some noise (cid:15) to obtain
S = UA ∗(cid:62) + (cid:15) .
The true alignment indices, I ∗ , are therefore given by I ∗ = arg max
A ∗ .
The argmax is taken over the columns.
For an example, see Figure 1B.
It may be tempting to ﬁnd an adjacency matrix, A , that minimises the difference between SA and
U .
However, ﬁrstly, 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 deﬁned, 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, E A ∗ log[ f θ ( U, S )] .
We choose to model f θ as a graph network applied to a bipartite graph 4 (Battaglia et al., 2018).
Graph networks take three inputs: nodes, edges and a global variable.
We represent the nodes as the concatenation of shufﬂed and unshufﬂed 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.
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, E output ∈ (cid:60) 2 N × 2 N .
We take the upper-right quarter of this matrix as the predicted adjacency matrix, A ∈ (cid:60) N × N and apply a softmax across the columns of A .
Recall that the predicted alignment indices are given by, I = arg max A .
3.1. L EARNING 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 shufﬂed, 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 ﬁrst 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 E input , E output ∈ (cid:60) (2 N +1) ×
(2 N +1) .
We take the upper-right quarter of the matrix as a prediction for A ∈ (cid:60) ( 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) (cid:80) i =0:
N A i,j = 1 for j = 0 , 1 , ..., N − 1 , this ensures that the columns of A are distributions over the shufﬂed slots, (2) 0 ≤ (cid:18) (cid:80) j =0: N − 1 A i,j (cid:19) ≤ 1 for i = 0 , 1 , ..., N − 1 , this term allows for the special case where if slot s i is empty, the sum across row i should be zero, and (3) (cid:80) j =0:
N − 1 A N,j = n dropped , for n dropped ≤ N , this ﬁnal term, says that if n dropped objects have been dropped from U , then the sum over the N th row will be n dropped .
3.1. L EARNING TO ACCOUNT FOR NEWLY APPEARING ( ADDED ) OBJECTS
Now we consider how to deal with newly appearing objects.
To learn to account for appearing (added) objects during self-supervised training, we can construct a U that is missing certain elements that are present in S .
We consider these to be the added elements.
Consider an example (illustrated in Figure 1D.)
U
=
[ A , − , C , − ] and S =
[ B , C , D , A ] .
In this case B and D are the added objects.
Again, accounting for A and C is easy; A has moved to slot 3 and C has moved to slot 1 .
However,
B and D could be assigned to either slot 0 or slot 2 .
Additionally, a single object may be assigned to both slots.
It is imperative that this does not happen, but is hard to enforce through regularisation alone.
Our solution, when adding objects, is to remove ambiguity by training the model to assign the ﬁrst new object that it encounters in S to the ﬁrst available empty slot in U .
By imposing an ordering, we ensure that there is only one solution.
For the model to be able to infer the order of the slots, we append slot indices to the object representations in S and U (after shufﬂing).
Finally, in order to add a new object to a new slot, is it necessary to ﬁrst determine whether the object is new.
To infer if an object is new, we hypothesise that it is not sufﬁcient to consider ﬁrst-order similarity between objects.
Knowing that s 0 is not similar to u i , ∀ i , is not sufﬁcient to say that, s 0 is a new object.
To be conﬁdent that s 0 is new object, we have to be sure that it is less similar to u i , ∀ i than any other element in S .
This reasoning requires the model to understand relative similarity , not just the similarity between independent pairs of entities.
Therefore, we expect that a single step of graph network computations will not be sufﬁcient to solve the alignment problem.
Rather, we expect that recurrent application of a graph network, often referred to as message passing , will be needed to infer information about relative similarity.
In summary our model the AlignNet (Figure 2) is able to deal with dropped objects, by appending a null slot, ∅ to S , allowing entities that have been dropped to be assigned to this slot.
Our model is also able to deal with objects that have been added, by imposing an order in which new objects are assigned to empty slots.
3. L EARNING TO ALIGN SEQUENCES OF SETS OF ENTITIES
The model presented so far does not account for objects that disappear and reappear at some point in the future, so here we describe how to incorporate that functionality.
Given a sequence of sets of objects, [ U t =0 , ..., U t =
T − 1 ] , where T is the number of time steps we would like to assign each unique object in ∪ τ =0: T − 1 U t = τ a slot index which the object maintains over time, t .
This is illustrated in
Figure 3A. The slot indices are analogous to the “sticky indices” referred to by Leslie et al.
(1998) and Pylyshyn (1989).
To achieve this, the model needs to not only account for the appearance of new objects or the disappearance of objects; the model also has to account for the reappearance of objects.
To deal with reappearing objects, the model requires a memory, M , to recall which slot (or index) each object belongs to.
We propose an iterative two-step process for aligning sequences of sets of entities over time.
In the ﬁrst step, the input is aligned with the memory, M , using a trained
AlignNet.
In the second step the memory is updated with the new objects.
More formally, let S t be the input and M t be the memory at the time, t , where M 0 has N empty slots.
At each time step, we compute an adjacency matrix, A t = f θ ( M t , S t ) and obtain a list of aligned objects, Y t , using hard alignment 5 .
The memory is updated, M t +1 =
wM t
+ (1 − w ) Y t where w is the last row of the predicted adjacency matrix and represents the probability, p ∅ ( Y t = ∅ ) , that a slot is empty.
Note that the parameters, θ , of the trained AlignNet, f θ , are held ﬁxed and there is no additional training, these steps are illustrated in Figure 3B.
5 Hard alignment: tf.gather( S t , argmax( A t )) .
[Table]
In sections 3.1.1 and 3.1.2 we only considered cases where objects are dropped from – or added to – a slot.
However, it is often the case that some slots in M t are empty and remain empty.
Therefore when training an AlignNet it is necessary to include cases where a slot remains empty from one time step to the next.
For example, U =
[ A , B , − , D ] and S =
[ A , D , B , − ] , the third slot in U is empty and there is no new object in S that could be assigned to the third slot.
E XPERIMENTS AND R ESULTS
We test our model in various different settings, isolating different behavioural properties and compare to baselines.
We separately show that the AlignNet is able to account for both dropped (Section 4.2) and added (Section 4.3) objects.
We then demonstrate how a trained AlignNet trained with both added and dropped objects may be used to align sequences of sets of objects (Section 4.4).
In our experiments, we report the alignment accuracy and the entropy over the columns of the predicted adjacency matrices.
In the experiments that follow we use a symbolic dataset that represents objects in a partially observable room that contains eight objects.
A random agent moves in the room collecting observations.
Each observation is a set of entities, U , where each u j is a six element vector representing the colour (three elements), shape and position (two elements) of the objects.
The elements in the vectors are normalised between [0 , 1] .
We add uniform noise to each object representation.
The level of noise varies depending on the experiment.
Details of the AlignNet architecture can be found in the appendix G.
4. B ASELINES AND RESULTS WHEN NO OBJECTS ARE ADDED OR DROPPED
When traditionally solving the alignment problem, a similarity matrix, also referred to as an adjacency matrix, is provided.
Construction of this matrix requires hand-crafting a similarity measure.
The AlignNet bypasses the need for this by implicitly learning a similarity measure suitable for solving the alignment problem.
We compare the AlignNet to a model that uses hand-crafted similarity measures to construct the adjacency matrix; the cosine and mean-squared-error, MSE, distances.
For the baselines we align two sets that have the same objects (no added or dropped objects) with noise, Uniform ( − 0 .
1 , 0 . 1) , added to the object representations.
Results in Table 1 show that while the hand-crafted methods achieve competitive accuracy, the entropy of the slot prediction is very high, meaning that differences between hard and soft alignment are large.
While hard alignment assigns a single element in U to a single element in S, for our system to remain differentiable, soft alignment is preferable.
We demonstrate that our model achieves similar performance for both hard and soft alignment 6 .
The baseline models are not designed to deal with cases where there are different numbers of elements in each set, U and S , or where there is no correspondence for some elements.
The AlignNet is able to cope with both of these cases.
In the next section, we demonstrate the AlignNet’s unique ability to deal with cases where elements in set U are not present in set S , we refer to this as dropping (dropping was explained in Section 3.1.1).
[Table]
R ESULTS FOR DISAPPEARING ( DROPPED ) OBJECTS
We train separate AlignNets to drop up to ﬁve objects (i.e. max num to drop ∈ { 1 , 2 , 3 , 4 , 5 } ), where at most max num to drop objects that were present in set U are replaced with empty slots in set S .
A small amount of uniform noise, (cid:15) ∼ Uniform ( − 0 .
1 , 0 . 1) , was added to each object representation in each set.
The AlignNet achieves 99 .
9% alignment accuracy when dropping up to ﬁve objects and the differences between hard and soft alignment are negligible (Figure 7 in the
Appendix).
We use the same hand-crafted baselines as before to give our results context.
Our model signiﬁcantly outperforms the baselines which are not able to handle cases where objects have been dropped (see Table 2 for a brief summary).
Additionally, in Figure 6 we show the effect on the AlignNet’s performance – when trained with at most two objects being dropped – with different amounts of noise (cid:15) ∼
Uniform ( − noise level, noise level ) added to the object representations.
We see that, the AlignNet learns to be robust to 0 .
01 ≤ noise level ≤ 0 .
2 , while the baselines perform poorly.
This is because the AlignNet learns its own similarity function, invariant to noise, rather than using a hand-crafted similarity function that is not invariant to noise.
Finally, we consider how well an AlignNet trained with max num to drop = 2 generalises to 2 < max num to drop ≤ 7 .
We found that an AlignNet trained on max num to drop =generalises with almost perfect accuracy when tested with 2 < max num to drop ≤ 7 (See Table 3).
This may be attributed, in part, to the graph network, which shares parameters across nodes and edges, applying the same model to all nodes, independent of their order.
Now we move on, demonstrating the AlignNet’s unique ability to cope with new objects appearing in set S (that are not present in set U ).
The AlignNet must learn to assign new objects in S indices that correspond to empty slots in U (adding new objects was described in Section 3.1.2).
4. R ESULTS FOR NEW OBJECTS APPEARING ( ADDED )
For each max num to add ∈ { 1 , 2 , 3 , 4 , 5 } we trained ﬁve AlignNets, for each max num to add at least one of the ﬁve models achieves an accuracy > 99% (see Figure 8).
Figure 9 demonstrates that the AlignNet is able to learn a similarity measure that is tolerant to noise; the AlignNet is able to achieve more than 99% accuracy when trained with different levels of noise.
One interesting ﬁnding was that, in order to achieve maximal accuracy and low entropy solutions, multiple message passing steps, in the form of recurrently applying the graph network to the input nodes, were needed.
Figure 12(b) in the Appendix shows that at least two message passing steps ( num reccurent steps = 2 ) are needed to achieve 99 .
9% accuracy.
Multiple message passing steps allow to the model to learn higher order similarities between entities.
The model does not just learn that two entities are similar, but that entities are more or less similar than other entities.
We found that two or more message passing steps also led to very low entropy (sharp attention) solutions, resulting in hard and soft alignment being indistinguishable from each other.
Again, this would be useful if the AlignNet were to be incorporated into a larger architecture since soft alignment is differentiable.
We also consider how well a model trained on examples where at most µ objects were added ( max num to add = µ ), generalises to cases where more than µ objects are added.
The results are shown in Table 4 of the Appendix.
We see that AlignNets trained with larger µ generalise better.
we trained ﬁve AlignNets and show the median performance with the standard deviations.
The
Aligner is trained with random numbers of objects being added, dropped and slots left empty.
The top ﬁve models all achieve 99.9% accuracy.
We have shown that the AlignNet can learn to add and drop a variable number of objects and is also tolerant to noise.
In the next section we look at training an AlignNet, which we will refer to as the
Aligner .
We then use the Aligner to write to an external slot-based memory, without any further training.
We begin by training the Aligner.
4. T RAINING A LIGNERS FOR WRITING TO MEMORY AND R ESULTS ALIGNING SEQUENCES OF SETS OF ENTITIES
When using a trained Aligner to write to memory, M t , the Aligner needs to be able to align an input,
S t , with the memory.
The Aligner must be able to deal with adding and dropping as well as dealing with empty slots.
The Aligner is trained, as before, to align shufﬂed lists, S , with unshufﬂed lists, U .
We must therefore ensure that pairs of S and U exhibit added objects, dropped objects and empty slots (empty slots were explained in Section 3.2).
We train our Aligner to add, drop and keep empty a random number of entities and slots.
We experimented with different numbers of recurrent, message passing steps (Figure 4) and found that the best models had three or four message passing steps, achieving 99 . 9% accuracy.
Now we explore how well our trained Aligner writes to memory.
Taking the best trained Aligner, we use the model to align sequences of up to seven steps with varying numbers of kept objects.
The number of kept objects, num keep slots , refers to the number of occupied slots in S t (Figure 3A. shows an example of num keep slots = 4 ).
The trained
Aligner writes to memory with 100% accuracy across seven time steps and for num keep slots =
{ 1 , 2 , 3 , 4 , 5 , 6 , 7 } , correctly adding all new elements to memory at each time step (see Figuresand 15 in the Appendix).
The AlignNet with a slot-wise object-based memory is an example of a persistent object memory.
In the next section we demonstrate how a persistent object memory can be useful for solving tasks in partially observable environments that contain objects.
4. E PISODIC PARTIALLY OBSERVABLE QUESTION ANSWERING
Finally, we explore question answering using the AlignNet for a domain containing objects that appear, disappear and reappear randomly within an episode (Figure 5).
We use a dataset called epsiodic-sort-of-clevr (see the Appendix H for further details) where the goal is to answer a question about the relationship between objects within an episode.
To solve this task, a model must learn to represent objects in memory and then use this memory to answer a question.
The AlignNet together with a relation network (Santoro et al., 2017), AlignNet-RN, can learn to answer questions from episodic-sort-of-clevr, reaching an accuracy of 93% for episodes containing six objects per scene, two scenes per episode, 100 observations per scene and two objects per observation.
The AlignNet can accurately align objects across multiple scenes within an episode.
Objects that disappear and reappear within a scene are successfully aligned with the memory and thus only new entities are written to empty slots in a ﬁxed sized memory.
We compare this to a relation network, RN, without the AlignNet, where all observed objects are added to the memory at each time step (Pritzel et al., 2017).
The computational complexity of the RN scales quadratically with the number of input objects, so the computational cost of an RN without an AlignNet is much greater than the AlignNet-RN model.
In this example, the AlignNet memory contains 12 object representations (six objects × two scenes), so the AlignNet-RN must compute 12 2 = 144 relations.
In comparison, the RN without the AlignNet has 400 object representations in memory (100 observations per scene × two scenes × two objects per observation).
Therefore the RN has to compute 400 2 = 160 , 000 relations which becomes increasingly infeasible.
[Figure]
The episodic-sort-of-clevr question answering dataset consists of a sequence of observations of objects within a scene, with objects appearing, disappearing and reappearing again across multiple scenes.
Each object is represented symbolically and each scene is given a unique label.
In this example, the episode contains two scenes, and each scene contains six objects, with four observations per scene and two objects per observation.
The AlignNet learns to write objects to an object-aligned memory.
This memory is then provided as input to a relation network, along with the question.
The relation network is trained to answer questions about the episode.
R ELATED W ORK
Human object individuation and identiﬁcation.
The AlignNet is inspired by work by Pylyshyn (1989) on “sticky indices”, a paradigm for explaining how humans keep track of entities as they go in and out of their visual ﬁeld.
When a new object appears, it is assigned an index.
The index then “sticks” to the object, so when the object moves, it retains the same index.
The exact mechanism by which the indices “stick” is not clear.
Noles et al.
(2005) suggest that new object states are compared to old object states and if they are similar enough the old states are updated and if they are not similar enough a new index is created.
This is similar to the function that the AlignNet learns to perform, using multiple message passing steps to learn when something is, or is not, similar enough .
Additionally, evidence suggests that we rely on both object properties and their locations to successfully individuate and identify objects (Leslie et al., 1998; Xu, 1999), which is why we use representation of both object properties and locations in our experiments.
Combinatorics: traditional and deep approaches.
The AlignNet is unique, compared to other traditional and deep learning approaches for solving combinatorics problems because it does not assume access to a similarity or adjacency matrix (Bello et al., 2016;
Vinyals et al., 2015; Milan et al., 2017), which is required for algorithms such as the Hungarian algorithm.
The AlignNet simultaneously learns an implicit similarity measure that helps it to solve the assignment problem.
Additionally, we consider a more general assignment problem where sets may have additional elements for which there is no match.
Interesting work by Lyu et al.
(2019) demonstrates that deep networks can learn permutation matrices that, like the AlignNet, have low entropy solutions.
However, their work does not focus on alignment and the adjacency matrices in the AlignNet are not permutation matrices because an additional row is appended to deal with dropped objects.
Andrychowicz & Kurach (2016) also develop an interesting model for sorting and merging sequences.
However their proposed method is nondifferentiable and does not focus on alignment.
Object tracking: deep learning techniques.
A key novel feature of the AlignNet is its ability to deal with both adding (appearing) and dropping (disappearing) objects.
He et al.
(2018) propose a model for tracking objects using a memory addressing mechanism.
The model terminates trackers when an object disappears, therefore the model cannot deal with reappearing objects, while the
AlignNet can.
Additionally, they assume that all the entities visible at t = 0 are all the objects visible across time, meaning that the model cannot deal with new objects appearing.
In the object tracking literature, when a new object appears, which is not sufﬁciently similar to the objects that have been seen before, it may be considered to be false a detection (Yoon et al., 2019).
Valmadre et al.
; Yang & Chan (2018) also include an external memory for tracking single objects.
Their memory is used for evolving templates for matching over time, while the AlignNet memory keeps track of all the entities that have been seen allowing the model to track multiple pre-extracted entities.
Our work is quite different to traditional object tracking, in that we focus on tracking pre-extracted entities, while most previous work focuses on tracking from pixels.
A move towards entities driving a need for alignment.
Several recent works in relational reasoning (Yi et al., 2018; Janner et al., 2018; Ferreira et al., 2019; Co-Reyes et al., 2019) and reinforcement learning (Zambaldi et al., 2018; Kulkarni et al., 2019) have demonstrated the beneﬁt of working at the entity, rather than pixel level.
In these applications, the environment is fully observable; objects do not appear or disappear from view.
While models exist for extracting entities (Burgess et al., 2019; Greff et al., 2019; Nash et al., 2017; Greff et al., 2016), as far as we know, there are no models suitable for keeping track of correspondence between entities over time in partially observable environments; this is the purpose of the AlignNet.
Learning permutation invariant functions on sets and graphs.
At the core of our model is a graph network (Battaglia et al., 2018; Xu et al., 2018) and in the Appendix we also show that a mutli-layer attention network (or transformer) (Vaswani et al., 2017; Wang et al., 2018) works well too (Figure 13).
We chose these models because they learn functions that are permutation invariant.
This is essential when dealing with sets, which have no order.
Many other models also learn permutation invariant functions such as relation networks (Santoro et al., 2017; Murphy et al., 2018; Battaglia et al., 2016), deep sets (Zaheer et al., 2017) and graph attention networks (Veliˇckovi´c et al., 2017).
Any of these models may have been suitable to place at the core of the AlignNet.
However, we found that the transformers scaled better and the graph networks were easy to train.
Finally, more similar to our own work, Zhang et al.
(2019) propose a model for predicting sets.
However when determining correspondences between elements in the sets they still use hand-crafted set losses and their model is not a simple feed forward model because they apply gradient descent on the set itself.
Reading from and writing to memory.
The AlignNet keeps track of multiple pre-extracted entities, over time, with a persistent slot-base memory, where each slot contains the representations for a single object.
When writing to memory, new observations should be placed in new slots and observations similar to those already in memory should either be updated or left as they are.
There are few examples of models that learn to write to memory (Graves et al., 2014).
This is often because learning what and where to write can be challenging.
In a special case, Parisotto & Salakhutdinov (2017) propose the Neural Map, a 2D spatial memory where memories are stored based on the location of the agent.
Often all observations, along with a look up key, are written to (an episodic) memory (Pritzel et al., 2017; Blundell et al., 2016) using a differentiable neural dictionary, DND.
In contrast, the AlignNet has a ﬁxed sized memory and learns to only write new entities to empty slots in memory.
To obtain sharp attention, we argue that knowing the similarity between two things is not sufﬁcient, we also need to know how similar a (key, query) pair is relative to other (key, query) pairs.
By incorporating multiple message passing steps, we demonstrate that our model may achieve low entropy (sharp) attention over slots in memory.
Another method for sharply attending to slots in memory is variational memory addressing (Bornschein et al., 2017).
However, Bornschein et al.
(2017) focus on reading from memory, rather than adding new entities to memory.
C ONCLUSION
Our AlignNet model is unique in several ways.
It is a model for dealing with alignment in the cases where two sets being aligned may have varying numbers of objects and each set may contain objects that are not present in the other.
Additionally, the AlignNet bypasses the need for hand-crafting a similarity measure, implicitly learning a similarity measure in order to solve the alignment problem.
Finally, we demonstrate the beneﬁts of having a persistent object-based memory when solving a question answering problem in a partially observable environment.
A CKNOWLEDGEMENTS Acknowledgements
A R ESULTS FOR DROPPING OBJECTS
Figure 6 shows the effect of different levels of corruption on alignment accuracy; the AlignNet is able to achieve more than 99% test accuracy at all noise levels when at most two objects are dropped.
Figure 7 shows test prediction accuracy and entropy for an AlignNet trained to drop up to ﬁve objects.
Table 3 demonstrates that an AlignNet trained to drop N objects, can drop more than N objects with almost perfect accuracy
How well does a model trained with at most two objects being dropped ( max num to drop
= 2 ) generalise to cases were more than two objects are dropped?
We trained an AlignNet with max num to drop = 2 and tested the model on examples where more than two objects were dropped.
We see that the model generalises well to more objects being dropped than those seen during training.
[Table]
B R ESULTS FOR ADDING OBJECTS
Figure 8 shows test prediction accuracy and entropy values for an AlignNet trained to add up to ﬁve objects.
Figure 9 shows the effect of different levels of corruption on alignment accuracy; the AlignNet is able to achieve more than 99% accuracy at all noise levels when at most two objects are added.
Table 4 demonstrates that an AlignNet trained to add at most µ objects, can add more than µ objects with high accuracy, particularly when µ ≥ 3 .
[Figure]
C R ESULTS FOR BOTH ADDING AND DROPPING OBJECTS
In most real world cases, it is likely that objects are both added to and dropped from U to give an S that has both empty slots and new entities.
We explore this case here.
We assess the AlignNet’s ability to deal with cases where objects are both added and dropped.
We show the effect of different levels of corruption ( noise level = { 0 . 01 , 0 .
05 , 0 . 1 , 0 . 2 } ) on alignment accuracy when at most two objects are added and at most two objects are dropped (See
Figure 10).
Again, we found that two recurrent, message passing steps were necessary and sufﬁcient to achieve maximal performance and low entropy (sharp attention) solutions.
The AlignNet achieves more than 99% accuracy at all noise levels when at most two objects are added, except at noise level = 0 . 2 , where the best model achieves 98.7% accuracy.
Visualisations of the results are shown in Figure 11.
[Figure]
D H OW MANY MESSAGE PASSING STEPS ARE NEEDED ?
Figure 12 shows that, while only one message passing step is needed when dropping objects, at least two message passing steps are needed when adding objects.
Additional message passing steps can make training less stable.
[Figure] [Figure]
[Table]
[Figure] [Figure]
Entities present in both the Input and the Added (and) or Dropped columns are placed in the same slot that they were found in, in the Input .
[Figure]
More than two message passing steps may results in instability and higher variability during training.
[Figure]
E T HE A LIGN N ET AS A GENERAL FRAMEWORK
The AlignNet may be more broadly thought of as a framework with the following key elements: (1) transforming sets in to lists of slots, with a ﬁxed order, sorting one list with respect to another, (2) a deterministic order for adding new objects, (3) appending an additional null slot for modelling dropped objects and (4) computing a loss between the predicted adjacency matrix and the true adjacency matrix.
The model that takes the shufﬂed and unshufﬂed lists as inputs and outputs the adjacency matrix is the core of our model.
In the experiments above, we used a graph network as the core, however other networks could be substituted.
For example, transformers (Vaswani et al., 2017; Wang et al., 2018) may be used.
Figure 13 shows that with two layers of transformers, the
AlignNet achieves 100% accuracy.
F A LIGNING SEQUENCES OF SETS OF ENTITIES OVER TIME .
Figure 14 shows the accuracy with which a trained Aligner writes to memory, for different numbers of kept slots, num keep slots , in the input, with no further training.
This is visualised in Figure 15.
Each symbolic object in each slot of U , S and M is rendered as an image in a different row (or slot).
We use the best trained Aligner from Section 4.4 and use it to align sequences – in which objects appear, disappear and reappear – with different numbers of kept objects.
The model achieves 100% accuracy at each alignment step for sequences with different numbers of kept objects.
G A LIGN N ET ARCHITECTURE AND TRAINING DETAILS .
For the results presented in this paper, the AlignNet was trained using an RMS-Prop optimiser with learning rate 2e-4, however Adam worked well too.
For the graph network, we used a four layer
MLP with sizes [128, 256, 512, 1024] for the edge model (Battaglia et al.
and used a three layer MLP for the node model with sizes [32, 32, 7].
We found that the model was robust to many different MLP conﬁgurations and worked with edge models with as few as two layers and 32 nodes in each layer.
However, training was faster with deeper networks and more nodes.
H E PISODIC QUESTION ANSWERING H. E PISODIC - SORT - OF - CLEVR DATASET
The episodic-sort-of-clevr question answering dataset contains episodic sequences of objects arranged together in a scene , along with a question about a scene and a corresponding answer.
The object shapes can be triangles or squares and each object has a unique colour.
There are several types of question, all asking about the relationship between objects in the scene.
For example: ‘What is the colour of the farthest object from the blue object?’
and ‘How many objects have the same shape as the red object?’
Questions are encoded as a binary string.
Speciﬁcally, a one-hot representation of an object colour is provided, along with a binary coding of the question type (‘farthest distance’, ‘nearest distance’ or ‘object count’) and a binary coding of an object shape (‘triangle’ or ‘square’).
The scene is presented episodically as a sequence of observations.
Each object can appear, disappear and reappear again, randomly, across the sequence of observations, just as objects in the natural world can appear, disappear and reappear as an observer changes view-point, for example.
To answer a question successfully, a model must keep track of all the objects across a sequence of observations, and then answer a question about the entire scene.
For example, suppose a red square appears in the ﬁrst observation of a scene, a blue square appears in the second observation and in the third observation, a green triangle appears along with a repeated appearance of the red square.
In order to answer a question about which object is nearest to (or furthest from) the green triangle, a model must retain a representation of the full scene in its memory, across all observations, before it can answer the question.
In the episodic-sort-of-clevr dataset, there can also be multiple scenes within an episode, each containing a different set of objects.
Each question is directed towards one scene only using a unique scene identiﬁcation label.
Each scene can consist of a sequence of observations.
These observations, each from a different scene can be interleaved amongst each-other.
Therefore, in order to successfully answer a question, a model must learn to represent each scene from multiple observations, without confusing objects from across scenes before it can then answer a question about the relationship between the objects within a particular scene.
As the number of scenes and observations increase, it will become increasingly difﬁcult for many standard models to scale sufﬁciently to support relational question answering.
The episodic-sort-of-clevr dataset is an episodic version of the Sort-of-Clevr dataset (Santoro et al., 2017), in which all objects are presented simultaneously within a single scene, which in turn, was inspired by the Clevr question answering dataset.
H. A LGIN N ET WITH R ELATION N ETWORK MODEL
We train an AlignNet together with a relation network, RN (Santoro et al., 2017), which we refer to as the AlignNet-RN, to solve the episodic-sort-of-clevr tasks.
We ﬁrst train the AlignNet to align objects within scenes and across observations.
After training, the AlignNet may be used to write entities, in a given episode, to a memory, M =
[ m 0 , m 1 , ...., m N ] .
At the end of an episode, when all entities have been written to memory, the memory can be treated as a set of objects, since the order is no longer important.
Next, we provide the AlignNet memory for a given episode as input to a RN.
The AlignNet memory can be naturally combined with the RN, because the RN treats its inputs as a set of objects, and because the RN has in-built relational structure that is matched to the task of reasoning about relationships.
The RN receives the question q and the memory, M , as input and is trained to produce an answer a = f φ (cid:16)(cid:80) i,j g ψ ( m i , m j , q ) (cid:17) as output.
For g ψ , we use a four-layer MLP with 256 units per layer and ReLu non-linearities.
For f φ we use a two-layer MLP, again with 256 units per layer and ReLu non-linearities, followed by a linear layer.
We use a learning rate of 0 . 0001 and a batch size of 8 .
We