\section{Introduction}
\vs{1}








\begin{figure*}[]
\vspace{0.2cm}
\begin{minipage}{1.0\textwidth}
    \centering
    \label{fig:model-diagram}
    \includegraphics[width=1.0\linewidth]{figures/topdownbottomup_v5.png}
    \caption{Overall layout of the proposed Bidirectional Recurrent Independent Mechanisms (BRIMs) model.  Information is
    passed forward in time using recurrent connections and information passed between layers using attention.  Modules attend to both the lower layer on the current time step as well as the higher layer on the previous time step, along with \emph{null}, which refers to a vector of zeros.  }
\end{minipage}
\end{figure*}


Deep learning emerged from the goal of learning representational hierarchies,
with the higher levels corresponding to
abstract features used for decision-making~\citep{Hinton06,Bengio-nips-2006-small,salakhutdinov2009deep}.  
Hierarchical models are often taken to imply that computation proceeds in a 
feedforward or \emph{bottom up} fashion, i.e., stage-wise information 
processing in which low-level (sensory) representations construct or modulate
high level (conceptual)  representations. However, they could equally well support the flow of information in a 
feedback or \emph{top down} fashion, i.e., information processing in which
high-level representations modulate lower-level representations. 
 Neuroscientists have noted that
reciprocity of connectivity among anatomically distinct areas 
of neocortex is common \citep{FellemanVanEssen1991,Rockland2015},
causing early visual areas to be modulated by later stages of processing~\citep{Bastos:2015eu}. The same is true for other sensory modalities as well~\citep{Manita:2015bp}.
Neuroimaging research has found evidence for distinct bidirectional activity flows with functional consequences
\citep{Dijkstra2017,Nielsen1999}, and neurophysiological and neuroanatomical studies indicate the important
role of top-down information for guiding processing resources to behaviorally relevant events \citep{baluch2011mechanisms,gilbert2013topdown}. 
The  Global Workspace theory~\citep{baars1997theatre,Dehaene-et-al-2017} posits that
top-down signals are necessary to broadcast information throughout neocortex, which
enables conscious states and verbal behavior.
Nonetheless, the neural mechanisms of interaction between bottom-up and
top-down pathways are still poorly understood. One goal of our research
is to conduct machine-learning experiments to explore the value of
top-down mechanisms for learning, achieving robustness to distributional shifts,
and guiding the development of novel and improved neural architectures.

In the cognitive science community, the relative contributions of bottom-up and 
top-down signals  have been an ongoing subject of debate for over 40 years
\citep{Kinchla1979,Rauss2013}.
The \citet{Mcclelland1981}
model of printed-word reading consisted of a hierarchy of
detectors, from letter fragments to letters to words, with bidirectional 
connectivity. The top-down connectivity 
imposed orthographic constraints of the vocabulary and helped to explain human
proficiency in reading. The model also accounted for puzzling behavioral 
phenomena, such as the fact that people are better at discriminating the letters \textsc{e} 
and \textsc{o} when a letter appears embedded in a word---such as \textsc{r\underline{e}ad}---than
when the same letter is presented in isolation.
In the model, partial activation from 
the representation of the word \textsc{read} provides top-down support for
the activation of the letter \textsc{e} in the context of the word, but not the \textsc{e} in isolation.
If top-down pathways provide expectations, they can
be considered as priors on subsequent inputs, and uncertainty
plays a key role in how information is combined from bottom-up
and top-down sources \cite{Weiss2002,Kersten:2004hd}. For instance, if one enters a familiar but
dark room, one's expectations will guide perception and behavior, 
whereas expectations are less relevant in the light of day.
Even without prior expectations, top-down pathways can leverage 
simultaneous contextual information to resolve local uncertainty,
e.g., the interpretation of the middle letters of these two words:
\includegraphics[]{images/THECAT.pdf}.
Thus, top-down processing leverages goals, temporal and concurrent context, and
expectations (conscious or otherwise) 
to steer or focus
the interpretation of sensory signals. Top-down processing helps to 
overcome intrinsically noisy or ambiguous sensory data, whereas 
bottom-up processing allows an agent to be reactive to unexpected or 
surprising stimuli.











\vspace{-4mm}
\paragraph{Information Flow in Deep Neural Networks.}
Models having within-layer recurrence are common (e.g., LSTM layers), but such recurrence involves a particular 
level of representation interacting with itself, as opposed to the strict top-down notion of higher levels acting on lower levels.
In contrast, a layered architecture with feedforward and feedback connections makes a structural distinction between higher and lower
levels of representations, and it requires a combination of information from two distinct sources.

Bidirectional layered models have long existed \cite{dayan1995helmholtz,Larochelle2008,salakhutdinov2009deep}, 
but these models have not received the same intense scrutiny and development as feedforward models.
The goal of this paper is to revisit bidirectional layered models with the aim of raising them to state-of-the-art in performance and more specifically, out of distribution generalization.
Simply taking existing models and incorporating bidirectional connectivity introduces challenges \cite{Iuzzolino2019}. Instead,
we explore the notion that bidirectional information flow must be coupled with additional mechanisms to 
effectively select from potential top-down signals, modulate
bottom-up signals and prevent a cacophony of intermixed signals.
We propose two specific mechanisms from the deep learning toolkit to dynamically route information flow:  \emph{modularity} and \emph{attention}.
With these two mechanisms, we outperform state-of-the-art feedforward models on challenging vision and language tasks.

