\part{Appendix}
\parttoc

\section{Desiderata for Recurrent Independent Mechanisms}
\label{sec:desiderata}

\iffalse
\begin{figure*}[t]
\centering
\begin{tikzpicture}[
    scale=1.0,
    black!50, text=black,
    font=\medium,

    node distance=1mm,
    dnode/.style={
        align=center,

        font=\small,
        rectangle,minimum size=15mm,rounded corners,

        inner sep=20pt},
    rnode/.style={
        align=center,
        font=\medium,

        rectangle,
        minimum width=5mm,
        minimum height=6mm,
        rounded corners,

        inner sep=3pt,
        thin, draw},
    tuplenode/.style={
        align=center,

        rectangle,minimum size=10mm,rounded corners,

        inner sep=15pt},
    darrow/.style={
        rounded corners,-latex,shorten <=5pt,shorten >=1pt,line width=2mm},
    mega thick/.style={line width=2pt}
    ]
\matrix[row sep=7.2mm, column sep=2mm] {
    \node (a_d) [rnode, minimum width=15mm, top color=greenfill!40, bottom color = greenfill!40] {Surfing RIM};
    &&
    &&
    \node (a_1) [rnode, minimum width=15mm, top color=greenfill!40, bottom color = greenfill!40] {$h_{t+2,2}$};
    &&
    &&
    \node (a_2) [rnode, minimum width=15mm, top color=greenfill!40, bottom color = greenfill!40] {$h_{t+4,2}$};
    \\
    \node (b_d) [rnode, minimum width=15mm, top color=redfill!40, bottom color = redfill!40] {Studying RIM};
    &&
    \node (b_1) [rnode, minimum width=15mm, top color=redfill!40, bottom color = redfill!40] {$h_{t+1,1}$};
    &&
    &&
    \node (b_2) [rnode, minimum width=15mm, top color=redfill!40, bottom color = redfill!40] {$h_{t+3,1}$};
    &&
    &&
    \node (b_3) [rnode, minimum width=15mm, top color=redfill!40, bottom color = redfill!40] {$h_{t+5,1}$};
    \\
    &&
    \node (x_1) [rnode, minimum width=15mm, top color=black!10, bottom color = black!10] {\makecell{Study\\for Test}};
    &&
    \node (x_2) [rnode, minimum width=15mm, top color=black!10, bottom color = black!10] {\makecell{Prepare\\for Surfing}};
    &&
    \node (x_3) [rnode, minimum width=15mm,top color=black!10, bottom color = black!10] {\makecell{Think about\\mathcal{D_{\mathrm{test}}} before wave}};
    &&
    \node (x_4) [rnode, minimum width=15mm, top color=black!10, bottom color = black!10] {\makecell{Start\\Surfing}};
    &&
    \node (x_5) [rnode, minimum width=15mm, top color=black!10, bottom color = black!10] {\makecell{Take\\the test}};
    \\
};

\draw[-latex,shorten <=1pt,shorten >=1pt,thick,color=blue!60]  (a_1) to node[]{} (a_2);
\draw[-latex,shorten <=1pt,shorten >=1pt,thick,color=red!60]  (b_1) to node[]{} (b_2);
\draw[-latex,shorten <=1pt,shorten >=1pt,thick,color=red!60]  (b_2) to node[]{} (b_3);
\draw[-latex,shorten <=1pt,shorten >=1pt,thick,color=black!30]  (x_1) to node[]{} (b_1);
\draw[-latex,shorten <=1pt,shorten >=1pt,thick,color=black!30]  (x_2) to node[]{} (a_1);
\draw[-latex,shorten <=1pt,shorten >=1pt,thick,color=black!30]  (x_3) to node[]{} (b_2);
\draw[-latex,shorten <=1pt,shorten >=1pt,thick,color=black!30]  (x_4) to node[]{} (a_2);
\draw[-latex,shorten <=1pt,shorten >=1pt,thick,color=black!30]  (x_5) to node[]{} (b_3);

\draw[-latex,shorten <=1pt,shorten >=1pt,thick,color=purple!60]  (a_1) to node[]{} (b_2);

\end{tikzpicture}

\caption{An illustration of the ideal behavior of the RIMs model, in which RIMs are sparsely activated and responsible for modeling different aspects of the world.  This helps prevent interference between unrelated aspects of the world.  }
\label{fig:blocks}

\end{figure*}
\fi

We have laid out a case for building models composed of modules which by default operate independently and can interact in a 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 as \textit{Recurrent Independent Mechanisms (RIMs)}, with distinct functions learned automatically from data.
Our model encourages sparse interaction, i.e.,  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.

\paragraph{Competitive Mechanisms:}
Inspired by the observations in the main paper, we propose that RIMs utilize competition to allocate representational and computational resources. As argued by \citep{ParKilRojSch18}, this tends to produce independence among learned mechanisms if the training data has been generated by independent physical mechanisms.

\paragraph{Top Down Attention:}

The points mentioned in Section~\ref{sec:rimsmodel} in principle pertain to synthetic and natural intelligent systems alike. Hence, it is not surprising that they also appear in neuroscience. For instance, suppose we are looking for a particular object in a large scene, using limited processing capacity.
The \emph{biased competition theory} of selective attention conceptualizes basic findings of experimental psychology and neuroscience \citep{Desimone1995}:
our capacity of parallel processing of and reasoning with high-level concepts is limited, and many brain systems representing visual information use competition to allocate resources.  Competitive interactions among multiple objects  occur automatically and operate in parallel across the visual field. Second, the principle of selectivity amounts to the idea that a perceiver has the ability to filter out unwanted information and selectively \emph{process} the rest of the information. Third,  \emph{top-down bias} originating from higher brain areas enables us to selectively devote resources to input information that may be of particular interest or relevance. This may be accomplished by units matching the internal model of an object or process of interest being pre-activated and thus gaining an advantage during the competition of brain mechanisms.

