Recurrent Independent Mechanisms
Authors: Anirudh Goyal
Alex Lamb
Jordan Hoffmann
Shagun Sodhani
Sergey Levine
Yoshua Bengio
Bernhard Schölkopf
Venue: ICLR 2021 Spotlight (2021)
Keywords: modular representations
better generalization
learning mechanisms

Abstract: We explore the hypothesis that learning modular structures which reflect the dynamics of the environment can lead to better generalization and robustness to changes that 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 compete with each other so they are updated only at time steps where they are most relevant.  We show that this leads to specialization amongst the RIMs, which in turn allows for remarkably improved generalization on tasks where some factors of variation differ systematically between training and evaluation.


Published as a conference paper at ICLR 2021
RECURRENT INDEPENDENT MECHANISMS
Anirudh Goyal', Alex Lamb', Jordan Hoffmann Shaguun Sodhani Sergey Levine*
Yoshua Bengio". "* --Benhaad Schölkopf
ABSTRACT
We explore the hypothesis that learning modular structures which reflect the dy- namics of the environment can lead to better generalization and robustness to changes that only affect 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 compete with each other so they are updated only at time steps where they are most relevant. We show that this leads to specialization amongst the RIMs, which in turn allows for remarkably improved generalization on tasks where some factors of variation differ systematically between training and evaluation.
INDEPENDENT MECHANISMS
Physical processes in the world often have modular structure which human cognition appears to exploit, 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 as statistical dependencies, they are ultimately due to dynamic processes governed by causal physical phenomena. These processes are mostly evolving independently 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 Earth as well as (weakly) to each other. Only occasionally will they strongly interact via collisions.
The notion of independent or autonomous mechanisms has been influential in the field of causal inference. A complex generative model, temporal or not, can be thought of as the composition of independent mechanisms or "causal" modules. In the causality community, this is often considered prerequisite for being able to perform localized interventions upon variables determined by such models (Pearl 2009). has been argued that the individual modules tend to remain robust or invariant even as other modules change, e.g., in the case of distribution shift (Schölkopf et al. 2012) Peters et al. 2017). This independence is not between the random variables being processed but between the description parametrization of the mechanisms: learning about one should not tell us anything about another, and adapting one should not require also adapting another. One may hypothesize that if brain is able solve multiple problems beyond single i.i.d. (independent and identically distributed) task, they may exploit the existence of this kind of structure by learning independent mechanisms that can flexibly be reused, composed and re-purposed.
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 forces and interventions. An agent needs not devote equal attention to all subsystems at all times: only those aspects that significantly interact need to be considered jointly when deciding or planning (Bengio) 2017). Such sparse interactions can reduce the difficulty of learning since few interactions need to be considered at time, reducing unnecessary interference when subsystem adapted. Models learned this way may better 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 al. 2017 The central question motivating our work is how gradient-bart deep learning approach can discover a representation of high-level variables which favour forming independent but
Mila, University of Montreal, ? Harvard University. MPT for Intelligent Systems, Tübingen, University of California, Berkeley," Eqqual advising," Equal Contribution. :ani rudhgoya 911 90 a com

---