\vspace{-4mm}
\paragraph{Dynamic Information Flow Using Attention.}
\citet{goyal2019recurrent} offer evidence that modularity, coupled with a flexible means of communication between modules, can 
dramatically improve generalization performance.
They describe a recurrent independent mechanisms
(RIMs) architecture that
dynamically controls information flow in a modular RNN. The recurrently connected modules compete
for external input and communicate sets of objects 
\emph{sparingly} via differentiable attention
mechanisms driven by the match between keys (from a source module)
and queries (from a destination module). Computation is sparse and modular, meaning that only a subset of the neural modules are active at any time. Composition of computations is dynamic
rather than static, in the sense that 
the attention mechanism uses context to select the subset of modules that are
activated and what is communicated to which module, with the input of one module dynamically selected, so the  values associated with selected objects are not always coming from the same modules or parts of the input.
The motivating intuition behind this increased flexibility is that if the relationship between modules changes between training and evaluation, a model which keeps the information in these modules sufficiently separate should still be able to correctly recombine information, even though the overall data distribution differs from what was seen during training. 







\vspace{-4mm}
\paragraph{Combining Top-down and Bottom-up Information with Attention over Modules.}
We now present our central proposal, which is that top-down and bottom-up signals should be combined explicitly and selectively by using attention.  On a given time step, only a fraction of the high-level concepts being detected are likely to be relevant for a particular stimulus.  For example, if a person is trying to classify an object in a dark room (which they are familiar with), the most relevant top-down signal is the person's prior knowledge of what that object looks like, as opposed to other knowledge about the room.  This motivates breaking up the hidden state into multiple modules, such that the top-down and bottom-up interactions can be appropriately focused.  
Thus we consider the use of modules at multiple levels of abstraction to be a key ingredient for our method.  We elect to build upon the Recurrent Independent Mechanisms (RIMs) framework for the modular building blocks, as their end-to-end training demonstrated successful specialization between modules \cite{goyal2019recurrent}. However, while network modularity revealed considerable computational advantages, the modules in the RIMs architecture are fully interconnected and thus there is no notion of layered or hierarchical structure.



To address this issue, we propose a new architecture, which we call \textit{Bidirectional Recurrent Independent Mechanisms (BRIMs)}. This new approach endows the attention-based, modular structure with a hierarchy of layers composed of competing modules, yielding organized contextual computations during training. The rationale is to provide an architectural backbone such that each layer of the hierarchy can send information in both a bottom-up direction as well as in a top-down direction.  We find that this  inductive bias, combined with the modularity mechanisms developed in~\citet{goyal2019recurrent}, provides considerable and further advantages for tasks where there is a change of input distribution from training to testing. 





\section{Preliminaries}




\paragraph{Multi-layer Stacked Recurrent Networks.}
The most common multi-layer RNN architecture is bottom-up and feed-forward, in the sense that higher layers are supplied with the states of the lower layers as inputs. An $L$-layer deep RNN is then concisely summarized as:
\begin{align}
    \textbf{y}_t &= D(\textbf{h}^L_t) \\
    \textbf{h}^l_t &= F^{l}(\textbf{h}^{l-1}_t, \textbf{h}^l_{t-1}) \label{rnn_hier}  \\
    \textbf{h}^0_t &= E(\textbf{x}_t)
\end{align}
with $l = 0,1,...,L$. For a given time $t$, $\textbf{y}_t$ denotes the model prediction, $\textbf{x}_t$ the input and $\textbf{h}^l_t$ the hidden state of the model at layer $l$. $D$ and $E$ denote the Decoder and Encoder for the model. $F^l$ represents the recurrent dynamics at the hierarchy level $l$ (e.g., an LSTM or GRU).






\paragraph{Key-Value Attention.} 
Key-value Attention (also sometimes called Scaled Dot Product attention), defines the backbone of updates to the hidden states in the proposed model. This form of attention is widely used in self-attention models and performs well on a wide array of tasks \citep{vaswani2017attention, santoro2018relational}. Given a set of queries $\textbf{Q}$, keys $\textbf{K}$ ($d$-dimensional) and values $\textbf{V}$, an attention score $\textbf{A}_S$ and an attention modulated result $\textbf{A}_R$ are computed as
\begin{align}
    \mathrm{\textbf{A}_S} &= \mathrm{Softmax} \left (\frac{\textbf{QK}^T}{\sqrt{d}} \right) \label{att_score}\\
    \mathrm{\textbf{A}_R} &= \mathrm{\textbf{A}_S}\; \textbf{V} \label{att_out}
\end{align}
\paragraph{Recurrent Independent Mechanisms.} RIMs \citep{goyal2019recurrent} consist of a single layered recurrent structure where 
the hidden state $\textbf{h}_t$ is decomposed into $n$ modules, $\textbf{h}_{t,k}$ for $k = 1, ... n$.
It also has the property that on a given time step, only a subset of modules is activated. 
In RIMs, the updates for the hidden state follow a three-step process. 
First, a subset of modules is selectively activated based on their determination of the relevance of their input. 
Second, the activated modules independently process the information made available to them.
Third, the active modules gather contextual information from all the other modules and consolidate this information in their hidden state.

\paragraph{Selective Activation.} Each module creates queries \mbox{$\bar{\textbf{Q}}={Q}_{inp}\,(\textbf{h}_{t-1})$} ($n \times d$ matrix) which are then combined with the keys $\bar{\textbf{K}}={K}_{inp}\,(\textbf{\o},\textbf{x}_t)$  and values $\bar{\textbf{V}}={V}_{inp}\,(\textbf{\o}, \textbf{x}_t)$ obtained from the input $\textbf{x}_t$ and zero vectors $\textbf{\o}$ to get both the attention score $\bar{\textbf{A}}_S$ and attention modulated input $\bar{\textbf{A}}_R$ as per equations \eqref{att_score} and \eqref{att_out}. Based on this attention score, a fixed number of modules $m$ are activated for which the input information is most relevant (where the null module, which provides no additional information, has low attention score). We refer to this activated set per time-step as $\mathcal{S}_t$.  

