AlignNet: Self-supervised Alignment Module
Authors: Antonia Creswell
Luis Piloto
David Barrett
Kyriacos Nikiforou
David Raposo
Marta Garnelo
Peter Battaglia
Murray Shanahan
Venue:  (2020)
Keywords: Graph networks
alignment
objects
relation networks

Abstract: 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 field 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 first 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.

Under review as a conference paper at ICLR 2020
ALIGNNET: SELF-SUPERVISED ALIGNMENT MODULE
Anonymous authors
Paper under double-blind review
ABSTRACT
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 field 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 as- sign unique indices to new objects when they first 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 graph network applied to a bipartite graph, in which the input nodes objects from two sets that we wish to align. The network is trained to predict the edges which con- nect 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 deal with the appearance, disappear- ance and reappearance of objects. Additionally, we demonstrate how a persistent object-based memory can help solve question-anae answering problems in partially observable environment.
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,
when they disappear and reappear after some delay (Leslic et al. 1998 Leslie & Kaldy, 2001: Noles et al., 2005; Pylyshyn| 1989) (Figure [TA).
The way 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, know that 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 often attributed "sticky indices" (Leslie et al. 1998 Leslie & Kaldy] 2001] Noles et al. 20055 (Pllyshyn| 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 represen- tations (Engelcke et al.) 2019, Greff al. 2016) from static visual scene, these include MONet (Burgess et al. 2019) and TODINE (Greff et al. 2019). These models present compelling unsuper- vised segmentation results. However, while MONet or IODINE can be used to extract objects from an agent's 's visual input over time, neither model determines the correspondence between step and the next, For example, if MONet or IODINE were used to extract entities from the frames shown in Figure TA the model would not know that the purple cone at 0 was the same as the purple cone at t 1.
time
In this paper, 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) persistent slot-wise object-based memory that stores unique objects in

---

Under review as a conference paper at ICLR 2020
unique slots, corresponding to their index, (2) the ability to assign 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 3 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 problen (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 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 persistent object-based memory.
Our contributions are as follows:
1. We are the first formulate the alignment problem such that neural network may learn similarity measure for aligning entities (Section 3). We propose the self-supervised Align- Net, 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 & 4.4
3. We show that the AlignNet implicitly learns similarity measure for aligning entities, tol- erant to noise, that outperforms hand-crafted similarity measures (Section 4.1 Figures 6 &
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 7cp. while remaining differentiable.
5. Our experiments, benefits of having persistent object-based memory (Section 4.5)
partially observable question-answering dataset, demonstrate the
We begin by introducing the general alignment problem.
2 THE GENERAL ALIGNMENT PROBLEM
Given two sets of entities, 24 = {uu. .... uN-1 }1} and S = fso.... .... SN-1), which contain represen- tations for instances of the same N unique objects, the alignment problem equivalent to finding the aligning indices, I = j....., -.- I). such that, U R si; for j 0, 1, ,..., N F The represen- tations 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 Ut contain similar entities, they may also be thought of as sets of objects observed in scene, by agent, at different points in time. As 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 [.A)) 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 account for:
Objects that disappear from one time step to the next, such an object leaving the field of view after head movement. We refer to objects that disappear as dropped objects.
New objects that appear from one time step to the next. We refer new objects as added objects.
Objects that reappear after being hidden for a number of time steps, which requires mem- ory.
'The Alignment problem may also he referred to as the assignment problem. Pn Numpy, Tensorflow and PyTorch this would be this would he implemented as S[ [f...erg (S, 1) and terch. gather cerch.gatherree (S, 1), respectively.
2

---

Under review as a conference paper at ICLR 2020
Reference Shuflld,, Shuille des Nign idk,
Adjacenay morkx.
Soft
Oxject disapears
1. lcolured squares)
Otjectt appwars
A Objects wapear ditapperr and change slote. Probleme
Data gerwortonn
Selution: apownd rut stot,
Adacendy
Ids,
KSk, P
3
0
Zor
1
d'aappariing (droppedi objects.
Prob non detorminssie solution
Solution: add new objects in endar
Adjcceryy matrkk..
t/
or
0
Learn to deal with new jaddedi objects.
Figure 1: Our solution to the alignment problem. A. A schematic demonstrating objects appearing and disappearing over time. B. An example of how an shuffled list, S' and aligning indices, are constructed from unshuffled list U, also in the case where objects are C. dropped and D. added. Note that each of the four columns of A is a one-hot vector corresponding to the four indices 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) dropped 3.1.I. We then propose to use the AlignNet to align sets of objects over time, using an object-based memory (Section 3.2
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 the term object to refer to a vector representing an object. A set of objects, U, may be converted to an wnshuffled list, U, by concatenating the elements in arbitrary order, U [uo:.....NNI]] To obtain a list of shuffled objects, 5 (9...., vo....,8--11. SN-1]. we can shuffle elements U and add noise. Formally, obtain S we sample random permutation matrix, A", and some noise e to obtain S = U + e. The true alignment indices, /", are therefore given by I = arg max A", The argmax is taken over the columns. For an example, see Figure
It may be tempting to find an adjacency matrix, A, that minimises the difference between SA and U. However, firstly, there no obvious way to quantify this similarity and secondly, it is often the case that an A that minimises this difference is not permutation matrix. Finally, depending the similarity defined, many different solutions may be found that minimise the difference, when in fact there is only one correct (alignment) solution. Rather, goal is to learn model, fe. such a categorical distribution over the alignment indices
how
that = (u, S), where each column of for each element in S, by maximising the expectation, E log[fo(U,
We choose to model fo as a graph network applied to a bipartite graph Battaglia et al. 2018). Graph networks take three inputs: nodes, edges and global variable. We represent the nodes as the concatenation of shuffled and unshuffled objects, nodes = [S, U]. where every node 24 is connected every node in S by two edges, going in each direction. There are connections
$[.] is the concatenation operation.
*orr details of graph network computation please refer Battaglia al. 2018
3

---

Under review as a conference paper at ICLR 2020
between nodes within U or S. We do not use the global variable. In our case [4] = NN and |S| N and the output edges, Eoutput E R2Nx2N. We take the upper-right quarter of this matrix as the predicted adjacency matrix, A E RNx and apply a softmax across the columns of A. Recall that the predicted alignment indices are given by, I = arg max A.
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 slot that can either contain an object or be empty. To learn to account for disappearing (dropped) objects during self 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 = AA, B, C. D] and 5 = [B, -. -, A], where " represents an empty slot. In this case the aligning indices, [" 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 ambiguity in the solution, creating problem illustrated in Figure
To avoid ambiguity in the solution when objects are dropped, we propose to append an additional null slot to S (see the Solution in (ggure[C.), 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, fe, are now nodes = [S, U, Ø), where Ø is the additional null slot, and the input and output edges are now Fimput, Fooutput +)(2N+)) We take the upper-right quarter of the matrix as a prediction for A E R( The adjacency matrix for the problem above is illustrated in Figure 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=o. A.j = 1 for j = 10.,...,NN N - 1, this ensures that
E5= Ane)) S for i = 1,.., N - 1. this term allows for the special case where if slot 8; is empty, the sum across row i should be zero, and (3) Σj= AN.j = Rdrropedd for Ridropped S. N, this final term, says that if ndropped objects have been dropped from U, then the sum over the N row will be ndropped-
the columns of A are distributions over the shuffled slots, (2) 0 S
3.1.2 LEARNING 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 [..) 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 slot 2. Additionally, 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 first new object that it encounters in S to the first 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 shuffling).
Finally, in order to add a new object to new slot, is it necessary to first determine whether the object new. To infer if an object is new, we hypothesise that it is not sufficient to consider first-order similarity between objects. Knowing that so is not similar to u; Vi, is not sufficient to say that, so is a new object. To be confident that so is new object, we have to be sure that it is less similar to u; Vi 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 sufficient to solve the alignment problem. Rather, expect that recurrent application of a graph network, often referred to as message passing, will be needed to infer information about relative similarity.
4

