Recurrent Independent Mechanisms
Authors: Anirudh Goyal
Alex Lamb
Shagun Sodhani
Jordan Hoffmann
Sergey Levine
Yoshua Bengio
Bernhard Scholkopf
Venue:  (2020)
Keywords: modular representations
better generalization
learning mechanisms

Abstract: Learning modular structures which reflect the dynamics of the environment can lead to better generalization and robustness to changes which only affect a few of the underlying causes. We propose Recurrent Independent Mechanisms (RIMs), a new recurrent architecture in which multiple groups of recurrent cells operate with nearly independent transition dynamics, communicate only sparingly through the bottleneck of attention, and are only updated at time steps where they are most relevant.  We show that this leads to specialization amongst the RIMs, which in turn allows for dramatically improved generalization on tasks where some factors of variation differ systematically between training and evaluation.

Under review as a conference paper at ICLR 2020
RECURRENT INDEPENDENT MECHANISMS
Anonymous authors
Paper under double-blind review
ABSTRACT
Learning modular structures which reflect the dynamics of the environment can lead better generalization and robustness to changes which only affect a few of the underlying causes. We propose Recurrent Independent Mechanisms (RIMs), a new recurrent architecture in which multiple groups of recurrent cells operate with nearly independent transition dynamics, communicate only sparingly through the bottleneck of attention, and are only updated at time steps where they are most relevant. We show that this leads to specialization amongst the RIMs, which in turn allows for dramatically improved generalization on tasks where some factors of variation differ systematically between training and evaluation.
INTRODUCTION
Physical processes in the world often have modular structure, with complexity emerging through combinations of simpler subsystems. Machine learning seeks to uncover and use regularities in the physical world. Although these regularities manifest themselves statistical dependencies, they ultimately due to dynamic processes governed by physics. These processes are often independent and only interact sparsely. For instance, we can model the motion of two balls as separate independent mechanisms even though they are both gravitationally coupled to Earth as well as (weakly) to each other. They may, however, occasionally strongly interact via collisions.
The notion of independent or autonomous mechanisms has been influential in the field of causal inference, where it upplied not only to dynamic processes but also to time independent datasets. For instance, it has been argued that the conditional distribution of the average annual temperature given the altitude of place is an abstraction a causal mechanism (subsuming complex physical processes involving air pressure, etc.) that is independent of the distribution of the altitudes of settlements (Peters et al. 2017), and will thus apply invariantly for, say, different countries in the same climate zone with different altitude distributions.
complex generative model, temporal or not, can thought of as the composition independent mechanisms "causal" modules. In the causality community, this is often considered prerequisite of being able to perform localized interventions upon variables determined by such models (Pearl 2009). It has been argued that the individual modules tend remain robust or invariant modules change, e.g., in the case of distribution shift (Schölkopf et al., 2012 Peters et al.) 2017). One may hypothesize that if a brain is able to solve multiple problems beyond single i.i.d. (independent and identically distributed) task, it would be economical to learn structures aligned with this, by learning independent mechanisms that can flexibly be reused, composed and re-purposed.
other
In the dynamic setting, we think of an overall system being assayed as composed of a number of fairly independent subsystems that evolve over time, responding to forces and interventions. learning agent then need not devote equal attention to all subsystems at all times: only those aspects that significantly interact need to be considered jointly when taking a decision or forming a plan (Bengio) 2017). Such sparse interactions can reduce the difficulty of learning since few interactions need to be considered at time, reducing unnecessary interference when a subsystem is adapted. Models learned this way may be more likely to capture the compositional generative (or causal) structure of the world, and thus better generalize across tasks where a (small) subset of mechanisms change while most of them remain invariant (Simon [991] Peters et al.) 2017: Parascandolo al. 2018). The central question motivating our work is how a machine learning approach can learn independent but sparsely interacting recurrent mechanisms in order benefit from such modularity.

---

Under review as a conference paper at ICLR 2020
Delint dymaniic
Spane Commonariiin
Defnnit Synunics
Spazie Commmication
dezaa albention
Campeity RINS
Beard unggeikn bave top dons. dkiona
lagan viteation!
Thea
he42
Query
Duving Uoudims
Inpet
Inprt
Disnin Goslica
Aelive H514
Figure I: Illustration of Recurrent Independent Mechanisms (RIMs). A single step under the proposed model occurs in four stages (left figure shows steps). In the first stage, individual RIMs produce query which is used to read from the current input. In the second stage, attention based competition mechanism used to select which RIMS to activate (right figure) based on encoded visual input (blue RIMs are active, based on attention score, white RIMs remain inactive). In the third stage, individual activated RIMs follow their own default transition dynamics while non-activated RIMs remain unchanged. In the fourth stage, the RIMs sparsely communicate information between themselves, also using attention.
2 RECURRENT INDEPENDENT MECHANISMS WITH SPARSE LNTERACTIONS
Our approach to modelling a dynamical system of interest divides the overall model into small subsystems (or modules), each of which is recurrent in order to be able to capture dynamics. We refer to these subsystems as Recurrent Independent Mechanisms (RIMs). where each RIM has distinct functions that are learned automatically from date] We refer RIM lc at time step as having state h4, where = 1,...,77 T. Each RIM has parameters Os. which are shared across all time steps.
At high level (see. Fig. 1. we want each RIM to have its own independent dynamics operating by default, and occasionally to interact with other relevant RIMs and with selected elements of the encoded input. The total number of parameters can be kept small since RIMs can specialize on simple sub-problems, similar to Parascandolo et al. (2018). This specialization and modularization not only has computational and statistical advantages (Baum & Haussler 1989 Bengio et al. 2019). but also prevents individual RIMs from dominating and modelling complex, composite mechanisms. We expect this to lead more robust systems than training one big homogeneous neural network (Schmidhuber] 2018). Moreover, modularity also has the desirable implication that a RIM should maintain its own independent functionality even as other RIMs are changed. more detailed account of the desiderata for the model is given in Appendix A
2.1 INDEPENDENT RIM DYNAMICS
Now, consider the default transition dynamics which we apply for each RIM independently and during which no information passes between RIMs. We use h for the hidden state after the independent dynamics are applied (and before attention is applied). First, for the RIMs which are not activated (we refer to the activated set as S,), the hidden state remains unchanged:
het1A
he.k Vk $ St.
Note that the gradient still flows through RIM on step where it is not activated. For the RIMs that are activated, we run a per-RIM independent transition dynamics. The form of this is somewhat flexible, but in this work opted to either GRU (Chung al., 2015) an L.STM (Hochreiter & Schmidhuber] 1997). We generically refer to these independent transition dynamics as Dç, and we emphasize that each RIM has its own separate parameters. Aside from being RIM-specific, the internal operation of the LSTM and GRU remain unchanged, and the active RIMs are updated by
h141, = Da(h,,)) = LST M (hr, A (in), ( e(P))
e S₁
(2)
as a function of the attention mechanism (())) applied on the current input, described in the next two subsections below, after explaining the key-value mechanism used to select arguments for this update.
Note that using the term mechanism both for the mechanisms that make up the world's dynamics as well for the computational modules that we leam model those mechanisms.
2

---

Under review as a conference paper at ICLR 2020
2.2 KEY- VALUE ATTENTION TO PROCESS SETS OF NAMED INTERCHANGEABLE VARIABLES
Each RIM should be activated and updated when the input is relevant to it. We thus utilize competition to allocate representational and computational resources. As argued by Parascandolo al. 20. this tends to produce independence among learned mechanisms, provided the training data has been generated by set of independent physical mechanisms. In contrast Parascandolo al.) (2018), we use an attention mechanism for this purpose. In doing so, we are inspired by findings from experimental psychology in the study of the interplay of top-down attention and bottom-up information flow, conceptualized in the biased competition theory of selective attention (Desimone & Duncan 1995): A brain's 's capacity for parallel processing of complex entities is limited, and many brain systems representing visual information use competition (operating in parallel across the visual field) to allocate resources, often biased by feedback from higher brain areas.
The introduction of content-based soft- attentien mechanisms (Bahdanau al. 2014) has opened the door to neural networks which operate on sets of typed interchangeable objects. This idea has been remarkably successful and widely applied to most recent Transformer-style multi-head dot product self attention models (Vaswani et al., 2017] Santoro et al., 2018), achieving new state-of-the-aa results in many tasks. Soft-attentoon uses the product of a query (or read key) Q of dimensionality Nr x d matrix Q, and d dimension of each key) to a set of N, objects each associated with a key (or write-key) matrix K" (N, x d), and after normalization with a softmax yields outputs in the convex hull of the values (or write-values) V; (row i of matrix V). Its result is computed as
QKT
Atlention(Q. (Q, K, V) softmax
V.
vd
where the softmax is applied to each row of its argument matrix, yielding set of convex weights. As a result, one obtains a convex combination of the values V. If the attention focused on one element for particular row (i.e., the softmax is saturated), this simply selects one of the objects and copies its value to row j of the result. Note that the dimensions in the key can be split into heads which then have their attention matrix and write values computed separately.
When the inputs and outputs of each RIM are a of objects or entities (each associated with a key and value vector), the RIM processing becomes generic object-processigg machine which can operate on "variables" in a sense analogous to variables in a programming language: as interchangeable arguments of functions. Because each object has a key embedding (which one understand both as a name and as a type), the same RIM processing can be applied to any variable which fits an expected "distributed type" (specified by a query vector). Each attention head then corresponds to a typed argument of the function computed by the RIM. When the key of an object matches the query, it can be used as input for the RIM. Whereas in regular neural networks (without attention) neurons operate on fixed variables (the neurons which are feeding them from the previous layer), the key-value attention mechanisms make it possible to select on the fly which variable instance (i.e. which entity or object) is going to be used as input for each of the arguments of the RIM dynamics, with a different set of query embeddings for each RIM. These inputs can come from the external input from the output of other RIMs. So, if the individual RIMs can represent these "functions with typed arguments," then they can "bind" to whatever input is currently available and best suited according to its attention score: the "input attention" mechanism would look at the candidate input object's 's key and evaluate its "type" matches with what this RIM expects (specified in the query).
2.3 SELECTIVE ACTIVATION OF RIMS AS A FORM OF TOP-DOwN MODULATION
The proposed model learns to dynamically select those RIMs for which the current input is relevant. We give each RIM the choice between attending to the actual input instances or special null input. The null input consists entirely of zeros and thus contains no information. At each step, we select the t (out of kr) RIMs in terms of their value of the softmax for the real input. Intuitively, the RIMs must compete on each step to read from the input, and only the RIMs that win this competition will be able read from the input and have their state updated.
In our use of key-value attention, the queries come from the RIMs, while the keys and values come from the current input. The mechanics of this attention mechanism follow from the Transformer (Vaswani et al.) 2017) and the RMC (Santoro et al. 2018), with the modification that the parameters of the attention mechanism ifself are separate for each RIM. The input attention for a particular RIM
3

