# Leveraging Low-Rank and Sparse
Recurrent Connectivity for Robust
Closed-Loop Control

Neehal Tumma1 Mathias Lechner2 Noel Loo2 Ramin Hasani2 Daniela Rus2
1Department of Computer Science, Harvard University
2Computer Science and Artificial Intelligence Lab, MIT
<EMAIL_ADDRESS>
###### Abstract

Developing autonomous agents that can interact with changing environments is
an open challenge in machine learning. Robustness is particularly important in
these settings as agents are often fit offline on expert demonstrations but
deployed online where they must generalize to the closed feedback loop within
the environment. In this work, we explore the application of recurrent neural
networks to tasks of this nature and understand how a parameterization of
their recurrent connectivity influences robustness in closed-loop settings.
Specifically, we represent the recurrent connectivity as a function of rank
and sparsity and show both theoretically and empirically that modulating these
two variables has desirable effects on network dynamics. The proposed low-
rank, sparse connectivity induces an interpretable prior on the network that
proves to be most amenable for a class of models known as closed-form
continuous-time neural networks (CfCs). We find that CfCs with fewer
parameters can outperform their full-rank, fully-connected counterparts in the
online setting under distribution shift. This yields memory-efficient and
robust agents while opening a new perspective on how we can modulate network
dynamics through connectivity.

##  1 Introduction

Building models that are robust under natural distribution shift has long been
a goal in artificial intelligence (Taori et al., 2020). Existing techniques
have sought to address this challenge through various approaches, including
domain adaptation (Farahani et al., 2020), transfer learning (Zhuang et al.,
2019) and data augmentation (Perez & Wang, 2017). However, these techniques
come with drawbacks: domain adaptation and transfer learning can be
computationally expensive, and data augmentation often suffers from robustness
gains that are not uniform across corruption types (Ford et al., 2019). More
generally, machine learning systems tend to perform poorly under distribution
shift because approaches like these only serve to ameliorate a problem that is
rooted in the model architecture itself.

Natural learning systems serve to address this problem by interacting with
their environment to understand the world. They make use of biologically-
inspired frameworks that account for distribution shifts by modeling temporal
structure (Hasani et al., 2020b). These models are particularly well-suited in
a paradigm where they are fit on offline, passive datasets using imitation
learning, but deployed in closed-loop, active testing settings (de Haan et
al., 2019). This domain is commonly referred to as the open-loop to closed-
loop causality gap (Lechner et al., 2022) in which generalization requires the
agent to learn a coherent representation of its world (Lechner et al., 2020).
One particularly effective framework known as closed-form continuous-time
neural networks (CfCs) were shown to outperform many state-of-the-art
recurrent models in closed-loop settings (Hasani et al., 2021). These models
leverage a sparse wiring structure induced at initialization, yielding compact
and interpretable networks (Chahine et al., 2023). However, the reason as to
why sparse connectivity is useful in closed-loop systems is poorly understood.
In this work, we will address this question by proposing and analyzing an
interpretable parameterization of connectivity in CfCs.

To do so, we turn to another class of models: low-rank recurrent neural
networks (low-rank RNNs). A low-rank RNN is a network whose _recurrent_
connectivity matrix is low-rank. Deriving inspiration from neural activity in
the brain, these networks have provably low-dimensional patterns of activity
which makes them successful in many simple neuroscience-based tasks
(Mastrogiuseppe & Ostojic, 2018). In this paper, we leverage ideas from
previous work on CfCs and low-rank RNNs in order to devise a novel
parameterization of recurrent connectivity that improves robustness in closed-
loop settings and offers interpretable measures for the network dynamics it
incites. In doing so, we show the following:

  * •

Low-rank and sparse recurrent neural networks can outperform their full-rank,
fully-connected counterparts under distribution shift in closed-loop
environments

  * •

Parameterizing connectivity as a function of rank and sparsity yields provable
and interpretable network dynamics

  * •

Pruning by inducing sparsity and pruning by enforcing a low-rank structure are
distinct in the types of network dynamics each induces

  * •

CfCs are more amenable to a low-rank and sparse connectivity prior than other
canonical recurrent architectures such as LSTMs

##  2 Related work

In this section, we describe previous works that are closely related to the
core findings of the paper.

Robustness in closed-loop settings. A class of continuous-time recurrent
models known as liquid neural networks (Hasani et al., 2020a) have been shown
to achieve state-of-the-art performance under distribution shift by explicitly
accounting for external interventions to environment conditions. Liquid neural
networks are a prime example of natural learning systems that learn robust,
and provably causal, representations (Vorbach et al., 2021). CfCs are a
specific instance of liquid neural networks that can leverage sparse
connectivity to learn robust representations in the closed-loop causality gap
setting.

Figure 1: Open-loop systems receive ground-truth observations $x_{i}$. Closed-
loop systems receive observations $x_{i}(a_{i-1})$ that are a function of the
previous actions the agent takes. There is no external feedback to correct the
agent in the closed-loop setting.

Outside of modifying the model itself, other approaches to training robust
models under an imitation learning framework include augmentation strategies,
human interventions (Ross et al., 2010), goal-conditioning (Codevilla et al.,
2017), reward conditioning (Srivastava et al., 2019), task-embedding (James et
al., 2018) and meta-learning (Finn et al., 2017). These advances fail to
consider the structure of the underlying policy; in contrast, liquid neural
networks improve the robustness of the decision-making process in autonomous
agents, leading to better generalization under the same training distribution
(Vorbach et al., 2021).

Relating connectivity and network dynamics. A longstanding goal in
computational neuroscience is to understand the relationship between structure
and function in the brain (Sporns, 2013). We can ask the analogous question in
the context of artificial neural networks: what role does model connectivity
at initialization play in learned network dynamics? Amongst recurrent neural
networks, echo state networks (Goodfellow et al., 2016) present one example of
a model that induces static sparsity in order to modulate the complexity of
dynamics along the recurrent dimension. However, unlike in our case, echo
state networks fix their recurrent connectivity during training. Other works
have found that sparse networks can yield provably more robust models (Guo et
al., 2018) that generalize across many tasks (Chen et al., 2022) outside the
recurrent setting. With respect to modulating network rank, the advent of the
low-rank RNN framework (Mastrogiuseppe & Ostojic, 2018) is one of the first
attempts at explicitly modeling the low-dimensional dynamics of the brain.
This is done by constructing a low-rank connectivity matrix with additive
random noise. Some follow-ups on this work include examining how low-rank
connectivity arises in full-rank settings (Schuessler et al., 2020) and an
analysis of how some tasks are more amenable to low-rank connectivity than
others (Dubreuil, 2022). This work borrows approaches and intuition from many
of these prior studies, but is distinct with respect to the proposed
connectivity, analysis of network dynamics and task setting.

Pruning at initialization. Neural network pruning typically either removes
connections in a costly iterative training-retraining paradigm (Frankle &
Carbin, 2018) or modifies the objective function to promote sparsity
(Goodfellow et al., 2016) which can lead to difficulties in optimization.
Recently, pruning at initialization has emerged as an approach to resolve
these issues. Pruning at initialization attempts to take a randomly
initialized network and remove weights before training (Wang et al., 2021).
Various approaches include using connectivity sensitivity (Lee et al., 2018),
maintaining dynamical isometry (Lee et al., 2019) and training on supermasks
(Ramanujan et al., 2019). Generally, these techniques outperform sparsity
induced randomly at initialization. Another approach to pruning is using rank
decomposition. One example is low-rank matrix factorization, which is good at
reducing size but at the cost of performance (Sainath et al., 2013). The most
effective low-rank compression techniques usually leverage a low-rank
approximation on the full-rank matrix either during or after training (Xu et
al., 2020). In contrast, in this work, we propose a parameterization of
connectivity that randomly induces sparsity and leverages a low-rank
decomposition of the recurrent weights at initialization – the most
straightforward and least costly form of pruning.

Alternative RNN parameterizations Existing works proposed other recurrent
architectures such as Lipschitz RNNs (Erichson et al., 2021), antisymmetric
RNNs (Chang et al., 2019) and Cayley RNNs (Helfrich et al., 2018) which, like
us, propose alternative parameterizations of recurrent connectivity. However,
unlike our study, these works introduce changes to the underlying network
architecture, whereas our work focuses on the applications of low-rank and
sparse recurrent connectivity in established RNN architectures. This is in
line with previous works in the closed-loop causality gap setting (Chahine et
al., 2023) which also restrict the scope of their models as such.

##  3 parameterization of connectivity

Consider a standard RNN whose input is denoted by
$\displaystyle{\bm{x}}_{t}\in\mathbb{R}^{n}$ and hidden state is given by
$\displaystyle{\bm{h}}_{t}\in\mathbb{R}^{h}$ for time step $t$, hidden size
$h$ and input size $n$. Then the functional form of the RNN is given by

$\displaystyle{\bm{h}}_{t}={\bm{W}}_{rec}\sigma(\displaystyle{\bm{h}}_{t-1})+{\bm{W}}_{inp}\displaystyle{\bm{x}}_{t}+\displaystyle{\bm{b}}$
(1)

such that ${\bm{W}}_{rec}\in\mathbb{R}^{h\times h}$ denotes the recurrent
weights, ${\bm{W}}_{inp}\in\mathbb{R}^{n\times h}$ denotes the input weights
and $\displaystyle{\bm{b}}$ denotes the bias. We now propose an alternative
parameterization of connectivity for the recurrent weights. Consider a
singular value decomposition of ${\bm{W}}_{rec}$ given by
${\bm{W}}_{rec}={\bm{U}}{\bm{\Sigma}}{\bm{V}}^{T}$. The canonical rank $r$
approximation of ${\bm{W}}_{rec}$ is
${\bm{U}}_{r}{\bm{\Sigma}}_{r}{\bm{V}}_{r}^{T}$ such that ${\bm{\Sigma}}_{r}$
consists of the top-$r$ singular values and ${\bm{U}}_{r}$ and
${\bm{V}}_{r}^{T}$ consist of the corresponding singular vectors. Using this
low-rank approximation, we can construct the following parameterization of the
recurrent weights for a given rank $r$ and sparsity level $s$:

${\bm{W}}_{rec}(r,s)={\bm{W}}_{1}(r){\bm{W}}_{2}(r)\odot{\bm{M}}(s)$ (2)

where ${\bm{W}}_{1}={\bm{U}}_{r}({\bm{\Sigma}}_{r})^{1/2}$,
${\bm{W}}_{2}=({\bm{\Sigma}}_{r})^{1/2}{\bm{V}}_{r}^{T}$ and
${\bm{M}}\in\mathbb{F}_{2}$ is a random binary mask such that each entry
${\bm{M}}_{ij}\sim 1-\textrm{Bernoulli}(s)$. This low-rank and sparse
parameterization of ${\bm{W}}_{rec}$ is a generalization of the
parameterization studied in Herbert & Ostojic (2022) to recurrent connectivity
of arbitrary rank and trainable weights ${\bm{W}}_{1},{\bm{W}}_{2}$ as opposed
to restricting ourselves to the setting of rank-1 matrices and fixed weights
in which dynamical analysis is more tractable, but the network is constrained
in expressivity. Note that the random mask ${\bm{M}}$ is fixed throughout
training.

###  3.1 Network dynamics at initialization

We propose this parameterization of recurrent connectivity as it allows us to
modulate aspects of the network that instill an inductive bias which proves to
be beneficial for performance in the causality gap setting. We care about two
measures of network dynamics in particular: the spectral radius of
${\bm{W}}_{rec}(r,s)$ and the spectral norm of ${\bm{W}}_{rec}(r,s)$.