---

Under review as a conference paper at ICLR 2020
Piwdiceed
Sharffeck
Graph (Coni
Shuill node
Cioie entoyy
For arironne makiing
For traling only
Aign
Figure 2: Schematic of the self-supervised AlignNet.
Inpuzt
Marnery a
Algnnd sequence,
slet
Seet
Slet
Skan
Siet
Blvt
beher Leaming aczou reapoewing otjecta, Pictua curpuze
Aagnert
You
Mast
Y, algrrmnt
Blet
@
@
NAC. the memory attana
Skat 1
Memery update:
Ma 19
Inita
Allgn nppu with
Uppate
Allgn input with approschh
Updite
Mign incut with memory
Update Memory
Figure 3: The Goal: Aligning sequence of sets of entities. A. Given input sequence of sets of objects (left), we would like to obtain a sequence of objects that are aligned over time (right). This means that once new object appears is assigned "sticky index" (Pylyshyn 1989) corresponding to slot index, any subsequent instances of that same object must be assigned to that same slot. B. To achieve this we propose two-step process: (1) Align inputs with respect to the memory, (2) update the memory with the new objects.
In summary, our model, the AlignNet (Figure(), 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 also able to deal with objects that have been added, by imposing an order in which new objects are assigned to empty slots.
3.2 LEARNING 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 incorporate that functionality. Given a sequence of sets of ob- jects. -4, 1). where T is the number of time steps, we would like to assign each unique object U+=0:T 14 slot index, which the object maintains over time, t. This is illustrated in Figure BA. 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 memory, M, to recall which slot (or index) each object belongs to. We propose iterative two-step process for aligning sequences of sets entities over time. In the first step, the input is aligned with the memory, M, using a trained AlignNet. In the second step the memory updated with the objects.
More formally, let S, be the input and M, be the memory at the time, f, where Mj has N empty slots. At each time step, we compute an adjacency matrix, = fo S.) and obtain list of aligned objects, Y, using hard alignment The memory is updated, M w., + (1 w,YY, where w the last row of the predicted adjacency matrix and represents the probability, (Y, 2). that slot is empty. Note that the parameters, 8, of the trained AlignNet, fe, are held fixed and there is no additional training, these steps are illustrated in Figure 3B.
Hard alignment: f. gather (St, a argmax (Aa).
5

---

Under review as a conference paper at ICLR 2020
Table 1: Comparing the AlignNet to hand-crafted baselines. Each experiment was run five times and the mean is shown here. (Training curves with standard deviations across runs are show in Figure 7
under mww.num.to
Accuracy
Entropy
Difference between hard and soft alignment MSE
AlignNet (Ours)
99.9%
0.0
0.0
MSE Cosine
99.2%
14.6
1.8
96.5%
15.8
2.6
In sections 3.1.1 and 3.1.2 we only considered cases where objects are dropped from or added to slot. However, often the case that some slots in M, are empty and remain empty. Therefore when training an AlignNet is necessary to include cases where a slot remains empty from one time step to the next. For example, = [A, B, -, D] and [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.
4 EXPERIMENTS AND RESULTS
We test our model various different settings, isolating different behavioural properties and com- pare 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 experiments, we report the alignment accuracy and the entropy over the columns of the predicted adjacency matrices.
In the experiments that follow we use symbolic dataset that represents objects in a partially observ- able room that contains eight objects. random agent moves in the room collecting observations. Each observation is a set of entities, 24 where each Uj is six element vector representing the colour (three elements), shape and position (two elements) of the objects. The elements 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.1 BASELINES AND RESULTS WHEN OBJECTS ARE ADDED OR DROPPED
When traditionally solving the alignment problem, similarity matrix, also referred to adja- cency 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 adjacency matrix; the cosine and mean geean-square-erron MSE, distances. For the baselines we align two sets that have the same objects (no added dropped objects) with noise, Uniiform(-01.0p added to the object representations.
model that uses hand-crafted similarity measures to construct the
Results in Table ! 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 single element in 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 alignmen
The baseline models are not designed to deal with cases where there are different numbers of ele- ments in each set, 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, demonstrate the AlignNet's 's unique ability to deal with cases where elements in set are not present in set S, we refer to this as dropping (dropping explained in Section 3.I.1.
Hard alignment: f .at.herreeeee (8, argmax (A) ) Soft alignment: tf.matmul (S,

---

Under review as a conference paper at ICLR 2020
Table 2: Comparing the AlignNet to hand-crafted baselines for adding and dropping. The AlignNet achieves higher accuracy when objects are added and dropped compared to baselines. (See Figures and 8 in the Appendix for more extensive results with more objects added and dropped.)
% accuracy
max mooz_num_ooaaee
0
0
2
max www.suuumtodddree
0
AlignNet (Ours)
99.9%
99.9%
99.9%
99.9%
MSE
99.2%
91.2%
93.7%
88.5%
Cosine
96.5%
87.9%
91.1%
86.2%
4.2 RESULTS FOR DISAPPEARING (DROPPED) OBJECTS
We train separate AlignNets drop up five objects (i.e. mar num to drop E (1, 5)), where at most maz maz-num to drop objects that were present in set U are replaced with empty slots in set S. small amount of uniform noise, c ~ Uni foiforrmO... 0.1, 0.1). was added to each object representation in each set. The AlignNet achieves 99.9% alignment accuracy when dropping up to five 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 results context. Our model significantly outperforms the baselines which are not able to handle cases where objects have been dropped (see Table 2 for 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 "niffor((nnoic level, noise_lvee)) added to the object representations. We see that, the Align- Net learns be robust to 0.01 < noise_leve < 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 www..num_oouue
= 2 generalises to to door
2 < mar_nuu nwm.te drop < 7. We found that an AlignNet trained on generalises with almost perfect accuracy when tested with 2 < maz_num commtoo.hae to drod < 7 (See Table B.. 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 L4). The AlignNet must learn to assign new objects in S indices that correspond to empty slots in (adding new objects was described in Section 3.1.2).
4.3 RESULTS FOR NEW OBJECTS APPEARING (ADDED)
For each mar num to add E (1,2,3, 5] we trained five AlignNets, for each mar num to add least one of the five models achieves an accuracy > 99% (see Figure g) Figure 9 demonstrates that the AlignNet is able to learn similarity measure that is tolerant to noise; the AlignNet is able achieve more than 99% accuracy when trained with different levels of noise.
One interesting finding 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-
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 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 larger architecture since soft alignment is differentiable.
are needed
We also consider how well model trained 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 of the Appendix. We see that AlignNets trained with larger l generalise better.

---

Under review as a conference paper at ICLR 2020
Figure How many message passing steps are needed? For each num.recuryent.ster we trained five 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 five models all achieve 99.9% accuracy.
1.00
30100010
3.35
0.30
3.35
330
425
We have shown that the AlignNet can learn to add and drop 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 external slot-based memory, without any further training. We begin by training the Aligner.
4.4 TRAINING ALIGNERS FOR WRITING TO MEMORY AND RESULTS ALIGNING SEQUENCES OF SETS OF ENTITIES
When using trained Aligner to write to memory, Mj, the Aligner needs to be able to align an input, S,, with the memory. The Aligner must be able deal with adding and dropping as well as dealing with empty slots. The Aligner is trained, as before, to align shuffled lists, S, with unshuffled lists, U.
We must therefore ensure that pairs of and 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 St (Figure 33A.. shows an example of num num_keep_ssots = 4). The trained Aligner writes to memory with 100% accuracy across seven time steps and for num_keep.ss {1, 2.3.4 4, 5, 6, 7}. correctly adding all new elements to memory at each time step (see Figures 14 and L5 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.5 EPISODIC PARTIALLY OBSERVABLE QUESTION ANSWERING
Finally, we explore question answering using the AlignNett for a domain containing objects that appear, disappear and reappear randomly within an episode (Figure 5) We use dataset called epsiodic-sort-eeeo -off---- (see the Appendix[H]for further details) where the goal is to answer a question about the relationship between objects within 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 relation network (Santoro al. 2017), AlignNet-RN, can learn to answer questions from episodic--sort-ofcclee, reaching an accuracy of 93% for episodes contain- ing 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 scene are successfully aligned with the memory and thus only new entities are written to empty slots in a fixed 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 x two scenes), so the AlignNet-RN must compute 12a = 144 relations. In comparison, the RN without the AlignNet has 400 object representations in memory (100 obser- vations per scene x two scenes x two objects per observation). Therefore the RN has to compute 400 = 160, 000 relations which becomes increasingly infeasible.
8

---

Under review as a conference paper at ICLR 2020
soneewent)
AeitgNNet
Relation
Queation: ropm 2. what colour etger
Network
Figure 5: Episodic-sorttof--eer question answering with the AlignNet and Relation Network The mrssodic-sort-of-llee question answering dataset consists of a sequence of observations of ob- jects within seene, with objects appearing, disappearing and reappearing again across multiple scenes. Each object is represented symbolically and each scene is given a unique label. this ex- ample, 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 input relation network, along with the question. The relation network is trained to answer questions about the episode.
5 RELATED WORK
Human object individuation and identification. 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 field. When new object appears, it is assigned index. The index then "sticks" to the object, so when the object moves, retains the same index. The exact mechanism by which the indices "stick" is not clear. Noles et al.] (2005) suggest that new object states compared to old object states and if they are similar enough the old states are updated and if they are not similar enough new index is created. This is similar 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.
Combinatories: traditional and deep approaches. The AlignNet is unique. compared to other traditional and deep learning approaches for solving combinatorics problems because does not assume access to similarity or adjacency matrix (Bello et al. 2016, Vinyals et al., [015]] Milan jet al. 2017), which is required for algorithms such as the Hungarian algorithm. The Align si- multancously 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 ele- ments for which there is no match.
Interesting work by Lyu et al. (2019) demonstrates that deep networks can learn permutation ma- trices that, like the AlignNet, have low entropy solutions. However, their work does not focus alignment and the adjacency matrices in the AlignNet are not permutation matrices because an addi- tional row appended to deal with dropped objects. [Andrychowicz & Kurach] (2016) also develop an interesting model for sorting and merging sequences. However their proposed method is non- differentiable 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 model for tracking objects using 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 / 0 are all the objects visible across time, meaning that the model cannot deal with new objects appearing. In the object tracking literature, when have been seen before, it may be considered to be false a detection (Yoon et al., 2019)).
object appears, which is not sufficiently similar to the objects that
Valmadre et al. 2017); Yang & Chan (2018) also include 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 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 tracking from pixels.

---

Under review as a conference paper at ICLR 2020
A move towards entities driving need for alignment. Several recent works in relational reason- ing et al. 2018 Janner et al. 20188: Ferreira et al. 2019 Co-Reyes et al. 2019) and reinforcement learning (Zambaldi al. [0018] Kulkarmi et al. 2019) have demonstrated the benefit of working at the entity, rather than pixel level. In these applications, the environment is fully observable; objects do not appear disappear from view.
While models exist for extracting entities (Burgess et al, 2019 Greff et al. 2019; Nash et al, 2017 Greff al.] 22016), 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 graph network (Battaglia al. 2018 Xu et al., 2018)) and in the Appendix we also show that mutli-layer attention network (or transformer) Vaaswani et al., 2017; Wang et al. (0188) 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 order. Many other models also learn permutation invariant functions such as relation networks (Santoro et al., 2017; Murphy et al., 20118] Battaglia et al., (016)), deep sets (Zaheer et al. 2017) and graph attention networks (Veličković 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 enti- ties, time, with 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 ob- servations similar those already in memory should either be updated 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 special case, Parisotto & Salakhutdinov (017)) propose the Neural Map. a 2D spatial memory where memories are stored based on the lo- cation 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. (016) using differentiable neural dictionary. DND. In contrast, the AlignNet has a fixed 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 sufficient, also need to know how similar (key, query) pair is relative 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 slots in memory is variational memory addressing (Bornschein et al.) 2017). However, [Borscchein et al. (017) focus reading from memory, rather than adding new entities to memory.
6 CONCLUSION
Our AlignNet model is unique several ways. It 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 benefits of having a persistent object-basce memory when solving a question answering problem in partially observable environment.
ACKNOWLEDGEMENTS
Acknowledgements
10

---

Under review as a conference paper at ICLR 2020
REFERENCES
Marcin Andrychowicz and Karol Kurach. Learning efficient algorithms with hierarchical attentive memory. arXiv preprint arXiv.1602.03218, 2016.
Renée Baillargeon. The acquisition physical knowledge in infancy: A summary in eight lessons. Blackwell handbook of childhood cognitive development, 1(46-83):1, 2002.
Peter Battaglia, Razvan Pascanu, Matthew Lai, Danilo Jimenez Rezende, et al. Interaction networks for learning about objects, relations and physics. In Advances in neural information processing systems, pp. 4502 4510, 2016.
Peter W Battaglia, Jessica B Hamrick, Victor Bapst, Alvaro Sanzhez-Gonzalez, Vinicius Zambaldi, Mateusz Malinowski, Andrea Tacchetti, David Raposo, Adam Santoro, Ryan Faulkner, et al. Relational inductive biases, deep learning, and graph networks. arXiv preprint arXiv:1806.0126. 2018.
Irwan Bello, Hieu Pham, Quoe V Le, Mohammad Norouzi, and Samy Bengio. Neural combinatorial optimization with reinforcement learning. arXiv preprint arXiv: 161 1eee:66111.9900 1.09940, 2016.
Charles Blundell, Benigno Uria, Alexander Pritzel, Yazhe Li, Avraham Ruderman, Joel Z Leibo, Jack Rae, Daan Wierstra, and Demis Hassabis. Model-free episodic control. arXiv preprint arXiv:1606.04460, 2016.
Jörg Bornschein, Andriy Mnih, Daniel Zoran, and Danilo Jimenez Rezende. Variational memory addressing in generative models. Advances in Neural Information Processing Systems, pp- 3920-3929) 2017.
Christopher P Burgess, Loic Matthey, Nicholas Watters, Rishabh Kabra, Irina Higgins, Matt Botvinick, and Alexander Lerchner. Monet: Unsupervised scene decomposition and represen- tation. arXiv preprint
2019.
John D. Co-Reyes, Rishi Veerapaneni, Michael Chang, Michael Janner, Chelsea Finn, Jiajun Wu, Josh Tenenbaum, and Sergey Levine. Discovering, predicting. and planning with objects. Inter- national Conference on Machine Learning Workshop, 2019.
Martin Engelcke. Adam R Kosiorek. Oiwi Parker Jones, and Ingmar Posner. Genesis: Gener- ative scene inference and sampling with object-centric latent representations. arXiv preprint 190.v 13052, 2019.
Fabio Ferreira, Lin Shao, Tamim Asfour, and Jeannette Bohg. Learning visual dynamics models of rigid objects using relational inductive biases. arXiv preprint wrXiv:1909.03745 2019.
Alex Graves, Greg Wayne, and Ivo Danihelka. Neural turing machines. arXiv:1410.5401, 2014.
arXiv preprint
Klaus Greff, Antti Rasmus, Mathias Berglund, Tele Hao, Harri Valpola, and Jürgen Schmidhuber. Tagger: Deep unsupervised perceptual grouping. Advances in Neural Information Processing Systems, pp. 4484-4492, 2016.
Klaus Greff, Raphaël Lopez Kaufmann, Rishab Kabra, Nick Watters, Chris Burgess, Daniel Zoran, Loic Matthey, Matthew Botvinick, and Alexander Lerchner. Multi-object representation learning with iterative variational inference. arXiv preprint wwXiv:/903.00440 2019.
Zhen He, Jian Li. Daxue Liu, Hangen He, and David Barber. Tracking by animation: Unsupervised learning of multi-object attentive trackers. arXiv preprint arXiv: 1809. 03137, 2018.
Michael Janner, Sergey Levine, William T Freeman, Joshua B TTnenbaum, Chelsea Finn, and Jiajun Wu. Reasoning about physical interactions with object-oriented prediction and planning. arXiv preprint arXiv: 10972, 2018
Tejas Kulkarni, Ankush Gupta, Catalin Ionescu, Sebastian Borgeaud, Malcolm Reynolds, Andrew Zisserman, and Volodymyr Mnih. Unsupervised learning of object keypoints for perception and control. arXiv preprint arXiv: 1906. 1883, 2019.