\paragraph{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 or distracting information, a 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 to out-of-distribution or distractor observations than training one big homogeneous neural network \citep{schmidhuber2018one}.

\paragraph{Modular Computation Flow and Modular Parameterization:} Each RIM should have its own dynamics operating by \textit{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 to \citep{ParKilRojSch18}.
This can speed up computation and improve the generalisation ability of the system \citep{baum1989size}. The individuals RIMs in the ensemble should be simple also to prevent individual RIMs from dominating and modelling complex, composite mechanisms. We refer to a parameterization as modular if most parameters are associated to individuals RIMs only. This has the desirable property that a RIM should maintain its own independent functionality even as other RIMs are changed (due to its behavior being determined by its own self-contained parameters).

\section{Extended Related Work}
\label{sec:extended_related_work}
\begin{table*}[bht]
{\renewcommand{\arraystretch}{1.3}
\caption{A concise comparison of recurrent models with modular memory.  }
\label{tb:related}
\scalebox{0.8}{
\begin{tabular}{|l|l|l|l|l|}
\hline
Method / Property         & \makecell{Modular\\Memory} & \makecell{Sparse \\Information Flow} & \makecell{Modular \\Computation Flow}     & \makecell{Modular\\Parameterization} \\ \hline
LSTM / RNN & \xmark & \xmark & \xmark & \xmark \\ \hline
Relational RNN \citep{santoro2018relational}  & \cmark   & \xmark & \cmark & \xmark \\ \hline
NTM \citep{graves2014ntm} & \cmark & \cmark & \xmark & \xmark \\ \hline
SAB\citep{ke2018sparse}  & \xmark & \cmark  & \xmark & \xmark  \\ \hline
IndRNN\citep{li2018independently} & \cmark & \xmark & \xmark & \cmark  \\ \hline
RIMs    & \cmark & \cmark  & \cmark & \cmark \\ \hline
\end{tabular}
}
}
\end{table*}

The present section provides further details on related work, thus extending Section~\ref{sec:related_work}.

\textbf{Neural Turing Machine (NTM)}. The NTM \citep{graves2014neural} has a Turing machine inspired memory with a sequence of independent memory cells, and uses an attention mechanism to move heads over the cells while performing targeted read and write operations.  This shares a key idea with RIMs: that input information should only impact a 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.

\textbf{Relational RNN}.  The Relational Models paper \citep{santoro2018relational} is based on the idea of using a multi-head attention mechanism to share information between multiple parts of memory.  It is related to our idea but a key difference is that we encourage the RIMs to remain separate as much as possible, whereas \citep{santoro2018relational} allows information between the parts to flow on each step (in effect making the part distribution only relevant to a particular step).  Additionally, RIMs has the notion of each RIM having its own independent transition dynamics which operate by default, whereas the Relational RNN only does computation and updating of the memory using attention.

\textbf{Sparse Attentive Backtracking (SAB)}.  The SAB architecture \citep{ke2018sparse} 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 RIMs model in that it considers separation between time steps rather than separation between RIMs

\textbf{Independently Recurrent Neural Network (IndRNN)}.  The IndRNN \citep{li2018independently} replaces the full transition matrix in a 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 an RNN to learn long-term dependencies without vanishing or exploding gradients. Analysis of the gradients shows that having smaller recurrent transition matrices mitigates the vanishing and exploding gradient issue. This may provide further explanation for why RIMs perform well on long sequences.

\textbf{Consciousness Prior} \citep{bengio2017consciousness}: 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 a subset of high-level variables to interact together at any particular time during inference.  This  motivates in RIMs the use of modules (corresponding to 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.e., heads, involved
for each). Our paper thus helps to validate the inductive bias of the consciousness prior.

\textbf{Recurrent Entity Networks}:   EnTNet \citep{henaff2016tracking} can be viewed as a set of separate recurrent models whose hidden states store the memory
slots. These hidden states are either fixed by the gates, or modified through a simple RNN-style
update. Moreover,  EntNet uses an independent gate for writing to each memory slot. Our work is related in the sense that we also have different recurrent models (i.e.,RIMs, though each RIM has different parameters), but we allow the RIMs to communicate with each other sparingly using an attention mechanism.

\textbf{Capsules and Dynamic Routing:} EM Capsules \citep{hinton2018matrix} and the preceding Dynamic Capsules \citep{sabour2017dynamic} use the poses of parts and learned part $\rightarrow$ 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.

\textbf{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 \citep{scarselli2008graph, bronstein2017geometric, watters2017visual, raposo2017discovering,santoro2017simple, gilmer2017neural, van2018relational, kipf2018neural, battaglia2018relational, tacchetti2018relational}. One can also view our proposed model as a relational graph neural network, where nodes are parameterized as individual RIM and edges are parameterized by the attention mechanism. Though, it 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 a meaning directly tied to the inputs and targets and the topology is fixed and given.

\textbf{Default Behaviour:} Our work is also related to work in behavioural research that deals with two modes of decision making \citep{Dickinson67,  botvinick2015motivation, kool2018mental}: an automatic system that relies on habits and a controlled system that uses some privileged 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 or habit system. There is RL work aiming to learn \textit{default policies}, which have shown to improve transfer and generalization in multi-task RL \citep{teh2017distral, goyal2019infobot}. RIMs are different in the sense that we are not trying to learn \textit{default policies} which affect the environment, instead we want to learn mechanisms, which are more about analyzing or understanding the environment. State-dependent activation of different primitive policies was also studied by \citet{goyal2019reinforcement}, and the authors showed that they can learn different primitives, but they also consider that only a single primitive can be active at a particular time-step. Also, note that primitive policies try to \textit{affect} the environment, whereas mechanisms try to \textit{process or understand} the environment.

\section{Model Setup Details and Hyperparameters}
\label{appendix:detailedexp}

\subsection{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 \citep{santoro2018relational} but with a 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 a 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 a mask which zeroes out the outputs from inactive RIMs. We compute the independent dynamics over all RIMs by using a separate LSTM for each RIM.  Following this, we compute the communication between RIMs as a multihead attention \citep{santoro2018relational}, with the earlier-discussed modification of having separate weight parameters for each block, and also that we added a skip-connection around the attention mechanism.  This attention mechanism used 4 heads and in general used a key size and value size of 32. We computed the updates for all RIMs but used the activated-block mask to selectively update only the activated subset of the RIMs.

The use of RIMs introduces two additional hyperparameters over an LSTM/GRU: the number of RIMs and the number of activated RIMs per step.  We also observed that having too few activated RIMs tends to 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.

\textbf{Multiple Heads:} Analogously to \citet{vaswani2017attention, santoro2018relational}, we use multiple heads both for communication between RIMs as well as input attention (as in Sec \ref{sec:selectiveactivationrims}) 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.

\textbf{Selective Activation.} In order to selectively decide which set of RIMs to activate, we also tried appending a vector of zeros to the input representation, and the RIMs, which pay least attention to the null input (in the input attention) are activated.

\subsection{Detailed Model Hyperparameters}

Table \ref{table::appendix::minigrid::hyperparameters} lists the different hyperparameters.
\begin{table}[tbh]
    \caption{Hyperparameters}
  \begin{center}
    \begin{tabular}{lr}
      \toprule
      Parameter & Value  \\
      \midrule

      Optimizer                                   & Adam\citep{Kingma2014}\\
      learning rate                                   & $7\cdot 10^{-4}$      \\
      batch size                                      & 64      \\
     Input keys &  64 \\
     Input Values & Size of individual RIM * 4 \\
     Input Heads & 4 \\
     Input Dropout & 0.1 \\
     Communication keys &  32 \\
     Communication Values & 32 \\
     Communication heads & 4 \\
     Communication Dropout & 0.1 \\
      \bottomrule
    \end{tabular}
  \end{center}
  \label{table::appendix::minigrid::hyperparameters}
\end{table}

\subsection{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 we have considered in the paper. We briefly mention a few architectural changes which we have studied:
\begin{itemize}
    \item On the output side, we concatenate the representations of the different RIMs, and use the concatenated representation for learning a policy (in RL experiments) or for predicting the input at the next time step (for bouncing balls as well as all other experiments). We empirically found that adding another layer of (multi-headed) key-value attention on the output seems to improve the results. We have not included this change in the RIMs implementation of this paper.
    \item In our experiments, we shared the same decoder for all the RIMs, i.e., we concatenate the representations of different RIMS, and feed the concatenated representations to the decoder. In the future it would be interesting to think of ways to allow a more ``structured'' decoder. The reason for this is that even if the RIMs generalize to new environments, the shared decoder can fail to do so. So changing the structure of decoder could be helpful.
    \item 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  \emph{efference copies} as proposed by \cite{von1867handbuch, vonHolst1950}, i.e., using copies of motor signals as inputs. Preliminary experiments shows that this improves the performance in Atari games.
\end{itemize}

\section{Experiment Details}

\subsection{Effect of Varying Number and Active RIMs on Copying Task}
\label{table:copying_ablation}
We used a learning rate of 0.001 with the Adam Optimizer and trained each model for 150 epochs (unless the model was stuck, we found that this was enough to bring the training error close to zero).  For the RIMs model we used 600 units split across 6 RIMs (100 units per block).  For the LSTM we used a total of 600 units.  We did not explore this extensively but we qualitatively found that the results on copying were not very sensitive to the exact number of units.

The sequences to be copied first have 10 random digits (from 0-8), then a span of zeros of some length, followed by a special indicator ``9'' in the input which instructs the model to begin outputting the copied sequence.

In our experiments, we trained the models with ``zero spans'' of length 50 and evaluated on 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. \ref{tb:copy_ablation} shows the effect of 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. \ref{tb:copy_ablation} shows that the proposed method is better than both of these baselines. For copy task, we used 1 head in input attention, and 4 heads for RIMs communication. We note that even with 1 RIM, its not exactly same as a LSTM, because each RIM can still reference itself.

\begin{table*}[htb!]
\centering
{\renewcommand{\arraystretch}{1.0}
\caption{\textbf{Error (CE for last 10 time steps) on the copying task}.  Note that while all of the methods are able to 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.  }
\label{tb:copy_ablation}
\centering
\scalebox{0.9}{
\begin{tabular}{lrrr}
\toprule
Approach & Train Length 50 & Test Length 200  \\
\midrule

\textbf{RIMs} & \textbf{0.00} & \textbf{0.00} \\

\midrule
With input Attention and No Communication & & \\
\midrule
\textbf{RIMs} ($k_T$ = 4, $k_A$ = 2, $h_{dim}$ = 600) & 2.3 & 1.6  \\
\textbf{RIMs} ($k_T$ = 4, $k_A$ = 3, $h_{dim}$ = 600) & 1.7 & 4.3  \\

\textbf{RIMs} ($k_T$ = 5, $k_A$ = 2, $h_{dim}$ = 600) & 2.5 & 4,7  \\
\textbf{RIMs} ($k_T$ = 5, $k_A$ = 3, $h_{dim}$ = 600) & 0.4 & 4.0\\
\textbf{RIMs} ($k_T$ = 5, $k_A$ = 4, $h_{dim}$ = 600) & 0.2 & 0.7 \\

\textbf{RIMs} ($k_T$ = 6, $k_A$ = 2, $h_{dim}$ = 600) & 3.3 & 2.4\\
\textbf{RIMs} ($k_T$ = 6, $k_A$ = 3, $h_{dim}$ = 600) & 1.2 & 1.0  \\
\textbf{RIMs} ($k_T$ = 6, $k_A$ = 4, $h_{dim}$ = 600) & 0.7 & 5.0\\
\textbf{RIMs}
($k_T$ = 6, $k_A$ = 5, $h_{dim}$ = 600) & 0.22 & 0.56 \\
\midrule
With No input Attention and Full Communication & & \\
\midrule
\textbf{RIMs} ($k_T$ = 6, $k_A$ = 6, $h_{dim}$ = 600) & 0.0 & 0.7  \\
\textbf{RIMs} ($k_T$ = 5, $k_A$ = 5, $h_{dim}$ = 500) & 0.0 & 1.7  \\
\textbf{RIMs} ($k_T$ = 2, $k_A$ = 2, $h_{dim}$ = 256) & 0.0 & 2.9  \\
\textbf{RIMs} ($k_T$ = 2, $k_A$ = 2, $h_{dim}$ = 512) & 0.0 & 1.8  \\
\textbf{RIMs} ($k_T$ = 1, $k_A$ = 1, $h_{dim}$ = 512) & 0.0 & 0.2  \\
\midrule
With input Attention and Full Communication & & \\
\midrule
\textbf{RIMs} ($k_T$ = 9, $k_A$ = 2, $h_{dim}$ = 900) & 0.24 & 0.15  \\
\textbf{RIMs} ($k_T$ = 9, $k_A$ = 3, $h_{dim}$ = 900) & 0.01 & 0.00  \\
\textbf{RIMs} ($k_T$ = 9, $k_A$ = 4, $h_{dim}$ = 900) & 0.01 & 0.00  \\
\textbf{RIMs} ($k_T$ = 9, $k_A$ = 5, $h_{dim}$ = 900) & 0.01 & 0.00  \\
\textbf{RIMs} ($k_T$ = 9, $k_A$ = 6, $h_{dim}$ = 900) & 0.01 & 0.01  \\
\textbf{RIMs} ($k_T$ = 9, $k_A$ = 7, $h_{dim}$ = 900) & 0.04 & 0.07 \\
\textbf{RIMs} ($k_T$ = 9, $k_A$ = 8, $h_{dim}$ = 900) & 0.00 & 0.10 \\
\textbf{RIMs} ($k_T$ = 9, $k_A$ = 9, $h_{dim}$ = 900) & 0.03 & 0.24 \\
\textbf{RIMs} ($k_T$ = 16, $k_A$ = 4, $h_{dim}$ = 1600) & 0.05 & 0.02 \\
\textbf{RIMs} ($k_T$ = 16, $k_A$ = 6, $h_{dim}$ = 1600) & 0.00 & 0.00 \\
\textbf{RIMs} ($k_T$ = 16, $k_A$ = 8, $h_{dim}$ = 1600) & 0.00 & 0.00 \\
\textbf{RIMs} ($k_T$ = 16, $k_A$ = 10, $h_{dim}$ = 1600) & 0.00 & 0.00 \\
\textbf{RIMs} ($k_T$ = 16, $k_A$ = 12, $h_{dim}$ = 1600) & 0.00 & 0.00 \\
\textbf{RIMs} ($k_T$ = 16, $k_A$ = 14, $h_{dim}$ = 1600) & 0.01 & 0.26 \\
\textbf{RIMs} ($k_T$ = 24, $k_A$ = 6, $h_{dim}$ = 2400) & 0.01 & 0.00 \\
\textbf{RIMs} ($k_T$ = 24, $k_A$ = 8, $h_{dim}$ = 2400) & 0.00 & 0.00 \\
\textbf{RIMs} ($k_T$ = 24, $k_A$ = 16, $h_{dim}$ = 2400) & 0.00 & 0.00 \\
\textbf{RIMs} ($k_T$ = 24, $k_A$ = 20, $h_{dim}$ = 2400) & 0.00 & 0.00 \\
\textbf{RIMs} ($k_T$ = 24, $k_A$ = 22, $h_{dim}$ = 2400) & 0.00 & 0.42\\
\bottomrule
\end{tabular}
}
}
\end{table*}

\newpage
\subsection{Effect of Varying Number and Active RIMs on Adding Task}
\label{table:adding_ablation}

In the adding task we consider a stream of numbers as inputs (given as real-values) and then indicate which two numbers should be added together as a 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 ability to ignore the numbers which it is not tasked with adding together.  We provide the results in Table \ref{tb:adding_ablation} which demonstrates that proposed model generalize better for longer testing sequences as well as adding multiple numbers. We ran the proposed model with different configurations like changing number of RIMs as well as number of active RIMs.

\begin{table*}[htb!]
\centering
{\renewcommand{\arraystretch}{1.0}
\caption{\textbf{Error CE on the adding task}.   }
\label{tb:adding_ablation}
\centering
\scalebox{0.9}{
\begin{tabular}{lrrr}
\toprule
Approach & Train Length 50 & Test Length 200  \\
\midrule
With input Attention and Full Communication & & \\
\midrule
\textbf{RIMs} ($k_T$ = 6, $k_A$ = 3, $h_{dim}$ = 300) & 0.0 & 0.0\\
\textbf{RIMs} ($k_T$ = 6, $k_A$ = 4, $h_{dim}$ = 300) & 0.0 & 0.0\\
\textbf{RIMs} ($k_T$ = 6, $k_A$ = 5, $h_{dim}$ = 300) & 0.0 & 0.0\\
\textbf{RIMs} ($k_T$ = 6, $k_A$ = 3, $h_{dim}$ = 600) & 0.0 & 0.0\\
\textbf{RIMs} ($k_T$ = 6, $k_A$ = 4, $h_{dim}$ = 600) & 0.0 & 0.0\\
\textbf{RIMs} ($k_T$ = 6, $k_A$ = 5, $h_{dim}$ = 600) & 0.0 & 0.0\\
\textbf{RIMs} ($k_T$ = 9, $k_A$ = 3, $h_{dim}$ = 900) & 0.0 & 0.0\\
\textbf{RIMs} ($k_T$ = 9, $k_A$ = 4, $h_{dim}$ = 900) & 0.0 & 0.0\\
\textbf{RIMs} ($k_T$ = 9, $k_A$ = 5, $h_{dim}$ = 900) & 0.0 & 0.0\\
\textbf{RIMs} ($k_T$ = 9, $k_A$ = 6, $h_{dim}$ = 900) & 0.0 & 0.0\\
\textbf{RIMs} ($k_T$ = 9, $k_A$ = 7, $h_{dim}$ = 900) & 0.0 & 0.0\\
\textbf{RIMs} ($k_T$ = 9, $k_A$ = 8, $h_{dim}$ = 900) & 0.0 & 0.0\\
\toprule
Approach & Train Length 500 & Test Length 1000  \\
\midrule
With input Attention and Full Communication & & \\
\midrule
\textbf{RIMs} ($k_T$ = 6, $k_A$ = 3, $h_{dim}$ = 600) & 0.0 & 0.0\\
\textbf{RIMs} ($k_T$ = 6, $k_A$ = 4, $h_{dim}$ = 600) & 0.0 & 0.0\\
\textbf{RIMs} ($k_T$ = 6, $k_A$ = 5, $h_{dim}$ = 600) & 0.0 & 0.0\\
\textbf{RIMs} ($k_T$ = 9, $k_A$ = 3, $h_{dim}$ = 900) & 0.0 & 0.0\\
\textbf{RIMs} ($k_T$ = 9, $k_A$ = 4, $h_{dim}$ = 900) & 0.0 & 0.0\\
\textbf{RIMs} ($k_T$ = 9, $k_A$ = 5, $h_{dim}$ = 900) & 0.0 & 0.0\\
\textbf{RIMs} ($k_T$ = 9, $k_A$ = 6, $h_{dim}$ = 900) & 0.0 & 0.0\\
\textbf{RIMs} ($k_T$ = 9, $k_A$ = 7, $h_{dim}$ = 900) & 0.0 & 0.0\\
\bottomrule
\end{tabular}
}
}
\end{table*}

\subsection{Sequential MNIST Resolution Task}
\label{sec:mnist_ablation}

In 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 a resolution of 14 x 14 (sequence length of 196).  We then evaluated on MNIST digits of different higher resolutions (16 x 16, 19 x 19, and 24 x 24).  When re-scaling the images, we used the nearest-neighbor based down-scaling and performed binarization after re-scaling.  We trained with a learning rate of 0.0001 and the Adam optimizer.  For RIMs we used a total of 600 hidden units split across 6 RIMs (100 units per 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 4 heads for RIMs communication.

\subsection{Bouncing Ball Environment}
We use the bouncing-ball dataset from \citep{van2018relational}. The dataset consists of 50,000 training examples and 10,000 test examples showing $\sim$50 frames of either 4 solid balls bouncing in a confined square geometry, 6-8 balls bouncing in a confined geometry, or 3 balls bouncing in a confined geometry with a random occluded region.
In all cases, the balls bounce off the wall as well as off one another.  We train baselines as well as proposed model for about 100 epochs using 0.0007 as learning rate and using Adam as optimizer \citep{Kingma2014}. We use the same architecture for encoder as well as decoder as in \citep{van2018relational}. 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, we highlight a few different results.

\subsubsection{Different RIMs attend to different balls}
\label{appendix:diffballs}

\begin{figure}[htb!]
    \centering

    \includegraphics[width=0.95\linewidth]{figures/geom2_jh.png}
    \caption{\textbf{Different RIMs attending to Different Balls}. For understanding what each RIM is actually doing, we associate each with a separate encoder, which are spatially masked.  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. 1.) Early in training, RIM activations correlated more strongly with the locations of the four different balls. Later in training, this correlation decreased and the active strips did not correlate as strongly with the location of balls. As the model got better at predicting the location, it needed to attend less 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.}
    \label{fig:geometric}
\end{figure}
In order to visualize what each RIM is doing, we associate each RIM with a different encoder.
By performing spatial masking on the input, we can control the possible spatial input to each RIM.
We use six non-overlapping horizontal strips and allow only 4 RIMs to be active at a time (shown in Fig.~\ref{fig:geometric}). The mask is fixed mask of zeros with a band of ones that is multiplied by the input to each encoder. Therefore, each of the 6 encoders gets 1/6th of the input. The goal was to see how the RIM activation patterns changed/correlated with the locations of the balls.
We find that early in training, the RIMs' activations are strongly correlated with the location of the 4 balls. However, after training has proceeded for some time this correlation deteriorates.
This is likely because the predictable dynamics of the system do not necessitate constant attention.

\subsubsection{Comparison with LSTM Baselines}
In Figures \ref{fig:lstm2}, \ref{fig:comparison}, \ref{fig:lstm}, and \ref{fig:err2} we highlight different baselines and how these compare to the proposed RIMs model.
\begin{figure}[htb!]
    \centering
    \includegraphics[width=0.95\linewidth]{figures/other.png}
    \caption{\textbf{Example of the other LSTM baselines}. For the 2 other experiments that we consider, here we show example outputs of our LSTM baselines. In each row, the top panel represents the ground truth and the bottom represents the prediction. All shown examples use an LSTM with 250 hidden units, as shown in Fig.~\ref{fig:err}. Frames are plotted every 3rd time step. The red line marks 10 rollout frames. This is marked because after this we  do not find BCE to be a reliable measure of dissimilarity.}
    \label{fig:lstm2}
\end{figure}
\vspace{-1mm}

\begin{figure}[htb!]
    \centering
    \includegraphics[width=0.95\linewidth]{figures/all_baseline.png}
    \caption{\textbf{Comparison of RIMs to LSTM baseline}. For 4 different experiments in the text, we compare RIMs to two different LSTM baselines. In all cases we find that during rollout, RIMs perform better than the LSTMs at accurately capturing the trajectories of the balls through time. Due 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 15 time steps, computing the binary cross entropy between the prediction and the true balls at each instant, as in \cite{van2018relational}.  In the predictions, the transparent blue shows the ground  truth, overlaid to help guide the eye.}
    \label{fig:comparison}
    \vspace{-1mm}
\end{figure}

\begin{figure}[htb!]
    \centering
    \includegraphics[width=0.95\linewidth]{figures/compare.png}
    \caption{\textbf{Comparison between RIMs and LSTM baseline}. For the 4 ball task and the 6-8 ball extrapolation task, here we show an example output of from our LSTM baseline and from RIMs. All shown examples use an LSTM with 250 hidden units, as shown in Fig.~\ref{fig:err}. Frames are plotted every 3rd time step. The red line marks 10 rollout frames. This is marked because after this we  do not find BCE to be a reliable measure of dissimilarity.}
    \label{fig:lstm}
\end{figure}
\begin{figure}[htb!]
    \centering
    \includegraphics[width=0.95\linewidth]{figures/err_all_2.png}
    \caption{\textbf{Comparison of RIMs to LSTM baseline}. For 4 different experiments in the text, we compare RIMs to two different LSTM baselines. In all cases we find that during rollout, RIMs perform better than the LSTMs at accurately capturing the trajectories of the balls through time. Due to the number of hard collisions, accurate modeling is very difficult.}
    \label{fig:err2}
\end{figure}

\subsubsection{Occlusion}
In Fig.~\ref{fig:curtain}, we show the performance of RIMs on the curtain dataset. We find RIMs are able to track balls through the occlusion without difficulty. Note that the LSTM baseline, is also able to track the ball through the ``invisible'' curtain.
\begin{figure}[htb!]
    \centering

    \includegraphics[width=0.95\linewidth]{figures/curtain.png}
    \caption{\textbf{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 as input.}
    \label{fig:curtain}
\end{figure}

\subsubsection{Study of Transfer}
It is interesting to ask how models trained on a dataset with 6-8 balls perform on a dataset with 4 balls. In Fig.~\ref{fig:interpolation} we show predictions during feed-in and rollout phases.
\begin{figure}[htb!]
    \centering

    \includegraphics[width=0.95\linewidth]{figures/interpolation.png}
    \caption{\textbf{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.

    }
    \label{fig:interpolation}
\end{figure}

\subsection{Video Prediction from Crops} \label{sec:bballs_appendix}

\begin{figure}
\centering
\includegraphics[width=0.9\textwidth]{video_crop/bbals.png}
\caption{Performance metrics on OOD one-step forward prediction task. \textbf{Gist:} RIMs outperforms all RNN baselines OOD.}
\label{fig:bb_ood}
\vspace{-10pt}
\end{figure}

\textbf{Dataset.} We train all models on a training dataset of $20$K video sequences with $100$ frames of $3$ balls bouncing in an arena of size $48 \times 48$. We also include an additional fixed ball in the center to make the task more challenging. We use another $1$K video sequences of the same length and the same number of balls as a held-out validation set. In addition, we also have $5$ out-of-distribution (OOD) test sets with various number of bouncing balls (ranging from $1$ to $6$) and each containing $1$K sequences of length $100$.

\textbf{Training.} We train all models until the validation loss is saturated, and select the best of three runs. During training, we automatically decay the learning rate by a factor of $2$ if the validation loss does not significantly decrease by at least $0.01\%$ for five consecutive epochs.

\textbf{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~\ref{fig:bb_ood}, we report the balanced accuracy (i.e. arithmetic mean of recall and specificity) and F1-scores (i.e. harmonic mean of precision and recall) to account for class-imbalance

\textbf{Results.} In Figure~\ref{fig:bb_ood}, we see that proposed method out-perform all non-oracle baselines OOD on the one-step forward prediction task and strike a good balance in regard to in-distribution and OOD performance.

\textbf{Baselines:} We compared the performance of the proposed method with a baseline LSTM model, as well as state of the art memory model, RMC \citep{santoro2018relational}. As a 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).

\subsection{Environment with Novel Distractors}

We evaluate the proposed framework using Adavantage Actor-Critic (A2C)  to learn a policy $\pi_\theta(a | s, g)$ conditioned on the goal. To evaluate the performance of proposed method, we use a range of maze multi-room tasks from the gym-minigrid framework \citep{gym_minigrid} and the A2C implementation from \citep{gym_minigrid}. For the maze tasks, we used agent's relative distance to the absolute goal position as "goal".

For the maze environments, we use A2C with 48 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 a neural network, which consists of 1 fully connected layer. We use RMSProp with an initial learning rate of $0.0007$ to train the models. Due to the partially observable nature of the environment, we further use a LSTM to encode the state and summarize the past observations.

\subsection{MiniGrid Environments for OpenAI Gym\label{sec:minigrid}}

The  MultiRoom environments used for this research are part of MiniGrid, which is an open source gridworld package\footnote{https://github.com/maximecb/gym-minigrid}. This package includes a family of reinforcement learning environments compatible with the OpenAI Gym framework. Many of these environments are parameterizable so that the difficulty of tasks can be adjusted (e.g., the size of rooms is often adjustable).

\subsubsection{The World}

In MiniGrid, the world is a grid of size NxN. Each tile in the grid contains exactly zero or one object. The possible object types
are wall, door, key, ball, box and goal. Each object has an associated discrete color, which can be one of red, green, blue, purple, yellow and grey. By default, walls are always grey and goal squares are always green.

\begin{figure}
    \centering
    \includegraphics[width=0.6\linewidth]{figures/Pickup.png}
    \caption{An example of the minigrid task.  }
    \label{fig:minigrid_pick}
\end{figure}

\subsubsection{Reward Function}

Rewards are sparse for all MiniGrid environments. In the MultiRoom environment, episodes are terminated with a positive reward when the agent reaches the green goal square. Otherwise, episodes are terminated with zero reward when a time step limit is reached. In the FindObj environment, the agent receives a positive reward if it reaches the object to be found, otherwise zero reward if the time step limit is reached.

The formula for calculating positive sparse rewards is $1 - 0.9 * (step\_count / max\_steps)$. That is, rewards are always between zero and one, and the quicker the agent can successfully complete an episode, the closer to $1$ the reward will be. The $max\_steps$ parameter is different for each environment, and varies depending on the size of each environment, with larger environments having a higher time step limit.

\subsubsection{Action Space}

There are seven actions in MiniGrid:  turn left, turn right, move forward, pick up an object, drop an object, toggle and done. For the purpose of this paper, the pick up, drop and done actions are irrelevant. The agent can use the turn left and turn right action to rotate and face one of 4 possible directions (north, south, east, west). The move 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.

\subsubsection{Observation Space}

Observations in MiniGrid are partial and egocentric. By default, the agent sees a square of 7x7 tiles in the direction it is facing. These
include the tile the agent is standing on. The agent cannot see through walls or closed doors. The observations are provided as a tensor of shape 7x7x3. However, note that these are  RGB images (which is different from original BabyAI paper).

\subsubsection{Level Generation}
\label{sec: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 at a random location in the first room. (3) Add the goal at a random location in the last room.
A neural network parameterized as CNN is used to process the visual observation.

We follow the same architecture as
\citep{gym_minigrid} but we replace the LSTM layer with BlockLSTM.

\subsubsection{Additional Ablation}
\label{sec:bouncing_balls_ablation_1head}
We present one ablation in addition to the ones in Section~\ref{sec:main_ablation}. In this experiment, we  study the effect on input attention (i.e top down attention) as well as the use of multi-headed head key-value attention. We compare the proposed model (with input attention as well as multi-headed key-value attention) with 2 baselines: (a) In which we remove the input attention (and force all the RIMs to communicate with each other (b) We use 1 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.~\ref{fig:ablate2}.

\begin{figure}[htb!]
    \centering

    \includegraphics[width=0.5\linewidth]{figures/ablate_histogram.png}
    \caption{\textbf{Ablation loss} For the normal, a 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.}
    \label{fig:ablate2}
\end{figure}
In Fig.~\ref{fig:onehead}, we show the predictions that result from the model with only one active head.
\begin{figure}[htb!]
    \centering

    \includegraphics[width=0.95\linewidth]{figures/one-head.png}
    \caption{\textbf{One head and no attention} Using one head and no attention models, we show the rollout predictions in blue. On top we show results on the 4 ball dataset and on the bottom we show results on the curtains dataset.}
    \label{fig:onehead}
\end{figure}

\subsection{Atari}
\label{appendix:atariexp}

We used open-source implementation of PPO from \citep{pytorchrl} with default parameters. We ran the proposed algorihtm with 6 RIMs, and kept the number of activated RIMs to 4/5. We have not done any hyper-parameter search for Atari experiments.

\begin{figure}[H]
    \centering
    \includegraphics[width=0.8\linewidth]{atari_results/rim5over_rim4.png}
    \caption{A comparison showing relative improvement of RIMs with $k_A=5$ over a $k_A=4$ baseline.  Using $k_A=5$ performs slightly worse than $k_A=4$ but still outperforms PPO, and has similar results across the majority of games.  }
    \label{fig:atari5vs4}
\end{figure}

\vspace{-1mm}
\begin{figure}[H]
    \centering
    \includegraphics[width=0.8\linewidth]{figures/rims5.png}
    \caption{RIMs-PPO relative score improvement over LSTM-PPO baseline \citep{schulman2017proximal} 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 $k_A=5$).}
    \label{fig:my_label}
\end{figure}

\section{Additional Experiments}

\iffalse
\subsection{Language Modeling}

\begin{table}[htb!]
\caption{Wikitext-2 results}
\label{tb:lm}
\centering
\begin{tabular}{lrrrr}
\toprule
Approach & Num. Parameters  & Train PPL & Valid PPL & Test PPL  \\
\midrule
LSTM (2-layer) & 21.2M & 39.78 & 109.25 & 102.53  \\
Relational Memory \citep{santoro2018relational} & 11M & n/a & 112.77 & 107.21 \\
\textbf{RIMs} (2-layer, $k_T$ = 6, $k_A$ = 6) & 23.7M & 41.27 & 103.60 & \textbf{98.66}  \\
\bottomrule
\end{tabular}
\end{table}

We investigate the task of word-based language modeling.
We ran experiments on the wikitext-2 dataset \citep{merity2016pointer}.  We ran each experiment for a fixed 100 epochs.  These results are in Table~\ref{tb:lm}.  Our goal in this experiment is to demonstrate the breadth of the approach by showing that RIMs performs well even on datasets which are noisy and drawn from the real-world.
\fi
\subsection{Imitation Learning: Robustness to Noise in State Distribution}
\label{sec:imitationmujoco}

Here, we consider  imitation learning where we have training trajectories generated from an expert (Table~\ref{tb:rl}). We evaluate our model on continuous control tasks in Mujoco (in our case, Half-Cheetah) \citep{todorov2012mujoco}.  We take the rendered images as input and compared the proposed model with recurrent policy (i.e., LSTM). Since, using rendered image of the input does not tell anything about the velocity of the Half-Cheetah, it makes the task partially observable.  In order to test how well the proposed model generalizes during test, we add some noise (in the joints of the half-cheetah body). As one can see, after adding noise LSTM baselines performs poorly. On the other hand, for the proposed model, there's also a drop in performance but not as bad as for the LSTM baseline.

\begin{table*}[htb!]
\caption{\textbf{Imitation Learning:} Results on the half-cheetah imitation learning task.  RIMs outperforms a baseline LSTM when we evaluate with perturbations not observed during training (left).  An example of an input image fed to the model (right).  }
\label{tb:rl}
\begin{minipage}{.5\textwidth}
{\renewcommand{\arraystretch}{1.2}
\begin{tabular}{lrr}
\hline
\toprule
\centering
Method / Setting & \makecell{Training \\Observed Reward} & \makecell{Perturbed States\\ Observed Reward} \\
\midrule
LSTM (Recurrent Policy) & 5400 $\pm$ 100                    & 2500 $\pm$ 300                          \\
\textbf{RIMs} ($k_T$ = 6, $k_A$ = 3) & 5300 $\pm$ 200   & 3800 $\pm$ 200 \\
\textbf{RIMs} ($k_T$ = 6, $k_A$ = 6) & 5500 $\pm$ 100  & 2700 $\pm$ 400 \\
\textbf{RIMs} (without Input attention) & 5400 $\pm$ 100  & 3200 $\pm$ 50 \\
\bottomrule
\end{tabular}
}
\end{minipage}%
\begin{minipage}{.4\textwidth}
\centering
\includegraphics[width=0.3\linewidth,right]{figures/cheetah.png}
\end{minipage}
\end{table*}

We use the convolutional network from \citep{ha2018world} as our encoder, a GRU \citep{chung2015recurrent} with 600
units as deterministic path in the dynamics model, and implement all other functions as two fully connected layers of size
256 with ReLU activations. Since, here we are using images as input, which makes the task, partially observable. Hence, 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_a = 3$. We follow the same setting as in \citep{hafner2018learning, sodhani2019learning}. We also compare the proposed method to the baseline where we dont include input attention (or top-down attention). AS \ref{tb:rl} shows, there's a decline in performance if we dont use input attention, hence justifying the importance

\begin{table*}[t]
\small

\centering
\begin{tabular}{ccc}
    \toprule
    Environment & LSTM-PPO & RIMs-PPO\\
    \midrule
    Alien & 1612  $\pm$ 44 & \textbf{2152} $\pm$ 81 \\
    Amidar & 1000  $\pm$ 58 & \textbf{1800} $\pm$ 43 \\
    Assault &  4000 $\pm$ 213 & \textbf{ 5400} $\pm$ 312 \\
    Asterix & 3090 $\pm$ 420 & \textbf{21040} $\pm$ 548 \\
    Asteroids & 1611.0 $\pm$ 200 & \textbf{ 3801 } $\pm$ 89 \\
    Atlantis & 3280000 $\pm$ 200000 & \textbf{3500000} $\pm$ 120000 \\
    BankHeist & 1153 $\pm$ 23 & \textbf{ 1195 } $\pm$ 4 \\
    BattleZone & 21000 $\pm$ 232.0 & \textbf{22000} $\pm$ 324 \\
    BeamRider & 698 $\pm$ 100 & \textbf{5320} $\pm$ 300 \\
    Bowling & 30 $\pm$ 5 & 42 $\pm$ 13 \\
    Boxing & 80  $\pm$ 3 & \textbf{95} $\pm$ 10 \\
    Breakout & 593 $\pm$ 90 & 590 $\pm$ 10 \\
    Centipede & 4600 $\pm$ 312 & \textbf{5534} $\pm$ 283 \\
    ChopperCommand & 11000 $\pm$ 790 & \textbf{12303} $\pm$ 412 \\
    CrazyClimber & \textbf{138000} $\pm$ 2412 & 132039 $\pm$ 1221 \\
    DemonAttack & 26320 $\pm$ 3234 & \textbf{230324} $\pm$ 4032\\
    DoubleDunk & \textbf{-3.0} $\pm$ 0.5 & -3.8  $\pm$ 0.3 \\
    Enduro & 1600 $\pm$ 200 & \textbf{2800} $\pm$ 232 \\
    FishingDerby & 20 $\pm$ 4 & \textbf{38} $\pm$ 8 \\
    Freeway & 29 $\pm$ 2 & \textbf{ 33 } $\pm$ 2 \\

    Gopher & 7000.0 $\pm$ 402 & \textbf{33000} $\pm$ 2210 \\
    Gravitar & 500 $\pm$ 100 & \textbf{ 1090 } $\pm$ 80 \\
    IceHockey & -5 $\pm$ 0.3 & -4 $\pm$ 1 \\
    Jamesbond & 425 $\pm$ 25 & \textbf{800} $\pm$ 100 \\
    Kangaroo & \textbf{ 13000 } $\pm$ 500 & 1800 $\pm$ 400 \\
    Krull & \textbf{10000} $\pm$ 500 & 7900 $\pm$ 200 \\
    KungFuMaster & 28000 $\pm$ 2000 & \textbf{51000} $\pm$ 800 \\
    NameThisGame & 4200 $\pm$ 400 & \textbf{6800} $\pm$ 300 \\
    Pong & 20 $\pm$ 1 & 20 $\pm$ 1 \\
    PrivateEye & 90 $\pm$ 3 & \textbf{ 100} $\pm$ 0 \\
    Qbert & 22000 $\pm$ 300 & \textbf{22500} $\pm$ 400 \\
    Riverraid & 7500 $\pm$ 300 & \textbf{12000} $\pm$ 100\\
    RoadRunner & 53000 $\pm$ 120 & \textbf{53430} $\pm$ 300 \\
    Robotank & 3 $\pm$ 1 & \textbf{ 11 } $\pm$ 2 \\
    SpaceInvaders & 1600 $\pm$ 40 & \textbf{ 2800} $\pm$ 80 \\
    StarGunner & 35000 $\pm$ 800 & \textbf{ 70000} $\pm$ 1200 \\
    TimePilot & 4000 $\pm$ 100 & \textbf{ 10000} $\pm$ 689 \\
    UpNDown & 70000 $\pm$ 6000 & \textbf{ 390000} $\pm$ 20000 \\
    VideoPinball & 90000 $\pm$ 5000 & \textbf{220000} $\pm$ 9000 \\
    WizardOfWor & 3833 $\pm$ 400 & \textbf{ 10800} $\pm$ 700 \\
    Zaxxon & 200 $\pm$ 100 & \textbf{ 15000} $\pm$ 600 \\
    \bottomrule
\end{tabular}
\caption{\small Scores obtained using PPO with the LSTM architecture and PPO with the RIMs architecture with $k_A = 5$.  }
\label{tbl:atari_large_quantative}
\end{table*}

\subsection{Transfer on Atari}

As a very preliminary result, we investigate feature transfer between randomly selected Atari games. In order to study this question, we follow the  experimental protocol of \cite{rusu2016progressive}.

We start by training RIMs on three source games (Pong, River Raid, and Seaquest) and test if the learned features transfer to a different subset of randomly selected target games (Alien, Asterix, Boxing, Centipede, Gopher, Hero, James Bond, Krull, Robotank, Road Runner, Star Gunner, and Wizard of Wor). We observe, that RIMs result in positive transfer in 9 out of 12 target games, with three cases of negative transfer. On the other hand progressive networks \citep{rusu2016progressive} result in positive transfer in 8 out of 12 target games, and two cases of negative transfer. We also compare to LSTM baseline, which yields positive transfer in  3 of 12 games.

\newpage

\subsubsection{Atari Results: Comparison with LSTM-PPO}

\begin{figure*}[htb!]
\begin{tabular}{ccccc}
\subfloat{\includegraphics[width = 0.95in]{atari_results/Alien.png}} &
\subfloat{\includegraphics[width = 0.95in]{atari_results/Amidar.png}} &
\subfloat{\includegraphics[width = 0.95in]{atari_results/Assault.png}} &
\subfloat{\includegraphics[width = 0.95in]{atari_results/Asteroids.png}}&
\subfloat{\includegraphics[width = 0.95in]{atari_results/Asterisk.png}}\\
\subfloat{\includegraphics[width = 0.95in]{atari_results/Atlantis.png}} &
\subfloat{\includegraphics[width = 0.95in]{atari_results/BankHeist.png}} &
\subfloat{\includegraphics[width = 0.95in]{atari_results/BattleZone.png}}&
\subfloat{\includegraphics[width = 0.95in]{atari_results/BeamRider.png}}&
\subfloat{\includegraphics[width = 0.95in]{atari_results/Bowling.png}} \\
\subfloat{\includegraphics[width = 0.95in]{atari_results/Boxing.png}} &
\subfloat{\includegraphics[width = 0.95in]{atari_results/Breakout.png}} &
\subfloat{\includegraphics[width = 0.95in]{atari_results/Centipede.png}}&
\subfloat{\includegraphics[width = 0.95in]{atari_results/Frostbite.png}} &
\subfloat{\includegraphics[width = 0.95in]{atari_results/ChopperCommand.png}} \\
\subfloat{\includegraphics[width = 0.95in]{atari_results/CrazyClimber.png}} &
\subfloat{\includegraphics[width = 0.95in]{atari_results/DemonAttack.png}}&
\subfloat{\includegraphics[width = 0.95in]{atari_results/DoubleDunk.png}} &
\subfloat{\includegraphics[width = 0.95in]{atari_results/Enduro.png}} &
\subfloat{\includegraphics[width = 0.95in]{atari_results/FishingDerby.png}} \\
\subfloat{\includegraphics[width = 0.95in]{atari_results/Freeway.png}}&
\subfloat{\includegraphics[width = 0.95in]{atari_results/Frostbite.png}}&
\subfloat{\includegraphics[width = 0.95in]{atari_results/Gopher.png}} &
\subfloat{\includegraphics[width = 0.95in]{atari_results/Gravitar.png}} &
\subfloat{\includegraphics[width = 0.95in]{atari_results/Hero.png}} \\
\subfloat{\includegraphics[width = 0.95in]{atari_results/IceHockey.png}}&
\subfloat{\includegraphics[width = 0.95in]{atari_results/james_bond.png}}&
\subfloat{\includegraphics[width = 0.95in]{atari_results/kangaroo.png}}&
\subfloat{\includegraphics[width = 0.95in]{atari_results/Krull.png}}&
\subfloat{\includegraphics[width = 0.95in]{atari_results/KungFuMaster.png}} \\
\subfloat{\includegraphics[width = 0.95in]{atari_results/MontezumaRevenge.png}} &
\subfloat{\includegraphics[width = 0.95in]{atari_results/NameThisGame.png}} &
\subfloat{\includegraphics[width = 0.95in]{atari_results/Pong.png}}&
\subfloat{\includegraphics[width = 0.95in]{atari_results/PrivateEye.png}}&
\subfloat{\includegraphics[width = 0.95in]{atari_results/Qbert.png}} \\
\subfloat{\includegraphics[width = 0.95in]{atari_results/Riverraid.png}} &
\subfloat{\includegraphics[width = 0.95in]{atari_results/RoadRunner.png}} &
\subfloat{\includegraphics[width = 0.95in]{atari_results/Robotank.png}}&
\subfloat{\includegraphics[width = 0.95in]{atari_results/SpaceInvaders.png}}&
\subfloat{\includegraphics[width = 0.95in]{atari_results/StarGunner.png}} \\
\subfloat{\includegraphics[width = 0.95in]{atari_results/Tennis.png}} &
\subfloat{\includegraphics[width = 0.95in]{atari_results/TimePilot.png}} &
\subfloat{\includegraphics[width = 0.95in]{atari_results/Tutankham.png}}&
\subfloat{\includegraphics[width = 0.95in]{atari_results/UpNDown.png}}&
\subfloat{\includegraphics[width = 0.95in]{atari_results/Venture.png}} \\
\subfloat{\includegraphics[width = 0.95in]{atari_results/VideoPinball.png}} &
\subfloat{\includegraphics[width = 0.95in]{atari_results/WizardOfWor.png}} &
\subfloat{\includegraphics[width = 0.95in]{atari_results/Zaxxon.png}}&
\end{tabular}
\caption{\textbf{Comparing RIMs-PPO with LSTM-PPO:} Learning curves for $k_A=4$, $k_A=5$ RIMs-PPO models and the LSTM-PPO baseline across all Atari games.}
\end{figure*}

\newpage

\subsubsection{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.

\begin{figure*}[htb!]
\begin{tabular}{ccccc}

\subfloat{\includegraphics[width = 1in]{atari_noinp/Amidar_noinp.png}} &
\subfloat{\includegraphics[width = 1in]{atari_noinp/Assault_noinp.png}} &
\subfloat{\includegraphics[width = 1in]{atari_noinp/asteroids_noinp.png}}&
\subfloat{\includegraphics[width = 1in]{atari_noinp/Asterix_noinp.png}}&
\subfloat{\includegraphics[width = 1in]{atari_noinp/Atlantis_noinp.png}}\\
\subfloat{\includegraphics[width = 1in]{atari_results/BankHeist.png}} &
\subfloat{\includegraphics[width = 1in]{atari_noinp/BattleZone_noinp.png}}&
\subfloat{\includegraphics[width = 1in]{atari_noinp/BeamRider_noinp.png}}&
\subfloat{\includegraphics[width = 1in]{atari_noinp/Bowling_noinp.png}} &
\subfloat{\includegraphics[width = 1in]{atari_noinp/Boxing_noinp.png}} \\
\subfloat{\includegraphics[width = 1in]{atari_noinp/Centipede_noinp.png}} &
\subfloat{\includegraphics[width = 1in]{atari_noinp/ChopperCommand_noinp.png}}&
\subfloat{\includegraphics[width = 1in]{atari_noinp/CrazyClimber_noinp.png}} &
\subfloat{\includegraphics[width = 1in]{atari_noinp/DoubleDunk_noinp.png}} &
\subfloat{\includegraphics[width = 1in]{atari_noinp/Enduro_noinp.png}} \\
\subfloat{\includegraphics[width = 1in]{atari_noinp/FishingDerby_noinp.png}}&
\subfloat{\includegraphics[width = 1in]{atari_noinp/Freeway_noinp.png}} &
\subfloat{\includegraphics[width = 1in]{atari_noinp/Frostbite_noinp.png}} &
\subfloat{\includegraphics[width = 1in]{atari_noinp/Gravitar_noinp.png}} &
\subfloat{\includegraphics[width = 1in]{atari_noinp/Hero_noinp.png}}\\
\subfloat{\includegraphics[width = 1in]{atari_noinp/Kangaroo_noinp.png}}&
\subfloat{\includegraphics[width = 1in]{atari_noinp/Krull_noinp.png}} &
\subfloat{\includegraphics[width = 1in]{atari_noinp/KungFuMaster_noinp.png}} &
\subfloat{\includegraphics[width = 1in]{atari_noinp/NameThisGame_noinp.png}} &
\subfloat{\includegraphics[width = 1in]{atari_noinp/Qbert_noinp.png}}\\
\subfloat{\includegraphics[width = 1in]{atari_noinp/SpaceInvaders_noinp.png}}&
\subfloat{\includegraphics[width = 1in]{atari_noinp/VideoPinball_noinp.png}}&
\subfloat{\includegraphics[width = 1in]{atari_noinp/WizardOfWor_noinp.png}}&
\subfloat{\includegraphics[width = 1in]{atari_noinp/Zaxxon_noinp.png}} \\
\end{tabular}
\caption{\textbf{Baseline agent with no input attention mechanism:} Here we compare the RIMs to the baseline, where their is no input attention (i.e., top down attention) as well as all the RIMs communicate with each other at all the time steps. Learning curves for RIMs-PPO models, Baseline Agent, the LSTM-PPO baseline across 30 Atari games.}
\label{fig:no_attention_rims_comm}
\end{figure*}

\newpage

\subsection{Bouncing MNIST: Dropping individual RIMs}
\label{appednix:sec:droprims}

We use the Stochastic Moving MNIST (SM-MNIST) \citep{denton2018stochastic} dataset which consists of sequences of frames of size $64\times 64$, containing one or two MNIST digits moving and bouncing off the walls. Training sequences are generated on the fly by sampling two different MNIST digits from the training set (60k total digits) and two distinct trajectories.

Here, we show the effect of masking out a particular RIM and study the effect of the masking on the ensemble of RIMs. Ideally, we would want different RIMs not to co-adapt with each other. So, masking out a particular RIM should not really effect the dynamics of the entire model. We show qualitative comparisons in Fig.\ \ref{fig:4_blocks_top_k_2}, \ref{fig:4_blocks_top_k_3}, \ref{fig:5_blocks_top_k_2}, \ref{fig:5_blocks_top_k_3}, \ref{fig:5_blocks_top_k_4}.
In each of these figures, the model gets the ground truth image as input for first 5 time steps, and then asked to simulate the dynamics for next 25 time-steps. We find that sparsity is needed otherwise different RIMs co-adapt with each other (for ex. see Fig. \ref{fig:4_blocks_top_k_3}, \ref{fig:5_blocks_top_k_3}, \ref{fig:5_blocks_top_k_4}).  We tried similar masking experiments for different models like RMC, Transformers, EntNet (which learns a mixture of experts), LSTMs, but all of them failed to do anything meaningful (after masking). We suspect this is partly due to learning a homogeneous network.

\begin{figure}[ht]
    \centering
    \includegraphics[width=0.6\linewidth, clip=true, trim = 0mm 0cm 0mm 0mm]{mnist_results_precrop/mnist_2_digit/400_4blocks_topk2/sample_75_0.png}

    \vspace{1cm}
    \includegraphics[width=0.6\linewidth, clip=true, trim = 0mm 0cm 0mm 0mm]{mnist_results_precrop/mnist_2_digit/400_4blocks_topk2/sample_75_1.png}

    \vspace{1cm}

    \includegraphics[width=0.6\linewidth, clip=true, trim = 0mm 0cm 0mm 0mm]{mnist_results_precrop/mnist_2_digit/400_4blocks_topk2/sample_75_2.png}
    \vspace{1cm}

    \includegraphics[width=0.6\linewidth, clip=true, trim = 0mm 0cm 0mm 0mm]{mnist_results_precrop/mnist_2_digit/400_4blocks_topk2/sample_75_3.png}
    \caption{4 RIMs, (top k = 2). Each sub-figure shows the effect of masking a particular RIM and studying the effect of masking on the other RIMs. For example, the top figure shows the effect of masking the first RIM, the second figure shows the effect of masking the second RIM etc.}
    \label{fig:4_blocks_top_k_2}
\end{figure}

\begin{figure}[ht]
    \centering
    \includegraphics[width=0.6\linewidth, clip=true, trim = 0mm 0cm 0mm 0mm]{mnist_results_precrop/mnist_2_digit/400_4blocks_topk3/sample_76_0.png}

    \vspace{1cm}
    \includegraphics[width=0.6\linewidth, clip=true, trim = 0mm 0cm 0mm 0mm]{mnist_results_precrop/mnist_2_digit/400_4blocks_topk3/sample_76_1.png}

    \vspace{1cm}
    \includegraphics[width=0.6\linewidth, clip=true, trim = 0mm 0cm 0mm 0mm]{mnist_results_precrop/mnist_2_digit/400_4blocks_topk3/sample_76_2.png}

    \vspace{1cm}
    \includegraphics[width=0.6\linewidth, clip=true, trim = 0mm 0cm 0mm 0mm]{mnist_results_precrop/mnist_2_digit/400_4blocks_topk3/sample_76_3.png}
    \vspace{1cm}
    \caption{4 RIMs, (top k = 3). 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.}
    \label{fig:4_blocks_top_k_3}
\end{figure}

\begin{figure}[ht]
    \centering
    \includegraphics[width=0.6\linewidth, clip=true, trim = 0mm 0cm 0mm 0mm]{mnist_results_precrop/mnist_2_digit/400_5blocks_topk2/sample_74_0.png}
    \vspace{1cm}
    \includegraphics[width=0.6\linewidth, clip=true, trim = 0mm 0cm 0mm 0mm]{mnist_results_precrop/mnist_2_digit/400_5blocks_topk2/sample_74_1.png}
    \vspace{0.5cm}
    \includegraphics[width=0.6\linewidth, clip=true, trim = 0mm 0cm 0mm 0mm]{mnist_results_precrop/mnist_2_digit/400_5blocks_topk2/sample_74_2.png}
    \vspace{0.5cm}
    \includegraphics[width=0.6\linewidth, clip=true, trim = 0mm 0cm 0mm 0mm]{mnist_results_precrop/mnist_2_digit/400_5blocks_topk2/sample_74_3.png}
    \vspace{0.5cm}
    \includegraphics[width=0.6\linewidth, clip=true, trim = 0mm 0cm 0mm 0mm]{mnist_results_precrop/mnist_2_digit/400_5blocks_topk2/sample_74_4.png}
    \caption{400dim, $k_T = 5$, $k_A = 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.}
    \label{fig:5_blocks_top_k_2}
\end{figure}

\begin{figure}[ht]
    \centering
    \includegraphics[width=0.6\linewidth, clip=true, trim = 0mm 0cm 0mm 0mm]{mnist_results_precrop/mnist_2_digit/400_5blocks_topk3/sample_79_0.png}
    \vspace{1cm}
    \includegraphics[width=0.6\linewidth, clip=true, trim = 0mm 0cm 0mm 0mm]{mnist_results_precrop/mnist_2_digit/400_5blocks_topk3/sample_79_1.png}
    \vspace{0.5cm}
    \includegraphics[width=0.6\linewidth, clip=true, trim = 0mm 0cm 0mm 0mm]{mnist_results_precrop/mnist_2_digit/400_5blocks_topk3/sample_79_2.png}
    \vspace{0.5cm}
    \includegraphics[width=0.6\linewidth, clip=true, trim = 0mm 0cm 0mm 0mm]{mnist_results_precrop/mnist_2_digit/400_5blocks_topk3/sample_79_3.png}
    \vspace{0.5cm}
    \includegraphics[width=0.6\linewidth, clip=true, trim = 0mm 0cm 0mm 0mm]{mnist_results_precrop/mnist_2_digit/400_5blocks_topk3/sample_79_4.png}
    \vspace{0.5cm}
    \caption{400dim, $k_T = 5$, $k_A = 3$. Each sub-figure shows the effect of masking a particular RIM and studying the effect of masking on 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 etc.}
    \label{fig:5_blocks_top_k_3}
\end{figure}

\begin{figure}[ht]
    \centering
    \includegraphics[width=0.6\linewidth, clip=true, trim = 0mm 0cm 0mm 0mm]{mnist_results_precrop/mnist_2_digit/400_5blocks_topk4/sample_68_0.png}
    \vspace{1cm}
    \includegraphics[width=0.6\linewidth, clip=true, trim = 0mm 0cm 0mm 0mm]{mnist_results_precrop/mnist_2_digit/400_5blocks_topk4/sample_68_1.png}
    \vspace{0.5cm}
    \includegraphics[width=0.6\linewidth, clip=true, trim = 0mm 0cm 0mm 0mm]{mnist_results_precrop/mnist_2_digit/400_5blocks_topk4/sample_68_2.png}
    \vspace{0.5cm}
    \includegraphics[width=0.6\linewidth, clip=true, trim = 0mm 0cm 0mm 0mm]{mnist_results_precrop/mnist_2_digit/400_5blocks_topk4/sample_68_3.png}
    \vspace{0.5cm}
    \includegraphics[width=0.6\linewidth, clip=true, trim = 0mm 0cm 0mm 0mm]{mnist_results_precrop/mnist_2_digit/400_5blocks_topk4/sample_68_4.png}
    \vspace{0.5cm}
    \caption{400dim, $k_T = 5$, $k_A = 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.}
    \label{fig:5_blocks_top_k_4}

\end{figure}

\section{Natural Language Processing: Language Modeling, Machine Translation, and Transfer Learning}

We performed some additional experiments to evaluate how well RIMs improve transfer learning on some widely studied NLP 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 two datasets are 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 to French (with shared encoders but separate decoders) or only on English to German.  The results are in Table~\ref{tb:nlp_transfer}.  The LSTM and RIMs models used have a 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 is about the same in both settings).

\begin{table}
\centering
\caption{Transfer from WMT to IWSLT (en $\rightarrow$ de) for different models.  Results in BLEU score (higher is better).  }
\begin{tabular}{lccc}
\hline
Training Data & (Vanilla) Transformer & LSTM & RIMs \\ \hline
en $\rightarrow$ de  & 22.89 & 21.32 & \textbf{23.71} \\
en $\rightarrow$ de, en $\rightarrow$ fr & 22.92 & 20.37 & \textbf{24.23} \\
\end{tabular}
\label{tb:nlp_transfer}
\end{table}

\part{}

\maketitle
\vspace{-5mm}
\begin{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.

\end{abstract}

\vspace{-2mm}
\section{Independent Mechanisms}
\vspace{-2mm}
\let\thefootnote\relax\footnotetext{\textsuperscript{1} Mila, University of Montreal,\textsuperscript{2}
Harvard University, \textsuperscript{3}
MPI for Intelligent Systems, T\"ubingen, \textsuperscript{4} University of California, Berkeley,
\textsuperscript{**} Equal advising, \textsuperscript{*} Equal Contribution.  :\texttt{anirudhgoyal9119@gmail.com}}

Physical processes in the world often have a 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 to 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 \emph{independent} mechanisms or ``causal'' modules. In the causality community, this is often considered a prerequisite for being able to perform localized interventions upon variables determined by such models \citep{Pearl2009}. It 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 \citep{SchJanPetSgoetal12,PetJanSch17}. This independence is not between the random variables being
processed but between the description or 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 a brain is able to solve multiple problems beyond a 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 to 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 \citep{bengio2017consciousness}.
Such sparse interactions can reduce the difficulty of learning since few interactions need to be considered at a time, reducing unnecessary interference when a subsystem is 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 \citep{simon1991architecture, PetJanSch17}.The central question motivating our work is how a gradient-based deep learning approach can discover a representation of high-level variables which favour forming independent but sparsely interacting recurrent mechanisms in order to benefit from the modularity and independent mechanisms assumption.

\begin{figure}
    \vspace{-18mm}
    \centering
    \subfloat{{\includegraphics[width=0.5\linewidth]{figures/RIMs2.png}}}
    \qquad
    \subfloat{{\includegraphics[width=0.3\linewidth]{figures/top_down_jh2.png}}}%
    \caption{\textbf{Illustration of Recurrent Independent Mechanisms (RIMs)}. A single step under the proposed model occurs in  four stages (left figure shows two steps). In the first stage, individual RIMs produce a query which is used to read from the current input.  In the second stage, an attention based competition mechanism is used to select which RIMs to activate (right figure) based on encoded visual input (blue RIMs are active, based on an attention score, white RIMs remain inactive).  In the third stage, individual activated RIMs follow their own default transition dynamics while non-activated RIMs remain unchanged.  In the fourth stage, the RIMs sparsely communicate information between themselves, also using key-value attention.
    \label{fig:rims_diagram}%
    \vspace{-3mm}
\end{figure}

\iffalse
\begin{figure*}[!ht]
    \centering
    \includegraphics[width=0.9\linewidth]{figures/rim_newfig.png}
        \caption{In \textbf{Recurrent Independent Mechanisms (RIMs)}, the recurrent state is divide into $k_T = 4$ modules, where $k_A=2$ are activated on each time step, based on which have the strongest attention on the input.  The activated RIMs update their state with independent recurrent dynamics and can read from other RIMs using attention, while the RIMs which are deactivated remain dormant.  }
        \label{fig:rims_diagram}
    \end{figure*}
\fi

\paragraph{Why do Models Succeed or Fail in Capturing Independent Mechanisms?}
While universal approximation 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 environments, and how to implement suitable inductive biases.   As the simplest case, we can consider training an RNN consisting of $k$ 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 to truly compartmentalize these different processes, a fraction $\frac{k-1}{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-distribution 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.

\paragraph{Assumptions on the joint distribution of high level variables.} The central question motivating our work is how a gradient-based 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~\citep{higgins2016beta, burgess2018understanding, chen2018isolating}, 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.

\section{RIMs with Sparse Interactions}
\label{sec:rimsmodel}

Our approach to modelling a dynamical system of interest divides the overall model into $k$ 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 \textit{Recurrent Independent Mechanisms (RIMs)}, where each RIM has distinct functions that are learned automatically from data\footnote{Note that we are overloading the term \textit{mechanism},
using it both for the mechanisms that make up the world's dynamics as well as for the computational modules that we learn to model those mechanisms. The distinction should be clear from context.}.
We refer to RIM $k$ at time step $t$ as having vector-valued state $h_{t,k}$, where $t=1,\dots, T$. Each RIM has parameters $\theta_k$, which are shared across all time steps.

At a high level (see Fig.~\ref{fig:rims_diagram}), 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-problems, and operate
on few key/value variables at a time
selected using an attention mechanism,
as suggested by the inductive bias from~\citet{bengio2017consciousness}. This specialization and modularization not only has computational and statistical advantages \citep{baum1989size}, 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 \citep{schmidhuber2018one}.
Moreover, modularity and the independent mechanisms hypothesis~\citep{PetJanSch17,bengio2019meta} also has the desirable implication that a 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~\ref{sec:desiderata}.

\subsection{Key-Value Attention to Process Sets of Named Interchangeable Variables}
\label{sec:kvattention}

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 a subset of the RIMs for each time step. As argued by \cite{ParKilRojSch18}, this tends to produce independence among learned mechanisms, provided the training data has been generated by a set of independent physical mechanisms.
In contrast to \cite{ParKilRojSch18}, we use an {\em attention mechanism} for this purpose.  The introduction of content-based soft-attention mechanisms~\citep{bahdanau2014neural} has opened the door to neural networks  which operate on {\em sets of typed interchangeable objects}.
This idea has been remarkably successful and widely applied to most recent Transformer-style multi-head dot product self attention  models \citep{vaswani2017attention, santoro2018relational}, achieving new state-of-the-art results in many tasks. Soft-attention uses the product of a \emph{query} (or \emph{read key}) represented as a matrix $Q$ of dimensionality $N_r \times d$, with $d$ the dimension of each key, with a set of $N_o$ objects each associated with a \emph{key} (or \emph{write-key}) as a row in matrix $K^T$ ($N_o \times d$), and after normalization with a softmax yields outputs in the convex hull of the \emph{values} (or \emph{write-values}) $V_i$ (row $i$ of matrix $V$).
The result is
$$\mathrm{Attention}(Q,K,V)=\mathrm{softmax} \left (\frac{QK^T}{\sqrt{d}} \right )V,$$
where the softmax is applied to each row of its argument matrix, yielding a set of convex weights. As a result, one obtains a 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 in the key can be split into \emph{heads} which then have their own attention matrix and write values computed separately.

When the inputs and outputs of each RIM are a set of objects or entities (each associated with a key and value vector), the RIM processing becomes a generic object-processing machine which can operate on subsymbolic ``variables'' in a sense analogous to variables in a programming language: as interchangeable arguments of functions, albeit with a distributed representation both for they name or type and for their value. Because each object has a key embedding (which one can understand both as a name and as a type), the same RIM processing can be applied to any variable which fits an expected "distributed type" (specified by a query vector). Each attention head then corresponds to a typed argument of the function computed by the RIM. When the key of an object matches the query 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 to 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 \textit{functions with typed arguments}, then they can \emph{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 key and evaluate if its ``type'' matches with what this RIM expects (specified with the corresponding query).

\subsection{Selective Activation of RIMs as a form of Top-Down Modulation}
\label{sec:selectiveactivationrims}

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 input information coming from the environment.  At each step, we select the top-$k_A$ (out of $k_{T}$) RIMs in terms of their attention score for the real input. Intuitively, the RIMs must compete on each step to read from the input, and only the RIMs that win this competition will be able 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 \citet{vaswani2017attention,santoro2018relational}, 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 a particular RIM is described as follows.

The input $x_t$ at time $t$ is seen as a set of elements, structured as rows of a matrix. We first concatenate a row full of zeros, to obtain
\begin{align}
    X &= \emptyset \oplus x_t.

\end{align}

As before, linear transformations are used to construct  keys ($K = XW^e$, one per input element and for the null element),  values ($V = XW^v$, again one per element), and  queries ($Q = h_{t} W_k^q$, one per RIM attention head). $W^v$ is a simple matrix mapping from an input element to the corresponding value vector for the weighted attention and $W^e$ is similarly a weight matrix which maps the input to the keys.  $W_k^q$ is a per-RIM weight matrix which maps from the RIM's hidden state to its queries.  $\oplus$ refers to the row-level concatenation operator.
The attention thus is
\begin{align}
    A^{(in)}_{k} &= \text{softmax} \left( \frac{h_{t} W_k^q(X W^e)^T}{\sqrt{d_e}} \right) X W^v , \text{ where } \theta^{(in)}_k =(W_k^q, W^e, W^v). \label{eq:non_spat_atten}
\end{align}

Based on the softmax values in (\ref{eq:non_spat_atten}), we select the top $k_A$ 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 $\mathcal{S}_t$.  Since the queries depend on the state of the RIMs, this enables individual RIMs to attend only to the part of the input that is relevant for that particular RIM, thus enabling selective attention based on a \textit{top-down attention} process (see. Fig \ref{fig:rims_diagram}). In practice, we use multiheaded attention, and multi-headed attention doesn't change the essential computation, but when we do use it for input-attention we compute RIM activation by averaging the attention scores over the heads.

\textit{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 a video) or not (such as language modeling, where each step has a 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 $x_t$ at time $t$ can be seen as an output of the encoder parameterized by a neural network (for ex. CNN in case of visual observations) i.e., $X = CNN(x_t)$. As before, linear transformations are used to construct position-specific input  keys ($K = XW^e$), position-specific values ($V = XW^v$), and RIM specific queries ($Q = h_t W_k^q$, one per RIM attention head).
The attention thus is

\vspace{-1mm}
\begin{align}
    A^{(in)}_{k} &= \text{softmax} \left( \frac{h_t W_k^q(X W^e)^T}{\sqrt{d_e}} \right) X W^v , \text{ where } & \theta^{(in)}_k =(W_k^q, W^e, W^v). \label{eq:spatial_atten}
\end{align}

In order for different RIMs to specialize on different spatial regions, we can use position-specific competition among the different RIMs. The contents of the attended positions are combined yielding a RIM-specific input. As before based on the softmax values in (\ref{eq:spatial_atten}), we can select the top $k_A$ RIMs (out of the total $k_{T}$ RIMs) to be activated for each spatial position, which have the highest attention on that spatial position.

\subsection{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 $\tilde{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 as $\mathcal{S}_t$) remain unchanged, acting like untouched memory elements, i.e.,
$    h_{t+1,k} = h_{t,k} \quad \forall k \notin \mathcal{S}_t.$
Note that the gradient still flows through a RIM on a step where it is not activated.  For the RIMs that are activated, we run a per-RIM independent transition dynamics.  The form of this is somewhat flexible, but we opted to use either a GRU \citep{chung2015recurrent} or an LSTM \citep{hochreiter1997long}.  We generically refer to these independent transition dynamics as $D_k$, and we emphasize that each RIM has its own separate parameters.  Aside from being RIM-specific, the internal operation of the LSTM and GRU remain unchanged, and active RIMs are updated by
\begin{equation*}
    \tilde{h}_{t,k} = D_k(h_{t,k}) = LSTM(h_{t,k}, A^{(in)}_{k}; \theta^{(D)}_k) \quad \forall k \in \mathcal{S}_t
\end{equation*}
as a function of the attention mechanism $A^{(in)}_{k}$ applied on the current input, described in the previous sub-section.

\subsection{Communication between RIMs}
\label{sec:communicationrims}

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 as in \cite{santoro2018relational} to prevent vanishing or exploding gradients over long sequences.
Using parameters $\theta^{(c)}_k =(\tilde{W}_k^q, \tilde{W}_k^e, \tilde{W}_k^v)$, we employ
\begin{align}
    Q_{t,k} = \tilde{W}^{q}_k \tilde{h}_{t,k}, \forall k \in \mathcal{S}_t \quad\quad
    K_{t,k} = \tilde{W}^{e}_k \tilde{h}_{t,k}, \forall k \quad\quad
    V_{t,k} = \tilde{W}^{v}_k \tilde{h}_{t,k}, \forall k
    \nonumber
\end{align}
\begin{align*}
    h_{t+1,k} &= \text{softmax} \left( \frac{Q_{t,k} (K_{t,:})^T}{\sqrt{d_e}} \right) V_{t,:} + \tilde{h}_{t,k}  \forall k \in \mathcal{S}_t .
\end{align*}

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.

\textit{Number of Parameters.}  RIMs can be used as a drop-in replacement for an LSTM/GRU layer. There is 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 are rather in small number.

\textit{Multiple Heads:} Analogously to \citet{vaswani2017attention, santoro2018relational}, we use multiple heads both for communication between RIMs as well as input attention (as in Sec \ref{sec:selectiveactivationrims}) 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.

\vspace{-3mm}

\section{Related Work\label{sec:related_work}}

\textbf{Neural Turing Machine (NTM) and Relational Memory Core (RMC):}  the NTM \citep{graves2014neural} 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 a sparse subset of the memory by default, while keeping most of the memory unaltered. RMC \citep{santoro2018relational}  uses a multi-head attention mechanism to share information between multiple memory elements. We encourage the RIMs to remain separate as much as possible, whereas \cite{santoro2018relational} 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.

\textbf{Separate Recurrent Models:} EntNet \citep{henaff2016tracking} and IndRNN \citep{li2018independently} can be viewed as a set of separate recurrent models. In IndRNN, each recurrent unit has completely independent dynamics, whereas EntNet uses an independent gate for writing to each memory slot. RIMs use  different recurrent models (with separate parameters), but we allow the RIMs to communicate with each other sparingly using an attention mechanism.

\textbf{Modularity and Neural Networks}: A network can be composed of several modules, each meant to perform a distinct function, and hence can be seen as a combination of experts \citep{jacobs1991adaptive, bottou1991framework,ronco1996modular, reed2015neural, andreas2016neural,ParKilRojSch18,rosenbaum2017routing, fernando2017pathnet, shazeer2017outrageously, kirsch2018modular, rosenbaum2019routing} routing information through a gated activation of modules. These works generally assume that only a single expert is active at a particular time step. In the proposed method, multiple RIMs can be active, interact and share information.

\textbf{Computation on demand:}  There are various  architectures \citep{el1996hierarchical, koutnik2014clockwork, chung2016hierarchical, neil2016phased, jernite2016variable, krueger2016zoneout} 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 dynamics of recurrent cells (using RIMs), and (b) we also control the inputs of each module (using transformer style attention), while many previous gating methods did not control the inputs of each module, but only whether they should be executed or not.

\vspace{-1mm}

\section{Experiments}

\vspace{-1mm}

The main goal of our experiments is to show that the use of RIMs improves generalization across changing environments and/or in modular tasks, and to explore how it does so.
Our goal is not to outperform highly optimized baselines; rather, we want to show the versatility of our approach by applying it to a range of diverse tasks, focusing on tasks that involve a changing environment.
We organize our results by the capabilities they illustrate: we address generalization based on temporal patterns, based on objects, and finally consider settings where both of these occur together.

\subsection{RIMs improve generalization by specializing over temporal patterns}
\vspace{-1mm}

We first show that when RIMs are presented with sequences containing distinct and generally independent temporal patterns, they are able to specialize so that different RIMs are activated on different patterns.  RIMs generalize well when we modify a subset of the patterns (especially those unrelated to the class label) while most recurrent models fail to generalize well to these variations.

\label{sec:copying}
\begin{figure}[h]
    \centering

    \includegraphics[width=0.79\linewidth,trim=8.0cm 0.0cm 4.0cm 0.0cm,clip]{figures/blocks_copying_fig.png}
    \caption{\textbf{Visualizing Activation Patterns}.

    For the copying task,  one can see that the RIM activation pattern is 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.}

    \label{fig:copying_pattern}

\end{figure}

\paragraph{Copying Task:}
First we turn our attention to the task of receiving a short sequence of characters, then receiving blank inputs for a large number of 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~\ref{tb:copy_mnist}), 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~\ref{table:copying_ablation}) 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.

\begin{table*}[t]
\footnotesize
\resizebox{0.36\columnwidth}{!}{
\begin{minipage}[t]{0.4\textwidth}

  \begin{tabular}{p{1.3ex}rrr|r|r}
  \toprule
  \multicolumn{4}{l|}{\textbf{Copying}} &
  \multicolumn{1}{c|}{\textbf{Train(50)}} &
  \multicolumn{1}{c}{\textbf{Test(200)}} \\
  & \parbox{2em}{$k_\textrm{T}$} & $k_\textrm{A}$ & $h_\textrm{size}$ &
  CE &
  CE \\
  \midrule
  \multirow{4}{*}{\rotatebox{0}{\parbox{2.5em}{\textbf{RIMs}}}}

  & 6 & 4 & 600 & \textbf{0.00} & \textbf{0.00} \\
  & 6 & 3 & 600 & \textbf{0.00} & \textbf{0.00} \\
  & 6 & 2 & 600 & \textbf{0.00} & \textbf{0.00} \\
  & 5 & 2 & 500 & \textbf{0.00} & \textbf{0.00} \\
  \midrule
  \multirow{2}{*}{\rotatebox{0}{\parbox{2.5em}{\textbf{LSTM}}}}
  & - & - & 300 &  0.00 & 4.32 \\
  & - & - & 600 &  0.00 & 3.56 \\
  \midrule
  \multirow{1}{*}{\rotatebox{0}{\parbox{2.5em}{\textbf{NTM}}}}
  & - & - & - &  0.00 & 2.54 \\
  \midrule
  \multirow{1}{*}{\rotatebox{0}{\parbox{2.5em}{\textbf{RMC}}}}
  & - & - & - & 0.00 & 0.13 \\
  \midrule
  \multirow{1}{*}{\rotatebox{0}{\parbox{2.5em}{\textbf{Transformers}}}}
  & - & - & - & 0.00 & 0.54 \\
  \bottomrule
  \end{tabular}
\end{minipage}}
\hspace{3.3em}
\resizebox{0.36\columnwidth}{!}{
\begin{minipage}[t]{0.4\textwidth}

  \begin{tabular}{p{1.3ex}rrr|r|r|r}
  \toprule
  \multicolumn{4}{l|}{\textbf{Sequential MNIST}} &
  \multicolumn{1}{c|}{\textbf{16 x 16}} &
  \multicolumn{1}{c}{\textbf{19 x 19}} &
  \multicolumn{1}{c}{\textbf{24 x 24}} \\
  & \parbox{2em}{$k_\textrm{T}$} & $k_\textrm{A}$ & $h_\textrm{size}$ &
  Accuracy &
  Accuracy &
  Accuracy \\
  \midrule
  \multirow{4}{*}{\rotatebox{0}{\parbox{2.5em}{\textbf{RIMs}}}}
  & 6 & 6 & 600 & 85.5 & 56.2 & 30.9\\
  & 6 & 5 & 600 & 88.3 & 43.1 & 22.1 \\
  & 6 & 4 & 600 & \textbf{90.0} & \textbf{73.4} & \textbf{38.1}  \\
  \midrule
  \multirow{2}{*}{\rotatebox{0}{\parbox{2.5em}{\textbf{LSTM}}}}
  & - & - & 300 &  86.8 & 42.3 & 25.2 \\
  & - & - & 600 &  84.5 & 52.2 & 21.9\\
  \midrule
  \multirow{1}{*}{\rotatebox{0}{\parbox{2.5em}{\textbf{EntNet}}}}
  & - & - & - &  89.2 & 52.4 & 23.5 \\
  \midrule
  \multirow{1}{*}{\rotatebox{0}{\parbox{2.5em}{\textbf{RMC}}}}
  & - & - & - &  89.58 & 54.23 & 27.75 \\
  \midrule
  \multirow{1}{*}{\rotatebox{0}{\parbox{2.5em}{\textbf{DNC}}}}
  & - & - & - &  87.2 & 44.1 & 19.8 \\
  \multirow{1}{*}{\rotatebox{0}{\parbox{2.5em}{\textbf{Transformers}}}}
  & - & - & - &  \textbf{91.2} & 51.6 & 22.9 \\
  \bottomrule
  \end{tabular}
\end{minipage}}
\caption{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).
}
\label{tb:copy_mnist}
\vspace{-4mm}
\end{table*}

\paragraph{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-distribution, we consider the task of classifying MNIST digits as sequences of pixels \citep{krueger2016zoneout} 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.

\textbf{Results:} Table \ref{tb:copy_mnist} 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 as a  more involved analogue of the copying task, as MNIST digits contain large empty regions.  It is essential that the model be able to store information and pass gradients through these regions.  The RIMs outperform strong baselines such as Transformers, EntNet, RMC, and (DNC) \citep{graves2016hybrid}.

\subsection{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.
\vspace{-3mm}
\paragraph{Bouncing Balls Environment:}
We consider a synthetic ``bouncing balls'' task in which multiple balls (of different masses and sizes) move using basic Newtonian physics \citep{van2018relational}.  What makes this task particularly suited to RIMs is 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 \citep{williams1989learning}.  We can then use this model to generate multi-step rollouts.  As a 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~\ref{fig:err}).

\begin{figure*}[htb!]
    \vspace{-5mm}
    \centering
    \includegraphics[width=0.75\linewidth]{figures/err_all_2.png}
    \includegraphics[width=0.24\textwidth,trim={0 0 17cm 0},clip]{video_crop/bbals.png}
    \vspace{-0.5cm}
    \caption{\textbf{Handling Novel Out-of-Distribution Variations}. We study the performance of 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 as reflected by the lower Cross Entropy (CE) [blue for RIMs, purple for LSTMs].  Notice the substantially better out-of-distribution generalization of RIMs when testing on a number of objects different from the one seen during training. (2nd to 4th plot). \textbf{ We also show (right plot) improved out-of-distribution generalization (F1 score) as compared to LSTM and RMC \citep{santoro2018relational} on another partial observation video prediction task.} X-axis = number of balls. For these experiments, the RIMs and baselines get an input image at each time step (see  Appendix \ref{sec:bballs_appendix}, figure.~\ref{fig:bb_ood} 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.}
    \label{fig:err}
\vspace{-4mm}

\end{figure*}

We take this further by evaluating RIMs on environments where the test setup is different from the training setup.  First we consider training with 4 balls and evaluating on an environment with 6-8 balls.  Second, we consider training with 6-8 balls and evaluating with just 4 balls. Robustness in these settings requires 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 on 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~\ref{fig:err}) 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.

\vspace{-2mm}

\paragraph{Environment with Novel Distractors:}
We next consider an object-picking reinforcement learning task from BabyAI \citep{chevalier2018babyai} in which an agent must retrieve a specific object in the presence of distractors.  We use a partially observed formulation of the task, where the agent only sees a small number of squares ahead of it. These tasks are difficult to solve \citep{chevalier2018babyai} 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.

\begin{wrapfigure}{r}{0.6\textwidth}
    \centering
    \vspace{-6mm}
    \includegraphics[width=0.45\linewidth,trim=0.5cm 0.0cm 1.0cm 1.5cm,clip]{figures/figure_normal_jh.png}
    \includegraphics[width=0.45\linewidth,trim=0.5cm 0.0cm 1.0cm 1.5cm,clip]{figures/figure_distractor_jh.png}
    \caption{\textbf{Robustness to Novel Distractors:}. Left: performance of the proposed method (blue) compared to an LSTM baseline (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.}
    \label{fig:rl_exp}
    \vspace{-5mm}
\end{wrapfigure}

Figure~\ref{fig:rl_exp} shows that RIMs outperform LSTMs on this task (details in appendix).  When evaluating with known distractors, the RIM model achieves perfect performance while the LSTM struggles.   When evaluating in an environment with novel unseen distractors the RIM doesn't 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 as what information they receive (as a function of their hidden state). This "top-down" attention results in a diminished representation of the distractor, not only enhancing the target visual information, but also suppressing irrelevant information.

\vspace{-3mm}
\subsection{RIMs improve generalization in complex environments}

We have investigated how RIMs use specialization to improve generalization to changing important factors of variation in the data.  While these improvements have often been striking, it raises a 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 environment changes as the agent learns and improves.  We conjecture that when applied to reinforcement learning, an agent using RIMs may be able to learn faster as its specialization leads to improved generalization to previously unseen aspects of the environment.  To investigate this we use an RL agent trained using Proximal Policy Optimization (PPO) \citep{schulman2017proximal} with a recurrent network producing the policy.  We employ an LSTM as a baseline, and compare results to the RIMs architecture.  This was a simple drop-in replacement and did not require changing any of the hyperparameters for PPO.  We experiment on the whole suite of Atari games and find that simply replacing the LSTM with RIMs greatly improves performance (Figure~\ref{fig:atari_rims4}).

\begin{figure*}[htb!]
    \centering
    \vspace{-3mm}
    \includegraphics[width=0.8\linewidth,trim=2.8cm 0.0cm 3.5cm 1.0cm,clip]{atari_results/rims4v3.png}

    \caption{RIMs-PPO relative score improvement over LSTM-PPO baseline \citep{schulman2017proximal} across all Atari games averaged over 3 trials per game.  In both cases, PPO was used with the exact same settings, and the only change is the choice of recurrent architecture.    More detailed experiments with learning curves as well as comparisons with external baselines are in Appendix~\ref{appendix:detailedexp}.  }
    \label{fig:atari_rims4}
    \vspace{-3mm}
\end{figure*}

There is also an intriguing connection between the selective activation in RIMs and the concept of affordances from cognitive psychology \citep{gibson1977theory, cisek2010neural}.  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 a contextual manner, building representations of potential actions that the environment currently affords.  For instance, in Demon Attack, one of the games where RIMs exhibit strong performance gains, the agent must quickly choose between targeting distant aliens to maximize points and avoiding fire from close-by aliens to avoid destruction (indeed both types of aliens are always present, but which is relevant depends on the player's position).  We hypothesize that in cases like this, selective activation of RIMs allows the agent to rapidly adapt its information processing to the types of actions relevant to the current context.

\vspace{-1mm}
\subsection{Discussion and Ablations}

\textbf{Sparse Activation is necessary, but works for a wide range of hyperparameters: } On the copying task, we tried a wide variety of sparsity levels for different numbers of RIMs, and found that using a sparsity level between 30\% to 70\% performed optimally, suggesting that the sparsity hyperparameter is fairly flexible (refer to Table~\ref{tb:copy_ablation}, ~\ref{tb:adding_ablation} in appendix).  On Atari we found that using $k_A=5$ slightly improved over results compared with $k_A=4$, but both had similar performance across the vast majority of games.

\textbf{Input-attention is 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~\ref{fig:no_attention_rims_comm}) in appendix for more details.

\textbf{Communication between RIMs improves performance:} For copying and sequential MNIST, we performed an ablation where we remove the communication between RIMs and varied the number of RIMs and the number of activated RIMs (Refer to Table~\ref{tb:copy_ablation} in appendix.).  We found that the communication between RIMs is essential for good performance.

\label{sec:main_ablation}
\iffalse
\begin{figure}[htb!]
    \centering
    \includegraphics[width=0.9\linewidth, height=0.9\linewidth]{figures/ablate.png}

    \caption{\textbf{RIM ablations on bouncing balls}. Using a model trained with 8 RIMs, 6 active, on the bouncing ball dataset, we randomly allow different numbers of RIMs to be active. We randomly select up to 25 possible permutations of between 1 and all 8 RIMs being active and report the mean and standard deviation of the 1-step prediction loss on the test and transfer dataset.}
    \label{fig:ablate}
    \vspace{-2mm}
\end{figure}
\fi

\vspace{-3mm}
\section{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 are independent by default, but interact sparingly. For the purposes of this paper, we note that the notion of RIMs is not limited to the particular architecture employed here. The latter is used as a vehicle to assay and validate our overall hypothesis (cf.\ Appendix~\ref{sec:desiderata}), but better architectures for the RIMs model can likely be found.

\section*{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 session{s}. 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. The authors are grateful to NSERC, CIFAR, Google, Samsung, Nuance, IBM, Canada Research Chairs, Canada Graduate Scholarship Program, Nvidia for funding, and Compute Canada for computing resources.  We are very grateful to  Google for giving Google Cloud credits used in this project. This project was also known by  the name ``Blocks'' internally at Mila.

\section{Introduction}

\Ani{Need to think of a strong teaser figure. }

Machine learning seeks to uncover and use regularities in the physical world. Although they manifest themselves as statistical dependencies, they are ultimately due to dynamic processes governed by physics. These processes are often independent and only interact sparsely. For instance, we can model the motion of two balls as separate independent mechanisms even though they are both gravitationally coupled to earth as well as (weakly) to each other. They may, however, occasionally strongly interact via collisions.

The notion of independent or autonomous mechanisms has been fruitful in the field of causal inference, where it is applied not only to dynamic processes but also to datasets independent of time. For instance, it has been argued that the conditional distribution of the average annual temperature given altitude (of a city) is an abstraction of a causal mechanism (subsuming complex physical processes involving air pressure, etc.) that is fairly independent of the distribution of the altitudes of cities \citep{PetJanSch17}, and will thus apply invariantly for, say, different countries in the same climate zone with different altitude distributions.

A complex generative model (temporal or not) can thus be thought of as the composition of independent causal modules or mechanisms that are independent. In the causality community, the latter is often considered a prerequisite to the possibility of being able to, in principle, intervene upon variables determined by such models. It has been argued that the individual modules tend to remain invariant even when other modules change, e.g., in the case of distribution shift \citep{PetJanSch17}. One may hypothesize that if a brain is able to solve multiple problems and multiple tasks beyond a single i.i.d.\ dataset, it would be economical to learn structures aligned with this by using independent components that can flexibly be reused and repurposed.

In the dynamic setting, we could think of the world as composed of a large number of fairly independent subsystems that change and evolve over time, responding to forces and interventions. For a learner interacting with such an environment, not all these subsystems will significantly interact with the learner most of the time.
A learner thus does not need to know the dynamics of all subsystems at each time-step to correctly model aspects of the world relevant to a current decision or plan.  Only some aspects of the world significantly interact and thus need to be considered jointly when taking a decision or forming a plan, as proposed with the consciousness prior~\citep{bengio2017consciousness}. We assay this in the context of multiple dynamical subsystems which can only interact with each other in a sparse way, i.e., with few subsystems influencing another subsystem at a particular time. Such sparse interactions reduce the difficulty of learning since few interactions need to be considered at a time, reducing unnecessary interference when a subsystem is adapted.
Nearly decomposable (or modular) structures \citep{simon1991architecture} are aggregates of mechanisms that interact together and can perform specific functions without affecting the remainder of the system.

Models learned this way may be more likely to capture the compositional causal structure of the world, and thus better generalize across tasks where a (small) subset of mechanisms change while most of them remain invariant \citep{PetJanSch17,ParRojKilSch17}.

The above points in principle pertain to synthetic and natural intelligent systems alike, and hence it is not surprizing that they also appear in neuroscience. For instance, suppose we are looking for a particular object in a large scene, using limited processing capacity.
The \emph{biased competition theory} of selective attention conceptualizes basic findings of experimental psychology and neuroscience \citep{Desimone1995}:
our capacity of parallel processing is limited, and many brain systems representing visual information use competition to allocate resources.  Competitive 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 to filter out unwanted information and selectively \emph{process} the remainder. Third,  \emph{top-down bias} originating from higher brain areas enables us to selectively devote resources to stimuli that may be of particular interest or relevance. This may be accomplished by units matching the internal model of an object or process of interest being pre-activated and thus gaining an advantage during competition of brain mechanisms.

The central question motivating our work is how a machine learning system can learn independent mechanisms in order to benefit from the hypothesized modular structure of the world. Many approaches in machine learning represent the world at a point in time with a single high-dimensional hidden state which attempts to describe the full world-state. It could thus be difficult for typical recurrent networks to describe a world which consists of nearly independent systems that only occasionally interact.

\Ani{For Anirudh, Write here a summary of contributions/experiments.}

\section{Recurrent Independent Mechanisms with Sparse Interactions}

\iffalse
\begin{figure*}[t]
\centering
\begin{tikzpicture}[
    scale=1.0,
    black!50, text=black,
    font=\medium,

    node distance=1mm,
    dnode/.style={
        align=center,

        font=\small,
        rectangle,minimum size=15mm,rounded corners,

        inner sep=20pt},
    rnode/.style={
        align=center,
        font=\medium,

        rectangle,
        minimum width=5mm,
        minimum height=6mm,
        rounded corners,

        inner sep=3pt,
        thin, draw},
    tuplenode/.style={
        align=center,

        rectangle,minimum size=10mm,rounded corners,

        inner sep=15pt},
    darrow/.style={
        rounded corners,-latex,shorten <=5pt,shorten >=1pt,line width=2mm},
    mega thick/.style={line width=2pt}
    ]
\matrix[row sep=7.2mm, column sep=2mm] {
    \node (a_d) [rnode, minimum width=15mm, top color=greenfill!40, bottom color = greenfill!40] {Surfing RIM};
    &&
    &&
    \node (a_1) [rnode, minimum width=15mm, top color=greenfill!40, bottom color = greenfill!40] {$h_{t+2,2}$};
    &&
    &&
    \node (a_2) [rnode, minimum width=15mm, top color=greenfill!40, bottom color = greenfill!40] {$h_{t+4,2}$};
    \\
    \node (b_d) [rnode, minimum width=15mm, top color=redfill!40, bottom color = redfill!40] {Studying RIM};
    &&
    \node (b_1) [rnode, minimum width=15mm, top color=redfill!40, bottom color = redfill!40] {$h_{t+1,1}$};
    &&
    &&
    \node (b_2) [rnode, minimum width=15mm, top color=redfill!40, bottom color = redfill!40] {$h_{t+3,1}$};
    &&
    &&
    \node (b_3) [rnode, minimum width=15mm, top color=redfill!40, bottom color = redfill!40] {$h_{t+5,1}$};
    \\
    &&
    \node (x_1) [rnode, minimum width=15mm, top color=black!10, bottom color = black!10] {\makecell{Study\\for Test}};
    &&
    \node (x_2) [rnode, minimum width=15mm, top color=black!10, bottom color = black!10] {\makecell{Prepare\\for Surfing}};
    &&
    \node (x_3) [rnode, minimum width=15mm,top color=black!10, bottom color = black!10] {\makecell{Think about\\mathcal{D_{\mathrm{test}}} before wave}};
    &&
    \node (x_4) [rnode, minimum width=15mm, top color=black!10, bottom color = black!10] {\makecell{Start\\Surfing}};
    &&
    \node (x_5) [rnode, minimum width=15mm, top color=black!10, bottom color = black!10] {\makecell{Take\\the test}};
    \\
};

\draw[-latex,shorten <=1pt,shorten >=1pt,thick,color=blue!60]  (a_1) to node[]{} (a_2);
\draw[-latex,shorten <=1pt,shorten >=1pt,thick,color=red!60]  (b_1) to node[]{} (b_2);
\draw[-latex,shorten <=1pt,shorten >=1pt,thick,color=red!60]  (b_2) to node[]{} (b_3);
\draw[-latex,shorten <=1pt,shorten >=1pt,thick,color=black!30]  (x_1) to node[]{} (b_1);
\draw[-latex,shorten <=1pt,shorten >=1pt,thick,color=black!30]  (x_2) to node[]{} (a_1);
\draw[-latex,shorten <=1pt,shorten >=1pt,thick,color=black!30]  (x_3) to node[]{} (b_2);
\draw[-latex,shorten <=1pt,shorten >=1pt,thick,color=black!30]  (x_4) to node[]{} (a_2);
\draw[-latex,shorten <=1pt,shorten >=1pt,thick,color=black!30]  (x_5) to node[]{} (b_3);

\draw[-latex,shorten <=1pt,shorten >=1pt,thick,color=purple!60]  (a_1) to node[]{} (b_2);

\end{tikzpicture}

\caption{An illustration of the ideal behavior of the RIMs model, in which RIMs are sparsely activated and responsible for modeling different aspects of the world.  This helps to prevent interference between unrelated aspects of the world.  }
\label{fig:blocks}

\end{figure*}
\fi

We have laid out a case for building models composed of modules which by default operate independently and can interact in a 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 as \textit{Recurrent Independent Mechanisms (RIMs)}, with distinct functions learned automatically from data.\footnote{Note that we are using the term \textit{mechanism} both for the mechanisms that make up the world's dynamics as well as for the computational modules that we learn to model those mechanisms.} Our model will encourage sparse interaction, i.e.,  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.

\paragraph{Competitive Mechanisms:}
Inspired by the observations in the previous section, we propose that RIMs utilize competition to allocate representational and computational resources. As argued by \cite{ParRojKilSch17}, this tends to produce independence among learned mechanisms if the training data has been generated by independent physical mechanisms in the first place.

\paragraph{Modular Memory:} Memory is divided into distinct modules, such that it is possible for the model to elect to read to or write from a subset of the modules.  This allows the model to easily keep different pieces of information well-separated over multiple time steps.

\paragraph{Sparse Information Flow:} Each RIMs' dynamics should only be affected by RIMs which are deemed relevant. The fundamental challenge is centered around establishing sensible communication between RIMs.  In the presence of noisy or distracting information, a large subset of RIMs should stay dormant, and not affected by the noise. This way, training an ensemble of these RIMs can be more robust to out-of-distribution or distractor observations than training one big homogeneous neural network \citep{schmidhuber2018one}.

\paragraph{Modular Computation Flow:} Each RIM should have its own dynamics operating by \textit{default}, in the absence of interaction with other RIMs.\Bernhard{btw, Neuroscience is getting increasingly interested in the brain's 'resting state' and 'default networks' ;) No need to mention it, but could be interesting to think about it some time.}
The total number of parameters (i.e., weights) can be reduced since the RIMs can specialize on simple sub-problems, similar to \citep{ParRojKilSch17}.
This can speed up computation and improve the generalisation ability of the system \citep{baum1989size}. The individuals RIMs in the ensemble should be simple also to prevent individual RIMs from dominating and modelling complex composite mechanisms.

\paragraph{Modular Parameterization:} We refer to a parameterization as modular if most parameters are associated to individuals RIMs only (with some exceptions, such as in the overall input and output). This has the desirable property that a RIM should maintain its own independent functionality even as other RIMs are changed (due to its behavior being determined by its own self-contained parameters).

\paragraph{Sparse Credit Assignment:} All parameters should be associated to individual RIMs.  The  model can then choose to propagate the error gradient into a sparse subset of the parameters only, implying that a RIM will be invariant (and thus remain functional) when exposed to data processed by unrelated RIMs.
Note that while sparse information flow is a property of the forward pass, sparse credit assignment is a property of how the model learns and updates RIMs.

In Table~\ref{tb:related} we discuss the extent to which related works have accomplished these desiderata and make the case that while several approaches have accomplished some of them, the RIMs model is the first to accomplish all of them together.

\section{RIMs: Proposed Model}

We previously laid out our motivation for the concept of RIMs and gave a high level description: a recurrent model is separated into RIMs which have independent dynamics and a RIM is only activated and updated when the input is relevant to that particular RIM.  However, this leaves many open questions regarding how we select which RIMs to activate and how we define independent dynamics.

At a high-level, our model divides the overall recurrent state into $K$ RIMs which each have a value across $T$ time steps.  Thus we refer to RIM $k$ and time step $t$ as having a state $h_{t,k}$.  Each RIM has parameters $\theta_k$, which are shared across all time steps.

\begin{figure*}[htb!]
\centering
\begin{tikzpicture}[
    scale=1.0,
    black!50, text=black,
    font=\medium,

    node distance=1mm,
    dnode/.style={
        align=center,

        font=\small,
        rectangle,minimum size=10mm,rounded corners,

        inner sep=10pt},
    rnode/.style={
        align=center,
        font=\medium,

        rectangle,
        minimum width=0m,
        minimum height=6mm,
        rounded corners,

        inner sep=1pt,
        thin, draw},
    tuplenode/.style={
        align=center,

        rectangle,minimum size=2mm,rounded corners,

        inner sep=5pt},
    darrow/.style={
        rounded corners,-latex,shorten <=5pt,shorten >=1pt,line width=2mm},
    mega thick/.style={line width=2pt}
    ]
\matrix[row sep=8.2mm, column sep=3mm] {
    &&
    &&
    &&
    &&
    &&
    &&
    &&
    &&
    \node (h1r) [rnode, minimum width=10mm, top color=red!40, bottom color = red!40] {$h_{t,1}$};
    &&
    \node (h2r) [rnode, minimum width=10mm, top color=blue!10, bottom color = blue!0] {$h_{t,2}$};
    \\
    \node (h1l) [rnode, minimum width=10mm, top color=red!40, bottom color = red!40] {$h_{t,1}$};
    &&
    \node (h2l) [rnode, minimum width=10mm, top color=blue!10, bottom color = blue!0] {$h_{t,2}$};
    &&
    \node (h3l) [rnode, minimum width=10mm, top color=red!40, bottom color = red!40] {$h_{t,3}$};
    &&
    &&
    \node (h2m1) [rnode, minimum width=10mm, top color=blue!10, bottom color = blue!0] {$h_{t,2}$};
    &&
    \node (h2m2) [rnode, minimum width=10mm, top color=blue!10, bottom color = blue!0] {$h_{t+1,2}$};
    &&
    &&
    \node (h2r1) [rnode, minimum width=10mm, top color=red!40, bottom color = red!40] {$h_{t,3}$};
    &&
    &&
    \node (h2r3) [rnode, minimum width=10mm, top color=red!40, bottom color = red!40] {$h_{t+1,3}$};
    \\
    \node (null_l) [rnode, minimum width=10mm, top color=black!20, bottom color = black!20] {$\emptyset$};
    &&
    \node (xt_l) [rnode, minimum width=10mm, top color=black!20, bottom color = black!20] {$x_t$};

    &&
    &&
    &&
    &&
    &&
    &&

    \node (null_r) [rnode, minimum width=10mm, top color=black!20, bottom color = black!20] {$\emptyset$};
    &&
    \node (xt_r) [rnode, minimum width=10mm, top color=black!20, bottom color = black!20] {$x_t$};

    \\
};

\draw[-latex,shorten <=1pt,shorten >=1pt,thick,color=black!20]  (null_l) to node[]{} (h1l);
\draw[-latex,shorten <=1pt,shorten >=1pt,thick,color=black!80]  (null_l) to node[]{} (h2l);
\draw[-latex,shorten <=1pt,shorten >=1pt,thick,color=black!20]  (null_l) to node[]{} (h3l);

\draw[-latex,shorten <=1pt,shorten >=1pt,thick,color=black!80]  (xt_l) to node[]{} (h1l);
\draw[-latex,shorten <=1pt,shorten >=1pt,thick,color=black!20]  (xt_l) to node[]{} (h2l);
\draw[-latex,shorten <=1pt,shorten >=1pt,thick,color=black!80]  (xt_l) to node[]{} (h3l);

\draw[-latex,shorten <=1pt,shorten >=1pt,thick,color=blue!80]  (h2m1) to node[label=above:id]{} (h2m2);
\draw[-latex,shorten <=1pt,shorten >=1pt,thick,color=red!80]  (h2r1) to node[label=above:$D_3$]{} (h2r3);

\draw[-latex,shorten <=1pt,shorten >=1pt,thick,color=red!30]  (h1r) to node[]{} (h2r3);
\draw[-latex,shorten <=1pt,shorten >=1pt,thick,color=blue!30]  (h2r) to node[]{} (h2r3);

\draw[-latex,shorten <=1pt,shorten >=1pt,thick,color=black!20]  (null_r) to node[]{} (h2r1);
\draw[-latex,shorten <=1pt,shorten >=1pt,thick,color=black!80]  (xt_r) to node[]{} (h2r1);

\end{tikzpicture}
\caption{An illustration of the mechanics of the Recurrent Independent Mechanisms (RIMs) model with ($K=3$ RIMs and $k_A=2$ RIMs active at any given point in time; in our applications, we generally use $K>5$ and $k_A=4$).  We use an attention mechanism to select the active RIMs (shown in red) based on the magnitude of their attention weights on the input $x_t$ (left).  This leads the non-activated RIM $\color{blue!70!black} h_{t,2}$ to remain unchanged on that step (middle).  It leads the activated RIMs to be updated using the input attention, then the independent RIM-specific dynamics $\color{red!80!black} D_{3}$, then the communication between RIMs (right, shown for a single activated RIM $\color{red!80!black} h_{t,3}$).  }
\label{fig:blocks_mechanics}
\end{figure*}

\begin{figure}
    \centering
    \includegraphics[width=0.7\linewidth]{figures/top_down_jh.png}
    \caption{Different mechanisms produce a query. The top-down inference enhances the target but also suppresses other stimuli, including the distractors.}
    \label{fig:minigrid_pick}
\end{figure}

\begin{figure}
    \centering
    \includegraphics[width=\linewidth]{figures/blocks.png}
    \caption{\JH{Updated figure. Let me know if you want any changes!} \Bernhard{why are some blocks white and some light blue? (apart from some being blue, which I assume means they are active)?} \Bernhard{from each block there is an arrow to the right, and then the blocks are printed again. That seems redundant? Does it illustrate that they all follow their default dynamics? but then I thought only the non-chosen ones follow the default dynamics, and others interact?}\Bernhard{please refer to the colors, e.g. what are the orange and pink boxes?} \Ani{Each block follows its own dynamics, and then the activated blocks, communicate with each other.} \Ani{I'm color blind, I'll change some light blue to white.}.An illustration of the mechanics of the Recurrent Independent Mechanisms (RIMs) model. A single step under the proposed model occurs in four steps.  In the first step, the RIMs produces a query vector using which they decide which portion of input to access.  In the second step, there's a competition process involved between different RIMs to select which RIM to activate. In the third step, each RIM follows its own default transition dynamics.  In the fourth step, the (active) RIMs communicate information between themselves using key-value attention process.}
    \label{fig:minigrid_pick}
\end{figure}

The RIMs architecture consists of the following elements:

\textbf{Sparse Activation of RIMs:} A system of RIMs could  modularize the recurrent dynamics, which has the advantage to localize the computation in the whole system. In case the input changes, the modularity should permit to modify only some part of the whole system i.e., states of only some RIMs should change. So, depending upon contextual information (information from the environment) we want only certain RIMs to modify their states i.e., we want to achieve sparse activation of RIMs.

\textbf{RIM Independent Dynamics:} A subset of the RIMs become activated, and these RIMs run their own independent dynamics, which are parameterized separately. \

\textbf{Communication Between RIMs:} The RIMs which are activated can communicate with other RIMs to exchange information.  Thus the RIMs will operate independently by default and only communicate amongst each other using some mechanism.  In next section, we show how we achieve this information exchange between different RIMs.

\subsection{Sparse Activation of RIMs}
The proposed model learns to dynamically select only a few RIMs for which the current input is relevant.  RIMs are triggered as a result of interaction between the current state of the RIM and input information coming from the environment.  The basic intuition behind the mechanism is that each RIM attends to both the input and a null input.  The null input consists entirely of zeros and thus contains no information.  Our criteria for activation is that we select the top-$k_A$ out of $k_{T}$ RIMs on each step which pay the most attention to the real input and pay the least attention to the null input.  Intuitively, the RIMs must compete on each step to read from the input, and only these RIMs that ``win'' this competition will be able to read from the input and have their state updated.

The attention mechanism that we use here is a key-value attention, where the read-key (or query) comes from the RIMs and the write-keys and write-values come from the input on the step and the null input.  The mechanics of this attention mechanism closely follow \citep{vaswani2017attention,santoro2018relational} but come with a major modification - which is that the parameters of the attention mechanism itself are separate per each RIM. The input attention for a particular RIM is described as follows.

\begin{align}
    \tilde{x}_t &= \emptyset \oplus x_t \\
    A^{(in)}_{\theta_k}(h_{t,k}) &= \text{softmax} \left( \frac{h_{t,k} W_k^q(\tilde{x}_t W^e)^T}{\sqrt{d_e}} \right) \tilde{x}_t W^v \label{eq:attention}, \text{ where } \theta_k =(W_k^q, W^e, W^v)
\end{align}

$W^v$ is a simple weight matrix mapping from the input to the write-values for the attention and $W^e$ is a weight matrix which maps the input to the write-keys.  $W_k^q$ is a per-RIM weight matrix which maps from the RIM's hidden state to its read-key.  $\oplus$ refers to the concatenation operator.

Based on this input attention, we select the $k_A$ RIMs (out of the total $K$ RIMs) to be activated for each step, and we call this set $\mathcal{S}_t$.  We do this by selecting the top-$k_A$ RIMs which have the highest attention over the input (and thus put the least attention on the null input).  $(A)_0$ refers to indexing the first input element of the attention matrix for each RIM (in this case attention on the null input).

\begin{equation}
    \mathcal{S}_t = \underset{k_A \text{ of } K}{\argmin} \left( \frac{h_{t,k} W_k^q(\tilde{x}_t W^e)^T}{\sqrt{d_e}} \right)_0
\end{equation}

This is the set $\mathcal{S}_t$ of $k_A$ elements per-step that it used to determine which RIMs are updated.

\subsection{RIM Dynamics}

Next, consider the default transition dynamics which we apply for each RIM independently and during which no information passes between RIMs.  We use $\tilde{h}$ for the hidden state after the independent dynamics are applied (and before the RIM-to-RIM communication has been applied).  First, for the RIMs which are not activated, the hidden state remains unchanged:
\vspace{-2mm}
\begin{equation}
    \tilde{h}_{t+1,k} = h_{t,k} \quad\quad \forall k \notin \mathcal{S}_t.
\end{equation}

An important note is that the gradient still flows through a RIM on a step where it is not activated.

For the RIMs that are activated, we run a per-RIM independent transition dynamics.  The form of this is somewhat flexible, but in this work we opted to use either a GRU \citep{chung2015recurrent} or an LSTM \citep{hochreiter1997long} for these transition dynamics.  We generically refer to these independent transition dynamics as $D_k$, and we emphasize that each RIM has its own separate parameters.  Aside from being RIM-specific, the internal operation of the LSTM and GRU remain unchanged, and they receive the result of the input-attention mechanism as input.

\begin{equation}
    \tilde{h}_{t+1,k} = D_k(h_{t,k}) = LSTM(h_{t,k}, A^{(in)}_{\theta_k}(h_{t,k})); \theta_k) \quad\quad \forall k \in \mathcal{S}_t
\end{equation}

\subsection{Communication between RIMs}

Although the RIMs operate independently by default, it is still the case that some time steps may require sharing information between the RIMs.  We accomplish this by using an attention mechanism controlling communication between the RIMs.  Only the activated RIMs are allowed to read from other RIMs, but the activated RIMs are allowed to read from any RIM, even those that are not activated.  The motivating intuition behind this is that non-activated RIMs are not related to the current input, so their value should not change.  However they may still store information that is relevant for other RIMs, so other RIMs should be able to read from them.  For this communication between RIMs, we also use a residual connection as in \citep{santoro2018relational} to prevent vanishing and exploding gradients over long sequences.

\begin{align}
    M_{t,k}^e &= W^{e_2}_k \tilde{h}_{t,k}, \quad \forall k \\
    M_{t,k}^v &= W^{v_2}_k \tilde{h}_{t,k}, \quad \forall k \\
    M_{t,k}^q &= W^{q_2}_k \tilde{h}_{t,k}, \quad \forall k \in \mathcal{S}_t \\
    h_{t+1,k} &= \text{softmax} \left( \frac{M_{t,k}^q (M_{t,:}^e)^T}{\sqrt{d_e}} \right) M_{t,:}^v + \tilde{h}_{t,k} \quad \quad \forall k \in \mathcal{S}_t
\end{align}

\section{Related Work}

A few papers discuss related ideas and thus we briefly outline in Tab. \ref{tb:related} how the proposed model is different from past work.

\begin{table}[htb!]
{\renewcommand{\arraystretch}{1.3}
\caption{A concise comparison of recurrent models with modular memory.  }
\label{tb:related}
\scalebox{0.8}{
\begin{tabular}{|l|l|l|l|l|}
\hline
Method / Property         & \makecell{Modular\\Memory} & \makecell{Sparse \\Information Flow} & \makecell{Modular \\Computation Flow}     & \makecell{Modular\\Parameterization} \\ \hline
LSTM / RNN & \xmark & \xmark & \xmark & \xmark \\ \hline
Relational RNN \citep{santoro2018relational}  & \cmark   & \xmark & \cmark & \xmark \\ \hline
NTM \citep{graves2014ntm} & \cmark & \cmark & \xmark & \xmark \\ \hline
SAB\citep{ke2018sparse}  & \xmark & \cmark  & \xmark & \xmark  \\ \hline
IndRNN\citep{li2018independently} & \cmark & \xmark & \xmark & \cmark  \\ \hline
RIMs    & \cmark & \cmark  & \cmark & \cmark \\ \hline
\end{tabular}
}
}
\end{table}

\textbf{Neural Turing Machine (NTM)}.  The Neural Turing Machine (NTM) has a Turing machine inspired memory with a sequence of independent memory cells, and uses an attention mechanism to move heads over the cells while performing targeted read and write operations.  This shares a key idea with RIMs: that input information should only impact a sparse subset of the memory by default, while keeping most of the memory unadultered.  However the RIMs model introduces the idea that each RIM has its own independent dynamics, whereas the memory cells remain dormant by default.

\textbf{Relational RNN}.  The Relational Models paper \citep{santoro2018relational} is based on the idea of using a multi-head attention mechanism to share information between multiple parts of memory.  It is related to our idea but a key difference is that we encourage the RIMs to remain separate as much as possible, whereas \citep{santoro2018relational} allow information between the parts to flow on each step (in effect making the part distribution only relevant to a particular step).  Additionally, RIMs has the notion of each RIM having its own independent transition dynamics which operate by default, whereas the Relational RNN only does computation and updating of the memory using attention.

\textbf{Sparse Attentive Backtracking (SAB)}.  The SAB architecture \citep{ke2018sparse} 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 in the backward pass is sparse.  It shares the motivation of using sparse attention to keep different pieces of information separated, but differs from the RIMs model in that it considers separation between time steps rather than separation between RIMs

\textbf{Independently Recurrent Neural Network (IndRNN)}.  The IndRNN \citep{li2018independently} uses something similar to a vanilla RNN but replaces the full transition matrix (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 an RNN to learn long-term dependencies without vanishing or exploding gradients, even when using an activation such as a relu.  They also produce a theoretical analysis of the gradients showing that having smaller recurrent transition matrices (in this case each unit can only look at itself) resolves the vanishing and exploding gradient issue.  This could also help to provide further explanation for why RIMs performs well.

\textbf{Modularity and Neural Networks}: In this framework, a neural network is composed using several neural modules, where each module is meant to perform a distinct function, and hence can be seen as a combination of experts \citep{jacobs1991adaptive, bottou1991framework,ronco1996modular, reed2015neural, andreas2016neural,ParRojKilSch17,rosenbaum2017routing, fernando2017pathnet, shazeer2017outrageously, kirsch2018modular} that route information through a gated activation of layers. However,
these works assume that only a single expert is active at a particular time step which limits the strength of these models.
Hence, our work is related in the sense that we model the recurrent dynamics using many modules, i.e., RIMs but in our framework. Different modules can interact and share information, and more than one RIM can be active at a particular time step.

\textbf{Consciousness Prior} \citep{bengio2017consciousness}: This is based on the assumption of a sparse graphical model describing the interactions between high-level variables, using gating mechanisms to select only a subset of high-level variables to interact at any particular time.  As a prior, it says that a variable can be predicted well from very few others.    This is closely related to our work in the sense high level abstract representation is based on the representations of the RIMs, which are activated sparsely and interact sparsely. Our paper thus helps to validate the consciousness prior idea.

\textbf{Recurrent Entity Networks}:   EnTNet \citep{henaff2016tracking} can be viewed as a set of separate recurrent models whose hidden states store the memory
slots. These hidden states are either fixed by the gates, or modified through a simple RNN-style
update. Moreover,  EntNet uses an independent gate for writing to each memory slot. Our work is related in the sense that we also have different recurrent models (i.e.,RIMs, though each RIM has different parameters), but we allow the RIMs to communicate with each other sparingly using an attention mechanism.

\Ani{Have to talk here how the proposed model is different from these things.}

\textbf{Phased LSTMs and Hierarchical RNNs.} We agree that the idea of having different computational modules which are executed on demand is indeed not new. There are important architectural differences though: (a) We modularize the dynamics of LSTM (using RIMs). (b) We also \textbf{control the inputs of each module (with transformer style attention)}, and many previous gating methods did not control the inputs of each module, only whether they should be executed or not. (c) Multiple RIMs can be active at the same time (unlike Mixture of experts, where you assume only one expert to be active).

\textbf{Computation on demand:} The central idea of the paper to separate a network into (functionally) different modules that perform computation on demand is not new. For the proposed LSTM variant, especially "Phased LSTM: Accelerating Recurrent Network Training for Long or Event-based Sequences" and "A Clockwork RNN" seem relevant, where also parts of the LSTM's hidden state are kept dormant for certain number of frames. While the proposed approach also provides a more sophisticated interaction between LSTM modules than "phased LSTMs", the described experiments do not underline or investigate this feature.

\Ani{Discuss all these}
\cite{el1996hierarchical, koutnik2014clockwork,  neil2016phased, chung2016hierarchical}

\Ani{Cite outrageously large neural networks paper.}

\section{Experiments}
We briefly outline the tasks on which we applied the RIMs model, and direct the reader to the appendix
for full details on each task and details on hyperparameter settings for the model. We have produced experimental results on four tasks: language modeling, sequence copying (with evaluation sequences longer than those seen during training), imitation learning (where states have high-level noise injected only during evaluation), and  Reinforcement Learning task with distractors.  Intriguingly, we found that RIMs led to the largest improvements when the training and test distributions were made to be systematically different.

Mention the list of things what our experiments actually show.

\begin{itemize}
    \item
    \item
\end{itemize}

\Ani{May be a good idea to structure the experiments is to say, that specialization can occur based on events as well as spatial information.}

\subsection{Baselines}

\begin{itemize}
    \item Relataional Memory
    \item Transformers and LSTM
    \item EntNet
    \item Visual Interaction Networks
    \item RIMs with no data-dependent competition.
    \item SKIP RNN: LEARNING TO SKIP STATE UPDATES IN
RECURRENT NEURAL NETWORKS
\end{itemize}

\subsection{Copying Task}

The copying task has been widely studied in the context of RNNs and neural sequence models.  It consists of receiving a short sequence of characters, then receiving blank inputs for a large number of steps, and then being forced to produce the original received sequence.  This task studies a simple property in isolation: is the model able to keep things in memory while no active input is being received.  LSTMs struggle to perform the copying task with very long sequences and especially struggle when evaluated on sequences which are longer than those seen during training.  The RIMs model may be able to address this by using one RIM to store the information at the start of the sequence, and then keep that RIM dormant until the information is retrieved.

\begin{table}[htb!]
\centering
{\renewcommand{\arraystretch}{1.0}
\caption{\textbf{Error (negative log likelihood) on the copying task}.  Note that while all of the methods are able to 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.  }
\label{tb:copy}
\centering
\scalebox{0.9}{
\begin{tabular}{lrrr}
\toprule
Approach & Train Length 50 & Test Length 75  \\
\midrule
LSTM (1-layer) & 0.01 & 6.13  \\
LSTM + Self. Attention & 0.00 & 4.60 \\
NTM \citep{graves2014neural}  & 0.00 & 1.30 \\
Relational Memory \citep{santoro2018relational} & 0.00 & 0.13 \\
\textbf{RIMs} ($k_T$ = 6, $k_A$ = 5) & 0.01 & 0.35 \\
\textbf{RIMs} ($k_T$ = 6, $k_A$ = 4) & \textbf{0.00} & \textbf{0.00}\\
\textbf{RIMs} ($k_T$ = 6, $k_A$ = 3) & \textbf{0.00} & \textbf{0.00}  \\
\textbf{RIMs} ($k_T$ = 6, $k_A$ = 2) & \textbf{0.00} & \textbf{0.00}  \\
\textbf{RIMs} ($k_T$ = 5, $k_A$ = 4) & \textbf{0.00} & \textbf{0.00} \\
\textbf{RIMs} ($k_T$ = 5, $k_A$ = 3) & \textbf{0.01} & \textbf{0.00}\\
\textbf{RIMs} ($k_T$ = 5, $k_A$ = 2) & \textbf{0.00} & \textbf{0.00}  \\
\textbf{RIMs} ($k_T$ = 4, $k_A$ = 3) & 0.14 & 0.10  \\
\textbf{RIMs} ($k_T$ = 4, $k_A$ = 2) & \textbf{0.00} & \textbf{0.00}  \\
\textbf{RIMs} ($k_T$ = 3, $k_A$ = 2) & \textbf{0.00} & \textbf{0.00}  \\
\midrule
With input Attention and No Communication & & \\
\midrule
\textbf{RIMs}
($k_T$ = 6, $k_A$ = 5) & 0.22 & 0.56 \\
\textbf{RIMs} ($k_T$ = 6, $k_A$ = 4) & 0.07 & 0.50\\
\textbf{RIMs} ($k_T$ = 6, $k_A$ = 3) & 0.12 & 0.10  \\
\textbf{RIMs} ($k_T$ = 6, $k_A$ = 2) & 0.33 & 0.24\\
\textbf{RIMs} ($k_T$ = 5, $k_A$ = 4) & 0.02 & 0.07 \\
\textbf{RIMs} ($k_T$ = 5, $k_A$ = 3) & 0.04 & 0.40\\
\textbf{RIMs} ($k_T$ = 5, $k_A$ = 2) & 0.25 & 0.47  \\
\textbf{RIMs} ($k_T$ = 4, $k_A$ = 3) & 0.17 & 0.43  \\
\textbf{RIMs} ($k_T$ = 4, $k_A$ = 2) & 0.23 & 0.16  \\
\textbf{RIMs} ($k_T$ = 3, $k_A$ = 2) & 0.06 & 0.27  \\

\bottomrule
\end{tabular}
}
}
\end{table}

\textbf{Results:} Table \ref{tb:copy} shows the result when the proposed model as well as baselines are trained on sequences of length 50, and evaluated on length 75. Due to the ability of RIMs, to store information for all times spans, the RIMs model outperform many strong baselines such as Neural Turing Machines \citep{graves2014ntm}, LSTMs combined with Self Attention, and Relational Memory architecture \citep{santoro2018relational}. We also conducted experiments where training was performed on sequence length 50 and evaluation on length 100, 150, 200. We observed that our proposed model achieved zero test error despite longer evaluation sequences.  Additionally we found that if we had 6 RIMs, the model generalized if we had top-k activation between $k_A=2$ and $k_A=4$.

\begin{figure}[htb!]
    \centering
    \includegraphics[width=0.9\linewidth,trim=8.0cm 0.0cm 4.0cm 0.0cm,clip]{figures/blocks_copying_fig.png}
    \caption{\textbf{Copying Task RIM Activation Pattern} for a model with $K=6$ RIMs and $K_A=3$ active RIMs per step (the activated RIMs are in black, non-activated in white).  We can see that the RIM activation pattern is distinct during the middle dormant part of the sequence.  }
    \label{fig:copying_pattern}
\end{figure}

\subsection{Sequential MNIST Resolution Task}

RIMs is motivated by the hypothesis that generalization performance can be improved by having modules which only activate on relevant parts of the sequence.  For further evidence that RIMs can achieve this, we considered the task of training to classify MNIST digits as sequences of pixels \citep{krueger2016zoneout} and studying if such a model could generalize to images of different resolutions from those seen during training.  Our intuition is that the RIMs model should have distinct subsets of the RIMs activated for pixels with the digit and empty pixels.  As a result, the RIMs model should generalize better to greater resolutions by keeping the RIMs which store pixel information dormant over the empty regions of the image.

\begin{table}[htb!]
\centering
\caption{\textbf{Sequential MNIST resolution:} Test Accuracy \% on the Sequential MNIST resolution generalization task.  Both the proposed and the Baseline model (LSTM) were trained on 14x14 resolution but evaluated at different resolutions.  }
\label{tb:mnist}
{\renewcommand{\arraystretch}{1.0}
\begin{tabular}{lrrr}
\toprule
\centering
Method / Evaluation & 16 x 16 & 19 x 19 & 24 x 24 \\
\midrule
Increase in Sequence Length (\%) & 30.6\% & 84.2\% & 196.9\% \\
LSTM & 56.34 & 13.44 & 8.80 \\
Relational Memory \citep{santoro2018relational} & 66.58 & 44.62  & 19.90 \\
\textbf{RIMs} ($k_T=6$, $k_A=4$) & 77.54 & 47.67 & 19.29  \\
\bottomrule
\end{tabular}
}
\end{table}

\subsection{MultiMNIST Task}

\subsection{Language Modeling}

We ran on the wikitext-2 dataset \citep{merity2016pointer}.  We ran each experiment for a fixed 100 epochs.  These results are in Table~\ref{tb:lm}.  Our goal in this experiment is to demonstrate the breadth of the approach by showing that RIMs performs well even on datasets which are noisy and drawn from the real-world.

\begin{table}[htb!]
\caption{Wikitext-2 results}
\label{tb:lm}
\centering
\begin{tabular}{lrrrr}
\toprule
Approach & Num. Parameters  & Train PPL & Valid PPL & Test PPL  \\
\midrule
LSTM (2-layer) & 21.2M & 39.78 & 109.25 & 102.53  \\
Relational Memory \citep{santoro2018relational} & 11M & n/a & 112.77 & 107.21 \\
\textbf{RIMs} (2-layer, $k_T$ = 6, $k_A$ = 6) & 23.7M & 41.27 & 103.60 & \textbf{98.66}  \\
\bottomrule
\end{tabular}
\end{table}

\subsection{Robustness to Novel Distractors}

An important property for machine learning systems is to be able to perform well even in the presence of distracting or irrelevant entities in the environment.  It is especially difficult and important to be resilient to these distractions when they are absent from the training data. For example, an LSTM with a single memory flow may struggle to keep the distracting elements separate from elements which are necessary for the task. Because the RIMs model uses attention to control which RIMs receive information at each step, it should be possible for the model to make a clear decision to not write distracting information into the important RIMs.

\textbf{Task Setup:} We setup an RL environment for object picking where we can directly study the impact of novel distractors, by training the model to pick one object during training (in the presence of distractors), and then introduce unseen distractors (i.e.,distractors which are not seen during training) during testing.  We use a partially observed formulation of the task, where the agent only observes a small number of squares ahead of it. These tasks are difficult to solve with standard RL algorithms, not only due to the partial observability of the environment but also the sparsity of the reward, given that the agent receives a reward only after reaching the goal.

\begin{figure}[htb!]
    \centering
    \includegraphics[width=0.45\linewidth,trim=0.5m 0.0cm 1.0cm 1.2cm,clip]{figures/figure_normal.png}
    \includegraphics[width=0.45\linewidth,trim=0.5cm 0.0cm 1.0cm 1.2cm,clip]{figures/figure_distractor.png}
    \caption{\textbf{Robustness to Novel Distractors:}. Left shows the performance of the proposed method as compared to the LSTM baseline to solve the object picking task in the presence of distractors.  Right shows the performance of the  proposed method and the baseline when novel distractors are added.  Success of an agent is measured, e.g. by the fraction of episodes where the agent was able to solve the task in 500 steps. (averaged over 500 examples}
    \label{fig:rl_exp}
\end{figure}

\textbf{Results:} Fig. \ref{fig:rl_exp} shows our result. The RIMs model is able to completely solve the task, on the other hand the baseline model (i.e., LSTM in this case) struggles to solve even the training task. Since, the LSTM has a single memory flow it may be difficult for LSTM to keep the information regarding distractors seperate from task relevant information. We also evaluate both the proposed model, as well as LSTM baseline when we add novel distractors (which are not seen during training). The performance of both the LSTM model as well as RIMs model degrade, though RIMs still performs better than the LSTM.

\section{Bouncing Ball experiment}
\JH{Adding temporary figures. Will update them.}
We assess the quality of RIMs by predicting the temporal dynamics of bouncing balls from \cite{van2018relational}.  The base dataset consists of  4 balls of two different masses, represented by their radii, bouncing off one another in a confined square geometry.
In all experiments, for 25 frames we feed in the ground truth solution to update the hidden state. For the next 25 frames, we use the previous output. This allows us to test the ability of the network to accurately predict future dynamics in a complex environment.

\subsection{Different RIMs attend to different balls (Geometric figures)}

\begin{figure}[htb!]
    \centering
    \includegraphics[width=0.95\linewidth]{figures/RGB_1.png}
    \caption{\textbf{RIMs on colored, bouncing balls:}. On top,  we show the prediction (with black background) and target (with white background) during the phase where at each time point, the true value is fed in. On the bottom, we show the rollout phase where the output from the previous time is the input to the next time-step.}
    \label{fig:rgb}
\end{figure}

\Sha{Add results of probing the RIM representations for decoding the object id to see if different blocks map to different objects}

\subsection{Balls of different colors}
\JH{Adding temporary figures. Will update them.}
Using the 4 ball dataset, we convert the input to an RGB image where each ball is assigned one of four colors at random.

\Sha{Add results of probing the RIM representations to predict the relative shape/size}

\Sha{Add results of probing the RIM representations to predict the relative shape/size}

\Sha{Add T-SNE representation of the RIM representations}

\Sha{Check if certain blocks are activated during collisions etc}

\subsection{Curtain Balls}
\JH{Adding temporary figures. Will update them.}
An additional dataset provided in \cite{van2018relational} has 3 balls though a small square region is excluded.
This allows us to assess the ability of balls to be tracked through the occluding region.

\subsection{Interpolation}

\subsection{Extrapolation}
We use a test-set with sequences containing 6-8 balls to evaluate the ability of each method to extrapolate their learned knowledge about physical interactions
between four balls to environments with more balls.

\subsection{Plug and Play}

\subsection{Effect of number of blocks}

\section{BabyAI - RL experiments}
As a final proof of concept we consider the task of learning
a proposed model capable of solving multiple RL problems at
the same time. The basic approach for this sort of task is to
train a model in a mixture of environments or equivalently
to train a shared model in multiple environments in parallel (Teh et al., 2017; Espeholt et al., 2018). However, this
sort of training can suffer from two drawbacks. First, it is
heavily reward scale dependent, and will be biased towards
high-reward environments. Second, environments that are
easy to train provide a lot of updates for the model and consequently can also bias the solution towards themselves.

\section{Atari: Specialization based on objects and events}

The study of vision must therefore include not
only the study of how to extract from images,
but also an inquiry into the nature of the internal
representations by which we capture this information and thus make it available as a basis for
decisions about our  actions.

\section{Discussions: Useful characteristic of proposed model}

\subsection{ The role of top-down influence}

RIMs use a soft key, query, and value type of attention similar to \cite{}, but instead of doing
“self"-attention where the queries come from the input (together with the keys and values) our queries come from each RIM which does not directly depend on the input. This enables our system to be both state/context dependent and input dependent.

It is important to emphasize several points about the proposed model. First, the model is fully
differentiable due to the use of soft-attention and can be trained using back-propagation. Second, the
query vectors are a function of the RIM state alone and not the observation — this allows for a
“top-down” mechanism where the RIM can actively query the input for task-relevant information rather than having to filter out large amounts of information. See section xyz in experiments for emperical evaluation.

a top-down bias is more directly responsible for the diminished distractor processing in high-load conditions than the competition among stimulus items. It is the application of this bias in favor of the target that results in diminished representation of the distractor. The top-down filter not only enhances the target but also suppresses other stimuli, including the distractor. The notion that enhancement of the target necessarily results in suppression of the other items in the display is fundamental to biased competition theory (Desimone and Duncan, 1995) and divisive inhibition more generally (Reynolds and Heeger, 2009). Because neurons representing individual items have mutually inhibitory connections, an increase in the firing rate of one results in a decrease in the other.

\subsection{Systematic generalization and modularity and "data dependent" activation of modules}

See section xyz in experiments for emperical evaluation.

\Ani{Ask Prof. Yoshua to see if it makes sense to say.}

\subsection{Multiagent and mechanism design}

\Ani{Ask Prof. bernhard to write a small description here.}

\subsection{Dynamic Binding}

\Ani{Say something:} Its important to emphasize that the topology of the graph induced where the nodes are RIMs, and the edges are determined by attention parameters is dynamic, while in most of the graph neural networks, the topology is fixed. See section xyz in experiments for emperical evaluation.

\subsection{Plug and Play Transfer}

The RIMs should be somewhat interchangeable, and not strongly co-adapted with other particular specific RIMs. In other words, they should be able to capture a certain type of dynamics even when other RIMs aren’t present. it possible to recompose RIMs learned on different tasks. Today’s neural networks are prone to fail to out of distribution samples. A hallmark of human intelligence by contrast is the ability to re-use previously acquired experience and expertise, to transfer it to radically different challenges. So, we should be able to plug and play.

\subsection{Adding More RIMs}

Adding capacity to an already trained deep network is generally quite tedious, as it involves changing the shape of the parameters, and it always requires retraining.  An intriguing possibility is that the RIMs model will allow for adding capacity to a model simply by adding new RIMs, as none of the existing parameters need to be changed.  In fact the new RIMs have no impact on the model's training unless the attention mechanism looks at the newly added RIMs, which we would expect to be relatively rare.  Progressively adding capacity has been explored for generative adversarial networks, where it achieved state of the art results \citep{karras2017progressive}.

\vspace{-1mm}

\section{Future Work}

What is the mechanism that incentivizes the RIMs to collectively optimize the global objective as a by-product of each RIM optimizing its own local objective?

What should the payoff structure be such that the Nash equilibrium would correspond to the solution for optimizing the above objective?

Is it possible to let the information about the state be the currency with which the agents communicate with each other? If so, what should this communication scheme be?

\section{Conclusion}

The world consists of many dynamic processes which primarily operate independently and only occasionally have meaningful interactions.  Despite this, the vast majority of machine learning models employ the opposite inductive bias: that all processes interact, and force the model to unlearn this bias.  This leads to inflexibility and often very poor generalization to changing task distributions.  We have proposed a new model, which we call Recurrent Independent Mechanisms (RIMs), in which we learn multiple recurrent modules which are independent by default, but interact sparingly. The learned RIMs can be flexibly recombined to solve more complex tasks. Our  positive experimental results validate the appropriateness of the previously proposed consciousness prior~\citep{bengio2017consciousness}. The proposed framework could also be very attractive for continual learning settings, where one could add more RIMs over time, thereby the already learned RIMs would keep their focus on particular aspect of the task, and newly added ones could specialize on novel aspects.

\newpage
\bibliographystyle{nips_2018}

\section{Introduction}

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 physics. These processes are often independent and only interact sparsely. For instance, we can model the motion of two balls as separate independent mechanisms even though they are both gravitationally coupled to Earth as well as (weakly) to each other. They may, however, occasionally strongly interact via collisions.

The notion of independent or autonomous mechanisms has been fruitful in the field of causal inference, where it is applied not only to dynamic processes but also to time independent datasets. For instance, it has been argued that the conditional distribution of the average annual temperature given altitude is an abstraction of a causal mechanism (subsuming complex physical processes involving air pressure, etc.) that is independent of the distribution of the altitudes of cities \citep{PetJanSch17}, and will thus apply invariantly for, say, different countries in the same climate zone with different altitude distributions.

A complex generative model (temporal or not) can thus be thought of as the composition of independently specified causal modules or mechanisms. In the causality community, the latter is often considered a prerequisite to the possibility of being able to, in principle, intervene upon variables determined by such models. It has been argued that the individual modules tend to remain invariant even when other modules change, e.g., in the case of distribution shift \citep{PetJanSch17}. One may hypothesize that if a brain is able to solve multiple problems and multiple tasks beyond a single i.i.d.\ dataset, it would be economical to learn structures aligned with this by using independent components that can flexibly be reused, composed and repurposed.

In the dynamic setting, we could think of the world as composed of a large number of fairly independent subsystems that change and evolve over time, responding to forces and interventions. For a learning agent, interacting with such an environment, not all these subsystems will significantly interact with the learner most of the time.
A learner thus does not need to know the dynamics of all subsystems at each time-step to correctly model aspects of the world relevant to a current decision or plan.  Only some aspects of the world significantly interact and thus need to be considered jointly when taking a decision or forming a plan \citep{bengio2017consciousness}. We assay this in the context of multiple dynamical subsystems which can only interact with each other in a sparse way, i.e., with few subsystems influencing other subsystems at a particular time. Such sparse interactions reduces the difficulty of learning since few interactions need to be considered at a time, reducing unnecessary interference when a subsystem is adapted.

Models learned this way may be more likely to capture the compositional causal structure of the world, and thus better generalize across tasks where a (small) subset of mechanisms change while most of them remain invariant \citep{simon1991architecture, PetJanSch17,ParRojKilSch17}. The central question motivating our work is how a machine learning system can learn
independent but weakly (or sparsely) interacting mechanisms in order to benefit from the hypothesized modular structure of the world. Many approaches in machine learning represent the world at a point in time with a single high-dimensional hidden state which attempts to describe the full world-state. It could thus be difficult for typical recurrent networks to describe a world which consists of nearly independent systems that only occasionally interact.

\section{Recurrent Independent Mechanisms with Sparse Interactions}

\iffalse
\begin{figure*}[t]
\centering
\begin{tikzpicture}[
    scale=1.0,
    black!50, text=black,
    font=\medium,

    node distance=1mm,
    dnode/.style={
        align=center,

        font=\small,
        rectangle,minimum size=15mm,rounded corners,

        inner sep=20pt},
    rnode/.style={
        align=center,
        font=\medium,

        rectangle,
        minimum width=5mm,
        minimum height=6mm,
        rounded corners,

        inner sep=3pt,
        thin, draw},
    tuplenode/.style={
        align=center,

        rectangle,minimum size=10mm,rounded corners,

        inner sep=15pt},
    darrow/.style={
        rounded corners,-latex,shorten <=5pt,shorten >=1pt,line width=2mm},
    mega thick/.style={line width=2pt}
    ]
\matrix[row sep=7.2mm, column sep=2mm] {
    \node (a_d) [rnode, minimum width=15mm, top color=greenfill!40, bottom color = greenfill!40] {Surfing RIM};
    &&
    &&
    \node (a_1) [rnode, minimum width=15mm, top color=greenfill!40, bottom color = greenfill!40] {$h_{t+2,2}$};
    &&
    &&
    \node (a_2) [rnode, minimum width=15mm, top color=greenfill!40, bottom color = greenfill!40] {$h_{t+4,2}$};
    \\
    \node (b_d) [rnode, minimum width=15mm, top color=redfill!40, bottom color = redfill!40] {Studying RIM};
    &&
    \node (b_1) [rnode, minimum width=15mm, top color=redfill!40, bottom color = redfill!40] {$h_{t+1,1}$};
    &&
    &&
    \node (b_2) [rnode, minimum width=15mm, top color=redfill!40, bottom color = redfill!40] {$h_{t+3,1}$};
    &&
    &&
    \node (b_3) [rnode, minimum width=15mm, top color=redfill!40, bottom color = redfill!40] {$h_{t+5,1}$};
    \\
    &&
    \node (x_1) [rnode, minimum width=15mm, top color=black!10, bottom color = black!10] {\makecell{Study\\for Test}};
    &&
    \node (x_2) [rnode, minimum width=15mm, top color=black!10, bottom color = black!10] {\makecell{Prepare\\for Surfing}};
    &&
    \node (x_3) [rnode, minimum width=15mm,top color=black!10, bottom color = black!10] {\makecell{Think about\\mathcal{D_{\mathrm{test}}} before wave}};
    &&
    \node (x_4) [rnode, minimum width=15mm, top color=black!10, bottom color = black!10] {\makecell{Start\\Surfing}};
    &&
    \node (x_5) [rnode, minimum width=15mm, top color=black!10, bottom color = black!10] {\makecell{Take\\the test}};
    \\
};

\draw[-latex,shorten <=1pt,shorten >=1pt,thick,color=blue!60]  (a_1) to node[]{} (a_2);
\draw[-latex,shorten <=1pt,shorten >=1pt,thick,color=red!60]  (b_1) to node[]{} (b_2);
\draw[-latex,shorten <=1pt,shorten >=1pt,thick,color=red!60]  (b_2) to node[]{} (b_3);
\draw[-latex,shorten <=1pt,shorten >=1pt,thick,color=black!30]  (x_1) to node[]{} (b_1);
\draw[-latex,shorten <=1pt,shorten >=1pt,thick,color=black!30]  (x_2) to node[]{} (a_1);
\draw[-latex,shorten <=1pt,shorten >=1pt,thick,color=black!30]  (x_3) to node[]{} (b_2);
\draw[-latex,shorten <=1pt,shorten >=1pt,thick,color=black!30]  (x_4) to node[]{} (a_2);
\draw[-latex,shorten <=1pt,shorten >=1pt,thick,color=black!30]  (x_5) to node[]{} (b_3);

\draw[-latex,shorten <=1pt,shorten >=1pt,thick,color=purple!60]  (a_1) to node[]{} (b_2);

\end{tikzpicture}

\caption{An illustration of the ideal behavior of the RIMs model, in which RIMs are sparsely activated and responsible for modeling different aspects of the world.  This helps prevent interference between unrelated aspects of the world.  }
\label{fig:blocks}

\end{figure*}
\fi

We have laid out a case for building models composed of modules which by default operate independently and can interact in a 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 as \textit{Recurrent Independent Mechanisms (RIMs)}, with distinct functions learned automatically from data.\footnote{Note that we are using the term \textit{mechanism} both for the mechanisms that make up the world's dynamics as well as for the computational modules that we learn to model those mechanisms.} Our model encourages sparse interaction, i.e.,  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.

\paragraph{Competitive Mechanisms:}
Inspired by the observations in the previous section, we propose that RIMs utilize competition to allocate representational and computational resources. As argued by \cite{ParRojKilSch17}, this tends to produce independence among learned mechanisms if the training data has been generated by independent physical mechanisms.

\paragraph{Top Down Attention and Dynamic Binding:}

The points mentioned in the previous section, in principle, pertain to synthetic and natural intelligent systems alike. Hence, it is not surprising that they also appear in neuroscience. For instance, suppose we are looking for a particular object in a large scene, using limited processing capacity.
The \emph{biased competition theory} of selective attention conceptualizes basic findings of experimental psychology and neuroscience \citep{Desimone1995}:
our capacity of parallel processing of and reasoning with high-level concepts is limited, and many brain systems representing visual information use competition to allocate resources.  Competitive interactions among multiple objects  occur automatically and operate in parallel across the visual field. Second, the principle of selectivity amounts to the idea that a perceiver has the ability to filter out unwanted information and selectively \emph{process} the rest of the information. Third,  \emph{top-down bias} originating from higher brain areas enables us to selectively devote resources to input information that may be of particular interest or relevance. This may be accomplished by units matching the internal model of an object or process of interest being pre-activated and thus gaining an advantage during the competition of brain mechanisms. Thus, in order for  different competing RIMs to be independent, they need to \textit{selectively} attend to the part of the input based on the top-down inference such that it leads that mechanism to ``win'' the competition, whenever the information present in that part of the input matches to what that particular RIM already expects.

\paragraph{Modular Memory:} Memory is divided into distinct modules, such that it is possible for the model to elect to read to or write from a subset of the modules.  This allows the model to easily keep different pieces of information well-separated over multiple time steps.

\paragraph{Sparse Information Flow:} Each RIMs' dynamics should only be affected by RIMs which are deemed relevant. The fundamental challenge is centered around establishing sensible communication between RIMs.  In the presence of noisy or distracting information, a large subset of RIMs should stay dormant, and not be affected by the noise. This way, training an ensemble of these RIMs can be more robust to out-of-distribution or distractor observations than training one big homogeneous neural network \citep{schmidhuber2018one}.

\paragraph{Modular Computation Flow and Modular Parameterization:} Each RIM should have its own dynamics operating by \textit{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 to \citep{ParRojKilSch17}.
This can speed up computation and improve the generalisation ability of the system \citep{baum1989size}. The individuals RIMs in the ensemble should be simple also to prevent individual RIMs from dominating and modelling complex, composite mechanisms. We refer to a parameterization as modular if most parameters are associated to individuals RIMs only. This has the desirable property that a RIM should maintain its own independent functionality even as other RIMs are changed (due to its behavior being determined by its own self-contained parameters).

\begin{figure}
    \centering
    \includegraphics[width=0.6\linewidth]{figures/blocks.png}
    \caption{An illustration of the mechanics of the Recurrent Independent Mechanisms (RIMs) model. A single step under the proposed model occurs in  four stages. In the first stage, individual RIMs produces a query which is used to read from the current input.  In the second stage, an attention based competition mechanism is used to select which RIMs to activate.  In the third stage, individual activated RIMs follow their own default transition dynamics while non-activated RIMs remain unchanged.  In the fourth stage, the RIMs communicate information between themselves using attention.}
    \label{fig:minigrid_pick}
\end{figure}

\section{RIMs: Proposed Model}

We previously laid out our motivation for the concept of RIMs and gave a high level description: a recurrent model that is separated into RIMs which have independent dynamics and parameters, with each RIM only being activated and updated when the input is relevant to that particular RIM.  However, this leaves many open questions regarding how we select which RIMs to activate and how we define independent dynamics.  At a high-level, our model divides the overall recurrent state into $K$ RIMs which each have a value across $T$ time steps.  Thus, we refer to RIM $k$ at time step $t$ as having state $h_{t,k}$.  Each RIM has parameters $\theta_k$, which are shared across all time steps.

The RIMs architecture consists of the following properties:

\textbf{Sparse Activation of RIMs:} A system of RIMs could  modularize the recurrent dynamics, which has the advantage of localizing the computation in the whole system. If the input changes, the modular nature of RIMs should result in only a subset of RIMs reacting to this change, rather than all of them.

\textbf{RIM Independent Dynamics:} A subset of the RIMs become activated, and these RIMs run their own independent dynamics, which are parameterized separately.

\textbf{Communication Between RIMs:} The RIMs which are activated can communicate with other RIMs to sparsely exchange information.  Thus, the RIMs will operate independently by default and only communicate with each other using key-value attention, described in the next section.

\subsection{Key-Value Attention to Process Sets of Named Interchangeable Variables}

The introduction of soft-attention mechanisms~\citep{bahdanau2014neural} has opened the door to neural networks  which operate on {\em sets of named interchangeable objects}.
This type of attention mechanism has been extremely successful and widely applied to most recent Transformer-style self-attention models \citep{vaswani2017attention, santoro2018relational}, achieving new state-of-the-art results in many tasks. The result of applying soft-attention with $N_h$  head queries $Q_j$ (row $j$ of $N_h \times d$ matrix $Q$, indicating attention head $j$) to a set of $N_o$ objects each associated with a key $K_i$ (row $i$ of $N_o \times d$ matrix $K$) and a value $V_i$ (row $i$ of matrix $V$) can be computed as
$$\mathrm{Attention}(Q,K,V)=\mathrm{softmax} \left (\frac{QK^T}{\sqrt{d}} \right )V$$
where the softmax is applied within each row of its argument matrix, yielding
for each attention head $j$, a set of convex weights. As a result one obtains a convex combination of the values $V_i$ for each head $j$. When the attention is focused on one element (the softmax is saturated) for a particular row, this just selects one of the objects and copies its value in row $j$ of the result.

When the inputs and outputs of each RIM is a set of objects (each associated with a key vector and a value vector), the RIM processing becomes a generic object-processing machine which can operate on ``variables'' in a sense analogous to variables in a programming language, as arguments of functions. Because each object has a key embedding (which one can understand both as a name and as a type), the same RIM processing can be applied to any variable which fits an expected type (or name) specified as a query for a particular head. Each attention head then corresponds to a typed argument of the function computed by the RIM. Whereas in regular neural networks (without attention) neurons operate on fixed variables (the neurons which are feeding them), the key-value attention mechanisms make it possible to select on the fly which variable is going to be used as input for each of the arguments of the RIM dynamics, with a different set of query embeddings for each RIM block. These inputs can come from the external input or from the output of other RIMs. So, if  the individual RIMs can represent these ``typed argument functions'', then they can ``bind'' to whatever input is currently available, then the ``input attention'' mechanism would look at the current input and try to understand if its ``type'' matches to what this RIM knows how to process as its input.

\subsection{Top Down Activation of RIMs}

\begin{wrapfigure}{r}{0.5\textwidth}
    \vspace{-13mm}
    \centering
    \includegraphics[width=0.5\linewidth]{figures/top_down_jh.png}
    \caption{Different RIMs produce a query (which is a function of the RIM state). The activated RIMs gets visual information depending upon their query. Hence,  The top-down inference (due to query) enhances the target (which could be different for each RIM) but also suppresses other visual information, including the distractors.}
    \label{fig:minigrid_pick}
\end{wrapfigure}

The proposed model learns to dynamically select only a few RIMs for which the current input is relevant.  RIMs are triggered as a result of interaction between the current state of the RIM and input information coming from the environment.  The basic intuition behind the mechanism is that each RIM attends to both the input and a null input.  The null input consists entirely of zeros and thus contains no information.  Our criteria for activation is that we select the top-$k_A$ out of $k_{T}$ RIMs on each step which pay the most attention to the real input and pay the least attention to the null input.  Intuitively, the RIMs must compete on each step to read from the input, and only these RIMs that ``win'' this competition will be able to read from the input and have their state updated.

The attention mechanism that we use here is a key-value attention, where the read-key (or query) comes from the RIMs and the write-keys and write-values come from the input on the step and the null input.  The mechanics of this attention mechanism closely follow \citep{vaswani2017attention,santoro2018relational} but come with a major modification - which is that the parameters of the attention mechanism itself are separate per each RIM. The input attention for a particular RIM is described as follows.

The input $x_t$ at time $t$ is seen as a set of element, structured as rows of a matrix. We first concatenate a row full of zeros, to obtain
\begin{align}
    X &= \emptyset \oplus x_t.

\end{align}

Then, linear transformations are used to construct  keys ($K = XW^k$, one per input element and for the null element),  values ($V = XW^v$, again one per element), and  queries ($Q = R W_k^q$, one per RIM block) where $R$ is a matrix with each row $r_i$ corresponding to the hidden state of an individual RIM (i.e $h_{t, k}$). $W^v$ is a simple matrix mapping from an input element to the corresponding value vector for the weighted attention and $W^k$ is similarly a weight matrix which maps the input to the keys.  $W_k^q$ is a per-RIM weight matrix which maps from the RIM's hidden state to its queries.  $\oplus$ refers to the row-level concatenation operator.

\begin{align}
    A^{(in)}_{\theta_k}(R) &= \text{softmax} \left( \frac{R W_k^q(X W^k)^T}{\sqrt{d_e}} \right) X W^v \label{eq:attention}, \text{ where } \theta_k =(W_k^q, W^e, W^v)
\end{align}

Based on this input attention, we select the $k_A$ RIMs (out of the total $K$ RIMs) to be activated for each step, and we call this set $\mathcal{S}_t$.  We do this by selecting the top-$k_A$ RIMs which have the highest attention over the input (and thus put the least attention on the null input).  $(A)_0$ refers to indexing the first input element of the attention matrix for each RIM (in this case attention on the null input).

\begin{equation}
    \mathcal{S}_t = \underset{k_A \text{ of } K}{\argmin} \left( \frac{R W_k^q(X W^k)^T}{\sqrt{d_e}} \right)_0
\end{equation}

This is the set $\mathcal{S}_t$ of $k_A$ elements per-step that it used to determine which RIMs are updated.  Since, RIMs use a  key-value attention where the queries come from  each RIM which depends on the state of the RIMs. This enables individual RIMs to attend only to the part of the input which is relevant for that particular RIM, thus enabling selective attention based on \textit{top-down inference} process.

\subsection{RIM Dynamics}

Next, consider the default transition dynamics which we apply for each RIM independently and during which no information passes between RIMs.  We use $\tilde{h}$ for the hidden state after the independent dynamics are applied (and before the RIM-to-RIM communication has been applied).  First, for the RIMs which are not activated, the hidden state remains unchanged:
\vspace{-2mm}
\begin{equation}
    \tilde{h}_{t+1,k} = h_{t,k} \quad\quad \forall k \notin \mathcal{S}_t.
\end{equation}

An important note is that the gradient still flows through a RIM on a step where it is not activated.  For the RIMs that are activated, we run a per-RIM independent transition dynamics.  The form of this is somewhat flexible, but in this work we opted to use either a GRU \citep{chung2015recurrent} or an LSTM \citep{hochreiter1997long} for these transition dynamics.  We generically refer to these independent transition dynamics as $D_k$, and we emphasize that each RIM has its own separate parameters.  Aside from being RIM-specific, the internal operation of the LSTM and GRU remain unchanged, and they receive the result of the input-attention mechanism as input.

\begin{equation}
    \tilde{h}_{t+1,k} = D_k(h_{t,k}) = LSTM(h_{t,k}, A^{(in)}_{\theta_k}(R_{k})); \theta_k) \quad\quad \forall k \in \mathcal{S}_t
\end{equation}

\text{where} $R_{k}$ refers to row $r_{k}$ of matrix R, corresponding to the hidden state of an individual RIM (i.e $h_{t,k}$).

\subsection{Communication between RIMs}

Although the RIMs operate independently by default, it is still the case that some time steps may require sharing information between the RIMs.  We accomplish this by using an attention mechanism controlling communication between the RIMs.  Only the activated RIMs are allowed to read from other RIMs, but the activated RIMs are allowed to read from any RIM, even those that are not activated.  The motivating intuition behind this is that non-activated RIMs are not related to the current input, so their value should not change.  However they may still store information that is relevant for other RIMs, so other RIMs should be able to read from them.  For this communication between RIMs, we also use a residual connection as in \citep{santoro2018relational} to prevent vanishing and exploding gradients over long sequences.

\begin{align}
    M_{t,k}^e &= W^{e_2}_k \tilde{h}_{t,k}, \quad \forall k \\
    M_{t,k}^v &= W^{v_2}_k \tilde{h}_{t,k}, \quad \forall k \\
    M_{t,k}^q &= W^{q_2}_k \tilde{h}_{t,k}, \quad \forall k \in \mathcal{S}_t \\
    h_{t+1,k} &= \text{softmax} \left( \frac{M_{t,k}^q (M_{t,:}^e)^T}{\sqrt{d_e}} \right) M_{t,:}^v + \tilde{h}_{t,k} \quad \quad \forall k \in \mathcal{S}_t
\end{align}

\section{Related Work}

A few papers discuss related ideas and thus we briefly outline how our proposed model is different from past work.

\textbf{Neural Turing Machine (NTM)}.  NTM \citep{graves2014neural} has a Turing machine inspired memory with a sequence of independent memory cells, and uses an attention mechanism to move heads over the cells while performing targeted read and write operations.  This shares a key idea with RIMs: that input information should only impact a sparse subset of the memory by default, while keeping most of the memory unaltered.  However the RIMs model introduces the idea that each RIM has its own independent dynamics, whereas the memory cells remain dormant by default.

\textbf{Relational RNN}.  The Relational Models paper \citep{santoro2018relational} is based on the idea of using a multi-head attention mechanism to share information between multiple parts of memory.  It is related to our idea but a key difference is that we encourage the RIMs to remain separate as much as possible, whereas \citep{santoro2018relational} allow information between the parts to flow on each step (in effect making the part distribution only relevant to a particular step).  Additionally, RIMs has the notion of each RIM having its own independent transition dynamics which operate by default, whereas the Relational RNN only does computation and updating of the memory using attention.

\textbf{Sparse Attentive Backtracking (SAB)}.  The SAB architecture \citep{ke2018sparse} 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 in the backward pass is sparse.  It shares the motivation of using sparse attention to keep different pieces of information separated, but differs from the RIMs model in that it considers separation between time steps rather than separation between RIMs

\textbf{Independently Recurrent Neural Network (IndRNN)}.  The IndRNN \citep{li2018independently} uses something similar to a vanilla RNN but replaces the full transition matrix (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 an RNN to learn long-term dependencies without vanishing or exploding gradients. Analysis of the gradients showing that having smaller recurrent transition matrices resolves the vanishing and exploding gradient issue. This may provide further explanation for why RIMs perform well.

\textbf{Modularity and Neural Networks}: Here, a neural network is composed of several neural modules, where each module is meant to perform a distinct function, and hence can be seen as a combination of experts \citep{jacobs1991adaptive, bottou1991framework,ronco1996modular, reed2015neural, andreas2016neural,ParRojKilSch17,rosenbaum2017routing, fernando2017pathnet, shazeer2017outrageously,fernando2017pathnet, kirsch2018modular, rosenbaum2019routing} that route information through a gated activation of layers. However,
these works assume that only a single expert is active at a particular time step which limits the strength of these models.
Hence, our work is related in the sense that we model the recurrent dynamics using many modules. Different modules can interact and share information, and more than one RIM can be active at a particular time step.

\textbf{Consciousness Prior} \citep{bengio2017consciousness}: This is based on the assumption of a sparse graphical model describing the interactions between high-level variables, using gating mechanisms to select only a subset of high-level variables to interact at any particular time.  As a prior, it says that a variable can be predicted well from very few others.    This is closely related to our work in the sense high level abstract representation is based on the representations of the RIMs, which are activated sparsely and interact sparsely. Our paper thus helps to validate the consciousness prior idea.

\textbf{Recurrent Entity Networks}:   EnTNet \citep{henaff2016tracking} can be viewed as a set of separate recurrent models whose hidden states store the memory
slots. These hidden states are either fixed by the gates, or modified through a simple RNN-style
update. Moreover,  EntNet uses an independent gate for writing to each memory slot. Our work is related in the sense that we also have different recurrent models (i.e.,RIMs, though each RIM has different parameters), but we allow the RIMs to communicate with each other sparingly using an attention mechanism.

\textbf{Computation on demand:} The central idea of the paper to separate a network into (functionally) different modules that perform computation on demand is not new. There are various different LSTM architectures \citep{el1996hierarchical, koutnik2014clockwork, chung2016hierarchical, neil2016phased, jernite2016variable} where  parts of the LSTM's hidden state are kept dormant for certain number of frames. Though, there are important architectural differences as compared to the proposed ar: (a) We modularize the dynamics of LSTM (using RIMs). (b) We also control the inputs of each module (with transformer style attention), and many previous gating methods did not control the inputs of each module, only whether they should be executed or not. (c) Multiple RIMs can be active at the same time (unlike Mixture of experts, where you assume only one expert to be active).

\textbf{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 \citep{scarselli2008graph, bronstein2017geometric, watters2017visual, raposo2017discovering,santoro2017simple, gilmer2017neural, van2018relational, kipf2018neural, battaglia2018relational, tacchetti2018relational}. One can also view our proposed model as a relational graph neural network, where nodes are parameterized as individual RIMs and edges are parameterized by the attention mechanism. Though, its important to emphasize that the topology of the graph induced in the proposed model is dynamic, while in most graph neural networks the topology is fixed.

\section{Experiments}

We briefly outline the tasks on which we applied the RIMs model, and direct the reader to the appendix for the full details of each task and details on hyperparameter settings for the model.
We organize our results by the model capabilities they illustrate. We address representation and generalization based on events, based on objects and affordances, as well as robustness, multi-task learning, and transfer.

\subsection{Specialization based on Objects: Bouncing Objects Experiment}

We consider a synthetic ``bouncing balls'' task in which multiple balls move using basic newtonian physics.  What makes this task particularly suited to RIMs is that the balls move independently on the majority of time steps, with the only exception being when the balls collide.  Thus we would hope to see that each RIM would model a subset of the balls and the RIMs would communicate most during collisions.  We show that the use of RIMs improves generalization to settings not observed during training.

\subsubsection{Different RIMs attend to different balls}
In the RIM model, we can associate each RIM with a different encoder.
By performing geometric masking we can control the possible spatial input to each block.
We use six non-overlapping horizontal strips and allow only 4 RIMs to be active at a time, shown in Fig.~\ref{fig:geometric}.
We find that early in training block activation is strongly correlated with the location of the 4 balls.
However, after training has proceeded for some time this correlation deteriorates.
This is likely because the predictable dynamics of the system do not necessitate constant attention.

\begin{figure}[htb!]
    \centering

    \includegraphics[width=0.95\linewidth]{figures/geom2_jh.png}
    \caption{\textbf{Geometric masking}. Each RIM is associated with a separate encoders, which are geometrically masked.  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. 1.) Early in training, RIM activations correlated more strongly with the locations of the four different balls. Later in training, this correlation decreased and the active strips did not correlate as strongly with the location of balls. As the model got better at predicting the location, it needed to attend less 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.}
    \label{fig:geometric}
\end{figure}

\subsubsection{4 Bouncing Balls}

\begin{figure}[htb!]
    \centering
    \includegraphics[width=0.95\linewidth]{figures/err_all.png}
    \caption{\textbf{Comparison of RIMs to LSTM baseline}. For 4 different experiments in the text, we compare RIMs to two different LSTM baselines. In all cases we find that during rollout, RIMs perform better than the LSTMs at accurately capturing the trajectories of the balls through time. Due to the number of hard collisions, accurate modeling is very difficult. In all cases, the first 15 frames of ground truth are fed in and then the system is rolled out for the next 10 time steps, computing the binary cross entropy between the prediction and the true balls at each instant, as in \cite{van2018relational}. See the Appendix for losses over the entire 35 frame rollout trajectory. Note that after some time, many balls no longer overlap making the error harder to interpret, as some of the amplitude is due to ``certainty'' in the incorrect ball.}
    \label{fig:err}
\end{figure}

\begin{figure}[htb!]
    \centering
    \includegraphics[width=0.95\linewidth]{figures/all_baseline.png}
    \caption{\textbf{Comparison of RIMs to LSTM baseline}. For 4 different experiments in the text, we compare RIMs to two different LSTM baselines. In all cases we find that during rollout, RIMs perform better than the LSTMs at accurately capturing the trajectories of the balls through time. Due 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 15 time steps, computing the binary cross entropy between the prediction and the true balls at each instant, as in \cite{van2018relational}. See the Appendix for losses over the entire 35 frame rollout trajectory. In the predictions, the transparent blue shows the ground  truth, overlaid to help guide the eye.}
    \label{fig:comparison}
\end{figure}

Here, show that the proposed model can capture the dynamics of various physical environments more accurately than other methods.
Since, different RIMs produce a query vector, dependent on which they get information from the input, which allows each RIM to use only  task-relevant information. Here, we show that the proposed architecture incorporates inductive biases about the existence of objects and interactions.

We assess the quality of RIMs by predicting the temporal dynamics of bouncing balls from \cite{van2018relational}.  The base dataset consists of  4 balls of two different masses, represented by their radii, bouncing off one another in a confined square geometry.
In Fig.~\ref{fig:err} we show how RIMs compare to two LSTM baselines (with 250 and with 512 hidden units).
Predictions from an LSTM baseline with 250 hidden units is shown in Fig.~\ref{fig:lstm}.
It is important to note that our baselines are able to capture the gestalt all of the tasks that we consider, with enough training: balls do move during rollout and for a dataset with an occlusion, balls do emerge from the occlusion after training.
During training, we predict only the next frame at each time step. Therefore, the learned dynamics are extremely complex, as the model is forced to extrapolate a complex system having been trained on very short steps.
In the experiments shown, we are more interested in the ability for the models to rollout the trajectory.
This allows us to test the ability of the network to accurately predict future dynamics in a complex environment.
For comparing rollout loss, we rollout for 10 frames (see Fig.~\ref{fig:err} and \ref{fig:comparison}).
After 10 frames, in the baselines and particularly in the 6-8 ball experiments, the collisions and inaccuracies in evolving trajectories results in a binary cross entropy becoming an inaccurate metric.

We consider three different experiments: (1) We consider the dynamics of the four bouncing balls.
(2) We consider the dynamics of three bouncing balls with an occlusion present.
(3) In the case of 4 bouncing balls, we give each ball a color attribute that needs to be tracked as well.

\subsubsection{Curtains Dataset}

An additional dataset provided in \cite{van2018relational} has 3 balls moving though a small square region is occluded, as shown in Fig.~\ref{fig:curtain}.
This allows us to assess the ability of balls to be tracked through the occluding region. Note that during the rollout phase balls are tracked through the occluding region. We therefore evaluate the capability of RIMs to handle occlusion by an invisible curtain. We find, as shown in Fig.~\ref{fig:curtain}, that RIMs do track balls as they move behind the occlusion.
An LSTM baseline is able to capture motion of balls behind the curtain as well, though typically quite inacurrately if the ball becomes fully occluded.

\subsubsection{Interpolation}
We train the RIMs model on the 6-8 ball dataset and then apply it, without training, on the 4 ball dataset (Fig.~\ref{fig:comparison} and an additional figure in the Supplementary Information).
Even though the model has never seen data with just 4 balls, and thus fewer active pixels, the 4 balls are conserved and reasonable trajectories are constructed.

\subsubsection{Balls of different colors}
\begin{figure}[htb!]
    \centering
    \includegraphics[width=0.95\linewidth]{figures/rgb_2.png}
    \caption{\textbf{RGB Rollout}. For the 4 ball dataset, we rollout 35 frames after feeding in the first 15 frames. We show to random samples from the test set, with the ground truth on top (white background) and the prediction on the bottom (black background).}
    \label{fig:rgb}
\end{figure}
Using the 4 ball dataset, we converted the input to an RGB image where each ball is assigned one of four colors at random (see Figs.~\ref{fig:err} and \ref{fig:rgb}).  The mean square error is used to compute the loss, otherwise training proceeds as before.
We find that during rollout, RIMs perform better than the LSTM baselines. Balls are not lost or created, nor do they change color.

\subsection{Keys as Affordances}

We also investigate the performance of RIMs on the Arcade Learning Environment. The study of vision must therefore include not only the study of how to extract from images, but also an inquiry into the nature of the internal
representations by which we capture this information and thus make it available as a basis for
decisions about our  actions.

\subsubsection{Transfer on Atari}
We next investigate feature transfer between randomly selected Atari games. This is an interesting question, because the visuals of Atari games are quite different from each other, as are the controls and required strategy. Though games like Pong and Breakout are conceptually similar (both involve hitting a ball with a paddle), Pong is vertically aligned while Breakout is horizontal: a potentially insurmountable feature-level difference. Other Atari game pairs have no discernible overlap, even at a conceptual level.

To this end we start by training RIMs on three source games (Pong, River Raid, and Seaquest)
and assess if the learned features transfer to a different subset of randomly selected target games (Alien, Asterix, Boxing, Centipede, Gopher, Hero, James Bond, Krull, Robotank, Road Runner, Star Gunner, and Wizard of Wor). Across all games, we observe, that RIMs result in positive transfer in 8 out of 12 target tasks, with only two cases of negative transfer. This compares favourably to baseline, which yields positive transfer in only 5 of 12 games.

Now consider the specific sequence Seaquest-to-Gopher, an example of two dissimilar games. Here,
the pretrain/finetune paradigm  exhibits negative transfer, unlike RIMs, perhaps because they are more able to ignore the irrelevant features.

\begin{figure}
\begin{tabular}{ccccc}
\subfloat[caption]{\includegraphics[width = 1.0in]{atari_results/Alien.png}} &
\subfloat[caption]{\includegraphics[width = 1.0in]{atari_results/Amidar.png}} &
\subfloat[caption]{\includegraphics[width = 1.0in]{atari_results/Assault.png}} &
\subfloat[caption]{\includegraphics[width = 1.0in]{atari_results/Asteroids.png}}&
\subfloat[caption]{\includegraphics[width = 1.0in]{atari_results/Atlantis.png}} \\
\subfloat[caption]{\includegraphics[width = 1.0in]{atari_results/BankHeist.png}} &
\subfloat[caption]{\includegraphics[width = 1.0in]{atari_results/BattleZone.png}} &
\subfloat[caption]{\includegraphics[width = 1.0in]{atari_results/BeamRider.png}}&
\subfloat[caption]{\includegraphics[width = 1.0in]{atari_results/Bowling.png}} &
\subfloat[caption]{\includegraphics[width = 1.0in]{atari_results/Boxing.png}} \\
\subfloat[caption]{\includegraphics[width = 1.0in]{atari_results/Breakout.png}} &
\subfloat[caption]{\includegraphics[width = 1.0in]{atari_results/Centipede.png}}&
\subfloat[caption]{\includegraphics[width = 1.0in]{atari_results/Frostbite.png}} &
\subfloat[caption]{\includegraphics[width = 1.0in]{atari_results/ChopperCommand.png}} &
\subfloat[caption]{\includegraphics[width = 1.0in]{atari_results/CrazyClimber.png}} \\
\subfloat[caption]{\includegraphics[width = 1.0in]{atari_results/DemonAttack.png}}&
\subfloat[caption]{\includegraphics[width = 1.0in]{atari_results/DoubleDunk.png}} &
\subfloat[caption]{\includegraphics[width = 1.0in]{atari_results/Enduro.png}} &
\subfloat[caption]{\includegraphics[width = 1.0in]{atari_results/FishingDerby.png}} &
\subfloat[caption]{\includegraphics[width = 1.0in]{atari_results/Freeway.png}}\\
\subfloat[caption]{\includegraphics[width = 1.0in]{atari_results/Frostbite.png}}&
\subfloat[caption]{\includegraphics[width = 1.0in]{atari_results/Gopher.png}} &
\subfloat[caption]{\includegraphics[width = 1.0in]{atari_results/Gravitar.png}} &
\subfloat[caption]{\includegraphics[width = 1.0in]{atari_results/Hero.png}} &
\subfloat[caption]{\includegraphics[width = 1.0in]{atari_results/IceHockey.png}}\\
\subfloat[caption]{\includegraphics[width = 1.0in]{atari_results/Krull.png}}&
\subfloat[caption]{\includegraphics[width = 1.0in]{atari_results/KungFuMaster.png}} &
\subfloat[caption]{\includegraphics[width = 1.0in]{atari_results/MontezumaRevenge.png}} &
\subfloat[caption]{\includegraphics[width = 1.0in]{atari_results/NameThisGame.png}} &
\subfloat[caption]{\includegraphics[width = 1.0in]{atari_results/Pong.png}}\\
\subfloat[caption]{\includegraphics[width = 1.0in]{atari_results/PrivateEye.png}}&
\subfloat[caption]{\includegraphics[width = 1.0in]{atari_results/Qbert.png}} &
\subfloat[caption]{\includegraphics[width = 1.0in]{atari_results/Riverraid.png}} &
\subfloat[caption]{\includegraphics[width = 1.0in]{atari_results/RoadRunner.png}} &
\subfloat[caption]{\includegraphics[width = 1.0in]{atari_results/Robotank.png}}\\
\subfloat[caption]{\includegraphics[width = 1.0in]{atari_results/SpaceInvaders.png}}&
\subfloat[caption]{\includegraphics[width = 1.0in]{atari_results/StarGunner.png}} &
\subfloat[caption]{\includegraphics[width = 1.0in]{atari_results/Tennis.png}} &
\subfloat[caption]{\includegraphics[width = 1.0in]{atari_results/TimePilot.png}} &
\subfloat[caption]{\includegraphics[width = 1.0in]{atari_results/Tutankham.png}}\\
\subfloat[caption]{\includegraphics[width = 1.0in]{atari_results/UpNDown.png}}&
\subfloat[caption]{\includegraphics[width = 1.0in]{atari_results/Venture.png}} &
\subfloat[caption]{\includegraphics[width = 1.0in]{atari_results/VideoPinball.png}} &
\subfloat[caption]{\includegraphics[width = 1.0in]{atari_results/WizardOfWor.png}} &
\subfloat[caption]{\includegraphics[width = 1.0in]{atari_results/Zaxxon.png}}\\
\end{tabular}
\caption{4 x 4}
\end{figure}
\

\subsection{Generalization based on events}

Due to the explicit, top down attention, (in the form of queries produce) \BS{Need to write why we expect generalization based on events here.}

\subsubsection{Copying Task}

The copying task has been widely studied in the context of RNNs and neural sequence models.  It consists of receiving a short sequence of characters, then receiving blank inputs for a large number of steps, and then being forced to produce the original received sequence.  This task studies a simple property in isolation: is the model able to keep things in memory while no active input is being received.  LSTMs struggle to perform the copying task with very long sequences and especially struggle when evaluated on sequences which are longer than those seen during training.  The RIMs model may be able to address this by using one RIM to store the information at the start of the sequence, and then keep that RIM dormant until the information is retrieved.

\begin{table}[htb!]
\centering
{\renewcommand{\arraystretch}{1.0}
\caption{\textbf{Error (negative log likelihood) on the copying task}.  Note that while all of the methods are able to 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.  }
\label{tb:copy}
\centering
\scalebox{0.9}{
\begin{tabular}{lrrr}
\toprule
Approach & Train Length 50 & Test Length 200  \\
\midrule
LSTM (1-layer) & 0.01 & 6.13  \\
LSTM + Self. Attention & 0.00 & 4.60 \\
NTM \citep{graves2014neural}  & 0.00 & 1.30 \\
Relational Memory \citep{santoro2018relational} & 0.00 & 0.13 \\

\textbf{RIMs} ($k_T$ = 4, $k_A$ = 2) & \textbf{0.00} & \textbf{0.00}  \\
\textbf{RIMs} ($k_T$ = 4, $k_A$ = 3) & 0.14 & 0.10  \\

\textbf{RIMs} ($k_T$ = 5, $k_A$ = 2) & \textbf{0.00} & \textbf{0.00}  \\
\textbf{RIMs} ($k_T$ = 5, $k_A$ = 3) & \textbf{0.01} & \textbf{0.00}\\
\textbf{RIMs} ($k_T$ = 5, $k_A$ = 4) & \textbf{0.00} & \textbf{0.00} \\

\textbf{RIMs} ($k_T$ = 6, $k_A$ = 2) & \textbf{0.00} & \textbf{0.00} \\
\textbf{RIMs} ($k_T$ = 6, $k_A$ = 3) & \textbf{0.00} & \textbf{0.00}  \\
\textbf{RIMs} ($k_T$ = 6, $k_A$ = 4) & \textbf{0.00} & \textbf{0.00}\\
\textbf{RIMs} ($k_T$ = 6, $k_A$ = 5) & 0.01 & 0.35 \\
\bottomrule
\end{tabular}
}
}
\end{table}

\textbf{Results:} Table \ref{tb:copy} shows the result when the proposed model as well as baselines are trained on sequences of length 50 and evaluated on sequences of length 75. Due to the ability of RIMs to store information for all length scales, the RIMs model outperform many strong baselines such as Neural Turing Machines \citep{graves2014ntm}, LSTMs combined with Self Attention, and Relational Memory architecture \citep{santoro2018relational}. We also conducted experiments where training was performed on sequence length 50 and evaluation on lengths 100, 150, 200. We observed that our proposed model achieved zero test error despite longer evaluation sequences.  Additionally we found that if we had 6 RIMs, the model generalized if we had top-k activation between $k_A=2$ and $k_A=4$.

\begin{figure}[htb!]
    \centering
    \includegraphics[width=0.6\linewidth,trim=8.0cm 0.0cm 4.0cm 0.0cm,clip]{figures/blocks_copying_fig.png}
    \caption{\textbf{Copying Task RIM Activation Pattern} for a model with $K=6$ RIMs and $K_A=3$ active RIMs per step (the activated RIMs are in black, non-activated in white).  We can see that the RIM activation pattern is distinct during the middle dormant part of the sequence.  }
    \label{fig:copying_pattern}
\end{figure}

\subsubsection{Sequential MNIST Resolution Task}

RIMs is motivated by the hypothesis that generalization performance can be improved by having modules which only activate on relevant parts of the sequence.  For further evidence that RIMs can achieve this, we consider the task of classifying MNIST digits as sequences of pixels \citep{krueger2016zoneout} and studying if such a model generalizes to images of different resolutions than those seen during training.  Our intuition is that the RIMs model should have distinct subsets of the RIMs activated for pixels with the digit and empty pixels.  As a result, the RIMs model should generalize better to greater resolutions by keeping the RIMs which store pixel information dormant over the empty regions of the image.

\begin{table}[htb!]
\centering
\caption{\textbf{Sequential MNIST resolution:} Test Accuracy \% on the Sequential MNIST resolution generalization task.  Both the proposed and the Baseline model (LSTM) were trained on 14x14 resolution but evaluated at different resolutions.  }
\label{tb:mnist}
{\renewcommand{\arraystretch}{1.0}
\begin{tabular}{lrrr}
\toprule
\centering
Method / Evaluation & 16 x 16 & 19 x 19 & 24 x 24 \\
\midrule
Increase in Sequence Length (\%) & 30.6\% & 84.2\% & 196.9\% \\
\midrule
LSTM(Hidden = 300) & 86.8 & 42.3 & 25.2 \\
LSTM (Hidden = 600) & 88.3 & 43.1 & 22.1 \\
\midrule
Relational Memory \citep{santoro2018relational} & 89.58 & 54.23  & 27.76 \\
\midrule
RIMs, Hidden = 300 ($k_T=6$, $k_A=4$) & 86.1 & 52.3 & 25.3  \\
RIMs, Hidden = 600 ($k_T=6$, $k_A=6$) & 85.5 & 56.2 & 30.9\\
RIMs, Hidden = 600 ($k_T=6$, $k_A=5$) & 84.5 & 52.2 & 21.9\\
RIMs, Hidden = 600 ($k_T=6$, $k_A=4$) & \textbf{90.0} & \textbf{73.4} & \textbf{38.1}  \\
\bottomrule
\end{tabular}
}
\end{table}

\textbf{Results:} Table. \ref{tb:mnist} shows the result of the proposed model on the Sequential MNIST Resolution Task.  When keeping the train and test sequence lengths the same, both models achieved comparable test set performance.  However, the RIMs model was relatively robust to changing the sequence length (by changing the resolution of the image), whereas the LSTM's performance degraded more catastrophically.  This can be seen as a slightly more involved analogue of the copying task, as MNIST digits have large empty regions with no information.  It is essential that the model be able to store information and pass gradients through these regions.

\subsection{Robustness to Distractors}

A top-down bias is more directly responsible for the diminished distractor processing in high-load conditions than the competition among stimulus items. It is the application of this bias in favor of the target that results in diminished representation of the distractor. The top-down filter not only enhances the target but also suppresses other stimuli, including the distractor. The notion that enhancement of the target necessarily results in suppression of the other items in the display is fundamental to biased competition theory (Desimone and Duncan, 1995) and divisive inhibition more generally (Reynolds and Heeger, 2009). Because neurons representing individual items have mutually inhibitory connections, an increase in the firing rate of one results in a decrease in the other.

\subsection{Multitask Learning Setup}

An important property for machine learning systems is the ability to perform well even in the presence of distracting or irrelevant information in the environment.  It is especially difficult and important to be resilient to these distractions when they are absent from the training data. For example, an LSTM with a single memory flow may struggle to keep the distracting elements separate from elements which are necessary for the task. Because the RIMs model uses attention to control which RIMs receive information at each step, it should be possible for the model to make a clear decision to not write distracting information into the important RIMs.

\textbf{Task Setup:} We setup an RL environment for object picking where we can directly study the impact of novel distractors by training the model to pick one object during training (in the presence of distractors), and then introducing new distractors  during testing.  We use a partially observed formulation of the task, where the agent only observes a small number of squares ahead of it. These tasks are difficult to solve with standard RL algorithms, not only due to the partial observability of the environment but also the sparsity of the reward, given that the agent receives a reward only after reaching the goal.

\begin{figure}[htb!]
    \centering
    \includegraphics[width=0.35\linewidth,trim=0.5cm 0.0cm 1.0cm 1.2cm,clip]{figures/figure_normal.png}
    \includegraphics[width=0.35\linewidth,trim=0.5cm 0.0cm 1.0cm 1.2cm,clip]{figures/figure_distractor.png}
    \caption{\textbf{Robustness to Novel Distractors:}. Left shows the performance of the proposed method as compared to the LSTM baseline to solve the object picking task in the presence of distractors.  Right shows the performance of the  proposed method and the baseline when novel distractors are added.  Success of an agent is measured, e.g. by the fraction of episodes where the agent was able to solve the task in 500 steps. (averaged over 500 examples}
    \label{fig:rl_exp}
\end{figure}

\textbf{Results:} Fig.~\ref{fig:rl_exp} shows our results. The RIMs model is able to completely solve the task though the baseline model (an LSTM in this case) struggles to solve even the training task. Since, the LSTM has a single memory flow it may be difficult for the LSTM to keep the information regarding distractors seperate from task relevant information. We also evaluate both the proposed model as well as LSTM baseline when we add novel distractors (which are not seen during training). The performance of both the LSTM model as well as the RIMs model degrade, though RIMs still perform better than the LSTM.

\section{Ablations}

\subsection{Importance of attention bottleneck for communication between RIMs}

\subsection{Importance of Sparsity of activation of the RIMs}

\subsection{No attention mechanisms to read the input}

\subsection{Importance of default behaviour}

\subsection{Importance of key value pairs}

\subsection{Effect of number of blocks}

\subsection{Effect of making the RIMs "dormant" ``Ablation studies''}

\begin{table}[htb!]
\centering
{\renewcommand{\arraystretch}{1.0}
\caption{\textbf{Error (negative log likelihood) on the copying task}.  Note that while all of the methods are able to 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.  }
\label{tb:copy}
\centering
\scalebox{0.9}{
\begin{tabular}{lrrr}
\toprule
Approach & Train Length 50 & Test Length 200  \\
\midrule
LSTM (1-layer) & 0.01 & 6.13  \\
LSTM + Self. Attention & 0.00 & 4.60 \\
NTM \citep{graves2014neural}  & 0.00 & 1.30 \\
Relational Memory \citep{santoro2018relational} & 0.00 & 0.13 \\

\textbf{RIMs} ($k_T$ = 6, $k_A$ = 2) & \textbf{0.00} & \textbf{0.00} \\
\textbf{RIMs} ($k_T$ = 6, $k_A$ = 3) & \textbf{0.00} & \textbf{0.00}  \\
\textbf{RIMs} ($k_T$ = 6, $k_A$ = 4) & \textbf{0.00} & \textbf{0.00}\\
\textbf{RIMs} ($k_T$ = 6, $k_A$ = 5) & 0.01 & 0.35 \\
\midrule
With input Attention and No Communication & & \\
\midrule
\textbf{RIMs} ($k_T$ = 4, $k_A$ = 2) & 0.23 & 0.16  \\
\textbf{RIMs} ($k_T$ = 4, $k_A$ = 3) & 0.17 & 0.43  \\

\textbf{RIMs} ($k_T$ = 5, $k_A$ = 2) & 0.25 & 0.47  \\
\textbf{RIMs} ($k_T$ = 5, $k_A$ = 3) & 0.04 & 0.40\\
\textbf{RIMs} ($k_T$ = 5, $k_A$ = 4) & 0.02 & 0.07 \\

\textbf{RIMs} ($k_T$ = 6, $k_A$ = 2) & 0.33 & 0.24\\
\textbf{RIMs} ($k_T$ = 6, $k_A$ = 3) & 0.12 & 0.10  \\
\textbf{RIMs} ($k_T$ = 6, $k_A$ = 4) & 0.07 & 0.50\\
\textbf{RIMs}
($k_T$ = 6, $k_A$ = 5) & 0.22 & 0.56 \\
\bottomrule
\end{tabular}
}
}
\end{table}

\section{Discussions}

\subsection{ The role of top-down influence}

\subsection{Adding More RIMs}

Adding capacity to an already trained deep network is generally quite tedious, as it involves changing the shape of the parameters, and it always requires retraining.  An intriguing possibility is that the RIMs model will allow for adding capacity to a model simply by adding new RIMs, as none of the existing parameters need to be changed.  In fact the new RIMs have no impact on the model's training unless the attention mechanism looks at the newly added RIMs, which we would expect to be relatively rare.  Progressively adding capacity has been explored for generative adversarial networks, where it achieved state of the art results \citep{karras2017progressive}.

\subsection{Hypothesis Testing}

These RIMs  have some kind of competition between them, so that they are competing to be maximally useful. Combined with a bit of randomness  in terms of how they are triggered, we might naturally get a lot of phenomena we associate with intelligence: competing “hypotheses” are actually competing RIMs, where multiple ways of solving a task might be present simultaneously until one of them very clearly wins out.

\vspace{-1mm}

\section{Future Work}

There are many exciting opportunities for future work. First, we have not conducted thorough
performance optimizations on RIMs. Future
work can also explore acceleration schemes that speed up training and inference.  Finally, we
look forward to seeing applications of RIMs like architecture in the case of generative models.

In the current paper, RIMs are collectively optimizing the single objective. In order to promote,What is the mechanism that incentivizes the RIMs to collectively optimize the global objective as a by-product of each RIM optimizing its own local objective?

\section{Conclusion}

The world consists of many dynamic processes which primarily operate independently and only occasionally have meaningful interactions.  Despite this, the vast majority of machine learning models employ the opposite inductive bias: that all processes interact, and force the model to unlearn this bias.  This leads to inflexibility and often very poor generalization to changing task distributions.  We have proposed a new model, which we call Recurrent Independent Mechanisms (RIMs), in which we learn multiple recurrent modules which are independent by default, but interact sparingly. The learned RIMs can be flexibly recombined to solve more complex tasks. Our  positive experimental results validate the appropriateness of the previously proposed consciousness prior~\citep{bengio2017consciousness}. The proposed framework could also be very attractive for continual learning settings, where one could add more RIMs over time, thereby the already learned RIMs would keep their focus on particular aspect of the task, and newly added ones could specialize on novel aspects.

\clearpage