The spectral radius $\lambda_{max}$ of ${\bm{W}}_{rec}(r,s)$ is widely
accepted as a proxy for the rate at which the gradient vanishes across the
time dimension (Pascanu et al., 2012). Namely, if $\lambda_{max}$ is less than
$1$, a recurrent network has a vanishing gradient. In many time series
applications, having a constant error flow is a desirable property, as
arbitrary data sequences may have long-term relations. However, in the case of
many closed-loop systems, learning long-term dependencies can be detrimental
in the online setting due to the short-term causality inherent to the task
(Lechner et al., 2020). To put this into context, in this work we consider
agents learning to play games in various Arcade Learning Environments (ALEs).
Agents that perform well in ALEs often use frame stacking which enforces a
strict short-horizon temporal prior (Horgan et al., 2018a) by leveraging a
look-back window of only a few input frames. Hence, networks like LSTMs or
GRUs may capture spurious long-term dependencies and thus learn inadequate
models. In contrast, the vanishing of gradients present in RNNs
counterintuitively enhances the performance of the agent as it places a short-
horizon prior on the temporal attention span of the network.

The spectral norm $\sigma_{max}$ of ${\bm{W}}_{rec}(r,s)$ is equivalent to its
Lipschitz constant which offers a proxy for the temporal robustness of the
network. This is important when considering performance under distribution
shift because perturbations applied to the input data will enter the hidden
state as shown in Equation (1). The recurrent activity that supervises the
agent is thus essential in the distribution shift paradigm as the agent can
not entirely rely upon the perturbed input to make decisions.

Leveraging results from random matrix theory on the asymptotic distributions
of eigenvalues and singular values of random matrices, we can prove the
following:

######  Theorem 1.

Given a recurrent neural network with an update function given by Equation
(1), parametrization in Equation (2) and a Glorot uniform spectral-
initialization scheme described in A.4, the following statements hold for the
network at initialization:

  * •

The spectral radius of ${\bm{W}}_{rec}(r,s)$ decreases as a function of $s$

  * •

The spectral radius of ${\bm{W}}_{rec}(r,s)$ increases as a function of $r$

  * •

The spectral norm of ${\bm{W}}_{rec}(r,s)$ decreases as a function of $s$

  * •

The spectral norm of ${\bm{W}}_{rec}(r,s)$ is constant as a function of $r$

For proofs of the above properties in the case of Glorot uniform
initialization and an empirical analysis extending to the case of orthogonal
initialization, refer to A.2.

Figure 2: Parameterization of recurrent connectivity as a function of rank and
sparsity.

##  4 Experiments

Experimental setup. We study the offline-online generalization gap of various
recurrent architectures parameterized by low-rank, sparse connectivity under
an imitation learning framework. In particular, we measure the performance of
our models in the Arcade Learning Environment (Bellemare et al., 2012) and
MuJoCo (Todorov et al., 2012). For ALEs, we run experiments in the Seaquest
and Alien environments and for MuJoCo we explore the HalfCheetah environment.
For each environment, we train a deep-Q network to generate expert
trajectories that we then use to fit our recurrent networks in an offline
setting. We evaluate the recurrent networks online, closed-loop when deployed
into the environment, both in-distribution and under distribution shift. For
more details on the experimental framework, refer to A.5.

Models. We examine the following recurrent architectures: CfCs, RNNs, LSTMs,
GRUs and temporal convolutional networks (TCNs). RNNs, LSTMs (Hochreiter &
Schmidhuber, 1997) and GRUs (Cho et al., 2014) refer to their canonical
implementations. TCNs (Lea et al., 2016) model the temporal structure of an
input sequence by performing a 1-dimensional convolution over the time
dimension and serve to address whether supervision along the recurrent
dimension is even necessary. CfCs (Hasani et al., 2021) are a modified version
of an RNN that admit the following functional form:

${\bm{h}}(t)=\sigma(-F({\bm{h}}_{t-1},{\bm{x}}_{t},\theta_{F}))\odot
G({\bm{h}}_{t-1},{\bm{x}}_{t},\theta_{G})+[1-\sigma(-F({\bm{h}}_{t-1},{\bm{x}}_{t},\theta_{f})]\odot
H({\bm{h}}_{t-1},{\bm{x}}_{t},\theta_{H})$

Here, $H$ and $G$ are vanilla RNNs and $F$ is what we refer to as a liquid
time constant network. In general, time constants are more thoroughly
motivated in the continuous-time setting in which a neural network is used to
model the derivative of the hidden state as opposed to the hidden state itself
(Chen et al., 2018). In that setting, a time constant represents a parameter
that characterizes the speed and coupling sensitivity of an ODE that models a
system. However, our work exists in the discrete-time setting in which we can
no longer interpret the time constant as a parameter of a continuous-time
system. Instead, we interpret $F$ as an adaptive gating mechanism that
interpolates between the state-space trajectories of $H$ and $G$ on a per-
element basis in the hidden vector ${\bm{h}}(t)$.

Finally, recall in Section 3.1 that we proposed a low-rank, sparse
parameterization of connectivity ${\bm{W}}_{rec}(r,s)$ for the recurrent
weights in a vanilla RNN. We generalize this parameterization to the other
recurrent architectures by simply parameterizing all recurrent weights in the
network in the form of ${\bm{W}}_{rec}(r,s)$. For specifics, refer to A.5.

###  4.1 Performance in closed-loop settings

Figure 3: Online performance of recurrent networks under different ranks and
sparsities in Seaquest environment. For offline performance, refer to A.3. a)
In-distribution rewards in the online, closed-loop setting. b) Rewards
averaged across 5 distribution shifts. We analyze these rewards in ranking
space since the rewards under each distribution shift are not on the same
scale.

We first explore the impact of our proposed connectivity parameterization for
various ranks $r$ and sparsities $s$ in the Seaquest environment (Figure 3).
We observe that the best performing models measured by online, closed-loop in-
distribution rewards tend to be low-sparsity, high-rank CfCs and LSTMs (Figure
3a). GRUs, on the other hand, tend to perform poorly for most $(r,s)$ which
aligns with previous work which also shows that GRUs are not particularly
well-suited for closed-loop settings (Chahine et al., 2023). Due to the poor
performance of GRUs relative to LSTMs, we will restrict most of our discussion
and analysis to CfCs, LSTMs and RNNs.

The most notable finding from the in-distribution results is that LSTMs and
CfCs are distinct with respect to their performance at lower sparsities; in
particular, we find that LSTMs tend to perform much better than CfCs at high
sparsities. Analogously, RNNs, which in general tend to perform worse than
CfCs, show similarly worse performance at high sparsities. We offer intuition
for these findings in Section 4.2 where we explore the relationship between
sparsity and the vanishing gradient.

Under distribution shift, we find that the best performing models are low-
rank, low-sparsity CfCs. More generally, across all recurrent architectures,
we find that the low-rank models tend to perform on par with, and in many
cases actually outperform, higher rank ones. This demonstrates that we can
construct models with fewer parameters at initialization yet still learn more
robust models. Interestingly, we find that the means through which we reduce
the parameter count matters: in particular for high-sparsity models, we don’t
achieve the same robustness that we observe in low-rank ones. We will provide
justification for the apparent disparity between sparsity and rank in Sections
4.3 and 4.4. For analogous results in the online setting in the Alien and
HalfCheetah environments as well the offline performance for all models, refer
to A.3.

###  4.2 Vanishing gradient analysis

Figure 4: Eigenspectra of recurrent weights ${\bm{W}}_{rec}(r,s)$ computed by
sorting $|\lambda_{i}|$ in trained CfC and LSTM networks. We measure these
eigenspectra across multiple ranks and sparsities.

Recall in Section 3.1 we motivated the parameterization ${\bm{W}}_{rec}(r,s)$
with respect to two measures, one of which was spectral radius which we
motivated from the perspective of the vanishing gradient. In particular, we
showed at initialization that the spectral radius of ${\bm{W}}_{rec}(r,s)$
decreases as a function of sparsity and increases as a function of rank. Here,
we aim to see if these properties persist throughout network training (for
details, refer to A.5).

In Figure 4, we find that they do: in particular, in both CfCs and LSTMs, the
low-rank and highly-sparse models tend to admit solutions with recurrent
weights of low spectral radius. Comparing the spectral radius in CfCs and
LSTMs, we note that CfCs tend to express the vanishing gradient even in the
baseline fully-connected, full-rank formulation as given by $\lambda_{max}<1$.
In contrast, LSTMs, which promote a more consistent error flow over time
(Jozefowicz et al., 2015), are more robust against expressing the vanishing
gradient. In the context of a closed-loop task, we know that it is beneficial
to limit the network’s temporal attention span to avoid learning spurious,
long-term relationships (Lechner et al., 2020). Since an LSTM in its baseline
form does not express the vanishing gradient, inducing sparsity at
initialization pushes the network into the vanishing gradient regime. This
result provides support as to why LSTMs are quite effective at navigating
environments in-distribution in spite of high-sparsity in their recurrent
weights (Figure 3a). In contrast, CfCs in their baseline form already lie in
the vanishing gradient regime: intuitively, inducing sparsity is less
effective when the baseline network inherently possesses an affinity to
selectively attend to past observations. This aligns with an analogous trend
observed in RNNs which are networks that, like CfCs, are amenable to learning
a vanishing gradient (refer to A.3).

###  4.3 Decomposing state-space trajectories

By analyzing the spectral radius of ${\bm{W}}_{rec}(r,s)$, we identified
sparsity as a modulator of the vanishing gradient. However, note that we also
observe, both at initialization and convergence, that the low-rank networks
tend to admit solutions with lower spectral radius as well. Yet, as shown
earlier, low-rank LSTMs tend to perform poorly in-distribution. To understand
why, we move beyond the leading eigenvalue, and examine the full recurrent
eigenspectrum.

Assume the recurrent weights are diagonalizable and can be equivalently
expressed by ${\bm{W}}_{rec}(r,s)={\bm{P}}{\bm{D}}{\bm{P}}^{-1}$. If we let
${\bm{k}}(t)={\bm{P}}^{-1}{\bm{h}}(t)$, then note that in the low-rank, no
sparsity case, exactly $n-r$ eigenvalues are $0$ which means that $n-r$
dimensions of the eigen-transformed ${\bm{h}}(t)$ are ignored by the network.
This is seemingly a harsh prior as it limits the expressivity of the network
significantly. In particular, because activity along the recurrent axis is
constrained to lie in the subspace spanned by the vectors comprising
${\bm{W}}_{rec}(r,s)$, the recurrent dynamics of low-rank recurrent networks
are much simpler. We can imagine that in the in-distribution setting, it is
not desirable to arbitrarily constrain the network’s capacity to learn
recurrently. This provides justification as to why, in spite of the low-rank
LSTMs also expressing the vanishing gradient, inducing sparsity into LSTMs
does not worsen the in-distribution performance of the network like
constraining rank does (Figure 3a).

In contrast, recall that under distribution shift, we observed the opposite:
inducing sparsity at initialization was unhelpful whereas constraining the
network to be low-rank improved robustness, particularly in CfCs. Constraining
the state-space trajectory of the network is one means of reducing its
variability in the presence of input perturbations. When artificial
distribution shifts are applied to the agent’s observations, they can corrupt
the hidden state. Supervision along the recurrent axis is pivotal when the
agent is faced with environmental occlusions as it needs to lean on some
notion of the past to make a decision in the present. By making the recurrent
state more robust, we are better able to generalize under distribution shifts.

Figure 5: Effective dimensionalities of recurrent state-space, input state-
space and full state-space trajectories collected during online testing as
measured by the explained variance of the top 5 principal components. a) CfC
state-space dynamics. b) LSTM state-space dynamics.

Inducing sparsity, on the other hand, makes the recurrent network dynamics
more complex. To measure this, we collected the state-space trajectories of
the agent during navigation and ran PCA on them to measure their effective
dimensionality. In addition to analyzing the full state-space trajectory, we
decomposed the trajectory into the activity driven along the recurrent axis
and the activity driven along the input axis. We find that with increasing
rank, in both CfCs and LSTMs, the dimensionality of the recurrent trajectories
increases; but, counterintuitively, with increasing sparsity, we find that it
also increases (Figure 5a, Figure 5b). To see why, we return to the
eigenspectra shown in Figure 4. Observe that as we increase sparsity, for a
fixed rank, the eigenvalues corresponding to the non-zero diagonal entries in
an eigen-decomposed ${\bm{W}}_{rec}(r,s)$ are pushed closer to zero; in
contrast, the eigenvalues corresponding to the zero diagonal entries are
lifted from zero. As the network stabilizes its recurrent eigenspectra, it
begins to balance its attention profile across each dimension of the hidden
state which complicates learned dynamics along the recurrent axis.