---

Under review as a conference paper at ICLR 2020
Alan M Leslie and Zsuzsa Kaldy. Indexing individual objects in infant working memory. Journal of Experimental Child Psychology, 78(1):61-74, 2001.
Alan M Leslie, Fei Xu, Patrice D Tremoulet, and Brian J Scholl. Indexing and the object concept: systems. Trends in cognitive sciences, 2(1):10-18, 1998.
"eeeelopngwhaa'adddweere andwhere systemss'ssstmss.
Jiancheng Lyu, Shuai Zhang. Yingyong Qi, and Jack Xin. Autoshufflenet Learning permutation matrices via an exact lipschitz continuous penalty in deep convolutional neural networks. arXiv preprint arXiv:I901.000 08624,, 2019.
Anton Milan, S Hamid Rezatofighi, Ravi Garg, Anthony Dick, and lan Reid. Data-driven approxi- mations to np-hard problems. In Thirty-FFrst AAAI Conference on Artificial Intelligence, 2017.
Ryan L Murphy, Balasubramaniam Srinivasan, Vinayak Rao, and Bruno Ribeiro. Janossy pool- ing: Learning deep permutation-invariant functions for variable-size inputs. arXiv preprint arXiv: 1.01900, 2018.
Charlie Nash, Ali Eslami, Chris Burgess. Irina Higgins, Daniel Zoran, Theophane Weber, and Peter Battaglia. The multi-entity variational autoencoder. In NIPS Workshops, 2017.
Nicholaus S Noles, Brian Scholl, and Stephen R Mitroff. The persistence of object file representa- tions. Perception & Psychophysics, 67(2):324-334, 2005.
Emilio Parisotto and Ruslan Salakhutdinov. Neural map: Structured memory for deep reinforcement learning. arXiv preprint vrXiv:1702.08366 2017.
Alexander Pritzel, Benigno Uria, Sriram Srinivasan, Adria Puigdomenech Badia, Oriol Vinyals, Demis Hassabis, Daan Wierstra, and Charles Blundell. Neural episodic control. In Proceedings of the 34th International Conference on Machine Learning-Vois Volume 70, pp. 2827-2836. JMLR. org, 2017.
Zenon W. Pylyshyn. The role of location indexes in spatial perception: sketch of the finst spatial- index model. Cognition, 32:65-99 97, 1989.
Adam Santoro, David Raposo, David G Barrett, Mateusz Malinowski. Razvan Pascanu, Peter Battaglia, and Timothy Lillicrap. A simple neural network module for relational reasoning. In Advances in neural information processing systems, pp. 4967-4976, 2017.
Elizabeth S Spelke and Katherine D Kinzler. Core knowledge. Developmental science, 10(1):89-96, 2007.
Jack Valmadre, Luca Bertinetto, João Henriques, Andrea Vedaldi, and Philip HS Torr. End-to- end representation learning for correlation filter based tracking. Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 2805-2813, 2017.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, łukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information processing systems, pp. 5998-6008, 2017.
Petar Veličković, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Lio, and Yoshua Bengio. Graph attention networks. arXiv preprint arXiv:1710.10903 2017.
Oriol Vinyals, Meire Fortunato, and Navdeep Jaitly. Pointer networks. In Advances in Neural Information Processing Systems, pp. 2692-2700, 2015.
Xiaolong Wang, Ross Girshick, Abhinav Gupta, and Kaiming He. Non-local neural networks. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 7794- 7803, 2018
Fei Xu. Object individuation and object identity in infancy: The role of spatiotemporal information, object property information, and language. Acta psychologica, 102(2-3):113-166 1999.
Keyulu Xu, Weihua Hu. Jure Leskovec, and Stefanie Jegelka. How powerful are graph neural networks? arXiv preprint arXiv: 1rxiv:1810.00826 2018.
12