---

Under review as a conference paper at ICLR 2020
is described as follows. The input 22 at time t is set of elements, structured as of matrix (for image data, it can be the output of the CNN). We first concatenate row full of zeros, to obtain
= 0 E Zt-
refers to the row-level concatenation operator. Then, linear transformations used construct keys (K = XW* one per input element and for the null element), values (V = XW", again one per element), and queries (Q RW7 k. one per RIM attention head) where R is matrix with each row corresponding to the hidden state of an individual RIM (i.e ht W is simple matrix mapping from an input element to the corresponding value vector for the weighted attention and W similarly a weight matrix which maps the input to the keys. W is a per-RIM weight matrix which maps from the RIM's ' hidden state to its queries. The attention thus is
R Vil.
A (
softmax
XW", where of!m) (W k W W
(4)
Based on the softmax values in 4 we select the top A RIMs (out of the total RIMs) to be activated for each step, which have the least attention on the null input (and thus put the highest attention on the input), and we call this set S, Since the queries depend on the state the RIMs, this enables individual RIMs to attend only to the part of the input that is relevant for that particular RIM, thus enabling selective attention based on top-down attention process (see. Fig In practice, we multiheaded attention, and multi-heade attention doesn't 't change the essential computation, but when we do use it for input-attention we compute RIM activation by averaging the attention scores over the heads.
2.4 COMMUNICATION BETWEEN RIMs
Although the RIMs operate independently by default, the attention mechanism allows sharing of information among the RIMs. Specifically, we allow the activated RIMs to read from all other RIMs (activated not). The intuition behind this is that non-activated RIMs not related the current input, so their value should not change. However they may still store contextual information that is relevant for activated RIMs. For this communication between RIMs,
residual connection as in (Santoro et al. 2018 to prevent vanishing or exploding gradients over long sequences.
Qr,K Whe
Vk E St Vk
(5)
K2, WeeeA
(6)
V1, = WiRe
Vk:
(Qun(Ke.)))) Ve, + he Ak E S,, where o(c) = =(Wg kt We, Wy).
1441. = softmax
(8)
Vde
As in the Transformer and RMC Vaswani al. (2017); Santoro et al. (2018). we multiple heads (as well as input attention (as in 2.3 by producing different sets of queries, keys, and values to compute linear transformation for each head (different heads have different parameters).
2.5 VARIATIONS ON THE RIMS ARCHITECTURE
The RIMs architecture that study is highly homogeneous and generally the only hyperparameters are the number of RIMs K and how many RIMs are activated on each time step KA. All of the datasets that we consider are temporal, yet there is distinction between datasets where the input on each time step is highly structured (such as a video, where each time step is an image) and where this not the case (such as language modeling, where each step is word or character). In the former case, we can get further improvements by making the activation of RIMs not just sparse across time but also sparse across the (spatial) structure.
4

---

Under review as a conference paper at ICLR 2020
3 RELATED WORK
Neural Turing Machine (NTM) and Relational Memory Core (RMC): the NTM (Graves et al. 2014a) consists of a sequence of independent memory cells, and uses an attention mechanism while performing targeted read and write operations. This shares key idea with RIMs: that input information should only impact a sparse subset of the memory by default, while keeping most of the memory unaltered. RMC (Santoro et al. 2018)) uses multi-head attention mechanism to share information between multiple memory elements. We encourage the RIMs to remain separate as much as possible, whereas Santoro al. (2188) allow information between elements to flow on each step in an unsconstrained way. Instead, each RIM has its own default dynamics, while in RMC, all the processes interact with each other.
Separate Recurrent Models: EnTNet (Henaff et al., 2016) and IndRNN (Li et al. 2018 can be viewed set of separate recurrent models. In IndRNN, each recurrent unit has completely independent dynamics, whereas EntNet uses an independent gate for writing to each memory slot. RIMs use different recurrent models (with separate parameters), but we allow the RIMs to communicate with each other sparingly using an attention mechanism.
Modularity and Neural Networks: A neural network is composed of several neural modules, where each module is meant to perform a distinct function, and hence can be seen combination of experts (Jacobs et al. 1991 Bottou Gallinari 1991 Ronco et al. 1997 Reed De Freitas 2015
Andreas et al., 2016 Parascandolo et al., [2018] 8 Rosenbaum et al., 2017 Fernando et al. 2017
Shazeer et al. 2017 Kirsch et al.) 2018 Rosenbaum et al. 2019) routing information through a gated activation of layers. These works generally assume that only single expert is active at particular time step. In the proposed method, multiple RIMs can be active, interact and share information.
Computation on demand: There are various architectures (El Hihi & Bengio 1996 Koutnik et al.
2014 Chung et al.) 2016 Neil et al. [016)] Jernite et al. 2016 Krueger et al. 2016) where parts of the LSTM's hidden state are kept dormant at times. The major differences as compared to the proposed architecture are that (a) we modularize the dynamics of recurrent cells (using RIMs), and (b) we also control the inputs of each module (using transformer style attention), while many previous gating methods did not control the inputs of each module, but only whether they should be executed or not.
EXPERIMENTS
The main goal of our experiments is to show that the use of RIMs improves generalization across changing environments and/or in modular tasks, and to explore how it does so. Our goal is not to outperform highly optimized baselines; rather, we want to show the versatility of our approach by applying it to range of diverse tasks, focusing on tasks that involve changing environment. We organize our results by the capabilities they illustrate: we address generalization based on temporal patterns, based on objects, and finally consider settings where both of these occur together.
4.1 RIMS IMPROVE GENERALIZATION BY SPECIALIZING OVER TEMPORAL PATTERNS
We first show that when RIMs are presented with sequences containing distinct temporal patterns, they are able to specialize so that different RIMs are activated on different patterns. As a result, RIMs are able to generalize well when we modify a subset of the patterns (especially those unrelated to the class label) while most recurrent models fail to generalize well to these variations.
4.1.1 COPYING TASK
First we turn our attention to the task of receiv- ing short sequence of characters, then receiv- ing blank inputs for large number of steps, and then being asked to reproduce the original se- quence. We can think of this as consisting of two temporal patterns which are independent: one where the sequence is received and another "dormant" pattern where input is provided.
Figure 2: Copying Task RIM Activation Pattern for a model with K = 6 RIMs and KA = 3 active RIMs per step (the activated RIMs are in black, non-activated white). We can that the RIM activation pattern distinct during the dormant part of the sequence.
5

---

Under review as a conference paper at ICLR 2020
Sequential MNIST
16
19 x
24 x Accuracy
Copying
Train(50)
Test(200)
kr
RA
huse
CE
CE
kr
kA
haiee
Accuracy
Accuracy
6
5
600
0.01
3.5
600
85.5
56.2
30.9
6
4
600
0.00
0.00
6
5
600
88.3
43.1
22.1
RIMS
600
RIMS
6
4
600
90.0
73.4
38.1
2
600
0.00
5
3
500
0.00
0.00
300
86.8
42.3
25,2
LSTM
600
84.5
52.2
21.9
300
0.00
2.28
LSTM
600
0.00
3.56
EntNet
89.2
52.4
23.5
NTM
0.00
2.54
RMC
89.58
54.23
27.75
RMC
0.00
0.13
DNC
87.2
19.8
Transformers -
0.00
0.54
Transformers-
91.2
51.6
22.9
Table I: Performance on the copying task (left) and sequential MNIST resolution task right). Error (CE on the last 10 time steps) the copying task. Note that while all of the methods able learn copy for length seen during training, the RIMs model generalizes sequences longer than those seen during training whereas the LSTM, RMC, and NTM degrade. Sequential MNIST resolution: Test Accuracy % on the Sequential MNIST resolution generalization task (see tex() after 100 epochs. Both the proposed and the Baseline model (LSTM) were trained 14x14 resolution but evaluated at different resolutions; results averaged over 3 different trials.
As an example of out-o--distribution generalization, we find that using RIMs, we can extend the length of this dormant phase from 50 during training to 200 during testing and retain perfect performance (Table D whereas baseline methods including LSTM, NTM, and RMC substantially degrade. In addition, find that this result is robust to the number of RIMs used well as to the number of RIMs activated per-step. Our results (Appendix C.5) show that communication between different RIMs well input attention is necessary to achieve good generalization. We consider this preliminary evidence that RIMs can specialize over distinct patterns in the data and improve generalization to settings where these patterns change.
4.1.2 SEQUENTIAL MNIST RESOLUTION TASK
RIMs are motivated by the hypothesis that generalization performance can be improved by having modules which only activate on relevant parts of the sequence. For further evidence that RIMs can achieve this fittof-ddstributtion we consider the task of classifying MNIST digits as sequences of pixels (Krueger et al.) (016)) and assay generalization to images of resolutions different from those seen during training. Our intuition is that the RIMs model should have distinct subsets of the RIMs activated for pixels with the digit and empty pixels. As result, RIMs should generalize better to greater resolutions by keeping the RIMs which store pixel information dormant over the empty regions of the image.
Table shows the result of the proposed model on the Sequential MNIST Resolution Task. the train and test sequence lengths agree, both models achieve comparable test set performance. However, the RIMs model relatively robust to changing the sequence length (by changing the image resolution), whereas the LSTM performance degraded more severely. This can be seen as more involved analogue of the copying task, MNIST digits contain large empty regions. It is essential that the model be able to store information and pass gradients through these regions. The RIMs outperform strong baselines such Transformers, EntNet, RMC, as well the Differentiable Neural Computer (DNC) (Graves et al., 2016).
Results:
4.2 RIMS LEARN TO SPECIALIZE OVER OBJECTS AND GENERALIZE BETWEEN THEM
We have presented evidence that RIMs specialize over temporal patterns. We now turn attention to showing that RIMs can specialize to objects, and show improved generalization to settings where add remove objects test time.

---

Under review as a conference paper at ICLR 2020
4.2.1 BOUNCING BALL ENVIRONMENT
We consider a synthetic "bouncing balls" task in which multiple balls (of different masses and sizes) move using basic Newtonian physics (Van Steenkiste et al., 2018). What makes this task particularly suited RIMs that the balls move independently most of the time, except when they collide. During training, we predict the next frame at each time step using teacher forcing (Williams & Zipser 1989). We can then use this model to generate multi-step rollouts.
As preliminary experiment, we train on sequences of length 51 (the previous standard), using binary cross entropy loss when predicting the next frame. We consider LSTM as baseline. We then produce rollouts, finding that RIMs are better able to predict future motion (examples in Figure 3 Figure 10 in Appendix and quantitative comparisons in Figure
1 Ohiests
3
+5
+6
-7
+9
10
+11
+12
+13
+1d
15
8
Truth
RIM
0
LSTM
Figure 3: Predicting Movement of Bouncing Balls. The first 15 frames of ground truth are given (last of those shown) and then the system rolled out for the next 15 time steps. We find that RIMs perform better than the LSTMs (predictions are in black, ground truth in bluc). Notice the blurring of LSTM predictions.
Tramndd on Objacta
Traised on f-8 Objents
Tasted on 6.8 Oblects
Tested on Ohieats
Curtain
Clos Enroop
1STM 513
200
Groand truth
Bollout
Figure 4: Handling Novel Duttof-Diitributtinn Variations. Here, we study the performance of our proposed model compared to an LSTM baseline. The first 15 frames of ground truth are fed in and then the system is rolled out for the next 10 time steps. During the rollout phase, RIMs perform better than the LSTMs in accurately predicting the dynamics of the balls as reflected by the lower Cross Entropy (CE) [see blue for RIMs, purple for LSTMJ. Notice the substantially better out-of of-distribution generalization of RIMs when testing on a different number of objects than during training.
We take this further by evaluating RIMs on environments where the setup is different from the training setup. First we consider training with 4 balls and evaluating on an environment with 6-8 balls. Second, we consider training with 6-8 balls and evaluating with just 4 balls. Robustness in these settings requires degree of invariance w.r.t. the number of balls.
In addition, we consider a task where we train on balls and then evaluate on sequences where part of the visual space is occluded by "curtain". This allows us to assess the ability of balls to be tracked (or remembered) through the occluding region. Our experimental results on these generalization tasks (Figure 4 show that RIMs substantially improve over an LSTM baseline. We found that increasing the capacity of the LSTM from 256 to 512 units did not substantially change the performance gap. suggesting that the improvement from RIMs is not primarily result of increased capacity.
4.2.2 ENVIRONMENT WITH NOVEL DISTRACTORS
We next consider an object-picking reinforcement learning task from BabyAl Booiseller-Boisverrt et al. 2018) in which an agent must retrieve a specific object in the presence of distractors. We use a partially observed formulation of the task, where the agent only sees a small number of squares ahead of it. These tasks are difficult to solve (Chevalier-Boiseee due to (1) the partial observability of the environment and (2) the sparsity of the reward, given that the agent receives a reward only after reaching the goal. During evaluation, we introduce new distractors to the environment which were not observed during training.
et al. 2018) with standard RL algorithms,
7

---

Under review as a conference paper at ICLR 2020
Figure 5 shows that RIMs outperform LSTMs on this task (details in appendix). When eval- uating with known distractors, the RIM model achieves perfect performance while the LSTM struggles. When evaluating in environment with novel unseen distractors the RIM doesn't achieve perfect performance but still outper- forms the LSTM. An LSTM with single mem- ory flow may struggle keep the distracting elements separate from elements which are nec- essary for the task, while the RIMs model uses attention to control which RIMs receive infor- mation at each step as well as what information they receive (as a function of their hidden state). This "top-down" bias results in diminished representation of the distractor, not only enhancing the target visual information, but also suppressing irrelevant information. The notion that enhancement of the relevant information necessarily results in suppression of irrelevant information is fundamental to biased competition theory (Desimone & Duncan 1995).
M
Figure 5: Robustness to Novel Distractors: Left: per-
formance of the proposed method compared to an LSTM baseline in solving the object picking task in the presence of distractors. Right: performance of proposed method the baseline when novel distractors are added.
4.3 RIMS IMPROVE GENERALIZATION IN COMPLEX ENVIRONMENTS
We have investigated how RIMs use specialization to improve generalization to changing important factors of variation in the data. While these improvements have often been striking, raises a question: what factors of variation should be changed between training and evaluation? One setting where factors of variation change naturally in reinforcement learning, as the data received from an environment changes reinforcement learning, an agent using RIMs may be able to learn faster as its specialization leads to improved generalization to previously unseen aspects of the environment.
the agent learns and improves. We conjecture that when applied to
To investigate this we use an RL agent trained using Proximal Policy Optimization (PPO) (Shhulman et al. (017) with recurrent network producing the policy. We employ an LSTM as a baseline, and compare results to the RIMs architecture. This was a simple drop-in replacement and did not require changing any of the hyperparameters for PPO. We experiment on the whole suite of Atari games and find that simply replacing the LSTM with RIMs greatly improves performance (Figure 6
There is also intriguing connection between the selective activation in RIMs and the concept of affordances from cognitive psychology (Gibson 1977] Cisek & Kalaska 2010). To perform well in environments with a dynamic combination of risks and opportunities, an agent should be ready to adapt immediately, releasing into execution actions which are at least partially prepared. This suggests agents should process sensory information in a contextual manner, building representations of potential actions that the environment currently affords. For instance, in Demon Attack, one of the games where RIMs exhibit strong performance gains, the agent must quickly choose between targeting distant aliens to maximize points and avoiding fire from close-by aliens to avoid destruction (indeed both types of aliens are always present, but which is relevant depends on the player's 'ss position). We hypothesize that in cases like this, selective activation of RIMs allows the agent to rapidly adapt its information processing to the types of actions relevant to the current context.
4.4 ABLATIONS
Role of Top-Down Modulation: Removing Input Attention We study the scenario where we remove the input attention process (Section 2.3 but still allow communication between RIMs (Section 2.4p. We train this agent on 30 ATARI games for 30M time-steps each and compare the performance of this agent with the normal RIMs-PPO agent. We find that the RIMs agent still outperform this agent on 11 out of 30 games, while on I game (Frostbite) we see the proposed baseline agent substantially improves the performance. For more details regarding the training curves, refer to Fig. 25 (in Appendix).
Importance of communication between RIMs: For copying, we performed an ablation where remove the communication between RIMs. We also varied the number of RIMs as well as the
8

---

Under review as a conference paper at ICLR 2020
0.50
0.25
0.00
-0.00
-0.7
11.00
Figure 6: RIMs-PPO relative score improvement over LSTM-PPO baseline (Schulman et al. 2017 across all Atari games averaged over trials per game. In hoth cases, PPO was used with the exact same settings, and the only change is the choice of recurrent architecture. More detailed experiments with learning curves in Appendix C
number of activated RIMs (T We found that the communication between RIMs is essential for good performance. We found similar results for the sequential MNIST resolution task.
Importance of sparsity of activation of the RIMs For the copying task, well as for the sequen- tial MNIST changed resolution task, we performed an ablation where we kept all RIMs active for all time steps (T. We found that we were not able to achieve strong generalization as compared to the best performing RIMs model. On Atari we found that using K4 = 5 slightly improved over results compared with A 4, but both had similar performance across the vast majority of games, suggesting that the A hyperparameter is reasonably flexible in practice.
Varying the number of attention heads for communication: Here, we study what happens the output of RIMs only has one 'object' rather than multiple ones (Section (.2). The intuition is
that RIM processing can be applied to any "head" which matches the query by an individual RIM. So, having more heads should help, different heads could be used by different RIMs, rather than every RIM competing for the same head. We study this in the context of bouncing balls. We found that using multiple heads improves the performance, thus validating our hypothesis (Sec. 2.2 See Appendix C.II for details.
Randomly Dropping Out RIMs: Modular structures are aggregates of mechanisms that can perform functions without affecting the remainder of the system, and interact needed. To what extent are trained RIMs able to model meaningful phenomena when other RIMs are removed? We performed experiment random RIM at test time. We found that in the absence of selective activation (i.e. when k = ky, Section C.13) the performance degraded very badly, but the performance degrades much less with selective activation. See Appendix C313 for details.
moving MNIST digits where train normally and "dropout"
CONCLUSION
Many systems of interest comprise multiple dynamical processes that operate relatively independently and only occasionally have meaningful interactions. Despite this, most machine learning models employ the opposite inductive bias. i.c., that all processes interact. This can lead to poor generalization (if data is limited) and lack of robustness to changing task distributions. We have proposed new architecture, Recurrent Independent Mechanisms (RIMs), in which we learn multiple recurrent modules that are independent by default, but interact sparingly. Our positive experimental results lend support to the consciousness prior (Bengio) 2017 7). i.c.. the importance of computational elements which focus few mechanisms at a time in order determine how high-level state evolves over time, with many aspects of the state not being affected by this attentive dynamics (i.e., following default dynamics). For the purposes of this paper, note that the notion of RIMs is not limited to the particular architecture employed here. The latter is used as a vehicle to assay and validate our overall hypothesis (cf. Appendix A), but better architectures for the RIMs model can likely be found.

---

Under review as a conference paper at ICLR 2020
REFERENCES
Jacob Andreas, Marcus Rohrbach, Trevor Darrell, and Dan Klein. Neural module networks. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 39-48, 2016.
Damitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to align and translate. arXiv preprint arXiv: 1rX.v::4090077. 2014.
Peter W Battaglia, Jessica B Hamrick, Vietor Bapst, Alvaro Gonzaez-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: 180..0106666666 2018.
Eric B Baum and David Haussler. What size net gives valid generalization? In Advances in neural information processing systems, pp. 81-90, 1989.
Yoshua Bengio. The consciousness prior. arXiv preprint arXiv: 2017.
Yoshua Bengio, Tristan Deleu, Nasim Rahaman, Rosemary Ke, Séhastien Lachapelle, Olexa Bilaniuk, Anirudh Goyal, and Christopher Pal. A meta-transfer objective for learning to disentangle causal mechanisms. arXiv:190..10112, 2019.
Léon Bottou and Patrick Gallinari. A framework for the cooperation of learning algorithms. In Advances in neural information processing systems, pp. 781-788, 1991.
Matthew Botvinick and Todd Braver. Motivation and cognitive control: from behavior to neural mechanism. Annual review of psychology, 66, 2015.
Michael M Bronstein, Joan Bruna, Yann LeCun, Arthur Szlam, and Pierre Vandergheynst. Geometric deep learning: going beyond euclidean data. IEEE Signal Processing Magazine, 34(4):18-42, 2017.
Maxime Chevalier-Boisver and Lucas Willems. Minimalistic gridworld environment for openai gym. https://gi /ttt:b/.//tttbb...../ ctt.huu.oooo////ommmmmm/o.oo. com/ ma x weaaaiimmeee i sareeeccc/yymmmee iitmmmmmiiii-... 2018.
Maxime Chevalier-Boiswe Thien Huu Nguyen, and Yoshua Bengio. Babyai: First steps towards grounded language learning with human in the loop. arXiv preprint wxxv:/8/0000822.. 2018.
Dzmitry Bahdanau, Salem Lahlou, Lucas Willems, Chitwan Saharia,
Junyoung Chung, Kyle Kastner, Laurent Dinh, Kratarth Goel, Aaron Courville, and Yoshua Bengio. A recurrent latent variable model for sequential data. In Advances in neural information processing systems, pp. 2980-2988, 2015.
Junyoung Chung. Sungjin Ahn, and Yoshua Bengio. Hierarchical multiscale recurrent neural networks. arXiv preprint 1rXiv:1609.0170000 2016.
Paul Cisek and John F Kalaska. Neural mechanisms for interacting with world full of action choices. Annual review of neuroscience, 33:269-298, 2010.
Emily Denton and Rob Fergus. Stochastic video generation with learned prior. arXiv preprint arXiv:1802.0768%. 2018.
Robert Desimone and Jody Duncan. Neural mechanisms of selective visual attention. Annual Review of Neuroscience, 18:193-221 222, 1995.
A. Dickinson. Actions and habits: the development of behavioural autonomy. Philosophical Transactions of the Royal Society B: Biological Sciences, 308(1135):67-78, 1985. ISSN 0080- 4622. doi: 10. 1098//ttt org/ htte//ootee///300//110
http:/ssst.119..00 1985.0010. URL
roya soci etypub www.hliishhiio iihiin
Salah El Hihi and Yoshua Bengio. Hierarchical recurrent neural networks for long-teme dependencies. In Advances in neural information processing systems, pp. 493-499, 1996.
10

---

Under review as a conference paper at ICLR 2020
Chrisantha Fernando, Dylan Banarse, Charles Blundell, Yori Zwols, David Ha. Andrei A Rusu, Alexander Pritzel, and Daan Wierstra. Pathnet: Evolution channels gradient descent in super neural networks. arXiv preprint arXiv:1011008344 2017.
James Gibson. The theory of affordances. Hilldale, USA, 1(2,, 1977.
Justin Gilmer, Samuel S Schoenholz, Patrick F Riley, Oriol Vinyals, and George E Dahl. Neural message passing for quantum chemistry. In Proceedings of the 34th International Conference on Machine Learning- Voolume 70, pp. 1263-1272. JMLR. org, 2017.
Anirudh Goyal, Riashat Islam, Daniel Strouse, Zafarali Ahmed, Matthew Botvinick, Hugo Larochelle, Sergey Levine, and Yoshua Bengio. Infobot: Transfer and exploration via the information bottleneck. arXiv preprint arXiv: 19011 2019a.
Anirudh Goyal, Shagun Sodhani, Jonathan Binas, Xue Bin Peng, Sergey Levine, and Yoshua Bengio. Reinforcement learning with competitive ensembles of ceeemation-constriinn primitives. arXiv preprint arXiv: 10667, 2019b.
Alex Graves, Greg Wayne, and Ivo Danihelka. Neural turing machines. arXiv:1410.5401, 2014a.
arXiv preprini
Alex Graves, Greg Wayne, and Ivo Danihelka. Neural turing machines. CoRR, URL
2b
Alex Graves, Greg Wayne, Malcolm Reynolds, Tim Harley, Ivo Danihelka, Agnieszka Grabska- Barwinska, Sergio Gómez Colmenarejo, Edward Grefenstette, Tiago Ramalho, John Agapiou, et al. Hybrid computing using a neural network with dynamic external memory. Nature, 538(7626)4471, 2016.
David Ha and Jürgen Schmidhuber. World models. arXiv preprint arXiv: 2018.
Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. arXiv preprint 1rXiv:181 1811,0555,, 2018.
Mikael Henaff, Jason Weston, Arthur Szlam, Antoine Bordes, and Yann LeCun. Tracking the world state with recurrent entity networks. arXiv preprint 1rXiv:16/2.03969) 2016.
Geoffrey E Hinton, Sara Sabour, and Nicholas Frosst. Matrix capsules with em routing. 2018.
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 9(8): 1735-1780, 1997.
Robert A Jacobs, Michael I Jordan, Steven J Nowlan, Geoffrey E Hinton, al. Adaptive mixtures of local experts. Neural computation, 3(1):79-87, 1991.
Yacine Jernite, Edouard Grave, Armand Joulin, and Tomas Mikolov. Variable computation in recurrent neural networks. arXiv preprint arXiv: 1.0688,, 2016.
Nan Rosemary Ke, Anirudh Goyal, Olexa Bilaniuk, Jonathan Binas, Michael C Mozer, Chris Pal, and Yoshua Bengio. Sparse attentive backtracking: Temporal credit assignment through reminding. In Advances in Neural Information Processing Systems, pp. 7640-7651, 2018.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv.1412.6988, 2014.
Thomas Kipf, Ethan Fetaya, Kuan-Chich Wang, Max Welling, and Richard Zemel. Neural relational inference for interacting systems. arXiv preprint arXiv: 1802 2018.
Louis Kirsch. Julius Kunze, and David Barber. Modular networks: computation. In Advances in Neural Information Processing Systems, pp. 2408-2418, 2018.
Learning decompose neural
Wouter Kool and Matthew Botvinick. Mental labour. Nature human behaviour; 2(12):899 -908899-00. 2018.

---

Under review as a conference paper at ICLR 2020
Ilya Kostrikov. Pytorch implementatiinoon of reinforcement learning algorithms. https://qitin htt.://ottikko///oroccc-------a-
-atooacccttt-aaaa 2018.
Jan Koutnik, Klaus Greff, Faustino Gomez, and Juergen Schmidhuber. A clockwork rnn. arXiv preprint arXiv: 1x0.v:140..3511, 2014.
David Krueger, Tegan Maharaj. János Kramár, Mohammad Pezeshki, Nicolas Ballas, Nan Rosemary Ke, Anirudh Goyal, Yoshua Bengio, Aaron Courville, and Chris Pal. Zoneout: Regularizing mns by randomly preserving hidden activations. arXiv preprint arXiv: 0/30, 2016.
Shuai Li, Wanqing Li, Chris Cook, Ce Zhu, and Yanbo Gao. Independently recurrent neural network (indran): Building a longer and deeper rnn. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 5457-5466, 2018.
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. arXiv preprint 1609. (09.000977,, 2016.
Daniel Neil, Michael Pfeiffer, and Shih-Chii Liu. Phased Istm: Accelerating recurrent network training for long event- based sequences. In Advances in neural information processing systems, pp. 30.38888299 2016.
Giambattista Parascandolo, Niki Kilbertus, Mateo Rojas-Carulla, and Berhhard Schölkopf. Learning independent causal mechanisms. In Proceedings the 35th International Conference on Machine Learning (ICML), pp. 4033 4041, 2018. URL http:://prcceeeiiggg.... maaps.mlr... m.... www....ressvvvv.
parascadollola. html
Judea Pearl. Causality: Models, Reasoning, and Inference. Cambridge University Press, New York, NY, 2nd edition, 2009.
Jonas Peters, Dominik Janzing, and Bernhard Schölkopf. Elements of Causal Inference Foundations and Learning Algorithms. MIT Press, Cambridge, MA, USA, 2017. ISBN 978-0-262-03731-0
David Raposo, Adam Santoro, David Barrett, Razvan Pascanu, Timothy Lillicrap, and Peter Battaglia. Discovering objects and their relations from entangled scene representations. arXiv preprint UrXiv:1702.05000 2017.
Scott Reed and Nando De Freitas. Neural programmerr-interpeeees. arXiv preprint 1..6:95511.06279 2015.
Eric Ronco, Henrik Gollee, and Peter Gawthrop. Modular neural networks and self-deccopoosttin. Technical Report CSC-96012, 1997.
Clemens Rosenbaum, Tim Klinger, and Matthew Riemer. Routing networks: Adaptive selection of non-linear functions for multi-task learning. arXiv preprint arXiv: 2017.
Clemens Rosenbaum, Ignacio Cases, Matthew Riemer, and Tim Klinger. Routing networks and the challenges of modular and compositional computation. arXiv preprint arXiv: /904. 92.9440..2747... 2019.
Andrei A Rusu, Neil C Rabinowitz, Guillaume Desjardins, Hubert Soyer, James Kirkpatrick. Koray Kavukcuoglu, Razvan Pascanu, and Raia Hadsell. Progressive neural networks. arXiv preprint 1rXx::1606.0407 04671, 2016.
Sara Sabour, Nicholas Frosst, and Geoffrey E Hinton. Dynamic routing between capsules. In Advances in neural information processing systems, pp. 9856-3866, 2017.
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.
Adam Santoro, Ryan Faulkner, David Raposo, Jack W. Rae, Mike Chrzanowski, Theophane Weber, Daan Wierstra, Oriol Vinyals, Razvan Pascanu, and Timothy P. Lillicrap Relational recurrent neural networks. CoRR, abs/1 806.01822, 2018. /tl./hccp://aaxttv.rg////////// 01822
Franco Scarselli, Marco Gori, Ah Chung Tsoi, Markus Hagenbuchner, and Gabriele Monfardini. The graph neural network model. IEEE Transactions on Neural Networks, 20(1):61-80 2008
12

---

Under review as a conference paper at ICLR 2020
Jürgen Schmidhuber. One big net for everything, arXiv preprint arXiv:1802.08864 2018.
Bernhard Schölkopf, Dominik Janzing, Jonas Peters, Eleni Sgouritsa, Kun Zhang, and Joris Mooij. On causal and anticausal learning. In J. Langford and J. Pineau (eds.), Proceedings of the 29th International Conference on Machine Learning (ICML), pp. 1255-1262, New York, NY, USA, 2012. Omnipress.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.063477 2017.
Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoe Lc, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. arXiv preprint arxiv: 100100630,,6888 2017.
Herbert A Simon. The architecture of complexity. In Facets of systems science, pp. 457-476. Springer, 1991.
Shagun Sodhani, Anirudh Goyal, Tristan Deleu, Yoshua Bengio, Sergey Levine, and Jian Tang. Learning powerful policies by using consistent dynamics model. arXiv preprint arXiv: 04355, 2019.
Andrea Tacchetti, Francis Song, Pedro AM Mediano, Vinicius Zambaldi, Neil C Rabinowitz, Thore Graepel, Matthew Botvinick, and Peter W Battaglia. Relational forward models for multi-agent learning. arXiv preprint arXiv: 1044, 2018.
Yee Teh, Victor Bapst, Wojciech M Czarnecki, John Quan, James Kirkpatrick, Raia Hadsell, Nicolas Heess, and Razvan Pascanu. Distral: Robust multitask reinforcement learning. In Advances in Neural Information Processing Systems, pp. 1496-4508 2017.
Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physics engine for model-based control. In 2012 IEEE/RSJ International Conference on Intelligent Robots and Systems, pp. 5026-5033. IEEE, 2012.
Sjoerd Van Steenkiste, Michael Chang, Klaus Greff, and Jürgen Schmidhuber. Relational neural expectation maximization: Unsupervised discovery of objects and their interactions. arXiv preprint 1rXiv:1802.10353 2018.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information processing systems, pp. 5998-6008, 2017.
H. L. F. von Helmholtz. Handbuch der physiologischen Optik, volume III. Voss, 1867.
Erich von Holst and Horst Mittelstaedt. Das reafferenzprinzip. Naturwissensxhaafeen 37(20):464-47, Jan 1950. doi: 10. O.1007/BF00622503..
Nicholas Watters, Daniel Zoran, Theophane Weber, Peter Battaglia, Razvan Pascanu, and Andrea Tacchetti. Visual interaction networks: Learning a physics simulator from video. In Advances in neural information processing systems, pp. 4539-4547, 2017.
Ronald Williams and David Zipser. A learning algorithm for continually running fully recurrent neural networks. Neural computation, 1(2):270-280) 1989.
13

---

Under review as a conference paper at ICLR 2020
A DESIDERATA FOR RECURRENT INDEPENDENT MECHANISMS
We have laid out case for building models composed of modules which by default operate independently and can interact in limited manner. Accordingly, our approach to modelling the dynamics of world starts by dividing the overall model into small subsystems (or modules), referred to as Recurrent Independent Mechanisms (RIMs), with distinct functions learned automatically from data.Our model encourages sparse interaction, i.c., we want most RIMs to operate independently and follow their default dynamics most of the time, only rarely sharing information. Below, we lay out desiderata for modules capture modular dynamics with sparse interactions.
Competitive Mechanisms: Inspired by the observations in the main paper, we propose that RIMs utilize competition to allocate representational and computational resources. As argued by (Parascandolo el al.) 2018). this tends to produce independence among learned mechanisms if the training data has been generaled by independent physical mechanisms.
The points mentioned in Section Din in principle pertain to synthetic and natural intelli- gent systems alike. Hence, it not surprising that they also appear neuroscience. For instance, suppose are looking particular object in a large scene, using limited processing capacity. The biased competition theory of selective attention conceptualizes basic findings of experimental psychology and neuroscience (Desimone) & Duncan 1995;: our capacity of parallel processing of and reasoning with high-level concepts is limited, and many brain systems representing visual information use competition to allocate resources. Competitive interactions among multiple objects occur automatically and operate parallel across the visual field. Second, the principle of selectivity amounts to the idea a perceiver has the ability to filter out unwanted information and selectively process the rest of the information. Third, top-down bias originating from higher brain areas enables us to selectively devote resources to input information that may be of particular interest or relevance, This may be accomplished by units matching the internal model of object process of interest being pre-activated and thus gaining an advantage during the competition brain mechanisms.
Top Down Attention:
Sparse Information Flow:
Each RIMs' dynamics should only affected by RIMs which are deemed relevant. The fundamental challenge is centered around establishing sensible communication between RIMs. In the presence of noisy or distracting information, a large subset of RIMs should stay dormant, and not be affected by the noise. This way, training ensemble of these RIMs can be more robust out-of-distribution distractor observations than training one big homogeneous neural network (Schmidhuber 2018).
Modular Computation Flow and Modular Parameterization: Fach RIM should have its dy-
namics operating by default, in the absence of interaction with other RIMs. The total number of parameters (i.c. weights) can be reduced since the RIMs can specialize on simple sub-problems, similar to (Parascandolo et al.) 2018). This can speed up computation and improve the generalisation ability of the system (Baum & Haussler 1989 The individuals RIMs in the ensemble should be simple also to prevent individual RIMs from dominating and modelling complex, composite mechanisms.
refer a parameterization as modular if most parameters are associated to individuals RIMs only. This has the desirable property that a RIM should maintain its own independent functionality even as other RIMs are changed (due to its behavior being determined by its own self- self-contained parameters).
B EXTENDED RELATED WORK
Table 2: A concise comparison of recurrent models with modular memory.
Modular Memory
Sparse
Modular Computation Flow
Modular Parameterization
Method / Property
Information Flow
C
L.TTM/ RNN
X
x
X
X
Relational RNN (Santoro al. 2018
/
x
NTM Grrves
2014b
/
/
X
SAB et al. 2018
8
/
X
X
IndRNN Li et al. 2018
/
x
X
/
RIMs
/
The present section provides further details on related work, thus extending Section
Neural Turing Machine (NTM). The NTM (Graves et al. 20144) has a Turing machine inspired memory with a sequence of independent memory cells, and uses an attention mechanism to move heads over the cells while performing targeted read and write operations. This shares a key idea with RIMs: that input information should only impact sparse subset of the memory by default, while keeping most of the memory unaltered. The RIM
14

---

Under review as a conference paper at ICLR 2020
model introduces the idea that each RIM has its own independent dynamics, whereas the mechanism for updating memory cells update is shared.
Relational RNN. The Relational Models paper (Santoro et al. 2018) is based on the idea of using a multi- attention mechanism to share information between multiple parts of memory. It is related to our idea but a key difference is that we encourage the RIMs to remain separate as much possible, whereas (Santoro et al. 2018) allows information between the parts to flow on each step (in effect making the part distribution only relevant to a particular step). Additionally, RIMs has the notion of each RIM having its own independent transition dynamics which operate by default, whereas the Relational RNN only does computation and updating of the memory using attention.
Sparse Attentive Backtracking (SAB). The SAB architecture (Ke et al. 2018 explores RNNs with self- attention across time steps well variants where the attention is sparse the forward pass and where the gradient sparse in the backward pass. It shares the motivation of using sparse attention to keep different pieces of information separated, but differs from the RIMs model in that it considers separation between time steps rather than separation between RIMs.
Independently Recurrent Neural Network (IndRNN). The IndRNN (Li et al. 2018 replaces the full transi- tion matrix in vanilla RNN (between time steps) a diagonal transition weight matrix. In other words, each recurrent unit has completely independent dynamics. Intriguingly they show that this gives much finer control over the gating of information, and allows for such an RNN to learn long-term dependencies without vanishing or exploding gradients. Analysis of the gradients shows that having smaller recurrent transition matrices mitigates the vanishing and exploding gradient issuc. This may provide further explanation for why RIMs perform well on long sequences.
Consciousness Prior (Bengio 2017): This is based on the assumption of sparse graphical model describing the interactions between high-level variables. using gating mechanisms to select only a subset of high-level variables to interact at any particular time. This is closely related to our work in the sense high level abstract representation is based on the representations of the RIMs, which are activated sparsely and interact sparsely, Our paper thus helps to validate the consciousness prior idca.
Recurrent Entity Networks: EnTNet (Henafi et al. 2016) can be viewed as set of separate recurrent models whose hidden states store the memory slots. These hidden states are either fixed by the gates, or modified through a simple RNN-style update. Moreover, EntNet uses an independent gate for writing to each memory slot. Our work is related in the sense that we also have different recurrent models (i.e.RIMs, though each RIM has different parameters), but we allow the RIMs to communicate with each other sparingly using an attention mechanism.
Capsules and Dynamic Routing: EM Capsules (Hinton et al. 2018) and the preceding Dynamic Capsules (Sabour et al. (2017) use the poses of parts and learned part + object relationships to vote for the poses of objects. When multiple parts cast very similar votes, the object assumed to be present, which is facilitated by an interactive inference (routing) algorithm.
Recent graph-based architectures have studied combinatorial generalization in the context of modeling dynamical systems like physics simulation, multi-object scenes, and motion-capture
Relational Graph Based Methods:
data, and multiagent systems (Scarselli et al. Bronstein et al. [017 Watters et al. [0017] Raposo et al.
2017
Santoro et al. 2017 Gilmer et al. [2017] 7 Van Steenkiste et al. [0018] Kipf et al. 2018 Battaglia al.
2018 Tacchetti et al. 2018 One can also view our proposed model as a relational graph neural network, where
nodes are parameterized as individual RIMs and edges are parameterized by the attention mechanism. Though, its important to emphasize that the topology of the graph induced in the proposed model is dynamic, while in most graph neural networks the topology is fixed.
Default Behaviour: Our work is also related to work in behavioural research that deals with modes of decision making (Dickinson) 1985; [Botinick & Braver 2015 Kool & Botvinick, 2018): an automatic systems that relics on habits and a controlled system that uses some privileged information for making decision making. The proposed model also has two modes of input processing, RIMs which activate uses some external sensory information, and hence analogous controlled system. RIMs which don't 't activate, they synonymous to habit based system. There some work done trying in Reinforcement learning, trying learn default policies, which have shown to improve transfer and generalization multi-task RI. (Tch et al. 2017 Goyal et al.) (201aa). The proposed method is different in the sense, we are not trying to learn default policies which effect the environment, instead we want to learn mechanisms. which try to understand the environment. State dependent activation of different primitive policies was also studied in (Goyal et al.) [2019)). and the authors showed that they can learn different primitives, but they also consider that only single primitive can active at a particular time step. Also, note that primitive policies try to effect the environment, whereas mechanism try to understand the enviormment.
15

---

Under review as a conference paper at ICLR 2020
C EXPERIMENTAL DETAILS AND HYPERPARAMETERS

---

Under review as a conference paper at ICLR 2020
would be interesting to think of ways allow a more "structured" decoder. The reason for this is that even if the RIMs generalize to new environments, the shared decoder can fail to do so. So changing the structure of decoder could be helpful.
For the RL experiments, we also tried providing the previous actions, rewards, language instruction as input to decide the activation of RIMs. This is consistent with the idea of efference copies as proposed by von Helmholtz (1867): yon Holst & Mittelstaedt (1950), i.e., using copies of motor signals as inputs. Preliminary experiments shows that this improves the performance in Atari games.
C.4 LANGUAGE MODELING
Table 4: Wikitext-2 results
Approach
Num. Parameters
Train PPL
Valid PPL
Test PPL
LSTM (2-layer)
21.2M
39.78
109.25
102.53
Relational Memory (Santoro et al. 2018
1 11M
n/a
112.77
107.21
RIMs (2-layer, ky 6, KA 6)
23.7M
41.27
103.60
98.66
We investigate the task of word-based language modeling. We ran experiments the wikitext-2 dataset (Merity al. 2016). We ran each experiment for fixed 100 epochs. These results are in Table Our goal in this experiment is to demonstrate the breadth of the approach by showing that RIMs performs well even on datasets which are noisy and drawn from the real-world.
C.5 COPYING TASK
We used a learning rate of 0.001 with the Adam Optimizer and trained each model for 150 epochs (unless the model was stuck, we found that this was enough to bring the training eтror close to zero). For the RIMs model we used 600 units split across 6 RIMs (100 units per block). For the LSTT we used total of 600 units. We did explore this extensively but we qualitatively found that the results copying were not very sensitive to the exact number of units.
The sequences to be copied first have 10 random digits (from 0-8), then a span of zeros of some length, followed by special indicator "9" in the input which instructs the model to begin outputting the copied sequence.
In our experiments, we trained the models with "zero spans" of length 50 and evaluated the model with "zero spans" of length 200. We note that all the ablations were run with the default parameters (i.c number of keys, values as for RIMs model) for 100 epochs. Tab. shows the effect two baselines as compared the RIMs model (a) When we allow the input attention for activation of different RIMs but we dont allow different RIMs to communicate. (b) No Input attention, but we allow different RIMs to communicate with each other. Tab, S shows that the proposed method is better than both of these baselines. For copy task, used head in input attention, and heads RIMs communication. We note that even with RIM, not exactly same a LSTM, because each RIM can still reference itself.
C.6 SEQUENTIAL MNIST TASK
this task we considered classifying binary MNIST digits by feeding the pixels to an RNN (in a fixed order scanning over the image). As the focus of this work is on generalization, we introduced a variant on this task
where the training digits are at resolution of 14 14 (sequence length of 196). We then evaluated on MNIST used the
digits of different higher resolutions (16 16. 19x 19, and 24 24). When re-scaling the images, nearest- neiget-- based down down- scaling and performed binarization after re-scaling. We trained with learning rate of 0.0001 and the Adam optimizer. For RIMs we used a total of 600 hidden units split across 6 RIMs (100 units per block). For the LSTM we used total of 600 units. We ran proposed model as well as baselines for 100 epochs. For sequential MNIST task, we used 1 head in input attention, and 4 heads for RIMs communication.
C.7 IMITATION LEARNING: ROBUSTNESS TO NOISE IN STATE DISTRIBUTION
Here, we consider imitation learning where we have training trajectories generated from an expert (Table G). We evaluate our model on continuous control tasks in Mujoco (in our case, Half-Cheetah) Todorov al. 2012). We take the rendered images as input and compared the proposed model with recurrent policy (i.e., LST LSTM). Since, using rendered image of the input does not tell anything about the velocity of the Half Cheetah, it makes the task partially observable. In order to test how well the proposed model generalizes during test, we add some noise (in the joints of the half-cheetah body). As one other hand, for the proposed model, there's also drop in performance but as bad as for the L.STM haseline.
after adding noise I.STM haselines performs poorly. On the
17

---

Under review as a conference paper at ICLR 2020
Table 5: Error (CE for last 10 time steps) the copying task. Note that while all of the methods are able learn to copy on the length seen during training. the RIMs model generalizes to sequences longer than those seen during training whereas the LSTM fails catastrophically.
Approach
Train Length 50
Test Length 200
RIMs
0.00
0.00
With input Attention and No Communication
RIMs (kT = Ky 2)
2.3
1.6
RIMs (KT k4 kA
1.7
4.3
KA = 2) RIMs (by- 55,
2.5
4,7
RIMs (KT =
0.4
4.0
RIMs (KT = =5. =A= 4)
0.2
0.7
RIMs (Ker =6. AAAA
3,3
2.4
RIMs (Ae =6. 33
1.2
1.0
RIMs (KT =6, K.
0.7
5.0
RIMs (KT = 6, AA = 5)
0.22
0.56
With No input Attention and Full Communication
RIMs (KT 6. A 6, hdii = 600)
0.0
0.7
RIMs (KT 5, A 5, hdim = =500)
0.0
1.7
RIMS (Ay 2, KA 2, Adim 2256)
0.0
2.9
RIMs (eT 2. KA hdim 512)
0,0
RIMs (ky I, A = haim 512)
0.0
0.2
Table 6: Imitation Learning: Results on the half-cheetah imitation learning task. RIMs outperforms a baseline LSTM when we evaluate with perturbations not observed during training (left). An example of an input image fed to the model (right).
Training
Perturbed States
Method Setting
Observed Reward
Observed Reward
LSTM (Recurrent Policy)
5400 + 100
2500 + 300
RIMs (kir 6, A = 3)
5300 200
3800 200
RIMs 6, A 6)
5500 + 100
2700 = 400
RIMs (without Input attention)
5400 + 100
3200 + 50
We use the convolutional network from Ha Schmidhuber 2018) as our encoder, GRU (Chung et al. 2015 with 600 units as deterministic path in the dynamics model, and implement all other functions as two fully connected layers of size 256 with ReLU activations. Since, here we are using images as input, which makes the task, partially observable. Hence, concatenate the past 4 observations, and then feed the concatenated observations input GRU (or our model). For our model, we use RIMs, each of size 100, and We follow the same setting as in (Hafner et al.] (2018] [Sodhani et al.] (019). We also compare the proposed method to the haseline where we dont include input attention (or top-down attention). shows, there's 's decline in performance if we dont use input attention, hence justifying the importance
ka
C.8 GENERALIZATION TO DISTRACTORS: ALGORITHM IMPLEMENTATION DETAILS
We evaluate the proposed framework using Adavantage Actor-Critic (A2C) to learn policy To(a|s, g) condi- tioned on the goal. To evaluate the performance of proposed method, we use a range of maze multi-room tasks from the gym minigrid framework (Chevalier- Boisvert & Willems 2018) and the A2C implementation from Boeiseitr-Boisvertt & Willems 2018). For the maze tasks, we used agent's relative distance to the absolute goal position "goal".
the maze environments, we use A2C with parallel workers. Our actor network and critic networks consist of two and three fully connected layers respectively, each of which have 128 hidden units. The encoder network is also parameterized as neural network. which consists of 1 fully connected layer. We use RMSProp with an initial fearning rate of 0.0007 to train the models. Due to the partially observable nature of the environment, we further use 1.STM to encode the state and summarize the past observations.
18

---

Under review as a conference paper at ICLR 2020
Figure 7: An example of the minigrid task.
C.9 MINIGRID ENVIRONMENTS FOR OPENAI GYM
The MultiRoom environments used for this research are part of MiniGrid, which is an open source gridworld packago This package includes family of reinforeement learning environments compatible with the OpenAl Gym framework. Many of these environments are parameterizable that the difficulty of tasks can be adjusted (e.g., the size rooms often adjustable).
C.9.1 THE WORLD
MiniGrid, the world grid size NxN. Each tile in the grid contains exactly zero or one object. The possible object types are wall, door, key, ball, box and goal. Each object has an associated discrete color, which can be one of red, green, blue, purple, yellow and grey. By default, walls are always grey and goal squares are always green.
C.9.2 REWARD FUNCTION
Rewards sparse for all MiniGrid environments. In the MultiRoom environment, episodes are terminated with a positive reward when the agent reaches the green goal square. Otherwise, episodes are terminated with zero reward when time step limit reached. In the FindObj environment, the agent receives positive reward if reaches the object to be found, otherwise zero reward if the time step limit is reached.
The formula for calculating positive sparse rewards 1 0.9 (slep cound /maz steps). That is, rewards are always between zero and one, and the quicker the agent can successfully complete an episode, the closer 1 the reward will be. The maz_ sttps parameter different for each environment, and varies depending on the size of each environment, with larger environments having higher time step limit.
C.9.3 ACTION SPACE
There are seven actions in MiniGrid: turn left, turn right, move forward, pick up an object, drop object, toggle and done. For the purpose of this paper, the pick up. drop and done actions irrelevant. The agent right action
the and face one of possible directions (north, south, cast, west). The move
turn left and forward action makes the agent move from its current tile onto the tile in the dircction is currently facing, provided there is nothing on that tile, or that the tile contains an open door. The agent can open doors if they are right in front of it by using the toggle action.
19

---

Under review as a conference paper at ICLR 2020
C.9.4 OBSERVATION SPACE
Observations in MiniGrid are partial and egocentric. By default, the agent sees a square of 7x7 tiles in the direction it is facing. These include the tile the agent standing on. The agent cannot through walls closed doors. The observations are provided as a tensor shape 7x7x3. However, note that these are not RGB images. Each tile is encoded using integer values; one describing the type of object contained the cell, one describing its color, and flag indicating whether doors are open or closed. This compact encoding was chosen for space efficiency and to enable faster training. The fully observable RGB image view of the environments shown in this paper is provided for human viewing.
C.9.5 LEVEL GENER
The level generation this task works follows: Generate the layout of the map number rooms with different sizes (at most size Y) and green goal) (2) Add the agent to the map at a random location the first room. (3) Add the goal at a random location in the last room. A neural network parameterized as CNN is used to process the visual observation.
We follow the same architecture BlockLSTM.
(Chvaalier Bheiseeer-Boisvertt & Willems 2018 but replace the LSTM layer with
C.10 BOUNCING BALLS
We use the bouncing-ball dataset from (Van Steenkiste al. 2018). The dataset consists of 50,000 training examples and 10,000 test examples showing ~00 frames of either 4 solid balls bouncing in confined square geometry, 6-8 balls bouncing in a confined geometry, or 3 balls bouncing in confined geometry with a random occluded region. In all cases, the balls bounce off the wall as well as off one another. We train baselines well proposed model for about 100 epochs using 0.0007 learning rate and using Adam optimizer (Kingmal Ba 2014). We use the same architecture for encoder as well decoder as in (Van Steenkiste et al. (018)) We train the proposed model as well as the baselines for 100 epochs. Our goal in this section is to give more thorough experimental results amitted from the main paper for the sake of brevity. Below, we highlight few different results.
C.10.1 DIFFERENT RIMS ATTEND TO DIFFERENT BALLS
Active region.
Innetive region
Ground truth
Prodletion (input foed) Prediction (rallout)
Figure 8: Different RIMs attending to Different Balls. For understanding what each RIM is actually doing, we associate each with separate encoder, which are spatially masked. Only encoders be active any particular instant and there are four different balls. did this to check if there would be the expected geometric activation of RIMs. 1.) Early in training. RIM activations correlated more strongly with the locations of the four different balls. Later in training, this correlation decreased and the active strips did not correlate as strongly with the location of balls. As the model got better at predicting the location, it needed to attend less the actual objects. The top row shows every 5th frame when the truth is fed in and the bottom shows the results during rollout. The gray region shows active block. In top row, the orange corresponds to the prediction and in the bottom, green corresponds to the prediction.
order visualize what each RIM doing, we associate each RIM with different encoder. By performing spatial masking on the input, we can control the possible spatial input to each RIM. We use six overlaaagg horizontal strips and allow only 4 RIMs to be active at time (shown in Fig. 88. The mask is fixed mask of zeros with a band of ones that is multiplied by the input each encoder. Therefore, each of the 6 encoders gets 1/6th of the input. The goal was balls. We find that early training. RIMs' activations are strongly correlated with the location the balls. However, after training has proceeded for some time this correlation deteriorates. This is likely because the predictable dynamics of the system do not necessitate constant attention.
how the RIM activation patterns changed/correlated with the locations of the
10.2 COMPARISON WITH LSTM BASELINES
In Figures 9 10
and 12 we highlight different baselines and how these compare to the proposed RIMs model.
20

---

Under review as a conference paper at ICLR 2020
Trained on 6-8 Objects Tested on 4 Objects
frames in rollout
Truth
LSTM
Curtain
Truth
LSTM
Figure 9: Example of the other LSTM baselines. For the other experiments that we consider, here we show example outputs of our LSTM baselines. In each row; the top panel represents the ground truth and the bottom represents the prediction. All shown examples use an LSTM with 250 hidden units, as shown in Fig. [ Frames are plotted every 3rd time step. The red line marks 10 rollout frames. This is marked because after this we do not find BCE to be reliable measure of dissimilarity.
1 Ohiects
33
+4
+5
+6
+7
+8
19
+10
+11
+12
+13
+1d
15
8*
Trith
RIM
1STM
Traiuel con Olriects
Teeted Lun G-S. Olviente
Truth
%
%
8.
8.
RIM
1.5TM
Trainel onn 6-8 Oboietts
Texted
Obiietta
Trith
RIM
1.TT
Curtains
Trith
TUM
LSTM
Figure 10: Comparison of RIMs to LSTM baseline. For 4 different experiments in the text, we compare RIMs two different LSTM baselines. In all cases we find that during rollout, RIMs perform better than the L.TMMs accurately capturing the trajectories the halls through time. Due to the number hard collisions, accurate modeling is very difficult. In all cases, the first 15 frames of ground truth are fed in (last 6 shown) and then the system is rolled out for the next 15 time steps, computing the binary cross entropy between the prediction and the true balls at each instant, in Van Steenkiste et al. (2018). See the Appendix for losses the entire 35 frame rollout trajectory. In the predictions, the transparent blue shows the ground truth, overlaid to help guide the eye.
1.10.3 OCCLUSION
In Fig. 13 we show the performance of RIMs on the curtain dataset. We find RIMs are able to track balls through the occi usion without difficulty. Note that the LSTM haseline, is also able to track the ball through the "invisible" curtain.
1.10.4 STUDY OF TRANSFER
It is interesting to ask how models trained on dataset with 6-8 balls perform on dataset with 4 balls. In Fig. 14 we show predictions during feed-in and rollout phases.
21

---

Under review as a conference paper at ICLR 2020
Objects
10 frames rollout
Truth
LSTM
RIM
Trained 4 Objects Tested on 6-8 Objects
Truth
LSTM
RIM
Figure 11: Comparison between RIMs and LSTM baseline. For the 4 ball task and the 6-8 ball extrapolation task, here we show example oulput of from LSTM baseline and from RIMs. All shown examples use an LSTM with 250 hidden units, as shown Fig. Frames are plotted every time step. The red line marks 10 rollout frames. This is marked because after this we do not find BCE to be a reliahle measure of dissimilarity.
Trined on Objocts
Trined on 6.8 Objocts
4 Obbeets
Tested on 6-8 Objects
Tested on 4 Obiects
Curtain
Entoopy
ESTM $12
ESTR vin
RDD
Ground truih
Rollont
30
23
90
12 20 4
Figure 12: Comparison of RIMs to LSTM baseline. For 4 different experiments in the text, we compare RIMs to two different LSTM baselines. In all cases we find that during rollout, RIMs perform better than the LSTMs accurately capturing the trajectories of the balls through time. Due the number of hard collisions, accurate modeling is very difficult.
Rodiout
Prome: 2
3
11
14
20
26
32
35
33
43
49
Prodic Thar
Figure 13: RIMs on dataset with an occlusion. We show two trajectories (top and bottom) of three balls, For the left frames, at cach step the true frame is used as input. On the right, outlined black, the previous output is used as input.
C.11 ABLATIONS
We present ablation in addition the ones in Section 4.4 In this experiment, study the effect input attention (i.e top down attention) well the use mulli-heade head key-value attention. compare the proposed model (with input attention as well as we remove the input attention (and force all the RIMs to communicate with each other (b) We use I head for key value attention compared to multi-headed key-value attention. Results comparing the proposed model, with these baselines shown Fig. 15
key value attention) with 2 baselines: (a) which
22

---

Under review as a conference paper at ICLR 2020
RIMs
Troining
Teasser
Podake
Figure 14: RIMs transferred on new data. We train the RIMs model on the 6-8 ball dataset (as shown in the top row). Then, we apply the model to the 4 ball dataset, as shown in the bottom.
Prodiction Loss
Binnsy Croes Enl sppy
0.25
RIMs
One-head
0:55
ingur
0.10
0.03
6.00
Frames 19+20
Figure 15: Ablation loss For the normal, one-head model, and without input attention, we show the loss during training and the loss for the 4th and 5th frame of rollout. We find that the one-head and without input attention models perform worse than the normal RIMs model during the rollout phase.
In Fig.
16 we show the predictions that result from the model with only active head.
Rollout
12
14
16
18
26
28
30
Truth
Prodiction
Truth
Prediction
Truuh
Prediction
Truth
Prediction
Figure 16: One head and no attention Using one head and no attention models, we show the rollout predictions in blue. On top we show results on the 4 ball dataset and on the bottom we show results on the curtains dataset,
23

---

Under review as a conference paper at ICLR 2020
Environment
LSTM-PPO
RIMs-PPO
Alien
1612 44
2152 1 81
Amidar
1000 58
1800 = 43
Assault
4000 = 213
5400 = 312
Asterix
3090 + 420
21040 = 548
Asteroids
1611.0 + 200
3801 = 89
Atlantis
3280000 + 200000
3500000 120000
BankHeist
1153 23
1195 4
BattleZone
21000 232.0
22000 + 324
BeamRider
698 + 100
5320 + 300
Bowling
+
42 1 13
Boxing
95 10
Breakout
593 + 90
590 10
Centipede
4600 + 312
5534 + 283
ChopperCommand
11000 790
12303 412
CrazyClimber Climber
138000 2412
132039 1 1221
DemonAttack
26320 1:: 3234
230324 1 4032
DoubleDunk
-3.0 1 0.5
-3.8 1 0.3
Enduro
1600 + 200
2800 + 232
FishingDerby
20 4
38 +
Freeway
+
33 =
Gopher
7000.0 + 402
33000 2210
Gravitar
500 + 100
1090 + 80
IceHockey
425 + +
-4 +
Jamesbond
800 100
Kangaroo
13000
1800
Krull
10000 + 500
7900 1 200
KungFuMaster
28000 + 2000
51000 = 800
NamcThisGame
4200 + 400
6800 + 300
Pong
+
20 +
PrivateEßye
+
0
Qbert
22000 + 300
22500 + 400
Riverraid
7500 + 300
12000 + 100
RoadRunner
53000 120
53430 300
Robotank
3
11 =
SpaceInvaders
1600 40
2800 80
StarGunner
35000 + 800
70000 + 1200
TimePilot
4000 + 100
10000 + 689
UpNDown
70000 + 6000
390000 20000
VideoPinball
90000 5000
220000 1 9000
WizardOrWor
3833 1 400
10800 700
Zaxxon
200 1 100
15000 1 600
Table 7: Scores obtained using PPO with the LSTM architecture and PPO with the RIMs architecture with
5.
C.12 ATARI
We used open-source implementation of PPO from (Kostriko)) 2018) with default parameters. We ran the proposed algorihtm with 6 RIMs, and kept the number of activated RIMs 4/5. We have not done any hyper-parameter search for Atari experiments.
24

---

Under review as a conference paper at ICLR 2020
Figure 17: A comparison showing relative improvement of RIMs with ka 5 over KA 4 bascline. Using A performs slightly than & 4 but still outperforms PPO, and has similar results across the majority of games.
0.75
0.25
Nomallzzd
-025
0050
Figure RIMs-PPO relative score improvement over L.STM-PPO baseline (Schulman et al. 2017) across all Atari games averaged over trials per game. In both cases PPO was used with the exact same settings with the only change being the choice of the recurrent architecture (RIMs with A 5).
C.12.1 TRANSFER ON ATARI
As very preliminary result, investigate feature transfer between randomly selected Atari games. In order to study this question, follow the experimental protocol of Rusu al. (0166).
We start by training RIMs three source games (Pong, River Raid, and Seaquest) and test if the Icarned features transfer to a different subset of randomly selected target games (Alien, Asterix, Boxing, Centipede, Gopher, Hero, James Bond, Krull, Robotank, Road Runner, Star Gunner, and Wizard of Wor). We observe, that RIMs result in positive transfer in 9 out of 12 target games, with three cases of negative transfer. On the other hand progressive networks (Rusu al. 2016) result in positive transfer in out of 12 target games, and cases of negative transfer. We also compare to LSTM baseline, which yields positive transfer in games.
C.13 BOUNCING MNIST: DROPPING OFF RIMS
We use the Stochastic Moving MNIST (SM-MNIST) (Denton & Fergus 2018) dataset which consists of sequenees of frames of size 64 x 64, containing one or two MNIST digits moving and bouncing off the walls. Training sequences are generated on the fly by sampling two different MNIST digits from the training set (60k total digits) and two distinct trajectories.
Here, show the effect of masking RIMs. Ideally, we would want different RIMs not co-adapt with each other. So, masking out a particular RIM should not really effect the dynamics of the entire model. We show qualitative comparisons Fig. 19 20
particular RIM and study the effect of the masking on the ensemble
21
23 In each of these figures, the model gets the ground truth image as input for first 5 time steps, and then asked to simulate the dynamics for next 25 time-steps. We find that sparsity is needed otherwise different
22
RIMs co-adapt with each other (for Fig.
23
We tried similar masking experiments for different
25

---

Under review as a conference paper at ICLR 2020
models like RMC, Transformers, EntNet (which learns mixture of experts), LSTMs, but all of them failed to do anything meaningful (after masking), We suspect this is partly due to learning homogeneous network.
26

---

Under review as a conference paper at ICLR 2020
5
5
6
6
5
45
4
4
6
6
OF
5
G
6
G
5
P
OF
&
&
&
C
%
R
P
A
F
8
7
7
4
E
E
A
A
A
T
7
7
R
E
EA
?
3
3//
3/
P
R
6
6
18
15
4
9
4
8
9
9
8
3
7
7
A
Y
4
"
16
+5
+
3
3
3/1
3/
3/
3/
3/
A
4
R
8
A/
GP
_
6
6
4
J
6
R
5
6
6
6
6
6
&
G
S
6
G
6
6
6
b
%
%
S
J
R
F
R
F
+
J
F
F
&
4
4
3
R
+5
8
&
6
C
69
tP
60
6
6
+
+
18
8
8
8
8
8
P
P
8
+5
3
9
8
3
6
9
3
9
6
8
R
R
6
OF
6
&
&
&
9
6
GC
&
OF
5
6
6
6
6
6
6F
66
GS
G 6
@ L
6
6
6PI
QC
6
GG
5
T
G
6
OF
S
S
6
6
G
G
G
S
S
R
S
ff
G
es
&s
&
&
BG
G
of
of
s
&
e
6
S
L
S
S
E
5
6
%
+
4
S
+
J
P
+
F
J
3
&
9
5
O
2
3
3
9
9
9
a
27
+
48
9
9
9
7
2
2D
9
#
7
7
7
7
7
4
3
2
2
*
+
$
b
3
3 3
9
7
7
8
A
&
7
Figure 19: 4 RIMs. (top k- = 2). Each sub- -ub-figure shows the effect of masking a particular RIM and studying the effect of masking on the other RIMs. For example, the top figure shows the effect of masking the first RIM, the second figure shows the effect of masking the second RIM etc.
27

---

Under review as a conference paper at ICLR 2020
S
5
3
5
4
3
4
Figure 20: 4 RIMs. (top k 3). Each sub-figure shows the effect of masking a particular RIM and studying the effect of masking the other RIMs. For example, the top figure shows the effect of masking the first RIM, the second figure shows effect masking second RIM etc.
28

---

Under review as a conference paper at ICLR 2020
%
40
6
P
6
P
%
P
P
2
8
8
8
9
8
Figure 21: 400dim, 5 RIMs, (top k = 2). Each sub-figure shows the effect of masking a particular RIM and studying the effect of masking on the other RIMs. For example, the top figure shows the effect of masking the first RIM, the second figure shows the effect of masking the sccond RIM etc.
29

---

Under review as a conference paper at ICLR 2020
Figure 22: 400dim, 5 blocks, (top k = 3). Each sub-figure shows the effect of masking particular RIM and studying the effect of masking the other RIMs. For examples, the top figure shows the effect of masking the first RIM, second figure shows the effect of masking the second RIM
30

---

Under review as a conference paper at ICLR 2020
Figure 23: 400dim, 5 blocks, (top k =). Each sub-fig sub-figure shows the effect of masking particular RIM and studying the effect of masking the other RIMs. For example, the top figure shows the effect of masking the first RIM, second figure shows the effect of masking the second RIM
31

---

Under review as a conference paper at ICLR 2020
C.13.1 ATARI RESULTS: COMPARISON WITH LSTM-PPO
Figure 24: Comparing RIMs-PPO with LSTM-PPO: Learning curves for KA 4, kA models and the LSTM- PPO baseline across all Atari games
RIMs-PPO
32

---

Under review as a conference paper at ICLR 2020
C.13.2 ATARI RESULTS: No INPUT ATTENTION
Here we compare the proposed method to the baseline. where we dont use input attention, and we force different RIMs to communicate with each at all the time steps.
Figure 25: Baseline agent with no input attention mechanism: Here we compare the RIMs to the baseline, where their is no input attention (i.e., top down attention) as well as at all the time steps. Learning curves for RIMs-PPO models, Baseline Agent, the LSTM-PPO baseline across 30 Atari games.
RIMs communicate with each other
33


Reviews:
[{"review_id": "SyxeOPO8tr", "reviewer": "AnonReviewer1", "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": "This work seems to propose an alternative to general RNN so that  the dynamics of sequential data can be better captured. The work is based on a hypothesis that a causal process can be modeled by \"independent\" modules and sparse interactions.\n \nThe paper is written in very fluent English, but the style is less technical. The impression is that the philosophical arguments and the machine learning realization have a gap in between. There is no much rigorous mathematical definition or derivation to back up the entire development. The mathematical symbols are a bit loosely defined. For example, it is unclear to the reviewer if h is a scalar or a vector.\n\nThe RIM idea seems to be derived from some ideas from the ``'causality community' . But the authors did not elaborate how significant will this structure change the state of art. In particular, reading the abstract or the introduction does not shed much light on what are the challenges of now the ML community is facing, and how this proposed RIM idea is going to help. This may have been obvious to the authors, but spelling them out may help the reviewer/readers to understand the contribution. The related work section helped a bit, but still unclear.\n\nThe reviewer feels that the paper stands at a high level in general, but lacks concrete examples/applications for general readers to appreciate the significance. Perhaps trying to re-organize this part could greatly help the readability.\n\nThe mathematical descriptions in 2.2, 2.3 and 2.4 are very hard to follow. There is no cost function (for training) around, but there are discussions of  'gradient'. Gradient of which function? ", "extra_scores": {"experience_assessment": "I do not know much about this area.", "review_assessment:_checking_correctness_of_experiments": "I did not assess the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I made a quick assessment of this paper.", "review_assessment:_checking_correctness_of_derivations_and_theory": "N/A"}, "extra_text": {}}, {"review_id": "SkxGMecLFB", "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": "This paper proposes a neural network architecture consisting of multiple independent recurrent modules that interact sparingly. These independent modules are not all used simultaneously, a subset of them is active at each time step. This subset of active modules is chosen through an attention mechanism. The idea behind this architecture is that it would allow the different modules to specialize in different mechanisms and that would allow compositionality. The empirical results suggest that the proposed approach is able to generalize better than traditional architectures (which all have the implicit assumption that all processes interact).\n\nThis paper is well-written and it provides a very thorough empirical analysis of the proposed idea. Because it is not in my area of expertise I’m not confident that I can assess its novelty or its relationship to other existing approaches.\n\nIn terms of presentation, I recommend the authors to enlarge some of the figures in the paper (e.g., I can’t read the small box in Figure 1) and to not use citations as nouns (e.g., “The mechanisms of this attention mechanism follow (Vaswani et al., 2017; Santoro et al., 2018), with the …”). I would also like to point out that although fairly different in how they tackle the problem, the work of Arjovsky et al. (2019) seems to be related to this one.\n\nThree questions I believe were not answered in the paper are:\n\n1) How is the performance related to the total number of subsystems (and the number of *active* ones). I can only see results related to that in Table 1, but the variation in the number of modules is pretty small (4-6). The results also don’t give any indication whether we want to have more modules active at each time step, if there’s a sweet spot, etc. It is said that the method seems to be robust to this choice but this claim is made because it performs similarly for the values 5 and 6 if I recall correctly.\n\n2) Is there any incentive in this architecture for a module to not simply “give up”? I mean, the modules are not necessarily incentivized to be used as often as possible, so could it be the case that a module learns to set its weights to zero?\n\n3) Would it make sense to present baseline results for an architecture that uses attention? It seems to me that LSTM was often the baseline of choice but RIMs have two important components: multiple LSTMs and an attention mechanism. Could the attention mechanism be explaining some of the results we are seeing?\n\nFinally, despite the very long appendix, I feel there are important details missing with respect to the empirical setup, at least in the Atari experiments which I’m more familiar with. Was stochasticity used, that is, sticky actions (Machado et al., 2018)? Moreover, for how long was PPO (and RIMs-PPO) trained in terms of number of frames? Finally, I’d recommend the authors to include a table with the actual average (and standard deviation) performance in each Atari games. It is really hard to know how well a method is doing by just squinting at learning curves. It is hard to know if the results are significant without a notion of variance.\n\n\nReferences:\n\nMartín Arjovsky, Léon Bottou, Ishaan Gulrajani, David Lopez-Paz: Invariant Risk Minimization. CoRR abs/1907.02893 (2019)\n\nMarlos C. Machado, Marc G. Bellemare, Erik Talvitie, Joel Veness, Matthew J. Hausknecht, Michael Bowling: Revisiting the Arcade Learning Environment: Evaluation Protocols and Open Problems for General Agents. J. Artif. Intell. Res. 61: 523-562 (2018)\n\n\n------\n\n\n>>> Update after rebuttal: I stand by my score after the rebuttal. \n\nUnfortunately I'm not an expert in this area and I don't feel confident in having a very strong opinion about this paper, willing to fight for its acceptance. I also agree with concerns raised by other reviewers. As I stated in the discussion with the authors, the clarifications and additional experiment does improve the paper a bit.", "extra_scores": {"experience_assessment": "I do not know much about this area.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper thoroughly.", "review_assessment:_checking_correctness_of_derivations_and_theory": "N/A"}, "extra_text": {}}, {"review_id": "B1lZO_ly9B", "reviewer": "AnonReviewer2", "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": "This paper draws inspiration from Physical world and considers an independent mechansim among recurrent modules. The authors apply the proposed RIM to several relatively simple tasks and show some advantages.\n\nIn general, I like the idea of making recurrent cells operate with nearly independent transition dynamics and interact only sparingly through the attention bottleneck. It is essentially to combine some environment prior into the model design. It makes senses to me that RIMs will work better in environments that objects and background are nearly independent and only interact with each other when collision happens. RIMs share similar to spirits with capsule networks, and its recurrent cells serve somewhat similar role to capsules. Such independent mechanism, selective activation and sparse communication is very inspiring and is indeed a potentially very useful way of modeling the physical world.\n\nFor the model itself, I appreciate its simplicity, but I also have some concerns.\n\n1) For the selective activation of RIMs, the number of activated RIMS is a hyperparameter and needs to be pre-defined. According to your experiments, I believe you need tune this hyperparameter a little bit in order to obtain the best performance. First of all, the design of a fixed number of activated RIMs does not seem to be reasonable and is also highly dependent on your task. I believe the framework will be more interesting if the model can determine this number automatically.\n\n2) I find it quite interesting that the top-down attention in selective RIM activation is corresponding to the states of these recurrent cells. I am wondering what if you do not select these top K activation and directly train it using the entire distribution of the soft attention output?\n\nFor the experiments, I think they all serve the purpose of showing the advanatges of RIMs quite well, except that thery are relatively easy task. However, it is still interesting to see that RIMs obtain significant gains over some baselines. Some of the details can be made more clear, such as loss function and evaluation metrics in every task. It is sometimes difficult to find what loss function you are using. I suggest the authors make the experiments more self-contained in the main paper, such that authors do not need frequently scroll down to the appendix and check the details.", "extra_scores": {"experience_assessment": "I do not know much about 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": "N/A"}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper has, at its core, a potential for constituting a valuable contribution. However, there was a shared belief among reviewers (that I also share) that the paper still has much room for improvement in terms of presentation and justification of the claims. I hope that the authors will be able to address the feedback they received to make this submission get where it should be.

Decision: Reject