The input-driven and full state-space trajectories, however, do not have clear
trends as a function of rank/sparsity in either LSTMs or CfCs. This implies
that the dimensionality of input-driven activity is little affected by changes
in the recurrent connectivity. However, if we analyze the full trajectory, we
find that for all pairs $(r,s)$, the dimensionality of LSTM trajectories is
much higher than that of CfC trajectories. This points to the inherent
complexity of LSTMs and recurrent networks in general that have intricate
gating mechanisms. In contrast, CfCs, which employ a simple, adaptive gating
mechanism in the form of the time constant network $F$, express dynamics that
are inherently simpler than dynamics in LSTMs. Interestingly, we find that RNN
dynamics tend to be less complex than LSTMs, but marginally more complex than
CfCs (refer to A.3). Since simpler state-space dynamics prove to be a
desirable property for agents navigating environments under distribution
shift, we find that CfCs, and to some extent RNNs as well, are demonstrably
more robust than LSTMs.

###  4.4 Robustness against distribution shift

Figure 6: Spectral norm of recurrent weights of trained networks as a proxy
for robustness under distribution shift.

Up to this point, we have demonstrated that LSTMs both vanish slower across
time and have more complex state-space dynamics than RNNs and CfCs which has
allowed us to justify their inferior performance under distribution shift.
Yet, we also observed that CfCs tend to outperform RNNs under distribution
shift, despite them not differing to the extent CfCs and LSTMs do with respect
to the network attributes we have explored so far.

To see why, we analyze the spectral norm of ${\bm{W}}_{rec}(r,s)$ which is
equivalent to the Lipschitz constant of the recurrent weights – a measure of
the network’s temporal robustness (for details, refer to A.5). We observe that
CfCs tend to learn recurrent weights with lower spectral norms than RNNs and
significantly lower spectral norms than LSTMs (Figure 6). Note that this holds
in spite of the spectral norm being equal to $1$ at initialization across all
of the networks (i.e., see A.4). This provides insight into the improved
robustness in CfCs relative to RNNs and also affirms why LSTMs are not
particularly robust under distribution shift. But this begs the question as to
why CfCs learn recurrent weights with lower spectral norm, despite the
architecture having two recurrent gates as opposed to the one in a vanilla
RNN. Recall the functional form of a CfC which includes a time constant
network $F({\bm{h}}_{t-1},{\bm{x}}_{t})$ that learns a vector of weights used
to interpolate between the $G$ and $H$ trajectories. Note that if
$F({\bm{h}}_{t-1},{\bm{x}}_{t})=[1,1,\dots,1]$, the network reduces to that of
an RNN by placing all of its weight on a single trajectory. Thus, we
hypothesize that the added expressivity conferred to the CfC via $F$ enables
the recurrent gates in $G$ and $H$ to learn functions that have lower
Lipschitz constants. And because the interpolation weights learned by $F$ are
a function of the input, the time constant network can dynamically adapt its
interpolation weights under distribution shift by changing how much it weighs
$G$ versus $H$. While this certainly warrants further exploration, the notion
of gating mechanisms improving performance in neural networks is not a new
one; recent work done in convolutional language models shows similar findings
(Poli et al., 2023). Given the apparent importance of the time-constant
network, we further explore the implications of the connectivity
parameterization on the learned interpolation weights. For this analysis,
refer to A.1.

###  4.5 Exploring the task dimension gap

We now have demonstrated two things: one is the efficacy of the proposed
connectivity parameterization in units of its interpretability as a modulator
of network dynamics. Second is the impact connectivity has on each of the
networks we examined: in particular, by disentangling the effects of rank and
sparsity we showed why LSTMs are more amenable to connectivity changes in-
distribution whereas CfCs tend to show the most promise when modulating
connectivity under distribution shift. Here, we further our intuition on the
performance of CfCs under distribution shift by understanding why they tend to
be the most amenable architecture to low-rank recurrent connectivity.

Figure 7: We individually consider the Frobenius norm of the weights at
initialization ${\bm{W}}_{0}$ and the change in the weights during training
$\Delta{\bm{W}}.$ Note that the models shown here have no sparsity.

As we have noted, the efficacy of this parameterization rests upon the ability
of the network to adhere to the prior throughout training. And while we find
that all networks converge to solutions with a spectral radius reasonably
close to where they began, the same does not hold true for spectral norm. In
particular, we found that the spectral norm of the recurrent weights in CfCs
remained much closer to 1 in the full-rank, fully-connected network than in
either LSTMs or RNNs (Figure 6). To better understand this, we borrow from an
analysis conducted by Schuessler et al. (2020) in which they decomposed the
recurrent weights as follows: ${\bm{W}}_{rec}={\bm{W}}_{0}+\Delta{\bm{W}}$
where ${\bm{W}}_{0}$ denotes the weights at initialization and
$\Delta{\bm{W}}$ denotes the change in the weights during training. In their
setting, the purpose of the decomposition was to demonstrate that in a set of
simple tasks, the Frobenius norm of the weights at convergence
$||{\bm{W}}_{rec}||_{F}$ is dominated by the norm of the weights at
initialization $||{\bm{W}}_{0}||_{F}$. In spite of the recurrent connectivity
they used being full-rank, they found that the changes in the weights learned
during training were in fact low-rank as measured by $||\Delta{\bm{W}}||_{F}$.
In our analysis, we find the opposite: $||{\bm{W}}_{0}||_{F}$ certainly does
not dominate the norm of the final weights and is in fact lower than
$||\Delta{\bm{W}}||_{F}$ (Figure 7). This reinforces the notion of task
dimension put forth by Schuessler et al. (2020) which describes the rank of
the training-induced connectivity changes as a function of the task the
network is trained on. In particular, in their work, they showed that the task
dimension of the simple tasks they examined was low and hence a network with
unconstrained, full-rank connectivity learned low-rank changes. In contrast,
in our setting we consider a significantly more complex task domain which
incites higher-rank changes in the recurrent connectivity. This brings forth
the notion of a task dimension gap between the offline, open-loop and online,
closed-loop settings: namely, the networks we examined are trained offline
without being exposed to distribution shifts and hence learn higher-rank
changes in their connectivity. In contrast, as we have shown, succeeding in
the closed-loop setting under distribution shift means learning lower-rank
dynamics. Thus, networks that are able to abide to our low-rank prior and
avoid learning high-rank changes in connectivity are better at generalizing
under distribution shift. This is precisely where CfCs supersede LSTMs and to
some extent RNNs as well. We find that despite each network starting at the
same $||{\bm{W}}_{0}||_{F}$, $||\Delta{\bm{W}}||_{F}$ is lowest in CfCs.

At a higher level, this notion of a task dimension gap which characterizes the
disconnect between agents learning passively in the offline setting and
learning actively in the online setting warrants further exploration and
formalization. Furthermore, the notion of task dimension put forth by
Schuessler et al. (2020) itself is more aptly decomposed into recurrent task
dimension and input task dimension. In particular, as we have shown, changes
in recurrent connectivity are not synonymous with changes in input
connectivity. A given task may necessitate learning connectivities of
different ranks along the input axis and recurrent axis: disentangling the two
is a necessity in understanding the dynamics of recurrent neural networks.

##  5 Conclusion

In this work, we investigated the use of a low-rank, sparse parameterization
of recurrent connectivity in various architectures as a means of improving
model robustness in closed-loop environments. We showed that this type of
connectivity was most amenable to CfCs and also showed promise in more
canonical networks like LSTMs and RNNs. Furthermore, we demonstrated the
interpretability of this prior by analyzing the network dynamics it induces as
a function of both rank and sparsity. Our results demonstrate an application
in pruning recurrent networks at initialization to improve performance under
distribution shifts.

## Reproducibility statement

To ensure the reproducibility of our work, we extensively detail the
experimental setup in the appendix as well as provide information regarding
the analyses we conducted. The bulk of these details can be found in appendix
A.5 which describes how we constructed the dataset, how the models were
trained (including details on hyperparameters), the evaluation metrics for the
models and the intuition/implementation regarding the analyses that we
performed on the trained models. A key portion of our results is driven by an
initialization scheme we proposed that deviates from default initializers
given in existing open-source implementations. We thoroughly describe how and
why our proposed initialization differs in appendix A.4, so the reader can
leverage it to reproduce our results. Regarding our theoretical work, we
provide proofs for the claims we made in the main portion of the paper which
can be found in appendix A.2. In that section, we clearly delineate the cases
in which we were unable to prove certain claims and had to resort to an
empirical analysis instead.