---

Under review as a conference paper at ICLR 2020
Tianyu Yang and Antoni Chan. Learning dynamic memory networks for object tracking. In Proceedings of the European Conference on Computer Vision (ECCV). pp. 152-167, 2018.
Kexin Yi, Jiajun Wu, Chuang Gan, Antonio Torralba, Pushmeet Kohli, and Josh Tenenbaum. Neural- symbolic vqa: Disentangling reasoning from vision and language understanding. In Advances in Neural Information Processing Systems, pp. 1031-1042, 2018.
Kwangjin Yoon, Du Yong Kim, Young-Chul Yoon, and Moongu Jeon. Data association for multi- object tracking via deep neural networks. Sensors, 19(3):559, 2019.
Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh, Barnabas Poczos, Ruslan R Salakhutdinov, and Alexander Smola. Deep sets. In Advances in neural information processing systems, pp- 3391-3401, 2017.
Vinicius Zambaldi, David Raposo, Adam Santoro, Victor Bapst, Yujia Li, Igor Babuschkin, Karl Tuyls, David Reichert, Timothy Lillicrap, Edward Lockhart, et al. Relational deep reinforcement learning. arXiv preprint arXiv: 01830, 2018.
Yan Zhang, Jonathon Hare, and Adam Prgel-Bennett Deep set prediction networks, 2019.
A
RESULTS 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 most two objects are dropped.
Figure 7 shows test prediction accuracy and entropy for an AlignNet trained to drop up to five objects.
Taable]jdemonstrann that an AlignNet trained to drop N objects, can drop more than N objects with almost perfect accuracy.
Table 3: How well does a model trained with at most two objects being dropped (mar_nwmmlld Lo_ropp 2) generalise to cases trained an AlignNet with max mas_num._o_drop = and tested the model on examples where more than two objects were dropped. We see that the model generalises well more objects being dropped than those seen during training.
than two objects dropped? We
mox mor drop
Test accuracy
0
100%
1
100%
2
99.9%
3
99.8%
99.8% 4
5
100%
6
99.6%
7
100%
B RESULTS FOR ADDING OBJECTS
Figure shows test prediction accuracy and entropy values for an AlignNet trained to add up to five objects.
Figure 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 added.
Jabblesdemnnttraln that an AlignNet trained to add at most μ objects, can add more than μ objects with high accuracy, particularly when μ 3.
13