Published as a conference paper at ICLR 2021
Debout
Sparee Communicttion
Delonit dynomles
Spanse Coammmlration
does
Compeing RDS
itsaed competiinn baoe cn top done ato0630
Sapart vitestion
h
hes
Botton -isal acxezazias
Iaput
Ingut
Actix
Inaetee RIM Vilan
Vauaal ingrat
Figure Illustration of Recurrent Independent Mechanisms (RIMs). single step under the proposed model occurs in four stages (left figure shows two 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 is used select which RIMs to activate (right figure) based 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 con-activated RIMs remain unchanged. In the fourth stage, the RIMs sparsely communicate information between themselves, also using key-value attention.
sparsely interacting recurrent mechanisms in order to benefit from the modularity and independent mechanisms assumption.
Why do Models Succeed or Fail in Capturing Independent Mechanisms? While universal ap- proximation theorems apply in the limit of large i.i.d. data sets, we are interested in the question of whether models can learn independent mechanisms from finite data in possibly changing envi- ronments, and how to implement suitable inductive biases. As the simplest case, we can consider training an RNN consisting of Ao completely independent mechanisms which operate on distinct time steps. How difficult would it be for an RNN (whether vanilla or LSTM or GRU) to correctly model that the true distribution has completely independent processes? For the hidden states truly compartmentalize these different processes, fraction k- of the connections would need to be set to exactly zero weight. This fraction approaches 100% as k approaches infinity. When sample complexity or out-of-distributtin generalization matter, we argue that having an inductive bias which favors this form of modularity and dynamic recombination could be greatly advantageous, compared to static fully connected monolithic architectures.
Assumptions on the joint distribution of high level variables. The central question motivating our work how gradient-bant basded deep learning approach can learn a representation of high level variables which favour learning independent but sparsely interacting recurrent mechanisms in order to benefit from such modularity assumption. The assumption about the joint distribution between the high-level variables is different from the assumption commonly found in many papers on disentangling factors of variation (Higgins et al. (2016; BBrresss et al.] 2018 Chen et al.) 2018). where the high level variables are assumed to be marginally independent of each other. We believe that these variables, (often named with words in language), have highly structured dependencies supporting independent mechanisms assumption.
RIMS WITH SPARSE INTERACTIONS
Our approach to modelling a dynamical system interest divides the overall model into small subsystems (or modules), each of which is recurrent in order to be able to capture the dynamics in the observed sequences. We refer to these subsystems as Recurrent Independent Mechanisms (RIMs), where each RIM has distinct functions that learned automatically from date] We refer to RIM k at time step t as having vector-valued state hy where 1,..., Each RIM has parameters 0g. which are shared across all time steps.
At a high level (see Fig. Dp. we want each RIM to have its own independent dynamics operating by default, and occasionally to interact with other relevant RIMs and selected elements of the encoded input. The total number of parameters can be kept small since RIMs can specialize on simple sub-
Note that we are overloading the term mechanism, using it both for the mechanisms that make up the world's dynamics well for the computational modules that we learn to model those mechanisms. The distinction should be clear from context.
2

---

Published as a conference paper at ICLR 2021
problems, and operate on few key/value variables at time selected using an attention mechanism, as suggested by the inductive bias from Bengio (2017). This specialization and modularization not only has computational and statistical advantages (Baum & Haussler 1989,, but also prevents individual RIMs from dominating the computation and thus facilitates factorizing the computation into easy to recombine but simpler elements. We expect this to lead to more robust systems than training one big homogeneous system (Schmidhuber, 2018)) Moreover, modularity and the independent mechanisms hypothesis (Peters et al. 2017; Bengio al. 2019) also has the desirable implication that RIM should maintain its own independent functionality even as other RIMs are changed. A more detailed account of the desiderata for the model is given in Appendix [
2.1 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, using an attention mechanism which selects and then activates only subset of the RIMs for each time step. As argued by Parascandolo et al.) (2018), this tends to produce independence among learned mechanisms, provided the training data has been generated by set of independent physical mechanisms. In contrast to Parascandolo et al. (2018), we use an attention mechanism for this purpose. The introduction of content- based soft- soft-attention mechanisms (Bahdanau al. 2014) has opened the door to neural networks which operate on sets of typed interchangeate objects. This idea has been remarkably successful and widely applied to most recent cuunsformer-ttlle multi-head dot product self attention models (Vaswani et al. 2017, Santoro et al., 2018), achieving new ofte-of-the-ass results in many tasks. Soft-attention uses the product of a query (or read key) represented as a matrix Q of dimensionality N, x d, with the dimension of each key, with a set of No objects each associated with a key (or write-key) as row in 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 of matrix V). The result is
QK?
Attention(Q. (Q, K, V) = softmax
V.
Vd
where the softmax is applied to each row of its argument matrix, yielding a set of convex weights. As result, one obtains convex combination of the values in the rows of V. If the attention is focused on one element for a 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 d dimensions the key can be split into heads which then have their own attention matrix and write values computed separately.
When the inputs and outputs of each RIM are set of objects or entities (each associated with key and value vector), the RIM processing becomes generic prroces-rooessing machine which can operate on subsymbolic "variables" in sense analogous variables in a programming language: as interchangeable arguments of functions, albeit with distributed representation both for they name or type and for their value. Because each object has a key embedding (which one can understand both 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 typed argument of the function computed by the RIM. When the key of an object matches the query of head k. it can be used as the k-th input vector argument 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 be used as input for each of the arguments of the RIM dynamics, with a different set of query embeddings for each RIM head. These inputs can come from the external input or 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 if its "type" matches with what this RIM expects (specified with the corresponding query).
2.2 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. RIMs are triggered as a result of interaction between the current state of the RIM and imput information coming from the environment. At each step, we select the top-k (out of kr) RIMs in terms of their
3

---

Published as a conference paper at ICLR 2021
attention score for the real input. Intuitively, the RIMs must compete each step to read from the input, and only the RIMs that win this competition will be able to 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. This differs from the mechanics of Vaswani et al. (017);; Santoro al. (2018), with the modification that the parameters of the attention mechanism itself are separate for each RIM rather than produced on the input side as in Transformers. The input attention for particular RIM described follows.
The input T2 at time / is seen as a set of elements, structured as rows of matrix. We first concatenate a row full of zeros, to obtain
X 0 ( Zt--
As before, linear transformations are used to construct keys =K= XW", one per input element and for the null element), values (V XW", again one per element), and queries (Q h2 W2 one per RIM attention head). W is simple matrix mapping from an input element to the corresponding value vector for the weighted attention and W is similarly 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. + refers to the row-level concatenation operator. The attention thus is
(XW XW", where o((m) (W2, W W²).
A (n) softmax
Vde
Based on the softmax values in (2. we select the top Ka RIMs (out of the total K 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 enables individual RIMs to attend only to the part of the input that is relevant for that particular RIM, thus enabling selective attention based use multiheaded attention, and multi-headed attention doesn't It 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.
the state of the RIMs, this
top-down attention process (see. Fig In practice,
For spatially structure input: All datasets we considered are temporal, yet there is a distinction between whether the input on each time step is highly structured (such as video) or not (such as language modeling, where each step has 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. The input F2 at time can be seen as an output of the encoder parameterized by neural network (for ex. CNN in case of visual observations) i.e., X CN N (zz). As before, linear transformations are used to construct position-speeifie input keys (K = XW*), values (V XW"), and RIM specific queries (Q attention thus is
hy W. one per RIM attention head). The
A softmax
XW",
where
e((h) = (W2, W W").
(3)
Vdl
In order for different RIMs specialize on different spatial regions, we competition among the different RIMs. The contents of the attended positions are combined yielding RIM-specific input. As before based the softmax values in (3). select the top k, RIMs (out of the total KT RIMs) to be activated for each spatial position, which have the highest attention on that spatial position.
position-specific
2.3 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. The hidden states of RIMs which are not activated (we refer to the activated set S₂) remain unchanged, acting like untouched memory elements, i.e., h441 h& Vk & S, Note that the gradient still flows through a RIM on a step where it is not activated. For the RIMs that are activated, we RIM independent transition dynamics. The form of this is
4

---

Published as a conference paper at ICLR 2021
somewhat flexible, but opted either GRU (Chung al. 2015)
LSTM (Hochreiter
& Schmidhuber 1997). We generically refer to these independent transition dynamics as Dg. and 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 active RIMs are updated by
h = (a = LSTM A/ ((0) ( Ak E St
function of the attention mechanism A ((r)) applied on the current input, described in the previous sub-section.
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 or not). The intuition behind this is that non-activated RIMs are not related to the current input, so their value needs not change. However they may still store contextual information relevant for activated RIMs later on. For this communication between RIMs, we use a residual connection in Santoro al. 2018) to prevent vanishing exploding gradients over long sequences. Using parameters %% (W2, Wg, WT), we employ
Qt,x = Whe Vk E St
Ki, = Wheee Vk:
V1,+ = WiT,,,,
(QeA(tii))) Vde
h = softmax
Ve, Ve,+ he E S,
We note that we can also consider sparsity in the communication attention such that a particular RIM only attends to sparse sub-set of other RIMs, and this sparsity is orthogonal to the kind used in input attention. In order to make the communication attention sparse, we can still use the same top-k attention.
Number of Parameters. RIMs can be used drop-in replacement for an LSTM/GRU layer. There a subtlety that must be considered for successful integration. If the total size of the hidden state is kept the same, integrating RIMs drastically reduces the total number of recurrent parameters in the model (because of having a block-sparse structure). RIMs also adds new parameters to the model through the addition of the attention mechanisms although these rather in small number.
Multiple Heads: Analogously to Vaswani et al. (2017); Santoro et al. (2018). we use multiple heads both for communication between RIMs well as input attention (as in Sec 2.2) by producing different sets of queries, keys, and values to compute a linear transformation for each head (different heads have different parameters), and then applying the attention operator for each head separately in order to select conditioning inputs for the RIMs.
3 RELATED WORK
Neural Turing Machine (NTM) and Relational Memory Core (RMC): the NTM (Graves et al. 2014a) updates independent memory cells using an attention mechanism to perform targeted read and write operations. RIMs share a key idea with NTMs: that input information should only impact 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 et al.) (2018) 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 as a set of separate recurrent models. IndRNN, each recurrent unit has completely independent dynamics, whereas EntNet uses 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 attention mechanism.
Modularity and Neural Networks: A network can be composed of several modules, each meant to perform distinct function, and hence can be seen as combination of experts (Jacobs al.
5

---

Published as a conference paper at ICLR 2021
1991 Bottou & Gallinari 1991 Ronco et al. 1997 Reed & De Freitas
2015 Andreas al. 2016
Parascandolo et al. 20188; Rosenbaum et al. 2017 Fernando et al. 2017 Shazeer et al. 2017 Kirsch et al., 2018 Rosenbaum et al., 2019) routing information through gated activation of modules. These works generally assume that only single expert is active at a 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 al. 2016 Neil et al. 2016 Jernite al. 2016 Krueger et al. 2016) where parts of the RNN's hidden state are kept dormant at times. The major differences to our architecture are that (a) we modularize the dynamies 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.
4 EXPERIMENTS
The main goal of our experiments is to show that the use of RIMs improves generalization across changing environments and/or modular tasks, and 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 to range diverse tasks, focusing on tasks that involve changing environment. We organize our results by the capabilities they illustrate: we address generalization based temporal patterns, based on objects, and finally consider settings where both 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 and generally independent temporal patterns, they able to specialize so that different RIMs activated on different patterns, RIMs generalize well when we modify a subset of the patterns (especially those unrelated the class label) while most recurrent models fail generalize well to these variations.
Figure 2: Visualizing Activation Patterns. For the copying task, one can see that the RIM activation pattern distinct during the dormant part of the sequence in the middle (activated RIMs black, non-activated white). X-axis=time, Y-axis=RIMs activation bit.
attention to the task of receiving short sequence of characters, then receiving blank inputs for a large number steps, and then being asked to reproduce the original sequence. 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 no input is provided. As an example of out-of-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, we find that this result is robust to the number of RIMs used as well as to the number of RIMs activated per-step. Our ablation results (Appendix (.T) show that all major components of the RIMs model are necessary to achieve this generalization. This is evidence that RIMs can specialize over distinct patterns in the data and improve generalization to settings where these patterns change.
Copying Task: First
Sequential MNIST Resolution Task: RIMs are motivated by the hypothesis that generalization performance can benefit from modules which only activate on relevant parts of the sequence. For further evidence that RIMs can achieve this out-of distrbbbbii,,,nnn we consider the task of classifying MNIST digits as sequences of pixels (Krueger al. 2016) 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. RIMs should generalize better to higher resolutions by keeping RIMs dormant which store pixel information over empty regions of the image.

---

Published as a conference paper at ICLR 2021
Copying
Train(50)
Test(200)
Sequential MNIST
16 16
19 19
24 x 24
kT
Rx
Duze
CE
CE
ky
AA
haze
Accuracy
Accuracy
Accuracy
6
4
600
0.00
0.00
600
85.5
56.2
30.9
600
0.00
0.00
600
88.3
43.1
22.1
RIMs
6
600
0.00
0.00
RIMS
6
600
90.0
73.4
38.1
5
2
500
0.00
0.00
300
86.8
42.3
25.2
LSTM
300
0.00
4.32
600
84.5
52.2
21.9
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.13
DNC
87.2
44.1
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 generalization (right). While all of the methods are able to learn to copy for the length seen during training. the RIMs model generalizes to sequences longer than those seen during training whereas the LSTM, RMC, and NTM degrade much more. On sequential MNIST, both the proposed and the Baseline models were trained on 14x14 resolution but evaluated at different resolutions (averaged over 3 trials).
Results: Table shows the result of the proposed model on the Sequential MNIST Resolution Task. If the train and test sequence lengths agree, both models achieve comparable test set performance. However, RIMs model is relatively robust to changing the sequence length (by changing the image resolution), whereas the LSTM performance degraded more severely. This can be seen more involved analogue of the copying task, as MNIST digits contain large empty regions. It is essential that the model be able store information and pass gradients through these regions. The RIMs outperform strong baselines such as Transformers, EntNet, RMC, and (DNC) (Graves et al. 2016).
4.2 RIMS LEARN TO SPECIALIZE OVER OBJECTS AND GENERALIZE BETWEEN THEM
We have shown that RIMs can specialize over temporal patterns. We now turn our attention to assaying whether RIMs can specialize to objects, and show improved generalization to cases where we add or remove objects at test time.
We consider a synthetic "bouncing balls" task in which multiple balls (of different masses and sizes) move using basic Newtonian physics (Van Steenkiste al. 2018). What makes this task particularly suited to 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 a binary cross entropy loss when predicting the next frame. We consider LSTMs as baselines. We then produce rollouts, finding that RIMs are better able to predict future motion (Figure 3).
Bouncing Balls Environment:
Triined on & Objpet
Trined on f&8 Ollyer
Testo Ohizota
Iestol Olijada
Saacr Cros baccps
LSTX 02
LSTH
1001
Cootol teath
Ralbot
Figure 3: Handling Novel Out-of-Diitriluttio Variations. We study the performance RIMs compared to an LSTM baseline (4 left plots). The first 15 frames of ground truth (yellow.orange) are fed in and then the system is rolled out for the next 35 time steps (blue.purple). During the rollout phase, RIMs perform better than the LSTMs in accurately predicting the dynamics of the balls reflected by the lower Cross Entropy (CE) [blue for RIMs, purple for L.S.Ms]. Notice the substantially better -ii-ttiiiitiittt generalization of RIMs when testing on a number of objects different from the one seen during training. (2nd to 4th plot). We also show (right plot) improved
-idt-ii-iitiitti generalization (F1 score) as compared to LSTM and RMC (Santoro al. (2018) on another partial observation video prediction task. X-axis = number of balls. For these experiments, the RIMs and baselines get input image each time slep (see Appendix (D3)] figure. [ for magnified image as well as more details). Here. TTO refers to the time travelling oracle upper bound baseline. that does not model the dynamics, and has access to true dynamics.
We take this further by evaluating RIMs on environments where the test setup is different from the training setup. First consider training with 4 balls and evaluating on environment with 6-8
7

---

Published as a conference paper at ICLR 2021
balls. Second, we consider training with 6-8 balls and evaluating with just 4 balls. Robustness in these settings requires a degree of invariance w.r.t. the number of balls.
In addition, we consider a task where we train on 4 balls and then evaluate sequences where part the visual space is occluded by a "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 (3) 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 related to capacity.
Environment with Novel Distractors: We next consider an object-pickinn reinforcement learning task from BabyAl (Chevalier-ooisvertee et al. 2018) in which an agent must retrieve specific object in the presence of distractors. We use a partially observed formulation of the task, where the agent small number of squares ahead of it. These tasks difficult to solve (Chevalier- cct al.) 2018) with standard RL algorithms, 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.
only
Figure shows that RIMs outperform LSTMs this task (details in ap- pendix). When evaluating with known distractors, the RIM model achieves perfect performance while the LSTM struggles. When evaluating an en- vironment with novel unseen distrac- tors the RIM doesn't it achieve perfect performance but strongly outperforms the LSTM. An LSTM with a single memory flow may struggle to keep the distracting elements separate from elements which are necessary for the task, while the RIMs model uses attention to control which RIMs receive information at each step as well what information they receive (as function of their hidden state). This "top-down" attention results in diminished representation of the distractor, not only enhancing the target visual information, but also suppressing irrelevant information.
Figure 4: Robustness to Novel Distractors: Left: performance of the proposed method (blue) compared to an LSTM bascline (red) in solving the object picking task in the presence of distractors. Right: performance of proposed method and the baseline when novel distractors are added.
4.3 RIMS IMPROVE GENERALIZATION IN COMPLEX ENVIRONMENTS
We have investigated how RIMs use specialization improve generalization to changing important factors of variation in the data. While these improvements have often been striking, it raises question: what factors of variation should be changed between training and evaluation? One setting where factors of variation change naturally is in reinforcement learning, as the data received from an enviromment changes as the agent learns and improves. We conjecture that when applied to reinforcement learning, agent using RIMs may be able to learn faster its specialization leads to improved generalization to previously unseen aspects of the environment. To investi gate this we use an RL agent trained using Proximal Policy Optimization (PPO) (Schulman et al.) 2017) with recurrent network producing the policy. We employ an LSTM as 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 5
There is also an 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, executing actions which are at least partially prepared. This suggests agents should process sensory information in 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 always present, but which is relevant depends the player's 's position). We
8

---

well as Published as a conference paper at ICLR 2021
088
025
616
4075
-18
Figure 5: 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 comparisons external baselimes are Appendix
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 DISCUSSION AND ABLATIONS
Sparse Activation is necessary, but works for wide range of hyperarrameters: On the copying task, we tried a wide variety of sparsity levels for different numbers of RIMs, and found that using sparsity level between 30% to 70% performed optimally, suggesting that the sparsity hyperparameter is fairly flexible (refer to Table 4 5 in appendix). On Atari we found that using l = 5 slightly improved over results compared with kΛ 4, but both had similar performance across the vast majority of games.
Input-attention necessary: We study the scenario where we remove the input attention process (i.e the top-down competition between different RIMs) but still allow the RIMs ot communicate with attention. We found that this degraded results substantially on Atari but still outperformed the LSTM baseline. See (Figure 20) in appendix for more details.
Communication between RIMs improves performance: For copying and sequential MNIST, we performed ablation where remove the communication between RIMs and varied the number of RIMs and the number of activated RIMs (Refer to Table i in appendix.). We found that the communication between RIMs is essential for good performance.
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.e., that all processes interact. This can lead to poor generalization and lack of robustness to changing task distributions. We have proposed a new architecture, Recurrent Independent Mechanisms (RIMs), in which we learn multiple recurrent modules that independent by default, but interact sparingly. For the purposes of this paper, we note that the notion of RIMs is not limited the particular architecture employed here. The latter is used as vehicle to assay and validate our overall hypothesis (cf. Appendix A). but better architectures for the RIMs model can likely be found.
ACKNOWLEDGEMENTS
The authors acknowledge the important role played by their colleagues at Mila throughout the duration of this work. The authors would like to thank Mike Mozer for brainstorming sessions. AG would like to thank Matthew Botvinick, Charles Blundell, Greg Wayne for their useful feedback. which has improved the work tremendously. The authors are grateful to Nasim Rahaman for letting us use video crop results. The authors would also like to thank Rosemary Nan Ke, Stefan Bauer, Jonathan Binas, Min Lin, Disha Srivastava, Ali Farshchian, Sarthak Mittal, Owen Xu, for useful discussions. The authors would also like to thank Shahab Bakhtiari, Kris Sankaran, Felix E. Leeb for proof reading the paper. This project was also known by the name "Blocks" internally at Mila.

---

Published as a conference paper at ICLR 2021
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: (rXiv::40.0077. 2014.
Peter W Battaglia, Jessica B Hamrick, Vietor Bapst, Alvaro Ganzaez-Gonzalez, Vinicius Zambaldi, Matcusz Malinowski, Andrea Tacchetti, David Raposo, Adam Santoro, Ryan Faulkner, et al. Relational inductive biases, deep learning, and graph networks. arXiv preprint arXiv: 18Xiv 188060026111 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ébastien Lachapelle, Olexa Bilaniuk, Anirudh Goyal, and Christopher Pal. A meta-transfer objective for learning to disentangle causal mechanisms. arXiv:1901.10912 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.
Christopher P Burgess, Irina Higgins, Arka Pal, Loic Matthey, Nick Watters, Guillaume Des- jardins, and Alexander Lerchner. Understanding disentangling in beta-vae. arXiv preprint arXiv:1804. 03599, 2018.
Ricky TQ Chen, Xuechen Li, Roger B Grosse, and David K Duvenaud. Isolating sources of disentanglement in variational autoencoders. In Advances in Neural Information Processing Systems, pp. 2610-2620, 2018.
Maxime Chevalier- and Lucas Willems. Minimalistic gridworld environment for openai gym. ht.tps://github..com cww.://mmccommmmmmmmmbbb.o gymmminigridddd 2018.
Maxime Chevalier- Dzmitry Bahdanau, Salem Lahlou, Lucas Willems, Chitwan Saharia, Thien Huu Nguyen, and Yoshua Bengio. Babyai: First steps towards grounded language learning with human in the loop. arXiv preprint arXiv:1810.0822 2018.
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 arXiv: (0t.::/60007000 2016.
Paul Cisck 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 arX1v.1802.07687, 2018.
Robert Desimone and Jody Duncan. Neural mechanisms of selective visual attention. Annual Review of Neuroscience, 18:193-222. 1995.
10

---

Published as a conference paper at ICLR 2021
A. Dickinson. Actions and habits: the development of behavioural autonomy. Philosophical Transactions of the Royal Society B: Biological Sciences, 308(1135),6778, 1985. ISSN 0080- 4622. doi: 10. 1tt98//sstb198.00000
Salah Hihi and Yoshua Bengio. Hierarchical recurrent neural networks long-term dependencies. In Advances in neural information processing systems, pp. 493-499, 1996.
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: 7701.08744 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- Volumee 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 2019a.
bottleneck. arXiv preprint
Anirudh Goyal, Shagun Sodhani, Jonathan Binas, Xue Bin Peng, Sergey Levine, and Yoshua Bengio. Reinforcement learning with competitive ensembles of informaation-constranned primitives. arXiv preprint arXiv:1906 (00,:/69900066677,, 2019b.
Alex Graves, Greg Wayne, and Ivo Danihelka. Neural turing machines. arXiv:1410.540) 2014a.
arXiv preprint
Alex Graves, Greg Wayne, and Ivo Danihelka. Neural turing machines. CoRR, abs/ 410.54010, 2014b. URL
abs/ 1410.5401.
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 arXiv: 1e:::81155055 1.04551, 2018.
Mikael Henaff, Jason Weston, Arthur Szlam, Antoine Bordes, and Yann LeCun. Tracking the world state with recurrent entity networks. arXiv preprint 112.0966, 2016.
Irina Higgins, Loie Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick, Shakir Mohamed, and Alexander Lerchner. beta-vae: Learning basic visual concepts with constrained variational framework. 2016.
Geoffrey E Hinton, Sara Sabour, and Nicholas Frosst. Matrix capsules with routing. 2018.
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 9()): 1735-1780, 1997.
Robert A Jacobs, Michael Jordan, Steven J Nowlan, Geoffrey E Hinton, al. Adaptive mixtures of local experts. Neural computation, 3(1):79-87) 1991.
Yacine Jemite, Edouard Grave, Armand Joulin, and Tomas Mikolov. Variable computation in recurrent neural networks. arXiv preprint 1rXiv:161.066188 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.
11

---

Published as a conference paper at ICLR 2021
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
Thomas Kipf, Ethan Fetaya, Kuan-Chich Wang, Max Welling, and Richard Zemel. Neural relational inference for interacting systems. arXiv preprint
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):89--90, 2018.
Ilya Kostrikov. implemchimmpe of reinforcement remtore learning algorithms. https://qithut
ccwp://kettilloo//rrorra-------- -atpoooactt---aaia 2018.
Jan Koutnik, Klaus Greff, Faustino Gomez, and Juergen Schmidhuber. A clockwork mn. arXiv
preprint arXiv: 1402.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:1606 01305, 2016.
Shuai Li, Wanqing Li, Chris Cook, Ce Zhu, and Yanbo Gao. Independently recurrent neural network (indrnn): Building longer and deeper mnn In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 5457-5466, 2018.
Daniel Neil, Michael Pfeiffer, and Shih-Chii Liu. Phased Istm: Accelerating recurrent network training for long or event-based sequences. In Advances in neural information processing systems, pp. 3882-3890, 2016.
Giambattista Parascandolo, Niki Kilbertus, Mateo Lojas-Carulla, and Bernhard Schölkopf. Learning independent causal mechanisms. Proceedings of the 35th International Conference on Machine Learning (ICML), pp.
2018. URL http: httt:://proceeeinneee ml www.gnwww....s.aasas.........e.e www.prressvvv. v
parascandolo 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 arXiv: 1702.00002000 2017.
Scott Reed and Nando De Freitas. Neural programmee--nterrreeess. arXiv preprint arXiv:1511.06279, 2015.
Eric Ronco, Henrik Gollee, and Peter Gawthrop. Modular neural networks and seff-decompositionn 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: 171100133, 2017.
Clemens Rosenbaum, Ignacio Cases, Matthew Riemer, and Tim Klinger. Routing networks and the challenges modular and compositional computation. arXiv preprint arXiv: /904. 12774, 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 17X1::1606604661,, 2016.
Sara Sabour, Nicholas Frosst, and Geoffrey E Hinton. Dynamic routing between capsules. In Advances in neural information processing systems, pp. 3856-3866, 2017.
12

---

Published as a conference paper at ICLR 2021
Adam Santoro, David Raposo, David G Barrett, Mateusz Malinowski, Razvan Pascanu, Peter Battaglia, and Timothy Lillicrap. 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//806.00 806.0182,, 2018. UtLhttpp//aaxxiv. wwtp://gaaab////ooeeeeeee/eg/ 01822
Franco Scarselli, Marco Gori, Ah Chung Tsoi, Markus Hagenbuchner, and Gabriele Monfardini. The graph neural network model. IEEE Transactions on Neural Networks, 20( (0(1):61-00 2008.
Jürgen Schmidhuber. One big net for everything. arXiv preprint FrXXiv:802.088860 2018.
Bernhard Schölkopf, Dominik Janzing, Jonas Peters, Eleni Sgouritsa, Kun Zhang, and Joris Mooij. On causal and anticausal learning. Langford and J. Pineau (eds.), Pruceedings of the 29th International Conference 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: 2017.
Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated nixture-of-experts layer. arXiv preprint arXiv:1701.cm 06538, 2017.
Herbert A Simon. The architecture of complexity. In Facets of systems science, pp. 4p..557-45 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: (eee: 2019.
Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdinov. Dropout: simple way prevent neural networks from overfitting. The journal of machine learning research, ((():1199-1195, 2014.
Andrea Tacchetti, H 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:1809 11044, 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. 4496 4506, 2017.
Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physics engine for model-based control. In 2012 JEEE/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 preprins 1xXi:18002.1353 2018.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N 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. Naturwissenschaffen 37(20):464 Jan 1950. doi: 10. O0.100//FF0622503.
Nicholas Watters, Daniel Zoran, Theophane Weber, Peter Battaglia, Razvan Pascanu, and Andrea Tacchetti. Visual interaction networks: Learning a physics simulator from video. Advances in neural information processing systems, pp. 4539-4544 2017.
Ronald Williams and David Zipser. A learning algorithm for continually running fully recurrent neural networks. Neural computation, 1(2):270-280, 1989.
13

---

Published as a conference paper at ICLR 2021
Appendix
Table of Contents
Desiderata for Recurrent Independent Mechanisms
14
Extended Related Work
15
Model Setup Details Hyperparameters
C.I
RIMs Implementation
16
Detailed Model Hyperparameters
................................
17
Other Architectraa Changes that Explored
17
D
Experiment Details
18
D.1
Effect of Varying Number and Active RIMs on Copying Task
18
D.2
Effect of Varying Number and Active RIMs on Adding Task
19
D.3
Sequential MNIST Resolution Task
19
D.4
Bouncing Ball Environment
20
D.5
Video Prediction from Crops
21
D.6
Environment Novel Distractors
23
MiniGrid Environments OpenAl Gym
24
Atari
25
E
Additional Experiments
27
E.1
Imitation Learning: Robustness Noise in State Distribution
27
E.2
Transfer on Atari
27
E.3
Bouncing MNIST: Dropping individual RIMs
31
F Natural Language Processing: Language Modeling, Machine Translation, and Trans-
fer Learning
31
G Pseudocode for RIMs Algorithm
37
A DESIDERATA FOR RECURRENT INDEPENDENT MECHANISMS
We have laid out a case for building models composed of modules which by default operate independently and interact in limited manner. Accordingly, our approach to modelling the dynamics of the world starts by dividing the overall model into small subsystems (or modules). referred to Recurrent Independent Mechanisms (RIMs), with distinct functions learned automatically from data. Our model encourages sparse interaction, i... 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 to capture modular dynamics with sparse interactions.
Competitive Mechanisms: Inspired by the observations in main paper. we propose that RIMs utilize competition to allocate representational and computational resources. As argued by (Parascandolo et al.) 2018). this tends to produce independence among learned mechanisms if the training data has heen generated by independent physical mechanisms.
The points mentioned in Section in principle pertain synthetic and natural intelli- gent systems alike. Hence, is not surprising that they also appear neuroscience, For instance, suppose are looking for 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 allocate resources. Competitive
Top Down Attention:
14

---

Published as a conference paper at ICLR 2021
interactions among multiple objects occur automatically and operate in parallel across the visual field. Second, the principle of selectivity amounts to the idea that a perceiver has the ability filter out unwanted information and selectively process the rest of the information. Third, top-down bias originating from higher brain arcas 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 an object process of interest being pre-activated and thus gaining advantage during the competition of brain mechanisms.
Sparse Information Flow: Each RIM's dynamics should only be affected by RIMs which are deemed relevant. The fundamental challenge is centered around establishing sensible communication between modules. In the presence of noisy distracting information, large subset of RIMs should stay dormant, and not be affected by the noise. This way, training an ensemble of these RIMs can be more robust out-of-distribution or distractor observations than training one big homogeneous neural network (Schmidhuber 2018).
Modular Computation Flow and Modular Parameterization: Each RIM should have its own dy- namics operating by default, in the absence of interaction with other RIMs. The total number of parameters (i.e. weights) can be reduced since the RIMs can specialize on simple sub-problems, similar (Parascandolo
2018). This can speed up computation and improve the gencralisation 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. We refer associated to individuals RIMs only. This has the desirable property that RIM should maintain its own independent functionality even as other RIMs are changed (due to its behavior being determined by its own self- contained parameters).
parameterization modular if most parameters
B EXTENDED RELATED WORK
Table 2: A concise comparison of recurrent models with modular memory.
Modular Mcmory
Sparse
Modular Computation Flow
Modular Parameterization
Method Property
Information Flow
LSTM RNN
X
x
Relational RNN (Santoro al.
2018
/
x
/
X
NTM (Graves et al. 2014b
/
/
X
X
SAB et al. 2018
x
/
X
X
IndRNN et al. 2018
/
x
X
RIMs
/
The present section provides further details on related work, thus extending Section 3
Neural Turing Machine (NTM). The NTM (Graves et al. 2014a) has a Turing machine inspired memory with sequence of independent memory cells, and uses attention mechanism 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 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 al. 2018) is based on the idea of using multi-head attention mcchanism share information between multiple parts of memory. It is related to our idca but a key difference is that we encourage the RIMs to remain scparate as much as possible, whereas (Santoro et al. 2018) allows information between the parts to flow on cach step (in effect making the part distribution only relevant to 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 (e et al. 2018 explores RNNs with self- attention across time steps as well as variants where the attention is sparse in the forward pass and where the gradient is sparse in the backward pass. It shares the motivation of using sparse attention to keep different pieces of information separated, but differs from the RTMs model in that considers separation between time steps rather than separation hetween RIMs.
Independently Recurrent Neural Network (IndRNN). The IndRNN (L al. 2018 replaces the full transi- tion matrix in vanilla RNN (between time steps) to 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 RNN to learn long-term dependencies without vanishing or
15

---

Published as a conference paper at ICLR 2021
exploding gradients. Analysis of the gradients shows that having smaller recurrent transition matrices mitigates the vanishing and exploding gradient issue. This may provide further explanation for why RIMs perform well on long sequences.
2017 This is based on the inductive bias of a sparse graphical model describing the interactions between high-level variables, each factor involving few variables and capaturing an independent mechanism, and using attention to select only subset of high-level variables to interact together at any particular time during inference. This motivates in RIMs the use of modules (corresponding factors) which dynamically select which variables (instances) they apply to, with modules activated sparsely (only those factors which need to interact) and communicating sparsely (only a few variables, i... heads, involved for each). Our paper thus helps to validate the inductive bias of the consciousness prior.
Consciousness Prior (Bengio
Recurrent Entity Networks: EnTNet (Henaff al. (016) can be viewed as 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.RIIM, though each RIM has different parameters), but we allow the RIMs to communicate with cach other sparingly using an attention mcchanism.
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 is assumed to be present, which is facilitated by an interactive inference (routing) algorithm.
Relational Graph Based Methods: Recent graph-based architectures have studied combinatorial generalization in the context of modeling dynamical systems like physics simulation, multi-object scenes, and motion-capture
data, and multiagent systems (Scarselli et al. 2008; [roontein et al.] 2017 Watters el al. [0117] Raposo el al.
2017 Santoro et al, 2017 Gilmer ct al. [201 7 Van Steenkiste et al. [2018] Kipt ct al. [2018] Battaglia et al.
2018 lacehetti et al. 2018). One also view our proposed model as relational graph neural network, where
nodes are parameterized as individual RIM and edges are parameterized by the attention mechanism. Though,
is important to emphasize that semantics of the nodes is learned and that the topology of the graph induced in
the proposed model is dynamic, while in most graph neural networks the nodes have meaning directly tied to the inputs and targets and the topology is fixed and given.
Our work is also related to work in behavioural research that deals with modes of decision making (Dickinson) 1985 Botvinick Braver 2015 Kool Botvinick 2018]: an automatic system that relies on hahits and a controlled system that uses some prvvleged information for decision-making. The proposed model also has two modes of input processing: activated RIMs rely on external sensory information, and hence seem analogous to the controlled system, while inactive RIMs may more loosely correspond to computation taking place outside of conscious processing and corresponding to the automatic hahit system. There is RL work aiming learn defaul pulicies, which have shown to improve transfer and generalization in multi-task RL (Teh et al.] 2017 Goyal el al. 2019a). RIMs are different in the sense that we are not trying to learn default policies which affect the environment, instead we want to learn mechanisms, which are more about analyzing understanding the environment. State-dependent activation of different primitive policies was also studied by Goyal al. (2019b), and the authors showed that they can learn different primitives, but they also consider that only single primitive can be active at particular time-step. Also, note that primitive policics try to affect the environment, whereas mechanisms try to process or understand the environment.
Default Behaviour:
Dropout (Srivastava et al. 2014) is regularizer which sets values of individual hidden units to zero randomly, and thus it is distinct trom RIMs which contains ensemble of interacting mechanisms (cach of themselves containing many units) compete to activate and share information using bottleneck of attention.
Dropout:
C MODEL SETUP DETAILS AND HYPERPARAMETERS
C.I RIMS IMPLEMENTATION
The RIMs model consists of three main components: the input attention, the process for selecting activated RIMs. and the communication between RIMs. The input attention closely follows the attention mechanism of (Santoro) et al. 2018) but with significant modification: that all of the weights within the attention mechanism are separate per-block. Thus we remove the normal linear layers and replace them with batch matrix multiplication over the RIMs (as each block has its own weight matrix). Note that the read-key (or query) is a function of the hidden state of each RIM.
For selecting activated RIMs, we compute the top-k attention weight over the RIMs (based on the query). We then select the activated RIMs, using mask which zerves out the outputs from inactive RIMs. We compute the independent dynamics over all RIMs by using separate LSTTM for each RIM. Following this, we compute
16

---

Published as a conference paper at ICLR 2021
the communication between RIMs multihead attention (Santoro al. 2018). with the earlier-diseussed modification having separate weight parameters for each block, and also that we added a skip-connection around the attention mechanism. This attention mechanism used heads and in general used a key size and value size of 32. We computed the updates for all RIMs but used the activated-blec mask to selectively update only the activated subset of the RIMs.
The use of RIMs introduces two additional hyperparameters over LSTM/GRU: the number of RIMs and the number of activated RIMs per step. We also observed that having too few activated RIMs tends hurt optimization and having too many activated RIMs attenuates the improvements to generalization. For the future it would be interesting to explore dynamic ways of controlling how many RIMs to activate.
Multiple Heads: Analogously to Vaswani et al. (2017 Santoro al. 2018 we multiple heads both for communication between RIMs well input attention in See by producing different of queries, keys, and values compute a linear transformation each head (different heads have different parameters), and then applying the attention operator for each head separately in order to select conditioning inputs for the
RIMs.
C.2 DETAILED MODEL HYPERPARAMETERS
Table 3 lists the different hyperparameters.
Table Hyperparameters
Parameter
Value
Optimizer
Adam(Kingma & Ba
2014
learning rate
10
batch size
64
Input keys
64
Input Values
Size of individual RIM 4
Input Heads
Input Dropout
0.1
Communication keys
32
Communication Values
32
Communication heads
4
Communication Dropout
0.1
C.3 OTHER ARCHITECTURAL CHANGES THAT WE EXPLORED
We have not conducted systematic optimizations of the proposed architecture. We believe that even principled hyperparameter tuning may significantly improve performance for many of the tasks have considered in the paper. We briefly mention a architectural changes which we have studied:
On the output side, we concatenate the representations of the different RIMs, and use the concatenated representation for learning policy (in RL experiments) or for predicting the input at the next time step (for bouncing balls well as all other experiments). We empirically found that adding another layer of mnuli--haded)) key-value attention on the output seems to improve the results. We have not included this change in the RIMs implementation of this paper.
In our experiments. we shared the same decoder for all the RIMs, i.e., we concatenate the represen- tations of different RIMS, and feed the concatenated representations to the decoder. In the future it 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 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 (yon Helmholtz) (1867): von Holst & Mittelstaedt (1950), i.e., using copies of motor signals as inputs. Preliminary experiments shows that this improves the performance in Atari games.
17

---

Published as a conference paper at ICLR 2021
D EXPERIMENT DETAILS
D.1 EFFECT OF VARYING NUMBER AND ACTIVE RIMS ON COPYING TASK
We used learning rate of 0.001 with the Adam Optimizer and trained cach model for 150 epochs (unless the model was stuck, we found that this was enough to bring the training enror close to zero). For the RIMs model used 600 units split across 6 RIMs (100 units per block). For the LSTM used total of 600 units. We did not explore this extensively we qualitatively found that results on copying were not very sensitive to exact number of units.
The sequenees 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 experiments, 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.e number of keys, values as for RIMs model) for 100 epochs. Tab. shows the effect two baselines as compared to 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, shows that the proposed method is better than both of these baselines. For copy task, used head in input attention, and heads RIMs communication. note even with RIM, not exactly same a LSTM, because each RIM can still reference itself.
18

---

Published as a conference paper at ICLR 2021
Table 4: 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 =4, KA 2. Itdim 600)
2.3
1.6
RIMs (KT 4. KA 3. lidim 600)
1.7
4.3
RIMs (by- =5, KA
haim 600)
2.5
4,7
RIMs (KT = 5, KA = 3, hdim = 600)
0.4
4.0
RIMs (KT = =5.kA= = 4, Iidim = 600)
0.2
0.7
RIMs (Ker =6, AA = 2, Adim 600)
3,3
2.4
RIMs (Ae =6, =6. KA 3, Adim 600)
1.2
1.0
RIMs (KT =6, KA
hdim 600)
0.7
5.0
RIMs (KT 6, AA = 5, haim 600)
0.22
0.56
With No input Attention and Full Communication
RIMs (KT = 6. A = 6, Iidi = 600)
0.0
0.7
RIMs (KT 5, A 5, hdim 500)
0.0
1.7
RIMs (Ay 2, KA 2, Adim 256)
0.0
2.9
RIMs (KT =2, KA
hdim 512)
0,0
RIMs (ky I, A = 1, haim = =12)
0.0
0.2
With input Attention and Full Communication
RIMs (by =9, A 2, Aai = =900)
0.24
0.15
RIMs (KT =9, KA = 3, Iddi 900)
0.01
0.00
RIMs (ky -A 4, Iddi 900)
0.01
0.00
RIMs (fy- =9, KA
haim 900)
0.01
0.00
RIMs (KT = =9, AA = ham 900)
0.01
0.01
RIMs (KT = =9, A = 7, hdim = 900)
0.04
0.07
RIMs (AT -9, KA = 8, Adim 900)
0,00
0.10
RIMs (KT KA 9, hdim 900)
0.03
0.24
RIMs (KT 16, A 4, Nati 1600)
0.05
0.02
RIMs (KT 16, A = 6, ham 1600)
0.00
0.00
RIMs (ky 16, KA = 8, Adim = 1600)
0.00
0.00
RIMs (ke 16, A 10, ham 1600)
0.00
0.00
RIMs (KT 16, A 12, haim 1600)
0.00
0.00
RIMs (KT 16, A 14, haim 1600)
0.01
0.26
RIMs (1y 24, A 6, haim 2400)
0.01
0.00
RIMS (KT 24, A = 8, Ado = 2400)
0.00
0.00
RIMs (KT 24, A = 16, hdmm 2400)
0.00
0.00
RIMs (KT 24, A 20, Adim 2400)
0.00
0.00
RIMS (AT 24, A 22, Aain 2400)
0.00
0.42
D.2 EFFECT OF VARYING NUMBER AND ACTIVE RIMS ON ADDING TASK
the adding task we consider a stream of numbers as inputs (given as real-values) and then indicate which two numbers should he added together as set of two input streams which varies randomly between examples. The length of the input sequence during testing is longer than during training. This is a simple test of the model's 's ability ignore the numbers which it is not tasked with adding together. We provide the results in Table Fallchhhich demonstrates that proposed model generalize better for longer testing sequences well adding multiple numbers. We ran the proposed model with different configurations like changing number RIMs as well as number of active RIMs.
D.3 SEQUENTIAL MNIST RESOLUTION 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 digits of different higher resolutions (16x 16, 19 19, and 24 x 24). When re-scaling the images, used the nearest-neighhor hased down-scaling and performed hinarization after re-scaling. We trained with learning
19

---

Published as a conference paper at ICLR 2021
Table 5: Error CE the adding task.
Approach
Train Length 50
Test Length 200
With imput Attention and Full Communication
RIMs (ky = 6. k4 3, Adim -300)
0.0
0.0
RIMs (kr =,, k. 4. Adim -300)
0.0
0.0
RIMs (by =6, k 5, dm =300)
0.0
0.0
RIMs (kr = ką 3. hdmm 600)
0.0
0.0
RIMs (kr = 6. kA 4, Adim 600)
0.0
0.0
RIMs (ky = =, k4 5, Adim 600)
0.0
0.0
RIMs (KT =9, k4 3, Adim 900)
0.0
0.0
RIMs (kr =9, A 4, hdim 900)
0.0
0.0
RIMs (kr =9, Æ4 5, Adm
0.0
0.0
RIMs (kr = 9. k 6, Adim 900)
0.0
0.0
RIMs (ky = 9, A 7, ddm 900)
0.0
0.0
RIMs (ky 9, A 8, ddm 900)
0.0
0.0
Approach
Train Length 500
Test Length 1000
With imput Attention and Full Communication
RIMs (by = 6, A 3, Aam 600)
0.0
0.0
RIMs (ky = 6, k 4, Adi 600)
0.0
0.0
RIMs (ky = 6, k. 5, haim 600)
0.0
0.0
RIMs (by = =9, k 3, Adim 900)
0.0
0.0
RIMs (kT =9. k 4. Adim 900)
0.0
0.0
RIMs (by =9, k 5. haim 900)
0.0
0.0
RIMs (by- = 9, A 6, Adi 900)
0.0
0.0
RIMs (ker = 9, A 7, Adim 900)
0.0
0.0
rate of 0.0001 and the Adam optimizer. For RIMs we used a total of 600 hidden units split across RIMs (100 units per RIM). For the LSTM we used a 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 heads for RIMs communication.
D.4 BOUNCING BALL ENVIRONMENT
We use the bouncing- ball dataset from (Van Steenkiste et al. 2018). The dataset consists of 50,000 training examples and 10,000 test examples showing 50 frames of either 4 solid balls bouncing in confined square geometry, 6-8 balls bouncing occluded region. In all cases, the balls bounce off the wall as well off another, We train baselines well as proposed model for about 100 epochs using 0.0007 as learning rate and using Adam as optimizer (Kingma) & Ba 2014). We use the same architecture for encoder as well as decoder as in (Van Steenkiste et al. 2018). 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 omitted from the main paper for the sake of brevity. Below, highlight few different results.
confined geometry. 3 balls bouncing in confined geometry with random
D.4.1 DIFFERENT RIMS ATTEND DIFFERENT BALLS
In order visualize what each RIM is doing, associate each RIM with different encoder. By performing spatial masking on input, we can control possible spatial input to each RIM. use six non-overlapping horizontal strips and allow only 4 RIMs to be active at time (shown in Fig. (). The mask fixed mask 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 to how the RIM activation patterns changed/correlated with the locations of the RIMs' activations are strongly correlated with the location the balls.
balls. We find that early training. However, after training has proceeded for some time this correlation deteriorates. This likely because the predictable dynamics of the system do not necessitate constant attention.
D.4.2 COMPARISON WITH LSTMM BASELINES
Figures and we highlight different baselines and how these compare to the proposed RIMs model.
20

---

Published as a conference paper at ICLR 2021
Active
Lunctäve region
Ground truth
Prediction (ipuut feed)
Prediction (rollout)
Figure 6: Different RIMs attending Different Balls. understanding what each RIM actually doing. we associate each with a separate encoder, which are spatially masked. Only encoders can active at any particular instant and there are four different balls. We 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 training. this correlation decreased and the active strips did not correlate strongly with the location balls. As the model got better at predicting the location, needed to attend less to the actual objects. The top row shows every Sth frame when the truth is fed in and the bottom shows the results during rollout. The gray region shows the active block. In the top row, the orange corresponds to the prediction and in the bottom, green corresponds the prediction.
Trained on 6-8 Objects Tested on 4 Objects
frames rollout
Truth
LSTM
Curtain
'Truhh
LSTM
Figure 7: Example of other LSTM baselines. For 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 LSTM with 250 hidden units, shown in Fig. Frames plotted every 3rd time step. The red line marks 10 rollout frames. This is marked because after this do not find BCE be reliable measure of dissimilarity.
D.4.3 OCCLUSION
In Fig. I we show the performance of RIMs on the curtain dataset. We find RIMs are able track balls through the occlusion without difficulty. Note that the LSTM baseline, also able to track the ball through the "invisible" curtain.
D.4.4 STUDY OF TRANSFER
is interesting to ask how models trained on a dataset with 6-8 balls perform on a dataset with 4 balls. Fig 12 we show predictions during feed-in and rollout phases.
D.5 VIDEO PREDICTION FROM CROPS
Dataset. We train all models on training dataset of 20K video sequences with 100 frames of 3 balls bouncing
of size 18 48. We also include additional fixed ball in the center make the task video sequences
challenging. use another
same length the same number balls as held-out validation set. In addition, we also have 5 of disttribbuttnno (OOD) test sets with various number of bouncing balls (ranging from 1 to 6) and each containing 1K sequences of length 100.
Training. We train all models until the validation loss is saturated, and select the best of three runs. During training, automatically decay the learning rate by factor of 2 if the validation loss does not significantly decrease by at least. 0.01% for five consecutive epochs.
Evaluation Criteria. After having trained on the training dataset with 3 bouncing balls, we evaluate the performance on all test datasets with 1 to 6 bouncing balls. In Figure we report the balanced accuracy (i.e. arithmetic mean of recall and specificity) and Fi-scores (i.e. harmonic mean of precision and recall) to account for class-imbalance
21

---

Published as a conference paper at ICLR 2021
O
18
16
88
18
110
111
122
14
15
Truth
RIM
18TM
Trniued
Olojects
Tested O.S. Obicents :
Truth
o-
%
8.
o 8.
1.TMM
s.
8.
Trined on R.R. Obiects Tested Obdiants
Truth
RIM
Curtains
Trith
1.5TM
Figure 8: Comparison of RIMs to LSTM baseline, For different experiments 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. Duc to the number of 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 time steps, computing the binary cross entropy between the prediction and the true balls at each instant, as in Van Steemkiste et al. (2018). In the predictions, the transparent blue shows the ground truth, overlaid to help guide the eye.
Objects
10 frames in rollout
Truth
LSTM
RIM
Trained on 4 Objects Tested on 6-8 Objects
Truth
LSTM
RIM
Figure 9: 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. 3 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 reliable measure of dissimilarity.
Results. In Figure 13.
that proposed method out-perform all non-oracle baselines OOD on the one-step forward prediction task and strike a good balance regard in-distribution and OOD performance.
Baselines: We compared the performance of the proposed method with a baseline LSTM model, as well as state of the art memory model, RMC (Santoro et al. 2018). As sanity check we also compare the performance of the proposed method to an oracle baseline that does not model the dynamics (refer to as TTO: Time Travelling Oracle).
22

---

Published as a conference paper at ICLR 2021
Trined on 4 Objects
Traned on 6.8 Objocis
4 O.jjects
Tosted on 6.8 Objects
Tested on 4 Objects
Chamii
Biuary Entrspy
LSTM 512
2500
RIM
Ground
Rollont
Figure 10: Comparison of RIMs to LSTM baseline. For 4 different experiments in the text, compare RIMs different LSTM baselines. In all cases find that during rollout, RIMs perform better than the LSTMs accurately capturing the trajectories the balls through time. to number hard collisions, accurate modeling is very difficult.
Reliont
Fruzu:: 2
11
14
20
23
26
85
38
43
4D
Tour
4
3
Figure 11: RIMs on dataset with an occlusion. We show two trajectories (top and bottom) of three balls. For the left frames, at each step the true frame is used as input. On the right, outlined in black, the previous output is used input.
RIMs
Troninng
Tree
Dransfrr
True
Figure 12: 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.
D.6 ENVIRONMENT WITH NOVEL DISTRACTORS
We evaluate the proposed framework using Adavantage Actor-Crite (A2C) learn policy To(als, 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 Cheislier-Boisvertt & Willems 2018). 9. 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, further I.STM encode the state and summarize the past observations.
23

---

Published as a conference paper at ICLR 2021
1.00
Balanced Accuracy
0.95
0.90
0.90
0.85
F3
Training Destr:
0.85
Traiiing Distr:
TTD
TTD
0.75
LSTM
RMC
0.80
LSTM
IMC
0.70
RJMS
Number Bouncing Balls
Number of Bouncing Balls
Figure 13: Performance metrics on OOD one-step forward prediction task. Gist: RIMs outperforms all RNN
baselines OOD.
Figure 14: An example of the minigrid task.
D.7 MINTGRID ENVIRONMENTS FOR OPENAI GYM
The MultiRoom environments used for this research part of MiniGrid, which is open source gridworld packagd] This package includes family of reinforeement learning environments compatible with the OpenAl Gym framework. Many these environments are parametcrizable so that the difficulty tasks can be adjusted (e.g.. the size of rooms is often adjustable).
D.7.1 THE WORLD
MiniGrid, the world grid size NxN. Each tilc in the grid contains exactly zero 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.
D.7.2 REWARD FUNCTION
MiniGrid environments. 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 is reached. In the FindObj environment, the agent receives positive reward if it reaches the object to be found, otherwise zero reward if the time step limit is reached.
Rewards are sparse
24

---

Published as a conference paper at ICLR 2021
The formula for calculating positive sparse rewards 1 0.9 (slep cound/ / /maw sleps). 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 max max_steps parameter is different for each environment, and varies depending on the size of each environment, with larger environments having higher time step limit.
D.7.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 turn left and forward action makes the agent move from its current tile onto the tile in the direction it 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.
the
right action rotate and face one of possible directions (north, south, cast, west). The move
D.7.4 OBSERVATION SPACE
Observations in MiniGrid are partial and egocentric. By default, the agent sees square of 7x7 tiles in the direction is facing. These include the tilc the agent standing on. The agent cannot see through walls closed doors. The observations are provided as a tensor of shape 7x7x3. However, note that these are RGB images (which is different from original BabyAl Al paper).
D.7.5 LEVEL GENERATION
The level generation in this task works as follows: (1) Generate the layout of the map (X number of rooms with different sizes (at most size Y) and green goal) (2) Add the agent to the map random location in the first. room. (3) Add the goal random location in the last room. neural network parameterized as CNN used to process visual observation.
We follow the same architecture as (Chevalier-Hoisee) & Willems 2018 we replace the LSTM layer with BLockLSTM.
D.7.6 ADDITIONAL ABLATION
We present one ablation in addition to the ones in Section 4.4 In this experiment, we study the effect on input attention (i.e top down attention) well the of muldi-headed head key-value attention. We compare the proposed model (with input attention well multi-headed key-value attention) with 2 haselines: (a) In which we remove the input attention (and force
RIMs communicate with each other (h) We use head for key-value attention as compared to multi-headed key-value attention. Results comparing the proposed model, with these two baselines is shown in Fig. 15
Profietion
Binuyy Chuss Eatrops
RIMs
0.20
One-heael
0.15
tuput atteution
0.10
6.05
Test. lass
Roilout
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 one active head.
D.8
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.
25

---

Published as a conference paper at ICLR 2021
Rollout
10
12
14
16
18
20
26
28
30
Truth
-
00 6
Prodiction
Truth
Prediction
Truth
Prediction
Truth
Prediction
Figure 16: One head and attention Using one head and attention models, show the rollout predictions blue. On top we show results on the 4 ball dataset and on bottom we show results on the curtains dataset.
4000
0.75
-1.0
Figure 17: A comparison showing relative improvement of RIMs with ka 5 over ka 4 baseline. Using ka 5 performs slightly worse than &a 4 but still outperforms PPO, and has similar results across the majority of games.
1.00
0.79
0200
000
220
Figure 18: RIMs-PPO relative score improvement over ASTM-PPO baseline (Shulman et al. 2017) across all Atari games averaged over 3 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 AA 5).
26

---

Published as a conference paper at ICLR 2021
E ADDITIONAL EXPERIMENTS
E
IMITATION LEARNING: ROBUSTNESS TO NOISE IN STATE DISTRIBUTION
Here, we consider imitation lcarning where we have training trajectories generated from an expert (Table We evaluate our model on continuous control tasks in Mujoco (in our case, Half-Cheetah) (Todorov et al. (012). We take the rendered images input and compared the proposed model with recurrent policy (i.e., LSTM). Since, using rendered image of input does not tell anything about the velocity of the Half-Cheetah, makes 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 can see, after adding noise LSTM baselines performs poorly. On the other hand, for the proposed model, there's also drop in performance but not as bad as for the LSTM baseline.
Table 6; Imitation Learning: Results on the half-cheetah imitation learning task. RIMs outperforms baseline LSTM when we evaluate with perturbations not observed during training (left). An example of input image fed to the model (right).
Training
Perturbed States
Method / Setting
Observed Reward
Observed Reward
LSTM (Recurrent Policy)
5400 + 100
2500 + 300
RIMs (AT = 6, ka= 3)
5300 + 200
3800 200
RIMs (ki = 6, k.a 6)
5500 + 100
2700 + 400
RIMs (without Input attention)
5400 + 100
3200 + 50
We use the convolutional network from (a & Schmidhuber 2018) as our encoder, GRU (Chung et al. 2015 with 600 units deterministic path in the dynamics model, and implement all other functions two fully connected layers of size 256 with ReI.U activations. Since, here we are using images as input, which makes the task, partially observable. Hence, we concatenate the past 4 observations, and then feed the concatenated observations input to GRU (or our model). For our model, we use 6 RIMs, each of size 100, and we set k. = 3. We follow the same setting in (Hafner et al. 2018 (oodhani al. 2019). We also compare the proposed method to the haseline where dont include input attention (or top-down attention). AS shooss there's decline performance if we dont use input attention, hence justifying the importance
E.2 TRANSFER ON ATARI
As a very preliminary result, we investigate feature transfer between randomly selected Atari games. order to study this question, we follow the experimental protocol Rusu et al. (2016).
We start by training RIMs on three source games (Pong, River Raid, and Seaquest) and test if the learned features transfer a different subset of randomly selected target games (Alien, Asterix, Boxing. Centipede, Gopher, Hero, James Bond, Krull, Robotank. Road Runner, Star Gunner, result in positive transfer 9 out 12 target games, with three cases of negative transfer. On the other hand progressive networks (Rusu et al. (016)) result in positive transfer in 8 out of 12 target games, and two cases of negative transfer. We also compare to L.S baseline, which yields positive transfer in 3 of 12 games.
Wizard of Wor). observe, that RIMs
27

---

Published as a conference paper at ICLR 2021
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 12211 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 + +.3
-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 +
Privatelß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
k 5.
E.2.1 ATARI RESULTS: COMPARISON WITH LSTM-PPO
28

---

Published as a conference paper at ICLR 2021
Figure 19: Comparing RIMs-PPO with LSTM-PPO: Learning curves for KA 4. KA models and the LSTM-PPO baseline across all Atari games.
RIMs-PPO
29

---

Published as a conference paper at ICLR 2021
E.2.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.
Bor
Figure 20: Baseline agent with no input attention mechanism: Here compare the RIMs to the baseline.
where their is input attention (i.e., top down atlention) well all the RIMs communicate with each other time steps. Learning curves for RIMs-PPO models, Baseline Agent, the LSTM-PPO baseline across 30 Atari games.
30

---

Published as a conference paper at ICLR 2021
Table 8: Transfer from WMT to IWSLT (en de) for different models. better).
Results in BLEU score (higher is
Training Data
(Vanilla) Transformer
LSTM
RIMs
en + de en de, en +
22.89
21.32
23.71
22.92
20.37
24.23
E.3 BOUNCING MNIST: DROPPING INDIVIDUAL RIMS
We use the Stochastic Moving MNIST (SM-MNIST) (Denton & Fergus 2018) dataset which consists of sequences frames of size 64 x 64, containing one or two MNIST digits moving and bouncing 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. masking out a particular RIM should not really effect the dynamics of the entire model. We show qualitative comparisons Fig. 2T [2
particular RIM and study the effect of the masking the ensemble
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 dynamies for next 25 time-steps. We find that sparsity is needed otherwise different
24
25
RIMs co-adapt with each other (for see Fig. 22 24 25). We tried similar masking experiments for different
models like RMC. Transformers, EntNet (which learns mixture of experts). LSTMs, but all them failed to do anything meaningful (after masking). We suspect this is partly due to learning a homogeneous network.
F NATURAL LANGUAGE PROCESSING: LANGUAGE MODELING, MACHINE TRANSLATION, AND TRANSFER LEARNING
We performed some additional experiments evaluate how well RIMs improve transfer leaminng on some widely studied NI.P datasets.
We performed an additional experiment where we trained a seq2seq model on the WMT machine translation dataset and evaluated on the IWSLT14 dataset (English to German). WMT consists of European Parliament text and news articles, whereas IWSLT14 consists of transcribed spoken text (for example, TED talks). Thus the datasets in the same general domain but have fairly distinct content distributions. Additionally, we considered either training the WMT model on both English to German and English French (with shared encoders but separate decoders) or only on English to German. The results are in Table 8 The LSTT and RIMs models used have comparable number of parameters in this experiment. We note that the multi-task training setup substantially hurts the performance of the LSTM baseline, but helps performance when using RIMs (the performance of the transformer about the same both settings).
31

---

Published as a conference paper at ICLR 2021
Figure 21: 4 RIMs. (top 2). Each sub- 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 second RIM etc.
32

---

Published as a conference paper at ICLR 2021
Figure 22: 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 the effect of masking the second RIM etc.
33

---

Published as a conference paper at ICLR 2021
4
%
4
4
46
96
%
5
+
5
6
&
8
8
8
5
O
6
2
6
B
40
&
Figure 23: 400dlm, kr 5, ka =
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 second RIM etc.
34

---

Published as a conference paper at ICLR 2021
29
69
Figure 24: 400dim, kr = 5, kA = the effect of masking the other RIMs. For examples, the top figure shows the effect of masking the first RIM, the second figure shows the effect of masking the second RIM
3. Each sub-figure shows the effect of masking a particular RIM and studying
35

---

Published as a conference paper at ICLR 2021
4. 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 second RIM etc.
Figure 25: 400dim, kr = 5, kA =
36

---

Published as a conference paper at ICLR 2021
G PSEUDOCODE FOR RIMS ALGORITHM
Assume these variables defined; rims_in (number of positions in the incoming input), rims_out (number RTMs current layer), rim_dim the number of hidden unite in each RIM, for simplicity is the aame on the input and output), baz (the batch size), & (number of RIMs to activate per step)
from torch import topk, like, like, arange,
def rims_core(ingg h):
at=_inp, att_acores = input_attention (inp, h)
mask = cin_oompetition (att_sccres)
h_next self.bl gro (att_inp, h)
next h_next rim commmmicmacct next, mask)
hent (mask) next (1-mask) sh return b_next
def input_ Input_attention (inp, h):
inp inp. reshape (baz, rims in, rim_dim))
att_inp, atto self.inp_att th, reshape (bsz, rims_out rim_dim) ) inp, inp)
att_inp - cet_iin.reshaaee ( (bsz, rim_dimvrims_out) return att__nnt attnl:,;,00
def cin_ccometiiionn (att_scores)
k_ind topk (att_scores, dim=1, sorted-Tree mask - ones ones_like (att (art_scores)
[1]
like
mask mask.index_pru (arange .usqueezee (1) 1,,k,,i, iiiin)), zeros_like (k (k_ind[0])
(arange
mask = mask. rwshapeeepeeeeee (bsz, rims rins_out,,) 1) repeat (1,1, 1, rim_dim) ) reshape ( (bsz,
rims_outarrm_ddm)
return inp, mask.detach()
def cin_commun.cation (b, mask):
h - h. rereaheee (bs2, rims_out, rim_dim)
h masaeed__pp grad #lock_gradiae (h, mask. reshape (bsz, rims__ottt out, rim (im_ddmm)) h_att,_ self mha gt..maakd_grddd..aaaa grad, masked_gradd
h_att hLatt. reshape (bsz, return h_att
Figure 26: The core RIMs module for single recurrent step using GRU independent dynamics. We show how the attention scores from input attention are used weight the input the RIMs (called att_ and construct a mask which controls which RIMs are allowed to update.
37


Reviews:
[{"review_id": "MOOiN4H6ZSR", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "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": "After rebuttal:\nI appreciate authors' detailed responses and an updated version of the paper. They mostly clear my concerns and doubt. I increase my rating to accept. \n--------------------------------------\nSummary:\nThis paper introduces a module that ensembles independent RNNs using a multi-head attention mechanism. This proposed recurrent independent mechanism (RIM) includes multi-head attention, top-k activation section, input attention, and communication modules. \nThe experiments on a range of diverse tasks show that RIMs generalizes better in many tasks than LSTMs. \n--------------------------------------\nPros:\n+ The paper is clearly written.\n+ The related works and the difference with the proposed model are explained in details. \n+ The experiments cover a wide range of scenarios from copying task to reinforcement learning. The additional experiments in Appendix are helpful. \n--------------------------------------\nConcerns:\n1. *Novelty:*\nIn my understanding, the core idea is essentially combining mutli-head top k attentions with RNNs. I appreciate that authors includes necessity of the proposed module and their insights. However, this paper simply combines existing works and thus lacks novelty. I ask authors to clarify it if I missed anything.\n\n2. *Model capacity:*\nAuthors claim that high performance with RIMs is not due to the increase of model capacity, and the model size is significantly reduced with RIMs when the comparing model has the same number hidden units. \nRelated to this, I have suggestions:\n    1. The model size of the proposed and comparing models should be reported. \n    2. Additionally, adding latency and FLOPs would be helpful. \n\n3. *Sparsity:*\nAuthors mention that sparsity is necessary in this model. How is the proposed model comparable with other sparse networks? Increase sparsity by adding an existing technique [1-4] in the standard LSTM can be another baseline for this model. Some previous works are listed here: \n    1. K-winner-take-all [1], local winner-take-all [2]\n    2. Dropout [3,4]\n\n4. *Missing references* regarding sparsity: [1-4]\n--------------------------------------\nMinor comments:\n- References of the model are missing in Table 1. \n- Page 8: 'allow the RIMs **ot** communicate with' -> 'allow the RIMs **to** communicate with'\n--------------------------------------\n[1] Majani, et al., On the K-Winners-Take-All Network, NeurIPS 1988.\n[2] Srivastava, et al., Compete to Compute, NeurIPS 2013.\n[3] Srivastava, et al., Dropout: A Simple Way to Prevent Neural Networks from Overfitting, JMLR 2014.\n[4] Molchanov, et al., Variational Dropout Sparsifies Deep Neural Networks, ICML 2017", "extra_scores": {}, "extra_text": {}}, {"review_id": "EBfHx0qq3E", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "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": "The authors argue that the world consists of largely independent causal mechanisms that sparsely interact. The authors propose a new kind of recurrent network (RIM) that presumably distills this world view into inductive biases. RIMs consist of largely independent recurrent modules that are sparsely activated and interact through soft attention. They tested RIMs on a number of supervised and reinforcement learning tasks, and showed that RIMs perform better than LSTMs and several other more recently proposed networks (e.g., Differential Neural Computers, Relational Memory Core, etc.) in these tasks. In particular, RIMs can generalize more naturally than other networks to out-of-distribution test set in presumably modular tasks.\n\nThe motivation for RIMs makes intuitive sense, even though it is perhaps debatable whether the largely independent causal mechanisms in the world should each be captured by a single RNN. The manuscript is easy to follow, the idea is quite interesting, and the model is empirically tested across a wide diversity of tasks.\n\n(1)\tMy main concern is that it is not clear to what extent the improved performance is due to the core concept of recurrent independent mechanisms, or due to other factors such as the use of attention and more hyperparameter tuning.\n\nI don’t believe the merit of this work should be judged exclusively in its improved performance over other architectures. However, since the authors focused most main figures on performance, it is worth better understanding the cause of that performance gain.\n\nIn some experiments, RIMs are only compared against LSTMs, and it is not clear whether the gain over LSTMs is due to the use of attention. In the authors’ submission last year, a similar question was raised by R3, and the authors correctly pointed out that in several other experiments, RIMs fared better than attention-based models such as Transformer and RMC. However, in these experiments--as far as I can tell—RIMs benefit substantially from more extensive hyperparameter tuning. From table 1 and 4, we see that the advantage of RIMs over competing architectures, especially attention-based ones, are similar in magnitude to the performance difference caused by reasonable hyperparameter variations.\n\nTo alleviate this concern, the authors could potentially show results on a few datasets studied in the papers of RMC, DNC, etc.\n\n(2)\tMy other related main concern is that the authors proposed RIMs as a response to a world full of largely independent objects (or variables), yet never showed that RIMs would break when the world substantially deviates from this ideal.\n\nTo convince readers that the sparsely activated nature of RIMs is suitable for our world, the authors could design a world with densely interacting objects, for example balls attached with springs, and shows the failure of RIMs, despite the same amount of hyperparameter tuning, in comparison to other architectures. We will have a much better understanding of RIMs when we know how to break it.\n\n(3)\tThe authors motivated RIMs as more natural for capturing the largely-independent components of the real world. However, it’s not clear from the authors’ results (Fig. 7, Figs. 22-26) that each recurrent module actually learns the dynamics of individual objects.\n\nFor example, Fig. 7 is titled “Different RIMs attending to different balls”, but as far as I can tell, this conclusion is not actually shown in the figure.\n\n---------------------\nEdit after author's responses\nMy first concern is addressed by the authors' response. My other two concerns were not really addressed, but I think these concerns should not preclude this manuscript from getting accepted. So I'm updating my score.\n\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "AqCUXvypV3p", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "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": "[Update after author's responses]\n\nI appreciate the responses provided by the authors. I think they answer my questions. In consequence, I update my review to favour accepting the paper.\n\n---\nThis paper introduces a new architecture composed of semi-independent recurrent networks that interact with each other, and with their environment, through attention. Furthermore, an attention mechanism is also used to select which networks are allowed to update their internal state at any given time. The authors present several experiment to demonstrate that this architecture outperforms simple LSTMs on various tasks.\n\nThe method seems interesting, and the experiments indicate some benefit. However, in its current state, the paper makes it very difficult to properly evaluate the method, due to a lack of explanations.\n\n- It is often unclear how many parameters/neurons the RIMs use.  As a result, we can't know whether the comparison to the LSTMs is fair. Please indicate *total* number of neurons and trainable parameters for both LSTMs and RIMs in all experiments.\n\n- Section D.4.1 and Figure 7 seem to contain important information about how the system works, but it is totally incomprehensible -e.g. how are the masks generated and moved around? Please rewrite this with an explanation of what exactly is going on.\n\n- I would appreciate some explanation about how exactly the system differentiates through the choice of which RNNs to activate or leave dormant. This seems to represent a hard all-or-nothing change, through which no smooth interpolation is possible, affecting all future time steps. How can we find a gradient over this step choice when backpropagating through time?\n\n- Related: how is the training actually performed? Do you backprop error through time over the whole history? I didn't see any description of the training process in the paper - though I may have missed it. There should certainly be one!\n\nIf these explanations are provided (and if the comparisons with LSTMs are fair) I think the paper would be acceptable. ", "extra_scores": {}, "extra_text": {}}, {"review_id": "3UL-oybJOQY", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 9, "confidence": 3, "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": "The authors propose to learn what they term recurrent independent mechanisms (RIMs), a new recurrent architecture with components with nearly independent transition dynamics.  RIMs exhibit excellent generalization on tasks in which the factors of variation differ systematically between the training and evaluation distributions.\n\nThe paper is very well organized and well written, with relatively simple and consistently clear explanations of key concepts. The authors are precise in this language use, noting for example that they use the term \"mechanism\" in two slightly different ways (footnote 1).  \n\nThe authors state that their central question is\" how a gradient-based deep learning approach can discover a representation of thigh-level variable which favour forming independent but sparsely interacting recurrent mechanisms in order to benefit from the modularity and independent mechanisms assumption.\"  The paper could be improved by more clearly explaining the unique benefits of a gradient-based approach to this task.\n\nThe experiments in section four are particularly thoughtful and well-designed.  Rather than merely comparing performance on some  set of benchmarks, the authors aim to elucidate key capabilities with specific experiments. ", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper proposes a novel recurrent network called RIMs for improving generalization and robustness to localized changes. The network consists of largely independent recurrent modules that are sparsely activated and interact through soft attention. The experiments on a range of diverse tasks show that RIMs generalizes better than LSTMs.

The overall feedback from reviewers is positive: the paper is well written, the idea is interesting, and the experiments cover a wide range of diverse tasks.

The main concerns of most reviewers are the fairness of comparison, the limited novelty, and lacking details on how and why the system works. The authors pointed out that RIMs are not a straightforward combination of attention and RNN, and it has fewer parameters than LSTMs. They also conducted ablation study to demonstrate the benefits of RIMs and provided the missing details in the revised version.

In summary, this paper presents an important research direction for systematic generalization using modularized network. The paper is well written, the idea is novel and interesting, and the experiments cover a wide range of diverse tasks. Hence, it makes a worthwhile contribution to ICLR and I am recommending acceptance of this paper.

Decision: Accept (Spotlight)