## References

  * Abadi et al. (2016) Martin Abadi, Paul Barham, Jianmin Chen, Zhifeng Chen, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Geoffrey Irving, Michael Isard, Manjunath Kudlur, Josh Levenberg, Rajat Monga, Sherry Moore, Derek G. Murray, Benoit Steiner, Paul Tucker, Vijay Vasudevan, Pete Warden, Martin Wicke, Yuan Yu, and Xiaoqiang Zheng.  Tensorflow: A system for large-scale machine learning.  In _12th USENIX Symposium on Operating Systems Design and Implementation (OSDI 16)_ , pp. 265–283, 2016.  URL https://www.usenix.org/system/files/conference/osdi16/osdi16-abadi.pdf.
  * Bellemare et al. (2012) Marc G. Bellemare, Yavar Naddaf, Joel Veness, and Michael Bowling.  The arcade learning environment: An evaluation platform for general agents.  _CoRR_ , abs/1207.4708, 2012.  URL http://arxiv.org/abs/1207.4708.
  * Brockman et al. (2016) Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba.  Openai gym.  _CoRR_ , abs/1606.01540, 2016.  URL http://arxiv.org/abs/1606.01540.
  * Chahine et al. (2023) Makram Chahine, Ramin Hasani, Patrick Kao, Aaron Ray, Ryan Shubert, Mathias Lechner, Alexander Amini, and Daniela Rus.  Robust flight navigation out of distribution with liquid neural networks.  _Science Robotics_ , 8(77):eadc8892, 2023.  doi: 10.1126/scirobotics.adc8892.  URL https://www.science.org/doi/abs/10.1126/scirobotics.adc8892.
  * Chang et al. (2019) Bo Chang, Minmin Chen, Eldad Haber, and Ed H. Chi.  Antisymmetricrnn: A dynamical system view on recurrent neural networks, 2019.
  * Chen et al. (2018) Tian Qi Chen, Yulia Rubanova, Jesse Bettencourt, and David Duvenaud.  Neural ordinary differential equations.  In Samy Bengio, Hanna M. Wallach, Hugo Larochelle, Kristen Grauman, Nicolò Cesa-Bianchi, and Roman Garnett (eds.), _NeurIPS_ , pp. 6572–6583, 2018.  URL http://dblp.uni-trier.de/db/conf/nips/nips2018.html#ChenRBD18.
  * Chen et al. (2022) Tianlong Chen, Zhenyu Zhang, Pengjun Wang, Santosh Balachandra, Haoyu Ma, Zehao Wang, and Zhangyang Wang.  Sparsity winning twice: Better robust generalization from more efficient training, 2022.
  * Cho et al. (2014) Kyunghyun Cho, Bart van Merrienboer, Çaglar Gülçehre, Fethi Bougares, Holger Schwenk, and Yoshua Bengio.  Learning phrase representations using RNN encoder-decoder for statistical machine translation.  _CoRR_ , abs/1406.1078, 2014.  URL http://arxiv.org/abs/1406.1078.
  * Codevilla et al. (2017) Felipe Codevilla, Matthias Müller, Alexey Dosovitskiy, Antonio M. López, and Vladlen Koltun.  End-to-end driving via conditional imitation learning.  _CoRR_ , abs/1710.02410, 2017.  URL http://arxiv.org/abs/1710.02410.
  * de Haan et al. (2019) Pim de Haan, Dinesh Jayaraman, and Sergey Levine.  Causal confusion in imitation learning.  _CoRR_ , abs/1905.11979, 2019.  URL http://arxiv.org/abs/1905.11979.
  * Dubreuil (2022) Valente A. Beiran M. et al. Dubreuil, A.  The role of population structure in computations through neural dynamics, 2022.
  * Erichson et al. (2021) N. Benjamin Erichson, Omri Azencot, Alejandro Queiruga, Liam Hodgkinson, and Michael W. Mahoney.  Lipschitz recurrent neural networks.  In _International Conference on Learning Representations_ , 2021.  URL https://openreview.net/forum?id=-N7PBXqOUJZ.
  * Farahani et al. (2020) Abolfazl Farahani, Sahar Voghoei, Khaled Rasheed, and Hamid R. Arabnia.  A brief review of domain adaptation, 2020.
  * Finn et al. (2017) Chelsea Finn, Tianhe Yu, Tianhao Zhang, Pieter Abbeel, and Sergey Levine.  One-shot visual imitation learning via meta-learning.  _CoRR_ , abs/1709.04905, 2017.  URL http://arxiv.org/abs/1709.04905.
  * Ford et al. (2019) Nic Ford, Justin Gilmer, Nicholas Carlini, and Ekin Dogus Cubuk.  Adversarial examples are a natural consequence of test error in noise.  _CoRR_ , abs/1901.10513, 2019.  URL http://arxiv.org/abs/1901.10513.
  * Frankle & Carbin (2018) Jonathan Frankle and Michael Carbin.  The lottery ticket hypothesis: Training pruned neural networks.  _CoRR_ , abs/1803.03635, 2018.  URL http://arxiv.org/abs/1803.03635.
  * Goodfellow et al. (2016) Ian Goodfellow, Yoshua Bengio, Aaron Courville, and Yoshua Bengio.  _Deep learning_ , volume 1.  MIT Press, 2016.
  * Guo et al. (2018) Yiwen Guo, Chao Zhang, Changshui Zhang, and Yurong Chen.  Sparse dnns with improved adversarial robustness.  _CoRR_ , abs/1810.09619, 2018.  URL http://arxiv.org/abs/1810.09619.
  * Hasani et al. (2020a) Ramin Hasani, Mathias Lechner, Alexander Amini, Daniela Rus, and Radu Grosu.  Liquid time-constant networks, 2020a.
  * Hasani et al. (2020b) Ramin Hasani, Mathias Lechner, Alexander Amini, Daniela Rus, and Radu Grosu.  A natural lottery ticket winner: Reinforcement learning with ordinary neural circuits.  In Hal Daumé III and Aarti Singh (eds.), _Proceedings of the 37th International Conference on Machine Learning_ , volume 119 of _Proceedings of Machine Learning Research_ , pp. 4082–4093. PMLR, 13–18 Jul 2020b.  URL https://proceedings.mlr.press/v119/hasani20a.html.
  * Hasani et al. (2021) Ramin M. Hasani, Mathias Lechner, Alexander Amini, Lucas Liebenwein, Max Tschaikowski, Gerald Teschl, and Daniela Rus.  Closed-form continuous-depth models.  _CoRR_ , abs/2106.13898, 2021.  URL https://arxiv.org/abs/2106.13898.
  * Helfrich et al. (2018) Kyle Helfrich, Devin Willmott, and Qiang Ye.  Orthogonal recurrent neural networks with scaled cayley transform, 2018.
  * Herbert & Ostojic (2022) Elizabeth Herbert and Srdjan Ostojic.  The impact of sparsity in low-rank recurrent neural networks.  _bioRxiv_ , 2022.  doi: 10.1101/2022.03.31.486515.  URL https://www.biorxiv.org/content/early/2022/04/01/2022.03.31.486515.
  * Hochreiter & Schmidhuber (1997) Sepp Hochreiter and Jürgen Schmidhuber.  Long short-term memory.  _Neural Computation_ , 9(8):1735–1780, 1997.
  * Horgan et al. (2018a) Dan Horgan, John Quan, David Budden, Gabriel Barth-Maron, Matteo Hessel, Hado van Hasselt, and David Silver.  Distributed prioritized experience replay.  _CoRR_ , abs/1803.00933, 2018a.  URL http://arxiv.org/abs/1803.00933.
  * Horgan et al. (2018b) Dan Horgan, John Quan, David Budden, Gabriel Barth-Maron, Matteo Hessel, Hado van Hasselt, and David Silver.  Distributed prioritized experience replay.  _CoRR_ , abs/1803.00933, 2018b.  URL http://arxiv.org/abs/1803.00933.
  * Huster et al. (2018) Todd Huster, Cho-Yu Jason Chiang, and Ritu Chadha.  Limitations of the lipschitz constant as a defense against adversarial examples.  _CoRR_ , abs/1807.09705, 2018.  URL http://arxiv.org/abs/1807.09705.
  * James et al. (2018) Stephen James, Michael Bloesch, and Andrew J. Davison.  Task-embedded control networks for few-shot imitation learning.  _CoRR_ , abs/1810.03237, 2018.  URL http://arxiv.org/abs/1810.03237.
  * Jozefowicz et al. (2015) Rafal Jozefowicz, Wojciech Zaremba, and Ilya Sutskever.  An empirical exploration of recurrent network architectures.  In Francis Bach and David Blei (eds.), _Proceedings of the 32nd International Conference on Machine Learning_ , volume 37 of _Proceedings of Machine Learning Research_ , pp. 2342–2350, Lille, France, 07–09 Jul 2015. PMLR.  URL https://proceedings.mlr.press/v37/jozefowicz15.html.
  * Lea et al. (2016) Colin Lea, Michael D. Flynn, René Vidal, Austin Reiter, and Gregory D. Hager.  Temporal convolutional networks for action segmentation and detection.  _CoRR_ , abs/1611.05267, 2016.  URL http://arxiv.org/abs/1611.05267.
  * Lechner et al. (2020) Mathias Lechner, Ramin Hasani, and Alexander Amini.  Neural circuit policies enabling auditable autonomy, 2020.
  * Lechner et al. (2022) Mathias Lechner, Ramin Hasani, Alexander Amini, Tsun-Hsuan Wang, Thomas A. Henzinger, and Daniela Rus.  Are all vision models created equal? a study of the open-loop to closed-loop causality gap, 2022.
  * Lee et al. (2018) Namhoon Lee, Thalaiyasingam Ajanthan, and Philip H. S. Torr.  SNIP: single-shot network pruning based on connection sensitivity.  _CoRR_ , abs/1810.02340, 2018.  URL http://arxiv.org/abs/1810.02340.
  * Lee et al. (2019) Namhoon Lee, Thalaiyasingam Ajanthan, Stephen Gould, and Philip H. S. Torr.  A signal propagation perspective for pruning neural networks at initialization.  _CoRR_ , abs/1906.06307, 2019.  URL http://arxiv.org/abs/1906.06307.
  * Liang et al. (2017) Eric Liang, Richard Liaw, Robert Nishihara, Philipp Moritz, Roy Fox, Joseph Gonzalez, Ken Goldberg, and Ion Stoica.  Ray rllib: A composable and scalable reinforcement learning library.  _CoRR_ , abs/1712.09381, 2017.  URL http://arxiv.org/abs/1712.09381.
  * Mastrogiuseppe & Ostojic (2018) Francesca Mastrogiuseppe and Srdjan Ostojic.  Linking connectivity, dynamics, and computations in low-rank recurrent neural networks.  _Neuron_ , 99(3):609–623.e29, 2018.  ISSN 0896-6273.  doi: https://doi.org/10.1016/j.neuron.2018.07.003.  URL https://www.sciencedirect.com/science/article/pii/S0896627318305439.
  * Pascanu et al. (2012) Razvan Pascanu, Tomás Mikolov, and Yoshua Bengio.  Understanding the exploding gradient problem.  _CoRR_ , abs/1211.5063, 2012.  URL http://arxiv.org/abs/1211.5063.
  * Perez & Wang (2017) Luis Perez and Jason Wang.  The effectiveness of data augmentation in image classification using deep learning.  _CoRR_ , abs/1712.04621, 2017.  URL http://arxiv.org/abs/1712.04621.
  * Poli et al. (2023) Michael Poli, Stefano Massaroli, Eric Nguyen, Daniel Y. Fu, Tri Dao, Stephen Baccus, Yoshua Bengio, Stefano Ermon, and Christopher Ré.  Hyena hierarchy: Towards larger convolutional language models, 2023.
  * Ramanujan et al. (2019) Vivek Ramanujan, Mitchell Wortsman, Aniruddha Kembhavi, Ali Farhadi, and Mohammad Rastegari.  What’s hidden in a randomly weighted neural network?  _CoRR_ , abs/1911.13299, 2019.  URL http://arxiv.org/abs/1911.13299.
  * Ross et al. (2010) Stéphane Ross, Geoffrey J. Gordon, and J. Andrew Bagnell.  No-regret reductions for imitation learning and structured prediction.  _CoRR_ , abs/1011.0686, 2010.  URL http://arxiv.org/abs/1011.0686.
  * Sainath et al. (2013) Tara N. Sainath, Brian Kingsbury, Vikas Sindhwani, Ebru Arisoy, and Bhuvana Ramabhadran.  Low-rank matrix factorization for deep neural network training with high-dimensional output targets.  In _2013 IEEE International Conference on Acoustics, Speech and Signal Processing_ , pp. 6655–6659, 2013.  doi: 10.1109/ICASSP.2013.6638949.
  * Scaman & Virmaux (2019) Kevin Scaman and Aladin Virmaux.  Lipschitz regularity of deep neural networks: analysis and efficient estimation, 2019.
  * Schuessler et al. (2020) Friedrich Schuessler, Francesca Mastrogiuseppe, Alexis Dubreuil, Srdjan Ostojic, and Omri Barak.  The interplay between randomness and structure during learning in rnns.  In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), _Advances in Neural Information Processing Systems_ , volume 33, pp. 13352–13362. Curran Associates, Inc., 2020.  URL https://proceedings.neurips.cc/paper_files/paper/2020/file/9ac1382fd8fc4b631594aa135d16ad75-Paper.pdf.
  * Schulman et al. (2017) John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov.  Proximal policy optimization algorithms.  _CoRR_ , abs/1707.06347, 2017.  URL http://arxiv.org/abs/1707.06347.
  * Sporns (2013) Olaf Sporns.  Structure and function of complex brain networks.  _Dialogues in Clinical Neuroscience_ , 15(3):247–262, 2013.  doi: 10.31887/DCNS.2013.15.3/osporns.  URL https://doi.org/10.31887/DCNS.2013.15.3/osporns.  PMID: 24174898.
  * Srivastava et al. (2019) Rupesh Kumar Srivastava, Pranav Shyam, Filipe Mutz, Wojciech Jaskowski, and Jürgen Schmidhuber.  Training agents using upside-down reinforcement learning.  _CoRR_ , abs/1912.02877, 2019.  URL http://arxiv.org/abs/1912.02877.
  * Taori et al. (2020) Rohan Taori, Achal Dave, Vaishaal Shankar, Nicholas Carlini, Benjamin Recht, and Ludwig Schmidt.  Measuring robustness to natural distribution shifts in image classification.  In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), _Advances in Neural Information Processing Systems_ , volume 33, pp. 18583–18599. Curran Associates, Inc., 2020.  URL https://proceedings.neurips.cc/paper_files/paper/2020/file/d8330f857a17c53d217014ee776bfd50-Paper.pdf.
  * Todorov et al. (2012) Emanuel Todorov, Tom Erez, and Yuval Tassa.  Mujoco: A physics engine for model-based control.  In _2012 IEEE/RSJ International Conference on Intelligent Robots and Systems_ , pp. 5026–5033, 2012.  doi: 10.1109/IROS.2012.6386109.
  * Vorbach et al. (2021) Charles J Vorbach, Ramin Hasani, Alexander Amini, Mathias Lechner, and Daniela Rus.  Causal navigation by continuous-time neural networks.  In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan (eds.), _Advances in Neural Information Processing Systems_ , 2021.  URL https://openreview.net/forum?id=ckVbQs5zD7_.
  * Wang et al. (2021) Huan Wang, Can Qin, Yulun Zhang, and Yun Fu.  Emerging paradigms of neural network pruning.  _CoRR_ , abs/2103.06460, 2021.  URL https://arxiv.org/abs/2103.06460.
  * Xu et al. (2020) Yuhui Xu, Yuxi Li, Shuai Zhang, Wei Wen, Botao Wang, Yingyong Qi, Yiran Chen, Weiyao Lin, and Hongkai Xiong.  TRP: trained rank pruning for efficient deep neural networks.  _CoRR_ , abs/2004.14566, 2020.  URL https://arxiv.org/abs/2004.14566.
  * Zhuang et al. (2019) Fuzhen Zhuang, Zhiyuan Qi, Keyu Duan, Dongbo Xi, Yongchun Zhu, Hengshu Zhu, Hui Xiong, and Qing He.  A comprehensive survey on transfer learning.  _CoRR_ , abs/1911.02685, 2019.  URL http://arxiv.org/abs/1911.02685.
  * Życzkowski & Sommers (1999) Karol Życzkowski and H. J. Sommers.  Truncations of random unitary matrices.  _Journal of Physics A_ , 33:2045–2057, 1999.  URL https://api.semanticscholar.org/CorpusID:27605911.