---

Under review as a conference paper at ICLR 2020
1911 Krann1
2013 Suce-12
13
20000
Accuracy different noise levels.
36
tetherthey
24
2aps
(b) Entropy at different noise levels.
Figure 6: Varying the noise level and comparing to baselines for mar mar.num www..to.mooommmmmmmmmmmmmmmm...m tropdop = 2. For each noise_lvee we trained five AlignNets with mar_num_toddroy 2 and show the median performance with the standard deviations. The AlignNet achieves > 99% accuracy at all noise levels and out-performs baselines.
C RESULTS FOR BOTH ADDING AND DROPPING OBJECTS
In most real world cases, it is likely that objects both added 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.02, 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 suffi- cient to achieve maximal performance and low entropy (sharp attention) solutions. The AlignNet achieves more than 99% accuracy all noise levels when most two objects added, except noise level = 0.2, where the best model achieves 98.7% 7% accuracy. Visualisations the results are shown in Figure 11
14

---

Under review as a conference paper at ICLR 2020
10
08-
04
10
Eihoeely cnker inteecsccroy exp-tertterree
9t95
2000
(a) Accuracy when dropping different numbers of objects.
(b) Entropy when dropping different numbers of objects.
map rum openz
4901
(c) Difference between hard and soft alignment is negligible.
Figure 7: Dropping up to five objects and comparing baselines. For each max fo we trained five AlignNets and show the median performance with the standard deviations. We also take the median across five runs of the baseline models. The AlignNet achieves 99.9% accuracy when dropping up to five objects.
D How MANY MESSAGE PASSING STEPS ARE NEEDED?
Figure shows that, while only 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.
15

