Title: Recurrent Independent Mechanisms
Authors: Anirudh Goyal
Alex Lamb
Jordan Hoffmann
Shagun Sodhani
Sergey Levine
Yoshua Bengio
Bernhard Schölkopf
Venue: ICLR 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.



Review discussion:
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.
The 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.
(1)	My 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.
I 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.
In 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.
To alleviate this concern, the authors could potentially show results on a few datasets studied in the papers of RMC, DNC, etc.
(2)	My 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.
To 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.
(3)	The 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.
For 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.
Edit after author's responses
My 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.
We thank the reviewer for their feedback and their generally positive assessment of our work.
> “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.”
We note that we have not done any hyper-parameter search. All the experiments in our paper use 6 RIM, and topk=4. What table 1 shows is that the models WITH sparsity (i.e., topk=4) perform better as compared to models without sparsity topk=6. We also note that the hyper-parameter (topk) is fairly flexible as shown in Table 4 in appendix.
> “To alleviate this concern, the authors could potentially show results on a few datasets studied in the papers of RMC, DNC, etc.”
RMC is a scalable version of the DNC, which uses attention as compared to a learned controller. We have in general found RMC to be sensitive to hyper-parameters. We have compared the proposed method to RMC for a sequential MNIST task as well as another partial observation video prediction task (Figure 4, right). On a partial observation video prediction task, we note that RMC indeed performs better than LSTM (confirming that memory is indeed useful while handling partial observability).
We also compared RMC on wiki-text2 (we took the open source code for RMC from here, https://github.com/L0SG/relational-rnn-pytorch 236 stars on the github repo, and 22 forks).  RMC gets the ppl of 107.21, (worse than LSTM), whereas RIMs gets a ppl of 98, whereas an unregularized LSTM gets a ppl of 105. We know this is not an impressive result, but this is to show that getting RMC "correct" is tricky.
Both RMC & LSTM have ~11M parameters whereas RIMs have around 8M parameters. We have also tried running RMC on the atari-ppo benchmark, but we have not been able to achieve better results as compared to LSTM baseline. It’s also difficult to tune hyper-parameters for RMC on atari tasks, as running on a particular game for 50M time-steps takes 3 days.
On the other hand, we experiment on the whole suite of Atari games and find that simply replacing the LSTM with RIMs greatly improves performance.
> “yet never showed that RIMs would break when the world substantially deviates from this ideal.”
One clear lower-bound to consider is that every individual RIM can itself act as an LSTM/GRU on its own, so in principle, a sufficiently large RIMs model should be at least as expressive as an LSTM or a GRU.
At the same time, RIMs should provide benefits from having multiple mechanisms whenever the data has multiple processes or factors of variation.  This was the case in all of the datasets and tasks that we studied, but it would be interesting to see if there are any exceptions.  One thing to consider is that the improvements from using RIMs are much larger when systematic generalization is required, and in the i.i.d. setting there is often still an improvement over LSTMs, but it is usually smaller.
> “To 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.”
When all the entities in the world are interacting with each other, RIMs won’t really “break” as such, since information can still be shared freely between the different mechanisms by using attention.  There might be a smaller benefit in the case reviewer mention (some kind of spring-world), but it’s not clear to us that it would hurt performance relative to something like an LSTM.
Thank you for the timely feedback (although mine took much longer)!
I am now more comfortable with the hyperparameter tuning part and the empirical comparison with other attention-based architectures part.
I'm still not convinced about the authors' response to the "understanding RIMs by making it break" part.
But I think that's fine. This can be left to the future.
Dear. Review,
We are happy to  know that the reviewer feel confident, and we thank the reviewer for increasing their score.
If there's any outstanding issue where we can improve further, we're happy to invest the time to improve the paper further.
Thanks for your time.
After rebuttal:
I appreciate authors' detailed responses and an updated version of the paper. They mostly clear my concerns and doubt. I increase my rating to accept.
Summary:
This 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.
The experiments on a range of diverse tasks show that RIMs generalizes better in many tasks than LSTMs.
Pros:
+ The paper is clearly written.
+ The related works and the difference with the proposed model are explained in details.
+ The experiments cover a wide range of scenarios from copying task to reinforcement learning. The additional experiments in Appendix are helpful.
Concerns:
1. *Novelty:*
In 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.
2. *Model capacity:*
Authors 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.
Related to this, I have suggestions:
1. The model size of the proposed and comparing models should be reported.
2.
Additionally, adding latency and FLOPs would be helpful.
3. *Sparsity:*
Authors 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:
1. K-winner-take-all [1], local winner-take-all [2]
2. Dropout [3,4]
4. *Missing references* regarding sparsity: [1-4]
Minor comments:
- References of the model are missing in Table 1.
- Page 8: 'allow the RIMs **ot** communicate with' -> 'allow the RIMs **to** communicate with'
[1] Majani, et al., On the K-Winners-Take-All Network, NeurIPS 1988.
[2] Srivastava, et al., Compete to Compute, NeurIPS 2013.
[3] Srivastava, et al., Dropout: A Simple Way to Prevent Neural Networks from Overfitting, JMLR 2014.
[4] Molchanov, et al., Variational Dropout Sparsifies Deep Neural Networks, ICML 2017
We thank the reviewer for their helpful feedback.
> “In my understanding, the core idea is essentially combining multi-head top k attentions with RNNs. I appreciate that authors includes necessity of the proposed module and their insights. ”
We think that this characterization is too reductionist.  In RIMs, we divide the recurrent state into multiple mechanisms with separate parameters and hidden states, which only communicate via bottleneck of attention.
Even if we ignore the competition between mechanisms, the RIMs technique is not a straightforward combination of multi-head attention with RNNs, because it also divides the model into multiple separated mechanisms.
An additional source of novelty is that modules compete to access input information. Each module only attends to that information which is relevant to it. None of the previous modular architectures have this inductive bias.  The activation of different RIM is input dependent and hence dynamic.  Because of the presence of input attention, a particular module mechanism if the information in the input matches to what that particular mechanism expects.
> “[Concern about] Model capacity: Authors claim that high performance with RIMs is not due to the increase of model capacity ... ”
RIMs can be used as a drop-in replacement for an LSTM/GRU layer. RIMs drastically reduces the total number of recurrent parameters in the model (because of having a block-sparse structure) but also adds new parameters to the model through the addition of the attention layers.  In all of our experiments RIMs had fewer parameters than the LSTM baseline.
Here are specific numbers on parameter counts, which we will add to the paper:
Task,              LSTM, RMC,   RIMs
Seq. CIFAR:   4.8M,  6.45M, 3.23M
Seq. MNIST: 4.2M,   4.7M,   2.8M
Copying:        2.2M,   2.8M,   1.6M
On Atari, the LSTM had 9M parameters and our RIMs model had 6.2M parameters.
For Bouncing balls, LSTM had 7.4M parameters and RIMs had 4.2M parameters.
As you can see, RIMs is not only fair, but uses significantly fewer parameters than the baselines we compare against.
> "Proposed model comparable with other sparse networks? Increase sparsity by adding an existing technique [1-4] in the standard LSTM can be another baseline ... (1) K-winner-take-all [1], (2) local winner-take-all [2], (3) Dropout [3,4]”
K-winner-take-all: We use a similar strategy as k-winner takes all, where different RIMs compute the activation score based on similarity to the input information. The top-k RIMs which have higher similarity score get to access input information, and update their hidden state, while the RIMs which lose the competition follow their default behaviour. An important distinction is that the attention score is a function of the hidden state of each RIM, and hence the activation can be “considered” as top-down as compared to scenario, where a critic decides which module to activate.
Comparison to Zoneout: We compared the proposed method to Zoneout, a state of the art dropout augmented LSTM (Krueger et. al, ICLR'17). On a conceptual level, RIMs and Zoneout are very different because RIMs is a new architecture aimed at improving systematic generalization whereas Zoneout is a regularizer which skips the recurrent state update randomly.
Nonetheless we compare with Zoneout on the Sequential MNIST task. We train the Zoneout augmented LSTM on a 14 x 14 MNIST digit, and test it on 16 x 16, 19 x 19, 21 x 21. LSTM gets (86.8, 42.3, 25.2), LSTM + Zoneout gets (88.4, 43.5, 25.5), whereas RIMs got (90, 73.4, 38.1). Zoneout improves the generalization performance of LSTM confirming the reviewer’s hypothesis, but RIMs performs much better.
[a] Zoneout, https://arxiv.org/abs/1606.01305
local winner-take-all: The key difference is that winner-take-all selectively activates individual units based on which has the largest value.
As such it is more like an activation function.
On the other hand, RIMs divides the recurrent state into multiple independent mechanisms, and an attention-based competition is used to selectively activate them (where each mechanism has many units).
Moreover, this competition drives which mechanisms can run their independent dynamics on a given step, which is a much stronger form of gating than setting an individual unit to zero if it has a smaller value.
Dropout: Dropout is a regularizer which sets values of individual hidden units to zero randomly, and thus it is distinct from RIMs which is a new architecture in which mechanisms (each of themselves containing many units) compete to activate and share information using attention.  In principle, RIMs and dropout should be complementary, since they have different purposes (systematic generalization vs. regularization).
> “4. Missing references regarding sparsity: [1-4]”
We thank the reviewer for pointing out these references. We will update the paper with the correct citations.
Hello,
We thank the reviewer for their feedback and valuable comments.
Since the first phase of response period is completed, if you have time and could indicate if there are any other concerns of yours which we have not addressed, we'd be happy to take a look.
Thanks for your time.
[Update after author's responses]
I appreciate the responses provided by the authors. I think they answer my questions. In consequence, I update my review to favour accepting the paper.
This 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.
The 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.
- 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.
- 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.
- 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?
- 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!
If these explanations are provided (and if the comparisons with LSTMs are fair) I think the paper would be acceptable.
We thank the reviewer for their feedback and their generally positive assessment of our work.
> “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.”
With RIMs, the number of hidden units is always the same as the LSTM baseline.  The number of parameters is generally decreased relative to the LSTM baseline (some parameters are removed due to the block-sparse structure of the parameters in RIMs, but some new parameters are added for the attention layers).
Here are specific numbers on parameter counts, which we will add to the paper:
Task,              LSTM, RMC,   RIMs
Seq. CIFAR:   4.8M,  6.45M, 3.23M
Seq. MNIST: 4.2M,   4.7M,   2.8M
Copying:        2.2M,   2.8M,   1.6M
On Atari, the LSTM had 9M parameters and our RIMs model had 6.2M parameters.  For Bouncing balls, we also note that we compared the RIMs to the LSTM baseline  which has almost twice as many parameters as compared to RIMs (RIMs: 4.2M, LSTM: 7.4M), and we found RIMs to perform better as compared to the LSTM baseline.
As we can see, RIMs is not only fair, but uses significantly fewer parameters than the baselines we compare against.
> “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?“
The reviewer raised the great point asking how the gradient flows when doing BPTT.
While it's true that no gradient flows through the non-selected RIM, gradient flows into the weighting $A_k$ over the selected RIM, which leads their $A_k$ values to be pushed up if they are more useful.  This competitive mechanism of top-$k$ selection with gradient flowing into the highest-valued entries optimizes quite easily in practice and was used successfully in BRIMs (Mittal et. al, ICML 2020) and Sparse Attentive Backtracking (Ke et. al, NeurIPS 2018).  We will update the text to make this clearer.
In practice, we’ll also add that we never had any difficulty optimizing RIMs, and have consistently found it as easy to train as a normal LSTM.
> “Explanation of Figure 7”
1. For understanding what each RIM is actually doing, we divide the input into 6 horizontal strips such that each part of the input is processed by a different encoder. We note that for actual results, we only used a single encoder. Only as an ablation result, we ran the experiment with 6 encoders, where each encoder processed a part of the input but this is not indicative of how the proposed architecture works.
2. We associate each horizontal strip with a separate encoder, which are spatially masked.
3. We consider a scenario where only 4 encoders can be 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. We find that, early in training, RIM activations correlated more strongly with the locations of the four different balls.
Later in training, this correlation decreased and the active strips did not correlate as strongly with the location of balls. As the model got better at predicting the location, it needed to attend less to the actual objects. The top row shows every 5th frame when the truth is fed in and the bottom shows the results during rollout. The gray region shows the active block. In the top row, the orange corresponds to the prediction and in the bottom, green corresponds to the prediction.
> " how is the training actually performed? Do you backprop error through time over the whole history?"
We note that we use RIMs as a drop in replacement of LSTMs, and everything else remains the same. For copying/sequential MNIST and bouncing ball task, we do backprop over the entire history (for the proposed method as well as all the baselines).
Hello,
We thank the reviewer for their feedback and valuable comments.
Since the first phase of response period is closing soon, if you have time and could indicate if there are any other concerns of yours which we have not addressed, we'd be happy to take a look.
Thanks for your time.
> "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."
We've updated the paper with a new Appendix G (end of the paper) which directly addresses your concern about "how the system works".  It uses PyTorch pseudocode and shows that the core RIMs algorithm can be described in just a few lines, while using three helper functions: input_attention, rim_competition, and rim_communication.
This change also should make it clearer how the gradient flows through the system.  The input is weighted in the input_attention function to produce (att_inp), while those attention scores (att_scores) are used to construct a hard-mask by top-k selection.  Basically, if a RIM has a strong drive to read from the input on a given step, it will trigger a hard-decision for the RIM to update its hidden state on that step.
Let us know if this makes things clearer, or if there's anything else you'd like to see here.  We feel like this makes the paper a lot better, and we're happy that your feedback helped us to make this improvement.
We also addressed your concern about the number of parameters and units in a RIMs model.  The number of hidden units is always the same, and the number of parameters is substantially reduced relative to the LSTM/GRU baseline.
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.
The 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).
The 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.
The 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.
We thank the reviewer for their feedback and their generally positive assessment of our work.
“The paper could be improved by more clearly explaining the unique benefits of a gradient-based approach to this task”
By gradient-based, we mostly meant that RIMs are learned end-to-end through gradient descent, so it can be applied very generally.

Meta-review:
role: PC
data: comment: 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)