##  Appendix A Appendix

###  A.1 Time constant analysis

Figure 8: Average absolute deviation from $0.5$ of time constant vectors in
trained CfC models.

In Section 4.4, we demonstrated that on average, CfCs learn lower spectral
norms in their recurrent weights than RNNs. Since CfCs and RNNs are distinct
only with respect to the time constant network $F$, here we provide an
analysis on what role modulating recurrent sparsity and rank has on the
learned interpolation weights. In particular, let ${\bm{w}}_{t}$ denote the
interpolation weights that are the output of $F$ at time $t$. We compute

$|{\bm{w}}_{t}-0.5|=\frac{\sum_{t=1}^{T}\sum_{i=1}^{h}|w_{ti}-0.5|}{hT}$

which measures the average deviation of the weights from $0.5$. We consider
this metric as a means of understanding how much the interpolation between $G$
and $H$ tends to rely upon only one of the trajectories (as we observe in the
case of an RNN) versus evenly combining them (as would be the case if $F$
learns to output $[\frac{1}{2},\frac{1}{2},\dots,\frac{1}{2}]$). We find that
as a function of increasing rank, the deviation of the time constant from
$0.5$ decreases (Figure 8). This means that low-rank CfCs tend to place more
of their weight on one trajectory than the other which demonstrates the
ability of the time constant network to align with the simpler recurrent
dynamics promoted in the low-rank setting.

###  A.2 Properties of low-rank and sparse connectivity at initialization

In this section we provide proofs for the spectral radius and spectral norm,
$\rho({\bm{W}}),||{\bm{W}}||$, respectively, for connectivity matrices
${\bm{W}}$ at initialization. For clarity, we iterate that
$\rho({\bm{W}})=\max_{i}|\lambda_{i}|$, i.e. the largest norm of eigenvalues
of ${\bm{W}}$, and $||{\bm{W}}||=\max_{i}\sigma_{i}$, the largest singular
value of ${\bm{W}}$. We consider sparse networks with Glorot uniform
initialization and orthogonal initialization in sections A.2.1 and A.2.2,
respectively, and orthogonal low-rank matrices in section A.2.3. Note that
practically, in our experiments, we consider networks with orthogonally
initialized recurrent weights which we motivate further in A.4. Extending
these results at initialization to fully trained networks at convergence is an
interesting direction for future work.

####  A.2.1 Glorot uniform and sparse

Figure 9: Spectral radius and spectral norm of uniform-sparse matrices as a
function of sparsity, $s$, for matrices initialized with Glorot uniform
initialization, with $n=512$. We additionally plot the theoretical predicted
value for large $n$.

We consider weight matrices ${\bm{W}}$ with dimension $n\times n$ and consider
the limit as $n\to\infty$. Recall that Glorot uniform initialization scheme
initializes weights with entries iid
${\bm{W}}_{ij}\sim\texttt{Unif}\left(-\frac{\sqrt{3}}{\sqrt{n}},+\frac{\sqrt{3}}{\sqrt{n}}\right)$,
resulting in an entry-wise variance of $\frac{1}{n}$. We generate a sparse
matrix by element-wise multiplying a matrix, ${\bm{W}}^{0}$ sampled from the
Glorot uniform initialization by a sparsity map ${\bm{M}}$, where
${\bm{M}}_{ij}\sim\texttt{Bernoulli}(p=1-s)$, with $s$ being the sparse
factor. Our final weight matrix is given by
${\bm{W}}={\bm{W}}^{0}\odot{\bm{M}}$. For sparsity $s$, the resulting variance
of entries is given by $\frac{1-s}{n}$. As the entries of ${\bm{W}}$ are all
iid, we can apply the Girko-Ginibri circular law for large $n$, which states
that the eigenvalues of ${\bm{W}}$ converge to a uniform disk in the complex
plane with radius $\sqrt{1-s}$, so we expect $\rho(s)=\sqrt{1-s}$. This
analysis follows closely with that of Herbert & Ostojic (2022).

For $||{\bm{W}}||$, we apply the Marchenko-Pastur law, which states that the
distribution of eigenvalues of $\frac{1}{n}{\bm{A}}{\bm{A}}^{T}$ values
converges to
$p_{{\bm{A}}{\bm{A}}^{T}}(\lambda)=\frac{1}{2\pi\sigma^{2}}\frac{\sqrt{(\lambda_{+}-\lambda)(\lambda-\lambda_{-})}}{\lambda}\mathbf{1}_{\lambda\in[\lambda_{-},\lambda_{+}]}$,
with $\lambda_{\pm}=\sigma^{2}(1\pm 1)^{2}$, if ${\bm{A}}$ has entries iid
from a distribution with zero mean and variance $\sigma^{2}$. For our setting
we let ${\bm{A}}=\sqrt{n}{\bm{W}}$, so $\sigma^{2}=1-s$. We see that the
maximal eigenvalue of ${\bm{A}}{\bm{A}}^{T}$ is thus $\lambda_{+}=4(1-s)$, and
so the upper bound for the largest singular value of ${\bm{W}}$ is
$||{\bm{W}}||=2\sqrt{(1-s)}$. We observe good empirical agreement with these
values in Figure 9.

####  A.2.2 Orthogonal and sparse

Figure 10: Spectral radius and spectral norm of orthogonal-sparse matrices as
a function of sparsity, $s$, with $n=512$. We additionally plot the
theoretical predicted value for large $n$ for the uniform-sparse
initialization. We see that orthogonal-sparse matrices behave like uniform-
sparse ones at high sparsities.

In the orthogonal-sparse case, we cannot apply the same arguments in section
A.2.1, as the entries of an orthogonal matrix are no longer iid. However, we
note that the entries of orthogonal matrices are approximately Gaussian when
considered individually (Życzkowski & Sommers, 1999). Thus, we expect with
high sparsity, the correlations between entries break down and the entries of
the matrix behave iid. Thus, we expect for large values of $s$,
$\rho({\bm{W}})$ and $||{\bm{W}}||$ to have the same behavior as in section
A.2.1. We verify this empirically in Figure 10, where we see that for large
values of $s$, $\rho({\bm{W}})\approx\sqrt{1-s}$ and $||{\bm{W}}||\approx
2\sqrt{1-s}$, as with section A.2.1. Interestingly, $||{\bm{W}}||$ is non-
monotonic in $s$. Calculating the exact forms of $\rho({\bm{W}})$ and
$||{\bm{W}}||$ is an interesting direction for future work.

####  A.2.3 Orthogonal and low rank

Figure 11: Spectral radius and spectral norm of orthogonal-low-rank matrices
as a function of rank ratio, $r$, with $n=512$. We additionally plot the
theoretical predicted value for large $n$.

Recall that to generate our low rank connectivity matrix ${\bm{W}}$, we first
generate an orthogonal matrix ${\bm{W}}^{0}$, then perform an SVD of it
${\bm{W}}={\bm{U}}{\bm{S}}{\bm{V}}^{T}=\sum_{i}^{n}\sum_{i}^{n}{\bm{u}}_{i}{\bm{v}}_{i}^{T}$,
and then truncate it at rank k to get
${\bm{W}}=\sum_{i}^{k}{\bm{u}}_{i}{\bm{v}}_{i}^{T}$ with $k\leq n$. For
$||{\bm{W}}||$, we note that because all the singular values are 1, truncating
the SVD does not change the maximum singular value so $||{\bm{W}}||=1$ for all
rank ratios $r=\frac{k}{n}$.

For $\rho({\bm{W}})$, we note that one valid SVD of ${\bm{W}}^{0}$ is
${\bm{W}}^{0}={\bm{W}}^{0}{\bm{I}}{\bm{I}}$, with ${\bm{U}}={\bm{W}}^{0}$,
${\bm{S}}={\bm{I}}$ and ${\bm{V}}={\bm{I}}$. It is sufficient to consider this
particular SVD, as other SVDs correspond to arbirtrary rotations of ${\bm{U}}$
and ${\bm{V}}$, which will not affect the final eigenvalue distribution in
expectation. For this particular SVD,
${\bm{W}}=[{\bm{W}}^{0}_{1},{\bm{W}}^{0}_{2},\ldots,{\bm{W}}^{0}_{k},0,\ldots,0]$,
i.e. we take the first $k$ columns of ${\bm{W}}^{0}$, and replace the
remaining entries with 0. Due to the large zero block in this matrix,
${\bm{W}}$ has the same non-zero eigenvectors and eigenvalues as the $k\times
k$ principal submatrix of ${\bm{W}}^{0}$: ${\bm{W}}^{0}_{1:k,1:k}$111for
eigenvectors we need to concatenate the remaining $n-k$ zeros to have the
correct dimension. Życzkowski & Sommers (1999) studies the eigenvalues of the
principal submatrix of random orthogonal matrices and shows that the
distribution of these eigenvalues has mean $\mathbb{E}[\lambda]=\sqrt{r}$ for
rank ratio $r=\frac{k}{n}$. Furthermore, $\mathbf{Var}[\lambda]\to 0$ as
$n\to\infty$ at fixed $r$ (Życzkowski & Sommers, 1999), so we have
$\rho({\bm{W}})=\sqrt{r}$ for large $n$. Empirical verification of this is
provided in Figure 11 for $n=512$.

###  A.3 Full results

####  A.3.1 Online and offline performance

Here, we provide offline validation/test loss, online in-distribution rewards
and online rewards under distribution shift for each Gym environment (Brockman
et al., 2016) that we ran experiments on: Seaquest, Alien and HalfCheetah. For
each environment, we considered 5 types of models: CfCs, RNNs, LSTMs, GRUs and
TCNs. For each of the recurrent architectures we considered models of various
5 different ranks and 4 different sparsities. Full details on the experimental
setup are given in A.5.

Figure 12: Online and offline performance of recurrent networks under
different ranks and sparsities in the Seaquest environment. Note that online
performance of Seaquest was presented in the main portion of the paper, but
here we include it alongside the offline losses evaluated on validation/test
set. a) In-distribution rewards in the online, closed-loop setting. b) Rewards
averaged across 5 distribution shifts. c) Loss of model evaluated on a
validation/test set.