---

Under review as a conference paper at ICLR 2020
Poad::
18G
led
00
0.4
20
63
dtpa
(a) Accuracy when adding different numbers of objects.
2tga
300g
Mtpa
Entropy when adding different numbers of objects.
32-4
20
311
27
30-4
30-
an
27
(c) Difference between hard soft alignment negligible.
Figure 8: Adding up to five objects and comparing to baselines. For each maz numn to add trained five AlignNets and show the median performance with the standard deviations. We also take the median across five runs of the baseline models. For each maz num. to add least of the five models achieves an accuracy > 99%.
16

---

Under review as a conference paper at ICLR 2020
185
145
800
02
04
6.
03
(a) Accuracy at different noise levels.
02
0. 2. 03
(b) Entropy at different noise levels.
Figure 9: Varying the noise level and comparing to baselines for maz_nam__ee each noise_lvel we trained five AlignNets with maz wwt.unumooodde = 2 and show the median performance with the standard deviations. We also take the median across five runs of the baseline models. The AlignNet achieves > 99% accuracy at all noise levels, out-performing baselines.
2. For
Table How well does an AlignNet trained with at mar num tow ao. = μ objects added, generalise to cases where more than objects are added (μ S maar_uum_to_r add S 7)? Test accuracy values shown in pink are those for which the model seeing more added objects than those seen during training.
Maximum number of objects added,
Test accuracy when μ objects are added using a model trained to add at most
two objects
three objects
four objects
five objects
H 0
100%
100%
100%
100%
J =
100%
100%
100%
100%
J 2
100%
100%
100%
100%
= 3
86%
96%
100%
100%
75%
90%
99%
99%
= 5
70%
80%
95%
98%
61%
70%
93%
98%
7
56%
63%
90%
96%
17

---

Under review as a conference paper at ICLR 2020
Upw. 30
14.
Deem dl caee cniuuin
70015
Falindt entanaDuioery Adh-ttoleeouees
35295
(a) Accuracy when adding and dropping at most two objects at different noise levels.
03
3.
0.6
1000 2xd--)
(b) Entropy when adding and dropping at most two objects at different noise levels.
01.4
(c) Difference between hard and soft alignment.
Figure 10: Varying the noise level and comparing to baselines for mar_nwm_looode mar mnar_num_to_drop = 2. For each noise_level we trained five AlignNets and show the median performance with the standard deviations. We also take the median across five runs of the baseline models. The best model at each noise level achieves > 99% accuracy at all noise levels, except at noise level 0.2, where the best model achieves 98.7% accuracy. All of our models out-perform baselines.
and
18