\paragraph{Independent Dynamics.} Given the attention modulated input obtained above, each activated module then undergoes an update in its hidden state:
\begin{align}
\bar{\textbf{h}}_{t,k} &= \begin{cases}F_k\, ( \bar{\textbf{A}}_{R_k}\,,\,\textbf{h}_{t-1,k}) \quad & k \in \mathcal{S}_t \\
\textbf{h}_{t-1,k} \quad & k \notin \mathcal{S}_t
\end{cases}
\end{align}
$F_k$ here stands for any update procedure, e.g. GRU or LSTM. 

\paragraph{Communication.} After an independent update step, each module then consolidates information from all the other modules. RIMs again utilize the attention mechanism to perform this consolidation. Active modules create queries $\hat{\textbf{Q}}=Q_{com}(\bar{\textbf{h}}_{t})$ which act with the keys $\hat{\textbf{K}}={K}_{com}(\bar{\textbf{h}}_{t})$ and values $\hat{\textbf{V}}={V}_{com}(\bar{\textbf{h}}_{t})$ generated by all modules and the result of attention $\hat{\textbf{A}}_R$ is added to the state for that time step:\blfootnote{Code is available at \href{https://github.com/sarthmit/BRIMs}{https://github.com/sarthmit/BRIMs}}

\begin{align}
    \textbf{h}_{t,k} &= \begin{cases}\bar{\textbf{h}}_{t,k} + \hat{\textbf{A}}_{R_k} \quad & k \in \mathcal{S}_t \\
    \bar{\textbf{h}}_{t,k} \quad & k \notin \mathcal{S}_t
    \end{cases}
\end{align}

\section{Proposed Method}

Our contribution is to extend RIMs to a multilayered bidirectional architecture, which we refer to as BRIMs.
Each layer is a modified version of the RIMs architecture \citep{goyal2019recurrent}. A concise representation of our architecture is depicted in \hyperref[fig:model-diagram]{Figure~\ref{fig:model-diagram}}. We now describe
the dynamic bottom-up and top-down flow of information in BRIMs.

\subsection{Composition of Modules}


We use the procedure provided in RIMs to decompose the hidden state $\textbf{h}^l_t$ on each layer $l$ and time $t$ into separate modules. Thus, instead of representing the state as just a fixed dimensional vector $\textbf{h}^l_t$, we choose to represent it as $\{((\textbf{h}^l_{t,k})_{k=1}^{n_l}, \mathcal{S}^l_t)\}$ where $n_l$ denotes the number of modules in layer $l$ and $\mathcal{S}^l_t$ is the set of modules that are active at time $t$ in layer $l$. $|\mathcal{S}^l_t| = m_l$, where $m_l$ is a hyperparameter specifying the number of modules  active in layer $l$ at any time. Each layer can potentially have different number of modules active. Typically, setting $m_l$ to be roughly half the value of $n_l$ works well.

\subsection{Communication Between Layers}
\label{within_hier}

We dynamically establish communication links between multiple layers using key-value attention, in
a way which differs radically from RIMs \citep{goyal2019recurrent}.  While many RNNs build a strictly bottom-up multi-layer dependency using \eqref{rnn_hier}, we instead build multi-layer dependency by considering queries $\bar{\textbf{Q}}={Q}_{lay}\,(\textbf{h}^l_{t-1})$ from modules in layer $l$ and keys $\bar{\textbf{K}}={K}_{lay}\,(\textbf{\o},\textbf{h}^{l-1}_{t}, \textbf{h}^{l+1}_{t-1})$ and values $\bar{\textbf{V}}={V}_{lay}\,(\textbf{\o},\textbf{h}^{l-1}_{t},\textbf{h}^{l+1}_{t-1})$ from all the modules in the lower and higher layers (blue and red arrows respectively; \hyperref[fig:model-diagram]{Figure~\ref{fig:model-diagram}}). From this three-way attention mechanism, we obtain the attention score $\bar{\textbf{A}}^l_S$ and output $\bar{\textbf{A}}^l_R$ (Equations \ref{att_score} and \ref{att_out} respectively).  Note that in the deepest layer, only the lower layer is used and on the first layer, the input's embedded state serves as the lower layer. Also note that the attention receiving information from the higher layer looks at the previous time step, whereas the attention receiving information from lower layer (or input) looks at the current time step. 
\begin{align*}
\small
    \text{Softmax} \left(\begin{bmatrix}
    \mathbf{Q}_{\;module\; 1} \\
    \mathbf{Q}_{\;module\; 2} \\
    \mathbf{Q}_{\;module\; 3} \\
    \end{bmatrix}
    \begin{bmatrix}
    \mathbf{\color{myred}K}_{\;\phi} & \mathbf{\color{myblue}K}_{\;l-1} & \mathbf{\color{mygreen}K}_{\;l+1}
    \end{bmatrix}
    \right)
    \begin{bmatrix}
    \mathbf{\color{myred}V}_{\;\phi} \\
    \mathbf{\color{myblue}V}_{\;l-1} \\
    \mathbf{\color{mygreen}V}_{\;l+1}
    \end{bmatrix} \\
    \begin{matrix}
    {\color{myred}\phi}: \text{Null} &
    {\color{myblue}l-1}: \text{Bottom Up} &
    {\color{mygreen}l+1}: \text{Top Down}
    \end{matrix}
\end{align*}
\subsection{Sparse Activation}
Based on the attention score $\bar{\textbf{A}}^l_S$, the set $S_t^l$ is constructed which comprises modules for which null information is least relevant. Every activated module gets its own separate version of input (as it's dependent on its query, which is a function of the hidden state of the module) which is obtained through the attention output $\bar{\textbf{A}}^l_R$. Concretely, for each activated module, this can be represented as:
\begin{align}
     \bar{\textbf{h}}^l_{t,k} &= F_k^l \, ( \bar{\textbf{A}}^l_{R_k}, \textbf{h}^l_{t-1, k}) \quad k \in \mathcal{S}^l_t
\end{align}
where $F_k^l$ denotes the recurrent update procedure.


\subsection{Communication Within Layers}

We also perform communication between the different modules within each layer (green arrows; \hyperref[fig:model-diagram]{Figure~\ref{fig:model-diagram}}). In order to enable this communication, we again make use of key-value attention. This communication between modules within a layer allows them to share information, albeit in a limited way through the bottleneck of attention. We create queries $\hat{\textbf{Q}}={Q}_{com}\,(\bar{\textbf{h}}^l_{t})$ from active modules and keys $\hat{\textbf{K}}={K}_{com}\,(\bar{\textbf{h}}^l_t)$ and values $\hat{\textbf{V}}={V}_{com}\,(\bar{\textbf{h}}^l_t)$ from all the modules to get the final update to the module state through residual attention $\hat{\textbf{A}}^l_R$ addition as follows:
\begin{align}
    \textbf{h}^l_{t,k} &= \begin{cases}\bar{\textbf{h}}^l_{t,k} + \bar{\textbf{A}}^l_{R_k} \quad&   k \in \mathcal{S}^l_t \\
    \textbf{h}^l_{t-1,k} \quad& k \notin \mathcal{S}^l_t
    \end{cases}
\end{align}
\begin{figure*}[]
\begin{minipage}{1.0\textwidth}
    \centering
    \includegraphics[width=1.0\linewidth]{plots/CIFAR_BRIMS.jpg}
   
    \includegraphics[scale=0.25]{plots/legend.pdf}
    \caption{On sequential CIFAR-10, we evaluate on 32x32 test images and change a certain number of the pixels to random values.  We show the \textbf{average activation weight (y-axis)} on Input (left), Null (middle) and Higher Layer (right) plotted against the \textbf{number of random pixels (x-axis)} (total number of pixels = 1024).  We see that as more pixels are set to random values, the model becomes increasingly reliant on the higher-level information.  }
    \label{fig:noise-cifar}
\end{minipage}
\vspace{-2mm}
\end{figure*}

\subsection{Training}

The architecture  proposed here doesn't rely on any additional losses and  thus can be used as a drop-in substitute for LSTMs and GRUs. For training we consider task-specific losses which range from classification and video prediction losses to RL losses depending on the problem.





\section{Related Work}

\textbf{Deep Boltzmann Machines:} Deep Boltzmann machines \citep{salakhutdinov2009deep} resemble our motivating intuitions, as they are undirected and thus have both top-down and bottom-up feedback through an energy function. Additionally, the model captures a probability distribution over the hidden states, such that the role of top-down feedback becomes most important when the bottom-up signal has the most uncertainty. However, a significant problem in deep Boltzmann machines, and with undirected graphical models in general, is that sampling and inference are both difficult. Sampling generally requires an iterative procedure, which may converge slowly.

\textbf{Helmholtz Machine:} Like autoencoders, the Helmholtz
machine \citep{dayan1995helmholtz} contains two separate networks: a generative network and a
discriminative network, and the wake-sleep
learning algorithm is applied to discover a latent representation
of the data in an unsupervised way. However, unlike in our proposal, the top-down path only
influences the learning of the bottom-up path, and not its actual
current computation.   



\textbf{Transformers:} The Transformer architecture \citep{vaswani2017attention} eschews the use of a recurrent state in favor of using attention to pass information between different positions. It lacks an explicit inductive bias towards a bottom-up followed by top-down flow, since at the lowest-layers of the network, the positions in the distant past can only undergo a small amount of processing before they can influence the lower levels of processing of later inputs.  

\textbf{Modulating Visual Processing with Language:} \cite{devries2017visual} proposes to modulate a visual image processing pipeline by adding in language representations at multiple layers, including early layers.  Since the language is already quite abstract, this can be seen as top-down modulation on the model's bottom-up processing.  While this improved results, it assumes that the raw data is associated with text which provides a top-down signal, as opposed to our proposed technique, in which the top-down signal is learned end to end from raw data and dynamically applied using top-down attention. 

\textbf{Hierarchical Multiscale Recurrent Neural Network:} This paper explored an architecture in which the higher layer at previous steps is given as additional inputs to an RNN's lower layers \citep{chung2016hierarchical}. Additionally the higher level information is flushed, copied, or updated, using a sparse gating operation. However a key difference from our work is that we use attention to determine whether and what to read from the higher or lower levels, as opposed to concatenating into the input. Additionally, our approach uses a modular recurrent state.

\textbf{Generative Classifiers:} These classify by selecting a class y to maximize $p(x|y)p(y)$.  This is a top-down form of classification, since a key factor is estimating the likelihood of a given sample $x$ under a generative model $p(x|y)$.  Notably, this purely top-down approach often struggles as learning a model of $p(x|y)$ is generally much harder and much more complex than a direct model of $p(y|x)$ \citep{genclass2002ng}.  

\textbf{Link to Global Workspace Theory:}
The GWT~\citep{baars1997theatre,Dehaene-et-al-2017} proposes that this high-level conscious content is broadcast across the brain via top-down connections
forming an active circuit with the bottom-up connections.
What has been missing from RIMs and which we propose
to study here is therefore a notion of hierarchy
of levels enabling this kind of bidirectional information flow
mediated by attention. One major difference between the current work, and GWT is that,  the kinds of modules that are usually considered in GWT are pretty high level, e.g., face recognition, gait recognition, object recognition, visual routines etc i.e., GWT considers \textit{macro} modules, while the current work focuses more on micro modules.













\section{Experiments}


The main goal of our experiments is to explore the effect of BRIMs on generalization.  In particular, we primarily focus on tasks where the training and testing distribution differ systematically, as we believe that the improved top-down modulation of BRIMs will help the model to adapt to variations unseen during training.  Additionally, we provide ablations and analysis which explore how the different components of the BRIMs architecture affect results.  

We show that BRIMs improve out-of-distribution generalization over sequence length, on sequential MNIST and CIFAR classification, and moving MNIST generation.  We then show that BRIMs better reason about physical movement in a synthetic bouncing balls dataset.  In all of these tasks, the test distribution differs dramatically from the training distribution.  We also show that BRIMs improve results on both language modeling and reinforcement learning. As the data distribution implicitly changes in reinforcement learning as a result of the policy changing during training, this provides further evidence that BRIMs improve out-of-distribution generalization in complex settings.  



\begin{table}
\vs{2}
\begin{center}
\renewcommand{\arraystretch}{1.0}
\resizebox{\columnwidth}{!}{
\begin{tabular}{ ccccc } 
 \hline
 \textit{Algorithm} & \textit{Properties} & \textit{16$\times$16} & \textit{19$\times$19} & \textit{24$\times$24}\\ 
 \hline
 LSTM & --- &86.8 & 42.3 & 25.2 \\
 LSTM & \textsc{h} & 87.2 & 43.5 & 22.9 \\
 LSTM & \textsc{h+b} & 83.2 & 44.4 & 25.3 \\
 LSTM & \textsc{h+a} & 84.3 & 47.5 & 31.0\\
 LSTM & \textsc{h+a+b} & 83.2 & 40.1 & 20.8\\
 \hline
 RMC & \textsc{a} & 89.6 & 54.2 & 27.8\\
 Transformers & \textsc{h+a+b} & \highlight{91.2} & 51.6 & 22.9\\
 RIMs & \textsc{a+m} & 88.9 & 67.1 & 38.1 \\
 \hline
 Hierarchical RIMs & \textsc{h+a+m} & 85.4 & 72.0 & 50.3 \\
 MLD-RIMs & \textsc{h+a+m} & 88.8 & 69.1 & 45.3 \\
 BRIMs (ours) & \textsc{h+a+b+m} & 88.6 & \highlight{74.2} & \highlight{51.4} \\
 \hline
\end{tabular}
}
\end{center}
\caption{Performance on the \textbf{Sequential MNIST resolution generalization:} Test Accuracy \% after 100 epochs.  All models were trained on 14x14 resolution but evaluated at different resolutions; results averaged over 3 different trials.}
\label{table:mnist}
\vspace{-6mm}
\end{table}

\subsection{Baselines}


In our experiments, we consider multiple baselines which use some aspects of BRIMs, such as stacked recurrent layers and modularity. We compare our architecture against various state-of-the-art models (Transformers \citep{vaswani2017attention}, Relational RNNs \citep{santoro2018relational}) while also providing ablation studies based on an LSTM backbone \citep{hochreiter1997long}.  By outperforming these baselines, we demonstrate the necessity of the contributions introduced in the BRIMs architecture. In particular, we try to disentangle the contributions of 
attention [\textsc{a}], hierarchy [\textsc{h}], modularity [\textsc{m}], and bidirectional [\textsc{b}] structure.



\textit{LSTM variants}: We consider variants of LSTMs based on different subsets of important properties. In particular, we experimented with standard LSTMs as well as hierarchical LSTMs without feedback [\textsc{h}], with feedback [\textsc{h+b}],
with attention [\textsc{h+a}], and with both [\textsc{h+a+b}].


\textit{Transformers} [\textsc{h+a}]: Self-attention based multi-layer architecture \citep{vaswani2017attention}. 

\textit{RMC, a relational RNN} [\textsc{a}]: Memory based recurrent model with attention communicating between saved memory and hidden states \citep{santoro2018relational}.  

\textit{Recurrent Independent Mechanisms (RIMs)} [\textsc{a+m}]: Modular memory based single layered recurrent model with attention modulated input and communication between modules \citep{goyal2019recurrent}. 

\textit{Hierarchical RIMs} [\textsc{h+a+m}]: Two layered modular architecture with attention based communication between the modules in the same layer and between different layers. Here the flow of information is unidirectional i.e information only flows from bottom to top, and hence no top-down information is being used. 

\textit{MLD-RIMs, RIMs with Multilayer Dynamics} [\textsc{h+a+m}]: The same as hierarchical RIMs except that the activation of RIMs on the higher layer is copied from the lower layer instead of being computed at the higher layer.  This more tightly couples the behavior of the lower and higher layers. 

\textit{BRIMs} [\textsc{h+a+m+b}]: Our model incorporates all these ingredients: layered structure, where each layer is composed of modules, sparingly interacting with the bottleneck of attention, and allowing top down information flow. 

For more detailed description about the baselines, we refer the reader to section \ref{appendix:baseline_desc} in Appendix.



\begin{table}
\vs{3}
\begin{center}
\renewcommand{\arraystretch}{1.0}
\resizebox{\columnwidth}{!}{
\scalebox{0.95}{\begin{tabular}{ ccccc } 
 \hline

 \textit{Algorithm} & \textit{Properties} & \textit{19$\times$19} & \textit{24$\times$24} & \textit{32$\times$32} \\ 
 \hline
 LSTM & ---  & 54.4 & 44.0 & 32.2\\
 LSTM & \textsc{h} & 57.0 & 46.8 & 33.2\\
 LSTM & \textsc{h+b} & 56.5 & 52.2 & 42.1 \\
LSTM & \textsc{h+a} & 56.7 & 51.5 & 40.0 \\
LSTM & \textsc{h+a+b} & 59.9 & 54.6 & 43.0\\
\hline
 RMC & \textsc{a} & 49.9 & 44.3 & 31.3 \\
 RIMs & \textsc{a+m} & 56.9 & 51.4 & 40.1 \\
 \hline
Hierarchical RIMs & \textsc{h+a+m} & 57.2 & 54.6 & 46.8 \\
MLD-RIMs & \textsc{h+a+m} & 56.8 & 53.1 & 44.5\\
BRIMs (ours) & \textsc{h+a+b+m} & \highlight{60.1} & \highlight{57.7} & \highlight{52.2} \\
 \hline
\end{tabular}}
\vspace{-4mm}
}
\end{center}
\caption{Performance on \textbf{Sequential CIFAR generalization:} Test Accuracy \% after 100 epochs.  Both the proposed and the Baseline model (LSTM) were trained on 16x16 resolution but evaluated at different resolutions; results averaged over 3 different trials.}
\label{table:cifar}
\vspace{-4mm}
\end{table}

\subsection{Model Setup}
We use a 2-layered setup with each layer consisting of a set of modules. The proposed architecture has 2 degrees of freedom:  the number of modules in each layer and number of active modules in each layer. For the supervised loss, the final state of the recurrent model is passed through a feed-forward network before making a prediction. For RL experiments, the concatenation of the state of lower level modules as well as higher level modules is used to compute the policy. For video prediction experiments, the state of the lower level module is fed into the decoder, and used for generating the image. Unless otherwise indicated we always  (a) learn an embedding for input tokens before feeding it to the RNNs, (b) use adam with a learning rate
of 0.001 and momentum of 0.9. We include more experimental results in the supplementary material.




\begin{figure}
  \centering
  {\includegraphics[width=0.75\linewidth]{images/Multiple_SSIM_89_final.png}}
 
  {\includegraphics[width=0.75\linewidth]{images/Multiple_MSE_89_final.png}}
 
  \caption{Comparison of BRIMs and SVG-LP
  \citep{denton2018stochastic} 
  on Stochastic Moving MNIST, evaluated using (a) SSIM (higher is better) and (b) MSE (lower is better).
  Models are trained on digits 0-7 and tested with digits 8-9.
  The models are conditioned on 5 time steps, and trained for 10 more time steps. SSIM is calculated by rolling out the models into future time steps.}
  \label{fig:mnist}
  \vspace{-5mm}
\end{figure}
\subsection{Sequential MNIST and CIFAR}
\vs{1}
\label{sPixel}

These tasks involve feeding a recurrent model the pixels of an image in a scan-line order and
producing a classification label at the end of the sequence. To study the generalization capacity of different recurrent models, we train on sMNIST at a resolution of (14,14) and test on the resolutions (16,16), (19,19) and (24,24). Similarly, we perform training for sCIFAR at resolution (16,16) and testing on (19,19), (24,24) and (32,32).  Because these images contain many distracting and uninformative regions, a model which is better able to ignore these patterns should generalize better to longer sequences.  Additionally many details in these images are difficult to understand without having good top-down priors.  These will be most salient when the specific character of these details changes as, for example, when we increase the resolution.  See Tables~\ref{table:cifar},~\ref{table:mnist} for results and appendix section~\ref{appendix:mnist} for more details.









\paragraph{Analysis of Top Down Attention.} We perform analysis on learned BRIMs models to demonstrate the nature of the learned bidirectional flow on sCIFAR task. We show that the model's reliance on top-down information increases when the bottom-up information is less reliable. Further, in general, the models query higher level top-down information only sparingly.

We experimentally verified that almost every example attends to the higher level, yet only a small fraction of steps attend to the higher level on typical examples.  Quantitatively, 95.1\% of images accessed the higher level at least five times (with a score $\geq$ 50\%), while on average only 2.86\% of the total attention (averaged over all modules) was to the higher level. This shows that the model uses top-down information consistently but sparsely.

We also analysed the influence of top-down connections when there is noise in the data. We show that if we train normally but make some of the pixels random (uniformly random noise) at test time, then the model puts more of its attention on the higher level, as opposed to the input (Figure~\ref{fig:noise-cifar}). This is evidence that the model learns to rely more heavily on expectations and prior knowledge when the input sequence is less reliable.


\begin{figure}
    \centering
    \includegraphics[width=0.9\linewidth,trim={0.5cm 0 0.5cm 2cm},clip]{figures/bouncing_balls_loss.png}
    \vspace{-4mm}
    \caption{We study the performance of the proposed model as compared to the LSTM as well as RIMs baseline. We feed the ground truth for the first 15 time steps, and then we ask the model to predict the next 35 time steps. During rollout phase, the proposed model performs better in terms of accurately predicting the dynamics of the balls for the 4Balls scenario, as well as for the more difficult occlusion scenario, as reflected by the lower BCE.}
    \label{fig:bouncing_balls_loss}
    \vspace{-4mm}
\end{figure}

\subsection{Adding Task}
\vs{1}

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 two additional input streams which vary randomly between examples.  The length of the input sequence during testing (2000) is longer than during training (1000).  This is a simple test of the model's ability to ignore the distractor numbers. While we train on two number additions, we also test it on cases where more than two numbers are to be added. We provide the results in Table \ref{table:adding}. We also point the readers to appendix section~\ref{appendix:adding} for additional experiments and details regarding this setup. Here, our Random Prediction baseline, unlike the models, assumes knowledge of the number of digits to add.

\vs{1}
\subsection{Moving MNIST: Video Prediction}
\vs{1}
\label{movingMNIST}

We use the Stochastic Moving MNIST (SM-MNIST)
dataset introduced by \citet{denton2018stochastic} which consists of sequences of frames of size 64 x 64,
containing one or two MNIST digits moving and bouncing
off the edge of the frame (walls). Training sequences were generated on the fly by sampling
two different MNIST digits from the training set (60k total
digits). We trained the proposed model on SM-MNIST by conditioning on 5 frames and training
the model to predict the next 10 frames in the sequence, and compared it to the SVG-LP model \citep{denton2018stochastic}.

To test out-of-distribution generalization capability, we trained models on image sequences from the SM-MNIST dataset containing only digits 0-7, and tested them on sequences containing only digits 8 \& 9. We compared the Structural Similarity index (SSIM) and the Mean Squared Error (MSE) of the best generated sequences from the baseline SVG-LP model, and those from ours. We followed the same evaluation protocol as SVG, and the results in \hyperref[fig:mnist]{Figure~\ref{fig:mnist}} demonstrate that our proposed method performs consistently better. For details about the experimental setup, we refer the reader to section \ref{appendix:moving_appendix} in Appendix.

\paragraph{Analysis of Top Down Attention.} We hypothesize that higher level encodes top-down beliefs (eg. location of digits) that drive the lower level to draw the digits for the next frame. To support this claim, we train a classifier to predict both the digits’ locations from the learned representation in 1 layered RIMs (82\% accuracy), higher layer of 2 layered LSTM (76\% accuracy), higher layer of BRIMs (88\% accuracy) and lower layer of BRIMs (58\% accuracy). This indicates that the information about the positions of the digits is stored in higher levels that could be useful for better modelling of the dynamics at a lower level especially when the digits overlap.

\begin{table}[t!]
\centering
\renewcommand{\arraystretch}{1.0}
\resizebox{\columnwidth}{!}{
\begin{tabular}{c|c|c|c} 
 \toprule
    \textit{Number of Values} & \textit{Random Prediction} & \textit{LSTM} & \textit{BRIMs (ours)} \\
    \midrule
        2 & 0.500 & 0.0842 & \highlight{0.0000} \\
        3 & 1.000 & 0.4601 & \highlight{0.0999} \\
        4 & 1.333 & 1.352 & \highlight{0.4564} \\
        5 & 2.500 & 2.738 & \highlight{1.232} \\
        10 & \highlight{9.161} & 17.246 & 11.90 \\
\bottomrule
\end{tabular}
}
\vs{3}
\caption{Results on Adding task with training and testing on 1000 and 2000 sequence lengths respectively. Number of values to be added are varied from 2 to 10 during test time.}
\label{table:adding}
\vspace{-6mm}
\end{table}


\subsection{Handling Occlusion in Bouncing Balls}
\label{bouncingballs}

We study the ability of the proposed model to model physical reasoning capabilities, on the bouncing balls task, a standard environment for evaluating physical reasoning capabilities exhibiting complex non-linear physical dynamics. We train BRIMs on sequences of $64 \times 64$ binary images over 51 time-steps that contain four bouncing balls with different masses corresponding to their radii. The balls are initialized with random initial positions, masses and velocities. Balls bounce elastically against each other and the image window. We also test the proposed method  in the scenario when an invisible ``curtain'' \citep{van2018relational} is present, which further tests the degree to which the model correctly understands the underlying dynamics of the scene.  For vizualizations about the predictions from the proposed model, refer to section \ref{appendix:bouncing_balls} in Appendix.

\paragraph{Analysis of Top Down Attention.} Achieving good performance on this task requires the model to track the ball from behind the curtain. We find that whenever the perceptual input is less reliable (i.e a ball is occluded) the perception is relied more by the information coming from higher level layers (higher top-down attention score).




\subsection{Language Modelling}

We consider word-level language modeling on the WikiText-103 dataset which has been used as a benchmark in previous work \citep{santoro2018relational}.  Language has an organic hierarchical and recurrent structure which is both noisy and flexible, making it an interesting and apt task to study for the proposed model.  Table~\ref{table:wiki} shows Validation and Test set perplexities for various models. We find that the proposed method improves over the standard LSTM as well as RMCs. 
\begin{table}
\begin{center}
\renewcommand{\arraystretch}{1.0}
\begin{tabular}{cccc} 
 \hline
 \textit{Algorithm} & \textit{Properties}  & \textit{Valid} & \textit{Test} \\ 
 \hline
 LSTM & --- & 38.2 & 41.8 \\
 RMC* (ours) & \textsc{h} & 36.2 & 38.3 \\
Hierarchical RIMs & \textsc{h+a+m} & 36.1 & 38.1 \\
BRIMs (ours) & \textsc{h+a+b+m} & \highlight{35.5} & \highlight{36.8} \\
 \hline
\end{tabular}
\end{center}
\caption{Perplexities on the WikiText-103 dataset. 
$(*)$ refers to our implementation of the model. The actual results from RMC are even better as compared to the proposed model, but the idea here is not to beat well tuned state of the art methods.}
\label{table:wiki}
\vspace{-5mm}
\end{table}

\vspace{-4mm}
\subsection{Reinforcement Learning}

Top-down and bottom-up signals are clearly differentiated in the case of reinforcement learning (or in active agents, more generally) in that they lead to different optimal behavior.  Since top-down signals change more slowly, they often suggest long-term and deliberative planning.  On the other hand, bottom-up signals which are surprising can often require an immediate reaction.  Thus, in selecting actions, points in which the top-down system dominates may benefit from activating a system with a long-term plan, whereas points where the bottom-up information dominates could benefit from a system which immediately takes action to resolve an unexpected threat.  On the whole, we expect that dynamically modulating activation based on bottom-up and top-down signal importance will be especially important in reinforcement learning.  Additionally, in reinforcement learning the data distribution implicitly changes as the policy evolves over the course of training, testing the model's ability to correctly handle distribution shift.  

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 well as RIMs architecture \citep{goyal2019recurrent} as baselines and compare results to the proposed architecture. This was a simple drop-in replacement and did not require changing any of the hyperparameters for PPO. We experimented on 10 Atari games (chosen alphabetically) and ran three independent trials for each game (reporting mean and standard deviation).  We found that simply using BRIMs for the recurrent policy improves performance (\hyperref[table:atari]{Table~\ref{table:atari}}).  We also found an improvement on Ms. Pacman, which requires careful  planning and is not purely reactive. 

\begin{table}[t!]
\large
\renewcommand{\arraystretch}{1.2}
\centering
\resizebox{\columnwidth}{!}{
\begin{tabular}{c|rm{1em}l|rm{1em}l|rm{1em}l}
    \toprule
    \textit{Environment} & \multicolumn{3}{c}{\textit{LSTM}} & \multicolumn{3}{c}{\textit{RIMs}} &  \multicolumn{3}{c}{\textit{BRIMs (ours)}}\\
    \midrule
    Alien & 1612  &$\pm$& 44 & 2152 &$\pm$& 81 & \highlight{4102} &$\pm$& 400 \\
    Amidar & 1000  &$\pm$& 58 & 1800 &$\pm$& 43 & \highlight{2454} &$\pm$& 100\\
    Assault &  4000 &$\pm$& 213 & 5400 &$\pm$& 312 & \highlight{5700} &$\pm$& 320 \\
    Asterix & 3090 &$\pm$& 420 & 21040 &$\pm$& 548 & \highlight{30700} &$\pm$& 3200\\
    Asteroids & 1611 &$\pm$& 200 & \highlight{3801} &$\pm$& 89 & 2000 &$\pm$& 300\\
    Atlantis & 3.28M &$\pm$& 0.20M & 3.5M &$\pm$& 0.12M & \highlight{3.9M} &$\pm$& 0.05M \\
    BankHeist & 1153 &$\pm$& 23 & \highlight{1195} &$\pm$& 4 & 1155 &$\pm$& 20 \\
    BattleZone & 21000 &$\pm$& 232 & 22000 &$\pm$& 324 & \highlight{25000} &$\pm$& 414 \\
    BeamRider & 698 &$\pm$& 100 & \highlight{5320} &$\pm$& 300 & 4000 &$\pm$& 323\\
    MsPacMan & 4598 &$\pm$& 100 & 3920 &$\pm$& 500 & \highlight{5900} &$\pm$& 1000 \\
\bottomrule
\end{tabular}
}
\caption{Results on the Atari Reinforcement learning task using PPO with different recurrent policy networks.}
\label{table:atari}
\vspace{-4mm}
\end{table}

\section{Ablations and Analysis}
We have demonstrated substantially improved results using BRIMs.  However, it is essential to understand which aspects of the newly proposed model are important. To this end we conduct ablation studies to demonstrate the necessity of bidirectional structure, attention, modularity, and hierarchy.  

\vspace{-1mm}
\paragraph{Role of Bidirectional Information Flow.}
\label{td}
We demonstrate that in different architectures, having top down information flow helps in generalization. Tables \ref{table:cifar},~\ref{table:mnist},~\ref{table:wiki} demonstrate the effectiveness of using top down connections in standard LSTMs, attention based LSTMs and hierarchical modules. We maintain that higher layers are able to filter out important information from lower layers and thus conditioning on it provides much more context as to what is relevant at any time in a sequence.  However, the results still fall short of BRIMs, indicating that this component alone is not sufficient. 



\vspace{-1mm}
\paragraph{Role of Attention.}
Tables~\ref{table:cifar},~\ref{table:mnist},~\ref{table:wiki},~\ref{table:atari} and Figure~\ref{fig:mnist} show the importance of having attention, as Transformers, RMC, attention-modulated LSTMs, RIMs and BRIMs lead to much better performance and generalization in domains ranging from supervised learning to reinforcement learning than their non attention counterparts. This suggests that key-value attention mechanism is better able to extract out important information and dependency as opposed to standard connections.  

\vspace{-1mm}
\paragraph{Role of Modularity.}
Tables~\ref{table:cifar},~\ref{table:mnist},~\ref{table:wiki},~\ref{table:atari} indicate that having layers as a composition of smaller modules with their own independent dynamics outperform architectures without modular decomposition. In particular, we note that RIMs, their hierarchical variants and BRIMs perform much better than their LSTM counterparts. This shows that having independent dynamics for modules as well as a communication channel between them are important ingredients for generalization.
\vspace{-1mm}
\paragraph{Role of Hierarchy.}
We consistently found that models with hierarchical structure perform much better than their shallow counterparts (Tables ~\ref{table:cifar},~\ref{table:mnist},~\ref{table:atari}). This relates to the discussion in Section \hyperref[td]{7.1} that higher layers are able to abstract and filter out important information and concepts learned in lower layers.

\vspace{-2mm}
\paragraph{Capacity of Network.}
While BRIMs introduce additional parameters through the attention mechanisms, they are roughly cancelled out by the block diagonal structure in the independent module dynamics. As an example, our LSTM and BRIMs models for sCIFAR both contain $\sim$ 1.5M parameters. Our additional experiments on larger LSTMs showed a drop in performance (2-3\%) on sMNIST and sCIFAR.






\section{Conclusion}

Top-down and bottom-up information are both critical to robust and accurate perception.  
Exploring the combination of these signals in network architectures has an important history in deep learning, and has been a central question in cognitive science for several decades.
Our work focuses on the idea that attention can be used to give models explicit dynamic and context dependent control over the combination of top-down and bottom-up signals.  Our simulation experiments have shown a critical role for network modularity: ensuring that specific top-down signals combine with specific bottom-up signals.  Using these insights we have proposed a new algorithm, \emph{Bidirectional Recurrent Independent Mechanisms (BRIMs)}, which achieves substantial improvements on sequence-length generalization tasks, language modeling, video generation, and reinforcement learning on Atari.  A detailed ablation study gives evidence that combining top-down and bottom-up signals through attention is a critical component for the improved results.  

\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 Sergey Levine and Nasim Rahaman for useful discussions. The authors would also like to thank Mark Pickett for feedback on the earlier version of 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. 