|  | Sparsity
---|---|---
Model |  Rank | 0 | 0.2 | 0.5 | 0.8
CfC | 1 |  $1860$ p m 70 |  $2000$ p m 90 |  $1900$ p m 70 |  $1720$ p m 110
| 5 |  $1760$ p m 120 |  $1860$ p m 100 |  $1710$ p m 130 |  $1530$ p m 130
| 16 |  $2190$ p m 120 |  $1980$ p m 120 |  $1700$ p m 110 |  $1940$ p m 100
| 27 |  $1960$ p m 130 |  $2100$ p m 80 |  $1900$ p m 140 |  $1800$ p m 130
| full |  $2100$ p m 80 |  $1640$ p m 170 |  $1520$ p m 110 |  $1650$ p m 120
LSTM | 1 |  $2140$ p m 90 |  $1790$ p m 80 |  $1950$ p m 140 |  $1980$ p m 100
| 5 |  $2030$ p m 70 |  $2030$ p m 100 |  $1890$ p m 130 |  $1910$ p m 110
| 16 |  $2030$ p m 90 |  $2040$ p m 150 |  $1940$ p m 130 |  $1940$ p m 150
| 27 |  $2150$ p m 70 |  $2140$ p m 100 |  $1940$ p m 120 |  $1900$ p m 130
| full |  $1940$ p m 110 |  $2120$ p m 70 |  $1940$ p m 130 |  $1940$ p m 130
RNN | 1 |  $2020$ p m 130 |  $1980$ p m 120 |  $1930$ p m 110 |  $1940$ p m 100
| 5 |  $2080$ p m 100 |  $1860$ p m 90 |  $2030$ p m 120 |  $1980$ p m 120
| 16 |  $2020$ p m 120 |  $1950$ p m 100 |  $1890$ p m 140 |  $2310$ p m 50
| 27 |  $1910$ p m 110 |  $1800$ p m 110 |  $2190$ p m 140 |  $1970$ p m 70
| full |  $1990$ p m 80 |  $1980$ p m 80 |  $1980$ p m 110 |  $1850$ p m 120
GRU | 1 |  $1760$ p m 140 |  $2240$ p m 60 |  $2160$ p m 100 |  $1940$ p m 90
| 5 |  $1830$ p m 80 |  $1820$ p m 110 |  $1830$ p m 120 |  $1960$ p m 90
| 16 |  $1980$ p m 120 |  $1980$ p m 100 |  $1690$ p m 100 |  $2030$ p m 100
| 27 |  $1710$ p m 110 |  $2050$ p m 120 |  $1700$ p m 130 |  $1810$ p m 100
| full |  $1770$ p m 90 |  $2060$ p m 120 |  $2000$ p m 130 |  $2030$ p m 110
TCN | 1 | — | — | — | —
| 5 | — | — | — | —
| 16 | — | — | — | —
| 27 | — | — | — | —
| full |  $1670$ p m 100 | — | — | —
Table 1: Mean episodic in-distribution rewards in Seaquest environment. Rewards are given $\pm 1$ SE which is shown in parentheses. |  | Sparsity
---|---|---
Model |  Rank | 0 | 0.2 | 0.5 | 0.8
CfC | 1 | 21.5(5.7) | 16.8(4.6) | 29.0(3.0) | 10.0(1.0)
| 5 | 4.1(1.2) | 4.0(1.4) | 37.3(2.5) | 66.5(2.6)
| 16 | 54.3(2.4) | 36.0(6.1) | 69.0(1.2) | 45.3(3.7)
| 27 | 22.3(3.8) | 49.5(4.7) | 50.5(7.0) | 53.8(2.7)
| full | 5.5(0.8) | 61.8(4.2) | 41.5(4.1) | 10.0(1.0)
LSTM | 1 | 33.0(4.5) | 50.0(4.1) | 43.8(5.2) | 35.5(1.9)
| 5 | 52.8(5.5) | 34.8(6.3) | 34.8(5.7) | 35.8(1.1)
| 16 | 38.0(4.8) | 61.0(6.0) | 50.5(4.9) | 50.5(5.7)
| 27 | 32.5(3.6) | 37.0(5.8) | 46.8(1.9) | 51.5(5.7)
| full | 49.8(5.9) | 45.5(4.2) | 56.7(4.1) | 34.5(3.1)
RNN | 1 | 19.8(3.5) | 34.0(6.5) | 38.8(4.7) | 40.0(6.6)
| 5 | 16.5(4.3) | 26.8(6.3) | 20.5(5.5) | 42.0(2.7)
| 16 | 38.5(3.1) | 52.0(1.7) | 54.8(5.7) | 39.3(5.2)
| 27 | 41.0(5.2) | 43.4(2.8) | 26.8(3.9) | 26.5(3.8)
| full | 33.8(6.9) | 29.8(2.9) | 32.3(3.3) | 58.3(2.1)
GRU | 1 | 42.3(2.8) | 40.3(3.5) | 41.3(5.8) | 38.3(4.7)
| 5 | 41.3(3.9) | 38.3(5.7) | 62.8(4.6) | 58.5(3.9)
| 16 | 29.8(4.6) | 38.5(4.2) | 52.5(5.6) | 48.3(5.2)
| 27 | 39.3(5.2) | 47.8(3.9) | 52.8(5.6) | 35.8(4.7)
| full | 72.0(1.6) | 49.0(6.2) | 55.8(3.2) | 13.5(0.9)
TCN | 1 | — | — | — | —
| 5 | — | — | — | —
| 16 | — | — | — | —
| 27 | — | — | — | —
| full | 71.9(3.2) | — | — | —
Table 2: Mean episodic reward rankings under distribution shift in Seaquest
environment. Reward rankings are given $\pm 1$ SE which is shown in
parentheses.

Figure 13: Online and offline performance of recurrent networks under
different ranks and sparsities in the Alien environment. a) In-distribution
rewards in the online, closed-loop setting. b) Rewards averaged across 5
distribution shifts. c) Loss of model evaluated on a validation/test set.

|  | Sparsity
---|---|---
Model |  Rank | 0 | 0.2 | 0.5 | 0.8
CfC | 1 |  $1449$ p m 169 |  $1996$ p m 346 |  $2332$ p m 268 |  $2045$ p m 460
| 5 |  $1274$ p m 256 |  $2054$ p m 346 |  $1627$ p m 330 |  $2402$ p m 347
| 16 |  $2309$ p m 339 |  $1099$ p m 132 |  $1041$ p m 120 |  $1148$ p m 108
| 27 |  $2686$ p m 429 |  $2310$ p m 470 |  $2463$ p m 413 |  $1425$ p m 198
| full |  $2930$ p m 525 |  $1564$ p m 285 |  $1163$ p m 179 |  $1717$ p m 262
LSTM | 1 |  $2249$ p m 283 |  $1951$ p m 287 |  $2257$ p m 304 |  $1992$ p m 301
| 5 |  $2212$ p m 290 |  $2277$ p m 369 |  $1967$ p m 364 |  $2684$ p m 437
| 16 |  $2138$ p m 328 |  $2212$ p m 292 |  $2665$ p m 312 |  $2056$ p m 404
| 27 |  $2722$ p m 340 |  $2736$ p m 356 |  $1955$ p m 336 |  $2257$ p m 293
| full |  $2461$ p m 338 |  $2334$ p m 344 |  $28$ p m 287 |  $2148$ p m 418
RNN | 1 |  $1897$ p m 287 |  $1788$ p m 280 |  $1968$ p m 243 |  $1950$ p m 271
| 5 |  $2345$ p m 265 |  $1566$ p m 280 |  $2004$ p m 255 |  $1900$ p m 341
| 16 |  $1328$ p m 254 |  $1864$ p m 271 |  $2368$ p m 151 |  $1855$ p m 260
| 27 |  $2262$ p m 263 |  $1938$ p m 288 |  $1830$ p m 247 |  $1746$ p m 269
| full |  $2369$ p m 225 |  $1741$ p m 323 |  $2371$ p m 278 |  $1386$ p m 251
GRU | 1 |  $1994$ p m 263 |  $1566$ p m 277 |  $1750$ p m 280 |  $1265$ p m 259
| 5 |  $1768$ p m 222 |  $2140$ p m 300 |  $2035$ p m 297 |  $1073$ p m 157
| 16 |  $1949$ p m 257 |  $2107$ p m 268 |  $1774$ p m 237 |  $1574$ p m 301
| 27 |  $1325$ p m 254 |  $1636$ p m 260 |  $1348$ p m 251 |  $1936$ p m 277
| full |  $1473$ p m 251 |  $1330$ p m 302 |  $1884$ p m 332 |  $2177$ p m 260
TCN | 1 | — | — | — | —
| 5 | — | — | — | —
| 16 | — | — | — | —
| 27 | — | — | — | —
| full |  $1580$ p m 230 | — | — | —
Table 3: Mean episodic in-distribution rewards in Alien environment. Rewards
are given $\pm 1$ SE which is shown in parentheses.

|  | Sparsity
---|---|---
Model |  Rank | 0 | 0.2 | 0.5 | 0.8
CfC | 1 | 25.5(3.6) | 52.8(3.0) | 43.8(5.3) | 36.0(5.8)
| 5 | 9.0(2.6) | 26.8(5.1) | 34.2(5.4) | 52.8(4.1)
| 16 | 24.0(4.4) | 9.2(2.6) | 68.0(2.7) | 54.2(1.6)
| 27 | 56.0(3.6) | 35.8(5.9) | 36.0(4.4) | 31.0(5.7)
| full | 58.2(3.3) | 56.2(4.9) | 44.8(2.3) | 39.5(4.1)
LSTM | 1 | 43.8(3.3) | 55.2(3.7) | 22.2(3.5) | 46.0(4.7)
| 5 | 26.2(5.3) | 76.5(2.2) | 32.0(6.1) | 37.0(5.1)
| 16 | 34.0(6.4) | 75.8(3.3) | 52.5(1.4) | 28.0(5.4)
| 27 | 36.2(5.8) | 51.8(3.3) | 49.2(3.9) | 48.2(3.7)
| full | 46.8(5.6) | 72.0(0.8) | 32.0(3.8) | 45.8(5.8)
RNN | 1 | 74.5(1.2) | 14.5(2.1) | 26.5(4.5) | 33.5(2.1)
| 5 | 13.2(1.8) | 24.5(5.1) | 31.5(3.2) | 30.0(4.8)
| 16 | 9.5(1.9) | 17.2(3.3) | 73.8(4.5) | 36.5(4.5)
| 27 | 31.8(4.6) | 24.0(3.7) | 43.8(4.1) | 30.5(5.2)
| full | 44.2(3.6) | 74.4(5.1)) | 29.8(4.2) | 19.2(1.7)
GRU | 1 | 35.5(2.3) | 29.5(6.0) | 41.0(1.9) | 45.2(3.7)
| 5 | 43.0(4.6) | 64.2(1.8) | 25.2(2.5) | 43.0(5.2)
| 16 | 42.5(3.5) | 44.2(5.3) | 33.2(6.1) | 54.0(1.9)
| 27 | 25.2(3.8) | 45.2(4.5) | 37.2(5.5) | 35.5(4.9)
| full | 54.5(2.6) | 34.2(3.9) | 49.5(5.8) | 54.8(4.8)
TCN | 1 | — | — | — | —
| 5 | — | — | — | —
| 16 | — | — | — | —
| 27 | — | — | — | —
| full | 50.8(5.2) | — | — | —
Table 4: Mean episodic reward rankings under distribution shift in Alien
environment. Reward rankings are given $\pm 1$ SE which is shown in
parentheses. Figure 14: Online and offline performance of recurrent networks
under different ranks and sparsities in the HalfCheetah environment. a) In-
distribution rewards in the online, closed-loop setting. b) Rewards averaged
across 5 distribution shifts. c) Loss of model evaluated on a validation/test
set.