---

Under review as a conference paper at ICLR 2020
(a) Model trained maz num to drop 4. Objects that dropped are assigned the null slot and left empty the Align column.
3<
(b) Model trained with mar num to add 4. New objects added in the order in which they appear.
(c) Model trained with mar_num_fooate 2 and tow.nnmoooorroo
Figure 11: Visualisations of AlignNet results. These examples are not cherry picked. The Imput column shows the list of objects, U, the Dropped column shows the list of objects, S, with entities dropped and the Add column shows the list of objects, S with additional entities not present in the Input. The Align column shows the aligned entities, slots that were empty in the Input may now be occupied by new objects and entities that were in the Input but dropped are replaced by empty slots. Entities present both the Input and the Added (and) or Dropped columns are placed the same slot that they were found in, in the Input.
19

---

Under review as a conference paper at ICLR 2020
1.000
0.998
test/accuracy
num num_recurree www.rreettsteeps.e.eeeeeesesssee
0.996
num securrent_step-=0
M
num, rww.rrectttttcee steps=3 3
0.994
sum_recurrent setesureent_steps=4 steps=4
num_rest jum_recurrent_see 5trrres
0.992
0.990
0.0
0.2
0.4
0.6
0.8
steps
1e6
(a) Dropping to two objects. One message passing step is sufficient,
1.00
0.98
test/accuracy
num recurrent_steps-1
0.96
num recurrent_stepsee
0.94
num juu_reecrreet_tt 3tepst33
num Suum_recurrent_sepaa
0.92
num, ruuuurettrentt_ 5
0.90
0.88
0.4
0.8
steps
1e6
(b) Adding up to two objects. At least two message passing steps are needed.
Figure 12: Comparing the number of message passing steps needed obtain high test accuracy when adding and dropping objects. For dropping, single message passing step is sufficient. However, for adding at least two message passing steps are needed to obtain good performance. More than two message passing steps may results in instability and higher variability during training.
20

---