|  | Sparsity
---|---|---
Model |  Rank | 0 | 0.2 | 0.5 | 0.8
CfC | 1 |  $2976$ p m 15 |  $2927$ p m 10 |  $2911$ p m 12 |  $2897$ p m 10
| 5 |  $2964$ p m 11 |  $2916$ p m 8 |  $2868$ p m 6 |  $2947$ p m 6
| 16 |  $2936$ p m 7 |  $2965$ p m 11 |  $2894$ p m 7 |  $3032$ p m 16
| 27 |  $2950$ p m 14 |  $2754$ p m 12 |  $2880$ p m 11 |  $2783$ p m 14
| full |  $2941$ p m 9 |  $2570$ p m 20 |  $2764$ p m 7 |  $2669$ p m 9
LSTM | 1 |  $2939$ p m 9 |  $2961$ p m 10 |  $2920$ p m 9 |  $2948$ p m 8
| 5 |  $2994$ p m 9 |  $2913$ p m 13 |  $2950$ p m 7 |  $3029$ p m 10
| 16 |  $2952$ p m 9 |  $3123$ p m 9 |  $3049$ p m 16 |  $2852$ p m 10
| 27 |  $2934$ p m 10 |  $3040$ p m 13 |  $2976$ p m 10 |  $2746$ p m 9
| full |  $2997$ p m 9 |  $2756$ p m 10 |  $2800$ p m 25 |  $2652$ p m 27
RNN | 1 |  $2969$ p m 13 |  $2887$ p m 22 |  $2929$ p m 15 |  $2846$ p m 10
| 5 |  $2963$ p m 10 |  $2981$ p m 10 |  $2789$ p m 9 |  $2902$ p m 11
| 16 |  $2965$ p m 7 |  $2967$ p m 10 |  $3005$ p m 9 |  $2904$ p m 12
| 27 |  $2974$ p m 8 |  $2887$ p m 8 |  $2836$ p m 91 |  $2711$ p m 26
| full |  $2955$ p m 9 |  $2897$ p m 7 |  $3058$ p m 11 |  $2911$ p m 66
GRU | 1 |  $2950$ p m 7 |  $2954$ p m 14 |  $2936$ p m 11 |  $2951$ p m 5
| 5 |  $2973$ p m 7 |  $2945$ p m 10 |  $2955$ p m 9 |  $2920$ p m 8
| 16 |  $2929$ p m 11 |  $2943$ p m 9 |  $2920$ p m 11 |  $2969$ p m 8
| 27 |  $2961$ p m 7 |  $2969$ p m 10 |  $2917$ p m 36 |  $2941$ p m 4
| full |  $2951$ p m 12 |  $2949$ p m 8 |  $2988$ p m 8 |  $2946$ p m 7
TCN | 1 | — | — | — | —
| 5 | — | — | — | —
| 16 | — | — | — | —
| 27 | — | — | — | —
| full |  $2850$ p m 11 | — | — | —
Table 5: Mean episodic in-distribution rewards in HalfCheetah environment.
Rewards are given $\pm 1$ SE which is shown in parentheses.

|  | Sparsity
---|---|---
Model |  Rank | 0 | 0.2 | 0.5 | 0.8
CfC | 1 | 41.0(7.8) | 7.0(1.2) | 29.5(5.6) | 46.5(0.6)
| 5 | 9.5(2.1) | 29.5(6.6) | 54.0(0.2) | 71.5(0.9)
| 16 | 35.0(6.8) | 53.5(3.9) | 26.0(2.8) | 8.0(0.2)
| 27 | 56.0(5.0) | 55.0(4.8) | 75.5(0.6) | 38.5(8.4)
| full | 13.5(0.9) | 41.0(9.5) | 76.0(0.2) | 75.5(1.1)
LSTM | 1 | 45.5(4.4) | 17.5(0.6) | 41.5(0.4) | 25.0(3.5)
| 5 | 36.5(5.9) | 40.0(6.0) | 40.0(5.2) | 27.0(6.5)
| 16 | 29.0(0.8) | 34.5(7.6) | 61.5(1.4) | 5.0(0.2)
| 27 | 39.5(4.4) | 14.5(2.4) | 33.5(2.9) | 27.5(5.4)
| full | 18.5(0.6) | 56.5(2.1) | 11.0(0.8) | 46.0(1.2)
RNN | 1 | 50.5(4.9) | 61.5(2.4) | 59.0(3.5) | 44.5(4.4)
| 5 | 28.0(1.8) | 38.5(4.9) | 32.5(2.4) | 54.5(0.9)
| 16 | 26.5(0.9) | 32.0(5.8) | 37.0(5.8) | 36.0(5.8)
| 27 | 54.0(0.4) | 50.5(4.9) | 39.0(1.8) | 34.0(2.5)
| full | 29.0(4.5) | 51.0(6.2) | 29.5(1.1) | 33.0(1.0)
GRU | 1 | 44.0(1.0) | 54.5(3.1) | 73.0(1.8) | 58.0(2.8)
| 5 | 28.0(2.5) | 45.5(4.6) | 35.5(2.6) | 58.5(3.9)
| 16 | 48.0(3.2) | 39.0(6.8) | 54.0(6.2) | 60.0(0.8)
| 27 | 45.5(1.1) | 35.0(8.2) | 46.0(3.2) | 23.5(1.6)
| full | 52.5(6.1) | 41.5(2.1) | 29.5(2.4) | 56.0(5.5)
TCN | 1 | — | — | — | —
| 5 | — | — | — | —
| 16 | — | — | — | —
| 27 | — | — | — | —
| full | 56.1(2.1) | — | — | —
Table 6: Mean episodic reward rankings under distribution shift in HalfCheetah
environment. Reward rankings are given $\pm 1$ SE which is shown in
parentheses.

####  A.3.2 Model analyses for RNNs

In this section, we provide plots for analyses we conducted on trained RNNs
that were not included in the main portion of the paper.

Figure 15: Eigenspectra of trained RNNs across various ranks and sparsities.
Note that like CfCs, RNNs tend to express the vanishing gradient. Figure 16:
Effective dimensionality of RNN state-space trajectories across various ranks
and sparsities. Note that the dynamics of RNNs are lower-dimensional than the
dynamics we observe in CfCs. This, alongside the disparity in spectral norm,
offers intuition as to why we find that CfCs tend to outperform RNNs under
distribution shift.

###  A.4 Initialization of recurrent weights

We leverage an initialization scheme consistent across each of the recurrent
networks in order to control for potential confounding effects that could
arise due to differences present in the default initializers for the recurrent
weights. We will first discuss what the default initialization scheme looks
like for the recurrent networks we considered and then motivate our proposed
method of initialization.

Default initialization schemes. We define the default initialization schemes
of recurrent weights in RNNs, LSTMs and GRUs as those implemented by
TensorFlow (Abadi et al., 2016) and the default initialization of recurrent
weights in a CfC as given by the open-source implementation presented in
Hasani et al. (2021). The default initializers for these models differ in a
couple key areas. For one, RNNs, LSTMs and GRUs are orthogonally initialized
whereas CfCs are initialized from a Glorot uniform distribution.

Figure 17: Spectral norm of random square matrices drawn from orthogonal and
Glorot uniform distributions for various sized matrices.

In Figure 17, we investigate the spectral norm of random matrices drawn from
orthogonal and Glorot uniform distributions. Due to the significant disparity
between the spectral norms inherent to the distributions, we opt to
standardize the initialization scheme across all recurrent models by drawing
the weights from an orthogonal distribution. We motivate this further in the
description of our intializer.

Another concern is that RNNs, LSTMs and GRUs are distinct with respect to how
they are orthogonally initialized: in particular, consider $[{\bm{W}}_{0}]$,
$[{\bm{W}}_{0},{\bm{W}}_{1},{\bm{W}}_{2},{\bm{W}}_{3}]$,
$[{\bm{W}}_{0},{\bm{W}}_{1},{\bm{W}}_{2}]$ which represent the concatenated
recurrent weights in RNNs, LSTMs and GRUs respectively. Rather than initialize
each ${\bm{W}}_{i}$ to be orthogonal, TensorFlow draws the concatenated set of
matrices from an orthogonal distribution. The following result demonstrates
why this is problematic in the context of our work.

######  Lemma 1.

Consider matrices ${\bm{A}}\in\mathbb{R}^{n\times p}$ and
${\bm{B}}\in\mathbb{R}^{n\times q}$. Let ${\bm{C}}=[{\bm{A}}\ {\bm{B}}]$
denote the concatenation of the two matrices. Then,
$||{\bm{A}}||\leq||{\bm{C}}||$ and $||{\bm{B}}||\leq||{\bm{C}}||$ where
$||\cdot||$ denotes the spectral norm.

###### Proof.

We will first consider the case of matrix ${\bm{A}}$. By definition,
$||{\bm{A}}||=\sup_{{\bm{x}}\neq 0}\frac{||{\bm{A}}{\bm{x}}||}{||{\bm{x}}||}$.
If we constrain ${\bm{x}}$ to be of unit norm, then
$||{\bm{A}}||=\sup_{{\bm{x}}\neq 0}||{\bm{A}}{\bm{x}}||$. Let ${\bm{z}}$
denote the unit norm vector that maximizes $||{\bm{A}}{\bm{x}}||$. Then, we
can construct another vector ${\bm{m}}=[{\bm{z}}\ 0\ 0\cdots\ 0]$ which
represents the concatenation of ${\bm{z}}$ which contains $p$ entries and a
zero-vector which contains $q$ entries. Note that by construction ${\bm{m}}$
is also unit norm. It follows that
$||{\bm{A}}{\bm{z}}||\leq||{\bm{C}}{\bm{m}}||$ which means that the maximum
attainable value of $||{\bm{C}}{\bm{x}}||$ is greater than or equal to the
maximum attainable value of $||{\bm{A}}{\bm{x}}||$ over all unit vectors
${\bm{x}}$. Thus, $||{\bm{A}}||\leq||{\bm{C}}||$. An analogous argument can be
made to show that $||{\bm{B}}||\leq||{\bm{C}}||$. ∎

By Lemma 1, it follows that the default initialization in LSTMs and GRUs
produces submatrices ${\bm{W}}_{i}$ with spectral norm less than the spectral
norm of the default initialized recurrent weights in RNNs. We will address
this discrepancy in our proposed initialization scheme.

Orthogonal spectral-initialization scheme. Orthogonal spectral-initialization
refers to the initialization scheme we utilized for the models in all the
results we presented. First and foremost, this standardizes the distribution
of the recurrent weights to be orthogonal. Note that we also formalize an
analogous initializer which we call Glorot uniform spectral-initialization
which is detailed in the next subsection.

In practice, we use the orthogonal distribution as it leans on the intuition
of our connectivity prior in which we show empirically that networks
initialized at lower spectral norms tends to converge to solutions with lower
spectral norms as well. Because we care about the robustness of the network as
measured by spectral norm at convergence and found that orthogonal weights
have lower spectral norm than Glorot uniform ones (Figure 17), we opt to draw
from the orthogonal distribution.

With respect to the concatenation issue discussed above, we address this by
instead drawing each recurrent weight matrix from an orthogonal distribution
as opposed to drawing the full, concatenated set of matrices from one. This
ensures equivalent spectral norms across all recurrent weights and models.

Finally, the spectral-initialization refers specifically to the rank-r SVD
performed on the full-rank weights ${\bm{W}}_{rec}$ to generate
${\bm{W}}_{1},{\bm{W}}_{2}$ (Section 3.1) for training. We opted to perform a
spectral decomposition as opposed to drawing ${\bm{W}}_{1}$ and ${\bm{W}}_{2}$
from orthogonal distributions individually because ${\bm{W}}_{1}{\bm{W}}_{2}$
is no longer orthogonally distributed. In addition, note that multiplying the
two matrices does not preserve the spectral norm of ${\bm{W}}_{rec}$. In
contrast, by the Eckhart-Young-Minsky theorem, a rank-r SVD is the most
efficient approximation of ${\bm{W}}_{rec}$ under the spectral norm.

Note that the drawback of this initialization scheme is that we are unable to
prove all the results posed in Theorem 1 due to the dependence between entries
in an orthogonal matrix (appendix A.2.2). In contrast, we are able to prove
Theorem 1 in the case of uniformly distributed matrices, which we motivate
next.

Glorot uniform spectral-initialization scheme. The Glorot uniform spectral-
initializer refers to the same construction as the orthogonal spectral-
initializer, except recurrent weights are drawn from Glorot uniform
distributions instead. Note that this distribution is valuable if one wants
theoretical guarantees for all the properties given in Theorem 1. However,
this comes at the cost of empirically less robust models as evidenced by the
learned weights possessing higher spectral norm relative to their orthogonally
initialized counterparts (results not shown).

###  A.5 Details of experimental setup

In this section, we extensively detail each of the experiments and analyses
that were presented in the main portion of the paper.

####  A.5.1 Generating expert trajectories

Arcade learning environments. Within the set of ALEs, we considered the
Seaquest and Alien environments. These environments in particular were chosen
as we were able to train reinforcement learning agents that achieved
performance competitive to state-of-the-art using out-of-the-box models
provided by RLlib (Liang et al., 2017). For the ALEs, we train an Ape-X DQN
(Horgan et al., 2018b) model using the Atari pre-processing framework detailed
in Horgan et al. (2018b). We employ the default implementation of the Ape-X
DQN model given by RLlib which is a tuned version of the model presented in
Horgan et al. (2018b). After the model is trained to perform sufficiently well
in the environment, we use it to generate expert trajectories. Specifically,
we generate 100 rollouts of the trained model’s observations and actions taken
in the environment which were later used to fit the recurrent models we
examine in an imitation learning framework.

MuJoCo. Within the set of MuJoCo environments, we considered the HalfCheetah
environment. Again, this environment was chosen based on its amenability to
models provided by RLlib. For the MuJoCo environment, we employ an
implementation of proximal policy optimization (PPO) given by RLlib (Schulman
et al., 2017). We also leverage the default set of hyperparameters given by
RLlib to train the model. As we did with the ALEs, we generate 100 rollouts of
the trained model’s observations and actions taken in the environment which
were later used to fit the recurrent models using imitation learning.

####  A.5.2 Imitation learning framework

For each recurrent architecture, we construct a model to be fit offline on the
generated expert trajectories. We have layers in the network that act to
preprocess the observations before entering the recurrent portion of the
network as well as output layers from the recurrent portion in accordance to
the action space specified by the environments. Since the observation and
action spaces in ALEs and MuJuCo environments are distinct, we construct
different network architectures for each.

First, we consider the network architectures used for ALEs. In their raw form,
observations are given by a $210\times 160\times 3$ dimensional image, but
they are preprocessed using the Atari preprocessing specified by Horgan et al.
(2018b). This yields new observations of dimension $84\times 84\times 3$ which
are used as inputs to the model.

layer type | activation
---|---
Conv(64, k=5, s=2) | relu
Conv(128, k=5, s=2) | relu
Conv(128, k=5, s=2) | relu
AveragePooling | none
TimeDistributed | none
recurrentNet | none
FC(numActions) | softmax
Table 7: ALE network architecture

In the table above, $k$ denotes the kernel size and $s$ denotes the stride of
the convolution. recurrentNet refers to the recurrent network module that was
varied and is discussed in further detail below. numActions refers to the
number of actions the agent can take in the environment. Since the action
space is a discrete categorical variable that can take on numActions different
values, the network is trained using categorical cross-entropy loss. We detail
the model hyperparameters and corresponding grid search in A.5.3.

Next, we will consider the network architecture employed for MuJoCo networks.

layer type | activation
---|---
FC(256) | relu
FC(256) | relu
TimeDistributed | none
recurrentNet | none
FC(numActions) | tanh
Table 8: MuJoCo network architecture

Since the action space for MuJuCo environments in continuous, the output is a
continuous variable that can take on numActions different values and the
network is trained using MSE loss. Furthermore, since the action space is
bounded within the interval $[-1,1]$, a tanh activation function is applied to
the output layer.

####  A.5.3 Model hyperparameters

Below is a table with the model hyperparameters that were used in both the ALE
networks and MuJoCo networks. Parameters in square brackets represent a grid
search over which the best performing model was chosen. Note that only offline
performance on the validation/test set was used to determine the best
performing model in the grid search. Parameters in curly braces are dimensions
across which the network was individually evaluated for comparison.

parameter | value
---|---
optimizer | Adam
hidden size | 64
learning rate | $[5*10^{-5},1*10^{-4},5*10^{-4}]$
epochs | 150
rank | $\\{1,5,16,27,\textrm{full}\\}$
sparsity | $\\{0,0.2,0.5,0.8\\}$
Table 9: Model hyperparameters

The architecture of the convolutional head in the ALE networks was used on the
basis of work presented in Lechner et al. (2022) which conducting extensive
grid searching across convolutional architectures over many ALEs. Analogously,
the dense head used in the MuJoCo architecture was drawn from Hasani et al.
(2021).

In all the environments, we considered models of rank
$1,5,16,27,\textrm{full}$ and sparsities of $0,0.2,0.5,0.8$. Note that the
recurrent weights in the full rank networks are not rank decomposed, as they
are for the low-rank networks. These ranks and sparsities were chosen such
that for each rank we examined, we also looked at a sparsity level for which
the number of recurrent parameters in each parameterization is roughly
equivalent. For example, since the networks we constructed had a hidden
dimension of size $64$, a recurrent matrix with a sparsity level of $0.5$ has
the same number of parameters as one with a rank of $16$. We did not present
results of networks with a very high recurrent sparsity like $0.95$ (which
would roughly be the sparse analog to rank-1 networks) due to optimization
difficulties encountered, particularly in LSTMs and GRUs. In order to ensure
that the performance of a given network wasn’t due to a favorable random
sparse mask, we trained each model $3$ times from different random seeds and
averaged over the model performance.

Finally, we also employed an initialization scheme for the recurrent weights
distinct from the default initialization schemes given in open-source
implementations. For details on this, refer to A.4.

####  A.5.4 Evaluation metrics

Trained models were evaluated offline on a validation/test set with respect to
their cross-entropy loss in the case of ALE networks and mean-squared error in
the case of MuJoCo networks. However, in practice, we care about how the agent
performs when deployed in the environment closed-loop. The online rewards were
computed by performing 10 simulations of the agent in the environments and
averaging over the rewards from each episode. The in-distribution setting
refers to the environments in which the observations are not modified aside
from preprocessing performed by RLlib.

In the distribution shift setting, for the ALE observations we perform 4 types
of perturbations to the agent’s observations: cutout, brighten, darken and
noise. In cutout, we remove some pixels from the center of the input image. In
brighten, we shift all the pixel values by a constant amount in the positive
direction and do the same for darken but shift the pixels in the negative
direction. To noise the image, we draw random noise from a uniform
distribution and add it to the input image. For distribution shifts in the
MuJoCo environments, we perform 3 types of perturbations: noise, dropout and
offset. To noise the observations, we draw random noise from a normal
distribution and add it to the observation vector. To perform dropout, we mask
a subset of the dimensions in the observation vector. To perform offset, we,
with equal probability, shift the observation vector in the positive or
negative direction by an amount constant across all dimensions.

To evaluate the models under distribution shift, since each of the
perturbations have different effects on the models which puts the rewards on
different scales, we instead rank the performance of the models within a given
distribution shift and then average across distribution shifts. Since we
examine $81$ models for a given environment (20 $(r,s)$ pairs across 4
recurrent networks and 1 TCN), the reward ranking of a model lies in the
interval $[1,81]$.

####  A.5.5 Analyses on trained models

In this section, we further elaborate on the chosen modes of analysis in the
trained recurrent networks and in cases where the metric computation is
potentially ambiguous we elaborate on the methodology used to compute it. All
the plots shown in the paper regarding these analyses were performed on models
trained in the Seaquest environment.

Eigenspectra analysis. We motivate the analysis of the eigenspectrum of
${\bm{W}}_{rec}(r,s)$ from two perspectives: the rate of the vanishing
gradient and the attention profile of the recurrent weights across dimensions
of the hidden state. The former follows from the analysis presented in Pascanu
et al. (2012) which demonstrates that RNNs with recurrent weights that have
spectral radius less than $1$ necessarily express the vanishing gradient
because $\lambda_{max}$ dictates the rate of decay/explosion of the gradients
with respect to the recurrent weights across time. While $\lambda_{max}<1$ is
not a sufficient condition for the vanishing gradient in more complex
recurrent architectures like CfCs, LSTMs and GRUs, we believe that it is a
reasonable approximation. Furthermore, from the perspective of
interpretability, this is a useful formulation given the provability of our
connectivity prior as a modulator of spectral radius. Hence, we use the
spectral radius of ${\bm{W}}_{rec}$ as a measure for the rate at which the
gradient vanishes: in particular, the closer $\lambda_{max}$ is to $0$, the
faster the time gradients decay across the temporal axis which inhibits
gradient flow from data points in the past.

The second motivation for examining the recurrent eigenspectra is to
understand the attention profile of the recurrent weights at a fixed point in
time. If we assume that ${\bm{W}}_{rec}$ is diagonalizable, we can
equivalently express it as ${\bm{W}}_{rec}={\bm{P}}{\bm{D}}{\bm{P}}^{-1}$
where ${\bm{P}}$ denotes a matrix of eigenvectors and ${\bm{D}}$ denotes a
diagonal matrix with the eigenvalues on the diagonal. If we denote the hidden
state of the network by ${\bm{h}}$, then we have that

${\bm{W}}_{rec}{\bm{h}}={\bm{P}}{\bm{D}}{\bm{P}}^{-1}{\bm{h}}={\bm{P}}{\bm{D}}{\bm{k}}$

where ${\bm{k}}$ is equal to ${\bm{h}}$ transformed by the eigenbasis. Note
that the matrix-vector product ${\bm{D}}{\bm{k}}$ dictates how individual
dimensions of ${\bm{k}}$ are stretched before being transformed back into the
original basis by ${\bm{P}}$. Hence, the closer the distribution of
eigenvalues is to uniform, the more balanced the attention profile of the
recurrent weights is across dimensions of the eigen-transformed ${\bm{h}}$
which has implications for the dimensionality of recurrent state-space
dynamics.

For recurrent networks that have more than one set of recurrent weights (CfCs,
LSTMs, GRUs), we compute and plot a single eigenspectrum that represents the
average over the eigenspectra of each of the individual recurrent weight
matrices.

Effective dimensionality of trajectories. To understand the complexity of the
dynamics across recurrent models, connectivity ranks and connectivity
sparsities, we collected the trajectories of the model during the simulation
of the agent in the online, closed-loop setting and fit a PCA on the entire
set of trajectories. We then computed the explained variance of the top 5
principal components (PCs) as a proxy for the effective dimensionality of the
trajectories. This is a canonical approach to analyzing state-space
trajectories in recurrent networks (Lechner et al., 2020).

However, we extended the canonical analysis beyond just the full state-space
trajectories by decomposing the trajectories into recurrently-driven activity
and input-driven activity. In particular, we considered
${\bm{W}}_{rec}{\bm{h}}(t-1)$ as the recurrent activity and
${\bm{W}}_{inp}{\bm{x}}(t)$ as the input activity. We analogously aggregated
these trajectories individually and computed the explained variance of the top
5 PCs. This gives us separate measurements for what the dimensionality of
activity looks like in the recurrent subspace versus the input subspace and
allows us to explicitly disentangle these two axes.

For recurrent networks that have more than one set of recurrent (and input)
weights (CfCs, LSTMs, GRUs), we computed the explained variance of the top 5
PCs individually for each set of weights and then averaged across them to
produce a single number summary for the complexity of network dynamics.

Lipschitzness of network. Recall in Section 4.4 we conducted an analysis on
the spectral norms of the recurrent weights as a proxy for network robustness,
leaning on the fact that the Lipschitz constant of a matrix is given by its
spectral norm. However, in practice, we perturb the input of the network which
is first transformed via convolutional layers before being fed into the
recurrent layer which means that in reality we care about the Lipschitz
constant of the entire network. A trivial upper bound for the global Lipschitz
constant of a neural network is computed by multiplying the spectral norms of
all the weights in the network. However, this has been shown in many cases to
be a poor proxy for network robustness due to the looseness of the bound
(Huster et al., 2018). Since computing tight bounds for network-wide Lipschitz
constants is NP-hard (Scaman & Virmaux, 2019), we decided it was reasonable to
consider only the spectral norm of recurrent weights in the network. To verify
this, we also examined the spectral norm of the input weights to see if they
were changed by modulations in the recurrent connectivity. We found that they
were not (results not shown) which enabled us to justify the use of only the
spectral norm from the recurrent weights as a means of analyzing how our
connectivity parameterization affects network robustness.