Under review as a conference paper at ICLR 2020
1.00
test/metrics _aligner accuracy
0.95
num num_recumee reruurerrrretteee recurrent_steps--
num num__ecuree recurrent recurrent_steps=2
num num__ecui recurrent recurrent_steps=3
steps=3
num cem__recreettt_e rer_errent_steps-aa
num tet__reuureettteeeeee
0.85
0.80
0.0
0.2
0.4
0.6
0.8
10
12
steps
1e6
The AlignNet as general framework with transformers at the core, rather than graph network. We found that least two layers transformers were needed achieve maximal performance.
Figure 13:
E THE ALIGNNET AS GENERAL FRAMEWORK
The AlignNet may be more broadly thought of as framework with the following key elements: (I) transforming sets in to lists of slots, with a fixed order, sorting one list with respect to another,
(2) deterministic order for adding new objects, (3) appending additional null slot for modelling dropped objects and (4) computing a loss between the predicted adjacency matrix and the true ad- jacency matrix. The model that takes the shuffled and unshuffled lists inputs and outputs the adjacency matrix is the core our model. In the experiments above, we used a graph network as the core, however other networks could be substituted. For example, transformers (Vaswaní al.) 2017 Wang et al. 2018) may be used. Figure 13 shows that with two layers of transformers, the AlignNet achieves 100% accuracy.
F ALIGNING SEQUENCES OF SETS OF ENTITIES OVER TIME.
Figure shows the accuracy with which trained Aligner writes memory, for different numbers of kept slots, num_ker num_keep keep_slots, the input, with no further training. This visualised in Figure 15 Each symbolic object in each slot of U, S and M is rendered as an image in different row (or slot).
1.06
1.04
Mean accuracy
num, keep stots=2
keep 33
num, keep, stots
num, keep sl
Time steps.
Figure 14: Accuracy with which sequences with different numbers of kept objects are aligned. 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 ALIGNNET ARCHITECTURE AND TRAINING DETAILS.
For the results presented in this paper, the AlignNet trained using RMS- optimiser with learning rate 2e-4, however Adam worked well too. For the graph network, we used a four layer ML.P with sizes [128, 256, 512, 1024] for the edge model (Battaglia et al. (2018) and used three
21

---

Under review as a conference paper at ICLR 2020
layer MLP for the node model with sizes [32, 32, 7]. We found that the model was robust to many different MLP configurations 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 EPISODIC QUESTION ANSWERING
H.1 EPISODIC-SORTT-OF-CII DATASET
The spisodic-sort-of-clee question answering dataset contains episodic sequences of objects ar- ranged together in a scene, along with question about a scene and a corresponding answer. The object shapes can be triangles or squares and each object has unique colour. There are several types of question, all asking about the relationship between objects in the scene. For example: 'What 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. Specifically, a one-hot representation of object colour is provided, along with 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 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 red square appears in the first observation of a scene, 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 question about which object is nearest to (or furthest from) the green triangle, model must retain a representation of the full scene in its memory, across all observations, before can answer the question.
sequence of observations. Each object can appear, disappear
In the episodic-sort- -ofl-of dataset, there can also be multiple scenes within an episode, each con- taining a different set of objects. Each question is directed towards one scene only using a unique scene identification label. Each scene can consist of sequence of observations. These observa- tions, 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 observa- tions, 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 observa- tions increase, it will become increasingly difficult for many standard models scale sufficiently to support relational question answering.
The episodic-sort-of--les dataset is an episodic version of the Sort-of-Clevi dataset (Santoro 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.2 ALGINNET WITH RELATION NETWORK MODEL
We train an AlignNet together with relation network, RN (Sntoroo et al., 2017). which we refer to as the AlignNet-RN. RNN to solve the cpisodic-sort-of--ley tasks. We first train the AlignNet to align objects within scenes and across observations. After training, the AlignNet may be used to write entities, in given episode, to a memory, M = [mo, Mi, m,,...,myyl v. At the end of an episode, when all entities have been written memory, the memory can be treated set of objects, since the order is no longer important.
Next, we provide the AlignNet memory for given episode as input a RN. The AlignNet mem- ory 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 fo (Ers (m,, my, as output. For Pg. we use a four-layer MI. with 256 units per layer and ReLu non-linearities. For fa we use a two-layer MLP, again with 256 units per layer and ReLu non-linearities, followed by linear layer. We use learning rate of 0.0001 and batch size of 8. We
22

---

Under review as a conference paper at ICLR 2020
train all of our models using the Adam optimiser with decay rate parameters 81 0.9, B2 0..99 and é = 10
23

---

Under review as a conference paper at ICLR 2020
Visualising shuffled sequence num keep slotsm].
Visualising the memory over time
Visualising the aiinedd sequerice
Visualising the shuffled sequence num keep plots=2.
Visualising the aligned sequence.
Visualising the memory time.
X
X
X
Visualising shuffled sequence slots-3
Visualising aligned sequence.
Visualising memory
Visualising the shutfled sequence num keep dlots
Visualising the memory over time
Visualising the alignod sequenco.
X
Visualising the shutfled sequence
num keep clottsm5
Visualising the memory time
Visualising the aligned sequence
X
Visualising
sequence
num
memory
aliined
x
X
Visualising the shuffled sequence
num keep slots-
Visualising the memory over time.
Visualising the aligned sequance.
X
x
Figure 15:
Visualisations of the aligned sequences over time. For each num num_keep_slots in the input sequence, show how train AlignNet, which we refer ory, M. We also show the resulting sequence aligned objects. We render each symbolic object representation, in each slot as image, to show all the objects in each slot.
Aligner, writes mem-
24


Reviews:
[{"review_id": "ByxlfiUntH", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "With the assumption of object persistence and inspired by the sticky indices, this paper proposed a novel object alignment method for matching arbitrary entities in different sets. The motivation is clear and easy to follow. The simulation experiment with the symbolic dataset gives impressive results. \n\nI like this idea and have one concern about the experiment. To be specific, all the reported results are obtained in the symbolic dataset to simulate the real-world case. Have you performed the entity alignment in the real-world data, such as the cross-lingual knowledge graphs, or cross-modal analysis? It's expected to experiment with more challenging datasets. By the way, have you released your code? This is important to an ICLR submission. \n", "extra_scores": {"experience_assessment": "I have published in this field for several years.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper."}, "extra_text": {}}, {"review_id": "B1x_gbizcr", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 1, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "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.\n\nI lean to reject this paper because in my opinion the proposed method is just a set of predefined rules to train a network to be able to perform object association between a new time step and a memory. Additionally the experimental evaluation is very weak in several points (see below). \n\n- Contribution: this approach proposes a set of rules to train a network to be able to learn the correct association  between two set of features. Additionally, the network is used together with a memory to keep track of previously seen objects. In my understanding the proposed work presents a set of hard-coded rules to train a network for tracking. However, the connection with tracking is presented only in related works (at the end of the paper) and no comparison with other tracking approach is presented. Additionally, the network model that is actually used for the experiments (graph network) is not presented in the main paper, which makes more difficult to understand how the network works for the given task.\n\n- Experimental evaluation: Experiments are performed on a very simple, simulated environment. Objects are very simple to recognize and their location (which is often  a difficult component of the problem) is given. As the task is quite straight forward, the model obtains almost 100% on all tests. Comparisons are made only with hand-crafted features. It is quite evident that a learned similarity measure will be better than hand-crafted distances. \n\nAuthors should explain more clearly that the proposed contribution is a set of rules used to learn a distance between features in order to associate object instances. \nA cleared connection with tracking should be provided from the introduction. In the evaluation there should be a comparison on multiple and more challenging tracking datasets. In this way we can compare the proposed training technique with other tracking approaches that also learn to a distance in order to associate object at different time steps. \n\n", "extra_scores": {"experience_assessment": "I have published one or two papers in this area.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}, {"review_id": "B1gLTsN2FB", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "\nThis paper proposes AlignNet, a bipartite graph network that learns to match to sets of objects. AlignNet has a slot-wise object-based memory that associates an index with each unique object and can discover new and re-appearing objects. Experiments are conducted on a symbolic dataset.\n\nI do not think the paper meets the acceptance threshold, and recommend for weak rejection. While the paper proposes an interesting architecture to address the alignment problem, it has noticeable flaws in its experimental designs.\n\nFirst, all the experiments are conducted on toy symbolic datasets, where the alignment problem is rather easy to solve. On the other hand, real-world scenarios can be far more complicated. For example, the appearance of the same object can change due to lighting and distance, and it is unreasonable to assume that their features would remain static (apart from simple uniform noises). In addition, the paper only compares against hand-crafted similarity measures (MSE and cosine). It is unfair to compare learned methods only to hand-crafted methods. As a reasonable and fair comparison, the paper should also compare AlignNet against learned similarity measures (such as a neural network supervised with ground-truth labels for alignments).\n\nThe toy dataset and simple baselines in this paper raise doubts on whether the proposed method is applicable to more complex scenarios (such as aligning two sets of objects in natural images through their appearance features).", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_experiments": "I carefully checked the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "N/A"}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper proposes a network architecture which labels object with an identifier that it is trained to retain across subsequent instances of that same object.

After discussion, the reviewers agree that the approach is interesting, well-motivated and written, and novel. However, there was unanimous concern about the experimental evaluation, so the paper does not appear to be ready for publication just yet, and I am recommending rejection.
Decision: Reject