Gated Neural ODEs: Trainability, Expressivity and Interpretability
Authors: Timothy Doyeon Kim
Tankut Can
Kamesh Krishnamurthy
Venue: Submitted to ICLR 2023 (2023)
Keywords: Computational Neuroscience
Dynamical Systems
Differential Equations
Neural ODEs
Gating
Interpretability

Abstract: Understanding how the dynamics in biological and artificial neural networks implement the computations required for a task is a salient open question in machine learning and neuroscience. In particular, computations requiring complex memory storage and retrieval pose significant challenge for these networks to implement or learn. Recently, a family of models described by neural ordinary differential equations (nODEs) has emerged as powerful dynamical neural network models capable of capturing complex dynamics. Here, we extend nODEs by endowing them with adaptive timescales using gating interactions. We refer to these as gated neural ODEs (gnODEs). Using a task that requires memory of continuous quantities, we demonstrate the inductive bias of the gnODEs to learn (approximate) continuous attractors. We further show how reduced-dimensional gnODEs retain their modeling power while greatly improving interpretability, even allowing explicit visualization of the structure of learned attractors. We introduce a novel measure of expressivity which probes the capacity of a neural network to generate complex trajectories. Using this measure, we explore how the phase-space dimension of the nODEs and the complexity of the function modeling the flow field contribute to expressivity. We see that a more complex function for modeling the flow field allows a lower-dimensional nODE to capture a given target dynamics. Finally, we demonstrate the benefit of gating in nODEs on several real-world tasks.

Under review as a conference paper at ICLR 2023
GATED NEURAL ODES: TRAINABILITY, EXPRESSIV- ITY AND INTERPRETABILITY
Anonymous authors Paper under double-blind review
ABSTRACT
Understanding how the dynamics in biological and artificial neural networks im- plement the computations required for task is salient open question in machine learning and neuroscience. In particular, computations requiring complex memory storage and retrieval pose significant challenge for these networks to implement learn. Recently, a family models described by neural ordinary differential equa- tions (nODEs) has emerged as powerful dynamical neural network models capable of capturing complex dynamics. Here, we extend nODEs by endowing them with adaptive timescales using gating interactions. We refer to these gated neural ODEs (gnODEs). Using a task that requires memory of continuous quantities, we demonstrate the inductive bias of the gnODEs to learn (approximate) contin- uous attractors. We further show how reduced-dimensionn modeling power while greatly improving interpretability. even allowing explicit visualization of the structure of learned attractors. We introduce novel measure of expressivity which probes the capacity of neural network to generate complex trajectories. Using this measure, we explore how the phase-space dimension of the nODEs and the complexity of the function modeling the flow field contribute to expressivity. We that more complex function for modeling the flow field allows a lower-dimensional nODE to capture a given target dynamics. Finally, we demonstrate the benefit of gating in nODEs several real-world tasks.
gnODEs retain their
INTRODUCTION
How can the dynamical motifs exhibited by artificial or biological network implement certain computations required for a task? This a long-standing question in computational neuroscience and machine learning (Vyas al., 2020: Khona & Fiete| 2022]. Recurrent neural networks (RNNs) have often been used to probe this question (Mante et al.) 2013 Vyas et al.] 2020; Driscoll et al., 2022), as they are flexible dynamical systems that be easily trained (Rumelhart al. 1986) to perform computational tasks. RNNs, particularly ones that incorporate gating interactions (Hochre- iter & Schmidhuber 1997 Cho al. 2014)) have been wildly successful in solving complex real- world tasks (Jozefowicz et al.) 2015).
While RNN models provide link between dynamics and computation, how their (typically) high- dimensional dynamics implement computation remains hard to interpret. On this note, we may turn neural ordinary differential equations (nODEs), class of dynamical models with velocity field parametrized by a deep neural network (DNN), which can potentially implement more com- plex computations in lower dimensions than classical RNNs (Chen et al., 22018) Kidger 2022 This increased complexity in lower latent/phase-spac dimensions subsequently helps in extracting interpretable, effective low-dimensional dynamics that may underlie a dataset or task (Kim et al, 2021
Despite their promise, nODEs remain under-explored in the following crucial aspects. Trainability: Can improve performance of nODEs by introducing gating interactions (Hochreiter & Schmid- huber 1997:
Cho et al.) 2014) to tame gradients in dynamical systems? Expressivity: How does
By classical RNNs, we mean the form of RNNs often considered the neuroscience, physics and Lognitive-science literature, where the interaction between units are additive, and the interaction strengths are represented by matrix (MeCulloch & Pitts] 1943 Sompolinsky et al. 1988 Elman 1990 Vogels et al. 2005
Sussillo Abboti 2009 Song et al. 2016 Yang et al. 2019).

---

Under review as a conference paper at ICLR 2023
the structure of the neural network modeling the velocity flow field influence nODE's capacity to model complex trajectories? Interpretability: Does the capability of low-dimensional nODEs to model complex data improve interpretability of the dynamical computation? We summarize below the main insights of our exploration of these questions.
Main Contributions
We leverage our understanding of gating interactions introduce the galed neural ODE (gnODE), We find that gating endows nODEs with adaptive timescales, and improves train- ability of nODEs on tasks involving long timescales or rich representations (Section 2 Appendix B).
We introduce a novel measure of expressivity related to the capacity of a neural network to store complex dynamical trajectories. nODEs and gnODEs are more expressive compared to RNNs in many parameter regimes (Sections 4 5.3 Appendices C F.
We demonstrate an inductive bias of gnODEs and other gated networks to utilize marginally-stable fixed points in a "flip-flop" task that requires storing continuous mem- ory. We further demonstrate the interpretability of the gnODEs solutions, which organize the marginally-staie fixed- points in an approximate continuous attractor (Section 5.2 Ap- pendix EE)
We show the advantage of gating in nODEs on real-world tasks (Sections 5.55. 5.5
Ap-
pendix G
We determine the critical initialization for nODEs using dynamical mean-field theory (Ap-
pendix
2 GATED NEURAL ODE
The gated neural ODE (gnODE) is described by
hh GGp (h, a)) @ l-hh- Fo(h, x)),
where 7 is the time constant, h E R is the hidden/latent state vector, and x(t) E RD is the input vector. The velocity vector field Fg RN RDD + R.V and the gating function Gp : IR RDD R.N are parameterized (via 9 and je, respectively) by neural networks. While Fg and G. in general can each be parametrized by any neural network, in this work, we restrict Fo and G. to fully- connected feedforward neural networks (FNN) Fo x) = s and (h, x) = s"x,, where
s' = (W..h U. b.), s = = (W! W!--- s b..). s.. d. (W... wh.-1g -igh. sh.-11 b
with * {ht 2). Here, W E NI.NN 8-- E R.N, b.: No.t1 and No Nr. N the phase-space (or latent) dimension. dy E (I, tanh] and d2
where I is the identity function and 1, Oa Φi. When > 1, we typically set da to be ReLU.
(x) = =-= When
Without the leak term - h and the gating interaction (i.e., setting G,, = 1). this reverts to form in which nODEs typically studied (Chen et al. 2018): Th Fo,h, (t) We include the leak term - -h our formulation because it allows us to initialize the weights of the (gated or non-gated) nODE in either the stable critical regime. Without the leak term, we show that the nODE is always dynamically unstable for any initialization, except for the zero initialization, and we expect this to hinder training (Abarbanel al. (2008); ; see Appendix for details).
When we set La L 1, Equation ( reduces to a "minimal gated recurrent unit" (mGRU; Ravanelli et al. (2018) ). which simplified version of the popularly used gated recurrent unit (GRU: Cho et al. (2014). When in addition the gating interaction is removed (G,, x)) = 1),
Equation reduces to widely studied class of models known as "Elman" (or "vanilla") RNNs
PIn Chen et al. 2018
and x<(t) 'Th - -h h W ox (h) + U + 6), is also popular in neuroscience models, where h can be interpreted
the internal voltage of neuron, and oh (h) as the output firing rate of the neuron; WDO h is the synaptic strength hetween neuron j and neuron i (Sompolinsky et al. 1988).
2

---

Under review as a conference paper at ICLR 2023
Krishnamurthy et al. (202) show that the mGRU exhibits manifold of marginally-stable fixed points in the limit of step-like gating function a for a wide range of parameters. This property is likely involved in shaping the inductive bias of gated networks, since it is useful in tasks requiring memory of continuous quantities (see Appendix B for an analysis of Jacobian spectrums of networks assuming different architectures, gated non-gated).
Can & Krishnamurthy (201);
3 RELATED WORK
Our work is closely related to neural controlled differential equations (nCDEs), developed in Kidger et al. (2020), which prescribes principled way to include inputs with nODEs: Th = Fg (h) da(() An important distinction between nODE and nCDE is that the nODE takes input æ (t), whereas nCDE uses the fime-derivtiive of the input dx/dl. /dl. Because the choice of the interpolation scheme used in nCDE also determines how the derivative is estimated, which scheme to use becomes criti-
cal (Morrill et al.) 2021). nODE avoids the complication of calculating the derivative, though it may not be as general as the nCDE (Kidger et al., 2020).
The primary motivation for introducing gating is its robust ability to generate long timescales and to address the exploding and vanishing gradients problem (EVGP) (Hochreiter & Schmidhuber, 1997]
Pascanu et al. 2013; Cho et al.) 20144] Can et al. 2020; Krishnamurthy et al.] 2022). Our work can be viewed as distilling the key elements of gating from GRUs and LSTMs responsible for long timescales and stable gradients (cf. Krishnamurthy et al. (2022); Can al. (2020)) and incorporat- ing them in MCDE-inspired models. Future work could incorporate these gating interactions more
directly in nCDEs, with their principled dealing of inputs and interpolation schemes.
Previous work explored improving the performance of nODEs by augmenting extra dimensions to the phase space (Dupont et al. 2019) or by regularizing nODEs to encourage simpler dynam- ics (Kelly et al.] 2020 Ghosh et 2020, Finlay et al.] 2020; Pal et al. (022). Gating can be applied in addition to these improvements, which we expect will make gnODE more powerful. We
also expect to see that gating will be beneficial for related model classes, such neural stochastic
differential equations (nSDEs) (Li et al. 2020).
Notable recent works have used RNNs based on discretized ODEs to deal with the EVGP, and achieve near state-of-the-art performance on various tasks. A RNN based on a system of coupled non-linear oscillators (coRNN) was introduced in Rusch & Mishra 2020), and this was extended to a Hamiltonian system with multiple (learned) time-scales in Rusch & Mishra (202T). In particular, the presence of the learned timescales was important for solving tasks with long timescales (Rusch) & Mishra 2021). In Rusch et al.] (02T), the authors introduced a RNN based on gated ODEs - the long expressive memory (LEM) that makes the timescales adaptive and effectively deals with the EVGP.
LEM in Rusch al. (2021) is in fact a special case of mGRU (i.e., La. - Log 1 in Equation (I) where the second half of the columns of W is constrained to be zero and W is constrained to an anti-diagonal block matrix. Moreover, based on our studies of the effects of gating, we suspect that the strength of the LEM in tasks involving long memory might partially stem from gating interactions (see Appendix B for discussion). Given the strong inductive bias conferred by gating on tasks requiring long memory, our work can also be considered as extending the ODEs considered in Rusch et al. (202T) to incorporate more flexible flow-fields as in nODEs and nCDEs. We include LEM in experiments real-world datasets for comparison (see Sections 5.5)
Finally, in addition to addressing the EVGP, we show in this work that gating introduces a powerful inductive bias for integrator-like behavior (see Section 5.2). It achieves this by forming a continuous manifold of marginally
fixed points, commonly referred to as continuous attractors (defined
in Appendix H for review, see Chaudhuri & Fiete (016)).
4 EXPRESSIVITY OF A NEURAL NETWORK
In order to compare architectures, it is useful to have a principled measure of expressivity in the dynamical setting. The metrie we use is inspired by the Gardner capacity (Gardner 1988 Engel
& Van den Broeck 200T), which measures the ability of an architecture to interpolate a random dataset, i.e., to fit noise. The Gardner capacity is also closely linked to the VC dimension (Abbaras)
3

---

Under review as a conference paper at ICLR 2023
et al.
2020
Engel & Van den Broeck 2001), ) and extended to temporal sequences in Bauer & Taylor (1991); Bressloff & Taylor (1992).
Krey
1991
We introduce the relevant concepts using discrete-time RNN of the form he+1 Fe (e()), assuming for simplicity that there is no input æ. The dataset we want to fit is a random time series & {Eo, L1, S2..... Gr). Assuming & are samples from some N-dimensional random process, perfect fit will require finding parameters 0 which satisfy the set of T equations fe+1 = Fe(Ee) 1. The space of solutions at a given T will occupy region of parameter space known as the Gardner volume, which is a function of T. The capacity is determined by the critical sequence length T at which the Gardner volume vanishes.
where / 0, 1, 2.,,
The longer the sequence a network can "memorize", the higher will its capacity/expressiity be. In typical systems, T scales with phase-space dimension N (see Appendix [ for worked-out canonical example). We suggest that an advantage of using a FNN Fg is that the capacity instead scales with total number number of parameters, which need not scale with phase-space dimension.
Based on this notion of expressivity, in Section 5.3, we train nODEs with a variety of architectures Fe on samples of an Shheein-Uhlenbeck process. our experiments, we measure instead the mean squared error between trajectories MSB(h,, E.).
5
EXPERIMENTAL RESULTS
5.1 PRELIMINARIES
In our experiments, we use libraries in Julia's (Bezanson et al. 2017) SciML ecosystem, Differ- entialEquations and FluxEl (Rackauckas & Nie 2017 Rackauckas et al. 2020), to imple- ment all network models presented the experiment, and choose to discretize dynamics of these networks using the canonical forward Euler method (except for the LEM, which is discretized with the wwrkwd Euler method, following Rusch et al. (202T)). We use the "discretize-then optimize" approach to obtain the gradient of the loss with respect to the network parameters (for more discussion on different choices of discretization and adjoint, see Appendix [D3] and D.4). Whenever there missing values in dataset, used natural cubic splines to interpolate the missing values, following Kidger et al. (2020).
5.2 N-BIT FLIP-FLOP TASK
We examine how a vanilla RNN, mGRU, GRU, nODE and gnODE implement the "n-bit flip-flop task" (Sussillo & Barak 2013). In the original n-bit flip-flop task (Sussillo & Barak 2013), the net- work is given a continuous stream of inputs coming from n independent channels. In each channel, transient pulse of value either +1 -1 is emitted random times. The network should contin- uously generate ni-dimensional outputs, where each dimension of the outputs should maintain the value of the most recent pulse in each channel (see Appendix [ illustration). Because each output channel of the network should take one of two values, the network should generate one of 2 outputs at each time point.
Consistent with previous findings (Sussillo & Barak 2013), when we trained our networks (see Appendix D.I for training details) on the 3-bit flip-flop task, we find that all networks we consider can reach validation mean squared error (MSE) < 0.01 on the task, for range of different phase space dimensions N, with appropriate hyperparameters. We also find that all networks use similar strategies to solve the task, with each of the 2 stable fixed points representing each output that the networks can take (see Appendix E. for details).
F.
Bariable-Amplitude Flip-Flop Task We then modified the task so that each pulse in each channel takes a real value sampled uniformly from -1 to [Figure[AA) We trained our networks from one of 27 different combinations of hyperparameters (i.e., learning rates, rates of weight decay and batch sizes; see Appendix E.2 for details). When we set the phase-space dimension our networks to be N 6, find that gnODE successfully reached validation MSE 0.01 with appropriate hyperparameters, while for other networks, all runs reached MSEs 2 0.025 (Figure BB). We verified that the validation MSEs of our networks converged after training (Figure [C). This suggests that
4

---

Under review as a conference paper at ICLR 2023
A
N
0.30
Channel 1
0.25
Validation MSE
0.15
Validation MSE
0.20
0.10
0.10
1
0.05
Channel 2
0.00
0.00
RNN
MGRU
GAU
nODE
gnODE
300
600
Epoch
E
20:
N - 100
- 100
10-
Validation MSE
0.15
maxx/RttAlon
0-
Channel 3
0.10
20
-10
0.05
99
-20
0.00
Time (5)
RNN
MGRU
GRU
NODE
gnODE
RNN
M3RU
GRU
NODE
gRODE
Figure 1: Networks assuming N 6 (A-C) and N = 100 (D-E) performing the variable-amplitude 3-bit flip-flop task. (A) An example validation trial with inputs in each channel shown in black, and the trained gnODE traces maintaining the previous pulse value shown in colors. (B) For each network, we tried 27 different hyperparameter configurations. Each circle represents the minimum validation MSE achieved during 600 epochs of training. Circles with black edges represent the minimum out of the 27 configurations.
(C) Validation loss traces as a function of epoch is shown for
the circles with black edges. Color codes are the same as in (B). (D) Same (B). (E) Each circle is the spectral abscissa of the Jacobian evaluated at detected fixed point. Bold horizontal lines indicate medians.
only gnODE is able to solve the task accurately when the phase-space dimension is low (i.c., N 6).
In contrast, when the phase-space dimension is high (N = 100), we find that the gnODE and GRU reached validation MSE < 0.01 and the vanilla RNN and mGRU reached validation MSE < 0.016 with appropriate hyperparameters (Figure [)). Thus, vanilla RNN, mGRU, GRU and gnODE can solve the task in high phase-space dimensions.
Structure of Solutions: Fixed-Points and Marginal Stability Following [Susilll & Barak (2013). to examine how these networks solve the task, we use Newton's 's method initialized from points in the trajectories taken by these networks, and find solutions that reach ||hl| < 0.01 (see Appendix E.3 for details on the fixed-point finding algorithm). For each 100-dimensional network that reached the minimum validation MSE among the 27 different hyperparameter configurations, we ran 10, 000 starting points to detect fixed points, and computed the maximum real component of the eigenvalues (i.e., spectral abscissa) of the numerical Jacobian obtained from each of the detected fixed points. The distribution of these spectral abscissas shows that the medians and the quartiles of the gated networks (mGRU, GRU, gnODE) are closer to zero, compared to those of the vanilla networks (vanilla RNN, nODE) (Figure (b.). This suggests that we detected more (effee- tively) marginally-stable fixed points for the gated networks compared to the vanilla networks. For the vanilla networks, we see that many of the detected fixed points are stable (i.e., spectral abscissas are much less than zero), in contrast to the gated networks. This suggests that a vanilla RNN may be reaching its solution using a combination of marginally-stable and stable fixed points, while the gated networks mostly rely on marginally-stable fixed points reach their solutions. We obtained similar results for networks assuming N = 6, although for these networks, only gnODE reached validation MSE < 0.01 (see Appendix E.4 for details).
Interpretability of gnODEs While analyses on the 100-dimensional vanilla RNN, mGRU, GRU and gnODE trained on the task can give useful insights, we found that when apply PCA on the trajectories taken by these networks, we needed at least 10 principal components to reach more than 0.9 variance explained, suggesting that the high-dimensional networks do not necessarily favor low-
5

---

Under review as a conference paper at ICLR 2023
A
D
<RODE traectory shatle food paint instable food point
Charnel
Ganee
Channel
Channel
Channel
Flow fields of gnODE with N = performing four different versions of the 2-bit flip-flop task. The input pulses C1 in Channel 1 and C2 in Channel 2 are given by (A) C1 E (-1,+1). The light gray lines indicate example trajectories taken by gnODE. (B) Square: real-valued input pulses in the interval c,, 02 E [-1,1]. Numerically-identin
Figure 2:
fixed points are omitted for better
visualization. (C) Rectangle: E [-2, 2] and e E [[-1,1]. (D) Annulus: S V + + 2
+
dimensional solutions in this setup. However, in principle, a dynamical system as simple as the one taking up 3 dimensions, which has cube filled with marginally-stable fixed points, can solve this task. Indeed, we find that when we set the phase-space dimension of gnODE to be N = 3, it can still achieve validation MSE < 0.01 with appropriate hyperparameters. We were not able to achieve this low MSE for other networks, suggesting gnODEs might be appropriate for studying the emergence of interpretable solutions to the citiaall-ampltude flip-flop task.
For simplicity, we turned to training 2-dimensional gnODE on the 2-bit flip-flop task and its vari- ants and plot the dimensional flow field such that the two axes describing this space are projected onto the axes that correspond to the outputs, Channel and Channel 2. For the fixed-amplitude task (where the pulse values can either be +1 or -1), we find stable fixed points, and find that each input perturbation moves the gnODE state from exactly one stable fixed point to another (Fig- ure (ee) We then trained a gnODE on a cetruble-amplitude 2-bit flip-flop task where the pulses can take values from 1 to 1. When we plot the flow field in the output space, we see that the velocity of the flows are close to zero, and this plane of fixed points roughly form a square between -1 and 1. Input perturbations try to move the gnODE state within the square, so that gnODE can hold onto the memory of the inputs (Figure 23. In summary, the gnODE learns a continuous attractor in the shape of square, and is solving the variable-amplitue flip-flop task in intuitively appealing way, by simply integrating the input.
The plane of fixed points show up not only for this particular task but also for other tasks. Instead of varying the values of the pulses from -1 to 1, we varied the values of the pulses in Channel I from -2 to 2 and find a rectangular attractor (Figure 2C). We also tried varying the statistics of the pulses so that pulses in the two channels are no longer independent, but appear at the same time, and the value taken by the pulse Channel 1, ct, and the value taken by the pulse in Channel 2, 02. satisfy S + < 2. We see a disk attractor with radius roughly of in this case. We do not a hole between radius 0 and 1 because crossing this region may be the fastest way from one state to another, and we did not explicitly penalize the network for crossing this region (Figure(p) Consistent with the flow field, we find that, even though the gnODE has not seen any inputs with pulse values satisfying 0 < V + oz < 1 during training, when it is given inputs with pulse values satisfying 0.5 < V + cg < 1, it generalizes well (MSE = 0.005; see Appendix E.5 for details). However, when it is given inputs with pulse values that are small (i.c., 0 < v c < 0.5), it does tend to mistake them as having no input at all, resulting in worse performance (MSE = 0.028). When the gnODE is given inputs with pulse values 2 < V + 2z < 4, does not generalize (MSE = 0.490).
We do not plot the flow fields for other networks assuming N 2, as all of the 27 runs with different hyperparameter configurations reached validation MSEs > 0.05 for vanilla RNN, mGRU and GRU, and validation MSEs > 0.02 for nODE (see Appendix E.5 for details).
These results together suggest that gnODEs might be flexible enough to learn more general manifold geometries, provided they are trained on an appropriate synthetic task. Furthermore, the geometry of the low-dimensional representations found by gnODEs directly inform their generalization capacities, thanks to their enhanced interpretability (see Figure g in Appendix E.5 for flow fields of gnODE trained more variants of the 2-bit flip-ftpp task).

---

Under review as a conference paper at ICLR 2023
5.3 FITTING AN ORNSTEIN- UHLENBECK TRAJECTORY
We introduce a task to measure the practical expressivity of a neural network. The task that the network has perform is perfectly fit a finite number of samples from an Omsttein Uhleen--k (OU) process,
Toudz = zou + ædt + doydw,
where w is a Wiener process. As long as TOU is sufficiently smaller than S/d where S is the total length of the trajectory and d is the distance between consecutive samples, we have samples that are reasonably uncorrelated. In our analysis, we set dim(z) = 30, TOU = 1s, λou = -1, (+t) = 1, and oou = 1, and sample at every is of this trajectory for 100s (therefore having a total of 100 samples). We train our networks on single trajectory of these 100 samples. For the vanilla RNN, mGRU and GRU, we systematically vary the phase-space dimension N and T of the model. For the nODE and gnODE, along with N and т, we also vary the number of hidden layers in F6 and the number of units N₃ in each hidden layer of Fg.
We generally see that, for all networks, when the model is closer to TOU, achieve lower training MSEs, confirming our intuition that networks perform best when their timescales match correlation time of the data (Figure B). We also confirmed that generally when we increase the number of units N₂ in each hidden layer, the networks become more expressive. Figure shows an example of this for gnODEs assuming TOU with nODEs and for different numbers of layers). The other side of this same coin is that hidden layers can act as a bottleneck for expressivity. We can see this in FigureB where for large phase- space dimension, small hidden layer can hurt expressivity. We also see that for various regimes, gnODE can be more expressive than other networks especially when N is low (Figure F F; see (ooo analyses not highlighted in the main text).
1s, and hidden layer in Fe (see Appendix F for results
By changing the model r on given dataset, we are effectively changing the difficulty of the task that the networks have to solve. Transients in the network will be relevant on timescales that scale as ~ т; therefore, for very large т, the velocity is suppressed and h evolves very slowly. This places greater burden on F6 to send small changes in the phase space into effectively orthogonal vectors in the OU time series. Therefore, we suspected that in the transient regime, the complexity of Fo becomes more important for fitting noise. Confirming our intuition, we that increases, the performance gap between networks that have more complex Fo (i.e., nODEs and gnODEs) and networks with simpler Fg (i.e., RNNs, mGRUs and GRUs) becomes larger (Figure D
5.4
REAL-WORLD TASKS
5.4.1 LATIN ALPHABET CHARACTER TRAJECTORY CLASSIFICATION
In this task, networks of different architectures were trained to classify 20 different Latin alphabet characters from cregularly-samplee time series consisting of the and y positions of the pen tip and the force on the tip. This dataset ("CharacterTraecttorie
is originally from the UEA time series classification archive (Bagnall et al. 2018). and we used the preprocessed data obtained from the o3 (see Appendix G.2 and Kidger et al. (2020) for details). We trained each
Neural CDE repository network by performing grid search over the hyperparameter space to find the of hyperparame- ters that minimizes the validation loss (see Appendix G.2 for details). Table TAA shows the accuracy of each network on the classification task.
5.4.2 WALKER2D KINEMATIC SIMULATION PREDICTION
The networks were given the task of predicting the dynamical evolution of the trajectories generated by the MuJoCo physics engine kinematic simulations (Todorov et al. 2012) The preprocessed
data for this task were obtained from the ODE-LSTM repository (see Appendix G.3 and Lechner & Hasani (2020) for details). While Lechner & Hasami (202); Xia et al. (2021) 1) did not choose to interpolate missing data with natural cubic splines, doing so helps with performances of the networks we show in Table B we generally MSFs that are lower than those reported in
CD8
http::///gtthbb..cmm///eeebbbb
1stms
7

---

Under review as a conference paper at ICLR 2023
A
Duta
B
GAV
gIODE. 1005
gHOCE 1/--1/111 -10
gIODE. 301
grod AA 6.17
groDe. 100
0.30
0.3
MSR
2-
0.25
0.20
0.2
0.15
0.10
0.0
10
10
Phase-space dimension
Phase-space dimension
D
E
F
0.40
0.35
AIRRU
0.4
GRU
0.30
0.25
AODE
grODE
Training MSE
0.3
Training MSE
0.35
Taining
0.2
0.30
0.10
316
1=
0.20
==
30s
0.00
10P
10
Phase-space dimension
Phase-spare dimension
Figure 3: Practical expressivities of vanilla RNN, mGRU, GRU and gnODE, and their dependence network timescale. (B-F) show means and standard deviations across 5 runs. (A) gnODE with hidden layer (N = N1 316) assuming E { {ls, 10s, 30s, 100s} fitting samples from the OU trajectory. (B) Training MSEs of gnODE in
(A). (C) Training MSEs of gnODE with 1 hidden layer, assuming т = 1s. (D) N 316 across all networks, N1 = 1000 for nODE and gnODE. (E) T 1s across all networks, N, 1000 for nODE and gnODE. (F) 30s across all networks, N, 1000 for nODE and gnODE.
Lechner & Hasani 2020); Xia et al. (021) (the lowest reported MSE on this task is 0.883 + 0.014, with an ODE-LSTM), Table IB shows the MSE of each network on the prediction task, with the hyperparameters that achieved the lowest MSE on the validation dataset (see Appendix G.3 for details).
5.4.3 SPEECH COMMANDS CLASSIFICATION
We trained the networks the task of classifying spoken words, such as "Stop" and "Go", based on audio recordings of these words. The dataset is originally from Warden (2018) and preprocessed using the pipeline the Neural CDE repository (see Appendix G.4 and Kidger al. (2020) for details). Table [C shows the test accuracy of each network on the classification task, with the hyperparameters that achieved the highest accuracy on the validation dataset (see Appendix G.4 for details).
5.5 CONCLUSIONS FROM THE REAL-WORLD TASKS
We observe in Table that, when the phase-space dimension N is fixed to be 100 for all networks, the gnODE performs competitively against other architectures considered in all of the three real- world tasks (Sections 1.4.1. 5.4.1 5.4.3. We also performed hyperparameter search over different Ns and generally find that gnODEs with lower Ns can perform competitively against other networks with higher Ns (see Appendices G.2 for for more details and discussion on the fairness of comparisons). This suggests a practical advantage in choosing gnODE over other networks for tasks involving sequences.
Consistent with results in Section 5.3 we show in Appendix G.4 how much the intrinsic time con- stant can influence the results of training. While nODE performs almost at chance level for the т considered Table TC, we find that changing т can significantly improve nODE performance. We also find that gnODE appears most robust to changes in т, suggesting that it would be reason- able default architecture when one is lacking a principled approach to choosing the time constant (A (Appendix (.4).
8

---

Under review as a conference paper at ICLR 2023
We additionally show that, for some datasets, the critical initialization determined in Appendix A when used together with Fo x) that has tanh as the final nonlinearity, can enhance performance of gnODE
those datasets (Appendices G.2 G.4.. Furthermore, show that gating function G (h. 2) with L > can give a boost in performance compared to a gating function Gpp (h, x)) with T_ 1 in gnODEs (Appendix G.4).
Table 1: Networks assuming N = 100 performing real-world tasks. (A) Test accuracy (mean + std, across 5 runs) on CharacterTrajectoriie (B) Test MSE on Walker2D kinematic simulation. (C)Test Test MSE on SpeechCommands.
(A) Character Prajettor/raeccooriie
(B) Walker2I)
(C) SpeachCommands
Model
Test Accuracy
Test MSE
Test Accuracy
mGRU
0.983 + 0.002
1.138 + 0.030
0.809 + 0.018
GRU
0.988 = 0.003
0.850 + 0.032
0.819 + 0.004
LSTM
0.987 = 0.002
0.865 = 0.009
0.768 + 0.012
LEM
0.988 = 0.003
0.709 = 0.009
0.794 = 0.007
nODE
0.898 = 0.089
0.707 = 0.023
0.140 = 0.012
gnODE
0.986 + 0.003
0.588 + 0.003
0.823 + 0.006
6 DISCUSSION
We introduced gated neural ordinary differential equations (gnODEs), novel nODE architecture which utilizes a gating interaction to dynamically and adaptively modulate the timescale. A syn- thetic n-bit flip-flop task (ef. [Sussillo & Barak (013)) was used to demonstrate the inductive bias of the gnODEs to learn continuous attractors. We also showed that, compared to other architectures, the gnODE can learn this task with lower phase-space dimension. This allows to inspect the nature of the solution learned in an intuitive and interpretable manner. We also formulated a princi- pled measure of expressivity for RNNs/nODES based on their ability fit random trajectories. We used this measure to investigate how the phase-space dimension and the complexity of the velocity field interact to shape the overall expressivity. We saw that when the phase-space dimension is low, the gnODE can be more expressive compared to the other architectures tested. Lastly, even though gating results in more parameters and slower per-iteration update of the network state, we empiri- cally showed that a gated network (whether it be gated RNN or gated nODE) can significantly improve performance compared to a vanilla network, both on carefully designed synthetic tasks and real-world tasks.
While we do not claim that each unit a gnODE can correspond to a biological neuron, there is evidence that biological neural networks utilize several of the mechanisms that are found in the gn- ODE. First, gating appears to be a generally observed phenomenon in biological neural networks. For example, gnODE can, similar neurons where gating corresponds to voltage-gated ion channels (Rusch et al. (022). In another example, negative-derivative feedback in an E-1 balanced network can be viewed as form of gating which dynamically changes the time constant (Lim & Goldman 2013). Furthermore, it is known that a gating mechanism allows network to robustly form continuous attractors (Can & Krish- namurthy 2021), which is thought to be prevalent in biological neural networks (Khona & Fiete, 2022]. Second, recent experiments show that neural population activities across a large number of brain regions and species can be described by a low-dimensional dynamical system (Churchland et al. 2012 Harvey et al., 2012; Mante et al., 20133 Kaufman et al. [2014] Nieh et al., (021). How- ever, our work shows that nigh-dimensional networks do not necessarily favor a low-dimensional solution to a low-dimensional task.
LEM, be mapped onto network of Hodgkin- Huxley
Among the networks that we considered in this work, gnODE is the only network that both uses a gating mechanism and is capable of learning complex dynamics even in low phase-space dimen- sions, consistent with the previous literature how biological neural networks work. These features make gnODE a uniquely powerful model for probing the connection between computation and dy- namics in artificial and biological neural networks.

---

Under review as a conference paper at ICLR 2023
REFERENCES
Henry D. I. Abarbanel, Daniel R. Creveling, and James M. Jeanne. Estimation of parameters in nonlinear systems using balanced synchronization. Phys. Rev. 77:016208, 2008.
Alia Abbaras, Benjamin Aubin, Florent Krzakala, and Lenka Zdeborová. Rademacher complexity and spin glasses: A link between the replica and statistical theories of learning. In Mathematical and Scientific Machine Learning, pp- 27-54. PMLR, 2020.
Anthony Bagnall, Hoang Anh Dau, Jason Lines, Michael Flynn, James Large, Aaron Bostrom, Paul Southam, and Eamonn Keogh. The uea multivariate time series classification archive, 2018. arXiv, 2018.
K Bauer and U Krey. On the storage capacity for temporal pattern sequences in networks with delays. Zeitschrifi für Physik B Condensed Matter, 84(1):131-141, 1991.
Jeff Bezanson, Alan Edelman, Stefan Karpinski, and Viral B Shah. Julia: A fresh approach to numerical computing. SIAM Review, 59(1):65 2017.
PC Bressloff and JG Taylor. Temporal sequence storage capacity of tume-aing neural net- works. Journal of Physics Mathematical and General, 25(4):833, 1992.
Nicolas Brunel. Is cortical connectivity optimized for storing information? Nature neuroscience, 19 (5):749-755, 2016.
Tankut Can and Kamesh Krishnamurthy. Emergence of memory manifolds. arXiv, 2021.
Tankut Can, Kamesh Krishnamurthy, and David J. Schwab. Gating creates slow modes and controls phase-space complexity in GRUs and LSTMs. In Proceedings of The First Mathematical and Scientific Machine Learning Conference, pp. 476-511, 2020.
Jack Carr. Applications of centre manifold theory, volume 35. Springer Science & Business Media, 1981.
R. Chaudhuri and I. Fiete. Computational principles of memory. Nature Neuroscience, 19:394-403, 2016.
Ricky T. Q. Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural Ordinary Dif- ferential Equations. Advances in Neural Information Processing Systems, 11.6571-6553 2018.
Kyunghyun Cho, Bart van Merrienboer, Caglar Gulcehre, Dzmitry Bahdanau, Fethi Bougares, Hol- ger Schwenk, and Yoshua Bengio. Learning phrase representations using rnn encoder-decoder for statistical machine translation, 2014.
M. M. Churchland,
Cunningham, M. T. Kaufman, J. Foster, P. Nuyujukian, Ryu, and K. V. Shenoy. Neural population dynamics during reaching. Nature, 487:51-56, 2012.
A Crisanti and H Sompolinsky. Path integral approach to random neural networks. Physical Review E, 9826),, 2018.
Edward De Brouwer, Jaak Simm, Adam Arany, and Yves Moreau. Gru-ode-bayes: Continuous modeling of sporadically-obsereed time series. arXiv, 2019.
Bernard Derrida and Yves Pomeau, Random networks of automata: simple annealed approxima- tion. EPL (Europhysics Letters), 1(2):45, 1986.
Darshil Doshi, Tianyu He, and Andrey Gromov. Critical initialization of wide and deep neural networks through partial jacobians: general theory and applications layernorm. arXiv preprint arXiv:2111.12145 2021.
Laura Driscoll, Krishna Shenoy, and David Sussillo. Flexible multitask computation in recurrent networks utilizes shared dynamical motifs, bioRxiv, 2022.
Emilien Dupont, Arnau Doucet, and Yee Whye Teh. Augmented Neural ODEs. In Advances in Neural Information Processing Systems, volume 32, 2019.
10

---

Under review as a conference paper at ICLR 2023
Jeffrey L. Elman. Finding structure in time. Cognitive Science, 14(2):179-21 1990.
Andreas Engel and Christian Van den Broeck. Statistical mechanics of learning. Cambridge Uni- versity Press, 2001.
Chris Finlay,
Jacobsen, Levon Nurbekyan, and Adam Oberman. How to train your neural ODE: the world of Jacobian and kinetic regularization. In Proceedings of the 37th Inter- national Conference on Machine Learning, volume 119, pp. 3154-3164, 2020.
Elizabeth Gardner. The space of interactions in neural network models. Journal of physics A: Mathematical and general, 21(1):257, 1988.
Amir Gholami, Kurt Keutzer, and George Biros. ANODE: Unconditionally Accurate Memory- Efficient Gradients for Neural ODEs. arXiv, 2019.
Arnab Ghosh, Harkirat Behl, Emilien Dupont, Philip Torr, and Vinay Namboodiri. STEER Sim- ple Temporal Regularization for Neural ODE. In Advances in Neural Information Processing Systems, volume 33, pp. 14831-14843, 2020.
Xavier Glorot and Yoshua Bengio. Understanding the difficulty of training deep feedforward neural networks. In Proceedings of the Thirteenth International Conference on Artificial Intelligence and Statistics, volume 9, pp. 249-256, 2010.
Danijar Hafner. Tips for training recurrent neural networks. Blog post, 2017.
URL
https:
//dani
C. D. Harvey, Coen, and D. W. Tank. Choice-specific sequences parietal cortex during a virtual-navigation decision task. Nature, 484:62-68, 2012.
Moritz Helias and David Dahmen. Statistical field theory for neural networks. Springer, 2020.
S. Hochreiter and J. Schmidhuber. Long short-term memory. Neural Computation, 9 1997.
Mike Innes Flux: Elegant machine learning with julia. Journal of Open Source Software, 3(25): 602, 2018.
lan D. Jordan, Piotr Aleksander Sokót, and Il Memming Park. Gated recurrent units viewed through the lens of continuous time dynamical systems. Frontiers in Computational Neuroscience, 15, 2021.
Rafal Jozefowicz, Wojeiech Zaremba, and 1lya Sutskever. An empirical exploration of recurrent network architectures. In International conference on machine learning, pp. 2342-2350. PMLR, 2015.
M. Т. Kaufman, M. M. Churchland, S. I. Ryu, and V. Shenoy. Cortical activity in the null space: permitting preparation without movement. Nature Neuroscience, 17(3):440-448, 2014.
Jacob Kelly, Jesse Bettencourt, Matthew James Johnson, and David Duvenaud. Learning differential equations that are easy to solve. In Neural Information Processing Systems, 2020.
Mikail Khona and Ila R. Fiete. Attractor and integrator networks in the brain. Nature Reviews Neuroxcience, 2022.
Patrick Kidger. On Neural Disferential Equations. PhD thesis, Oxford, February 2022. arXiv: 2202.02435.
Patrick Kidger, James Morrill, James Foster, and Terry Lyons. Neural Controlled Differential Equa- tions for Irregular Time Series. Advances in Newral Information Processing Systems, 2020.
Timothy Doyeon Kim, Thomas Zhihao Luo, Jonathan W. Pillow, and Carlos D. Brody. Inferring la- tent dynamics underlying neural population activity via neural differential equations. Proceedings of the 38th International Conference on Machine Learning, 2021.

---

Under review as a conference paper at ICLR 2023
Kamesh Krishnamurthy, Tankut Can, and David J. Schwab. Theory of gating in recurrent neural networks. Phys. Rev. X, 12:011011, 2022.
Mathias Lechner and Ramin Hasani. Learning Long-Term dependencies in Peggularly-Samplle time series. arXiv, 2020.
Jaehoon Lee, Yasaman Bahri, Roman Novak, Samuel S Schoenholz, Jefffrr Pennington, and Jascha Sohl-Dickstein. Deep neural networks as gaussian processes. arXiv preprint 2017.
Xuechen Li, Ting-Kam Leonard Wong, Ricky T. Q. Chen, and David Duvenaud. Scalable gradients for stochastic differential equations. arXiv, 2020.
Sukbin Lim and Mark S Goldman. Balanced cortical microcircuitry for maintaining information in working memory. Nature neuroscience, 16(9):1306 1314, 2013.
Iyy Loshchilov and Frank Hutter. Decoupled weight decay regularization. International Confer- ence on Learning Representations, 2019.
V. Mante, D. Sussillo, K. V. Shenoy, and W. T. Newsome, Context-dependent computation by recurrent dynamics prefrontal cortex. Nature, 503(7474):78- 84, 2013.
Warren S. McCulloch and Walter Pitts. logical calculus of the ideas immanent in nervous activity. The bullletin of mathematical biophysics, 5(4): 115-133, 1943.
Patrick Kofod Mogensen and Asbjørn Nilsen Riseth. Optim: A mathematical optimization package for julia. Journal of Open Source Software, 3(24), 2018.
James Morrill, Patrick Kidger, Lingyi Yang. and Terry Lyons. Neural controlled diffcrential equa- tions for online prediction tasks. arXiv, 2021.
Jcan-Pierre Nadal. Neural networks that learn temporal sequences. In Measures of Complexity, pp- 54-61. Springer, 1988.
E. H. Nieh, M. Schottdorf, N. W. Freeman, R. Low, S. Lewallen, S. A. Koay, L. Pinto, J. L. Gau- thier, C. D. Brody, and D. W. Tank. Geometry of abstract learned knowledge in the hippocampus. Nature, 2021.
Derek Onken and Lars Ruthotto. Diseretize-optimize vs. optimize-dissretizz for time-series regres- sion and continuous normalizing flows. arXiv, 2020.
Avik Pal, Yingho Ma, Viral Shah, and Christopher V Rackauckas. Opening the blackbox: Acceler- ating neural differential equations by regularizing internal solver heuristics. In Proceedings of the 38th International Conference on Machine Learning, volume 139, pp. 8325-8335, 2021.
Razvan Pascanu, Tomas Mikolov, and Yoshua Bengio. On the difficulty of training recurrent neural networks. In International conference on machine learning, pp. 1310-1318. PMLR, 2013.
Christopher Rackauckas and Qing Nie. Differentialleuuaionn:]] jl- performant and feature-rich ecosystem for solving differential equations in julia. Journal of Open Research Sofwware 5(1), 2017.
Christopher Rackauckas, Yingbo Ma, Julius Martensen, Collin Warner, Kirill Zubov, Rohit Supekar, Dominic Skinner, and Ali Ramadhan. Universal differential equations for scientific machine learning, arXiv preprini arXiv:2001.04:: 2020.
Adityanarayanan Radhakrishnan, Mikhail Belkin, and Caroline Uhler. Wide and deep neural net- works achieve optimality for classification. arXiv, 2022.
Mirco Ravanelli, Philemon Brakel, Maurizio Omologo, and Yoshua Bengio. Light gated recurrett units for speech recognition. IEEE Transactions on Emerging Topics in Computational Intelli- gence, 102, 2018.
D. E. Rumelhart, G. E. Hinton, and R. J. Williams. Learning representations by back back-propagating errors. Nature, 323(6088):533-536, 1986.
12

---

Under review as a conference paper at ICLR 2023
T Konstantin Rusch and Siddhartha Mishra. Coupled oscillatory recurrent neural network (corn)): An accurate and (gradient) stable architecture for learning long time dependencies. arXiv preprint 0rXiv:2010000951, 2020.
T Konstantin Rusch and Siddhartha Mishra. Unicornn: A recurrent model for learning very long time dependencies. In International Conference on Machine Learning, pp. 9168-9178. PMLR, 2021.
Konstantin Rusch, Siddhartha Mishra, N Benjamin Erichson, and Michael W Mahoney. Long expressive memory for sequence modeling, arXiv preprint arXiv:2110.04744 2021.
Samuel S Schoenholz, Justin Gilmer, Surya Ganguli, and Jascha Sohl-Dickstein. Deep information propagation. ICLR, 2017.
Jannis Schuecker, Sven Goedeke, and Moritz Helias. Optimal sequence memory in driven random networks. Physical Review 844)004,, 2018.
H. Sompolinsky, A. Crisanti, and H. Sommers. Chaos in random neural networks. Phys. Rev. Lett. 61:259-262, 1988.
Haim Sompolinsky and Ido Kanter. Temporal association in asymmetric neural networks. Physical review letters, 57(22):2861, 1986.
H. Francis Song, Guangyu R. Yang, and Xiao-Jing Wang. Training excitatory-inhibitrry recurrent neural networks for cognitive tasks: A simple and flexible framework. PLOS Computational Biology, 12:1 1-30, 2016.
David Sussillo and L.F. Abbott. Generating coherent patterns of activity from chaotic neural net- works. Neuron, 6444)5544-557, 2009.
David Sussillo and Omri Barak. Opening the Black Box: Dw---imensiinll Dynamics in High-
Dimensional Recurrent Neural Networks. Neural Computation, 25()):626-649, 2013.
John G Taylor. Neural network capacity for temporal sequence storage. International Journal of Neural Systems, 2(01n02):47-54, 1991.
Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physics engine for model-based con- trol. 2012 IEEE/RSJ International Conference on Intelligent Robots and Systems, pp. 5026-5033, 2012.
Tim P. Vogels, Kanaka Rajan, and L.F. Abbott. Neural network dynamics. Annual Review of Neu- roscience, 28(1):357-374 2005.
Saurabh Vyas, Matthew D Golub, David Sussillo, and Krishna V Shenoy. Computation through neural population dynamies. Annual Review of Neuroscience, 43:249-275, 2020.
Pete Warden. Speech commands: dataset for limited-vocabulary speech recognition, 2018.
Christopher Williams. Computing with infinite networks. Advances neural information process- ing systems, 9, 1996.
Hedi Xia, Vai Suliafu, Hangjie Ji, Tan Minh Nguyen, Andrea Bertozzi, Stanley Osher, and Bao Wang. Heavy Ball Neural Ordinary Differential Equations. In Advances in Neural Information Processing Systems, 2021.
Guangyu Yang, Madhura Joglekar, H. Song, William Newsome, and Xiao-Jing Wang. Task repre- sentations in neural networks trained to perform many cognitive tasks. Nature Neuroscience, 22, 2019.
13

---

Under review as a conference paper at ICLR 2023
A CRITICAL INITIALIZATION FOR NEURAL ODES
In this Appendix, we will determine the critical initialization for neural ODEs. First, defining the model
h - + Fo 2:)
(2)
We will consider a multi-layer perceptron (MLP) network function Fe defined according to the equations
Fo(h, x) del at a f+1 = -W'd(a") bb, for t= -- a' = - + Ux 6
(3)
(4)
(5)
W E R-N-LIXN b' ERN,, a BR, hER", ZERP.
(6)
This is equivalent to the feedforward neural networks (FNN) defined in the main text under the identification s² a² and s = a(a²) for 1, 1. We have also separated W from because they should be scaled differently.
Jacobian A useful quantity in studying the dynamics and assessing stability is the instantaneous Jacobian D. This will be related to the input-output Jacobian I of the MLP, where
0 (Fo), -. Jh,
To
(a"3)) W (e'(a')) )
Using this, the instantaneous Jacobian of the nODE
Dij = -ōi Jij-
(8)
A.1 MEAN-FIELD THEORY
Initialization and Mean-Field Scaling We consider two choices of scaling which lead to a mean- field theory, each informed by popular initialization schemes in machine learning. The first the Kaiming scaling of the weights:
W! N (0,
oe N₂
Kaiming scaling
(9)
with No = Nr = N being the dimension of the phase space in which h lives. We also naturally would like U3 (0, o D), in order for the input to not be unnecessarily suppressed by No- This is only a problem if and are significantly mismatched.
Alternatively, take inspiration from the popular Glorot initialization and
W N (0, N₂ Near) )
Glorot scaling
(10)
The mean-field theory then requires taking N, + oo (including No) while keeping their ratios fixed. Defining the aspect ratio
a4+1 = Net1/Ne, we will develop the results below assuming the following initialization scheme
(11)
W
of
o/= + a1+1)
(12)
Keeping ag makes this equivalent Glorot scaling, whereas setting all o₂ recovers Kaiming scaling.
By keeping or unspecified, we have actually introduced more flexibility to what is typically under- stood by these initialization schemes. In fact, what is usually called Kaiming/Gloroot initialization has o. = 2. We will keep to this convention, and refer to Kaiming/Glorot scaling when o is not explicitly fixed.
14

---

Under review as a conference paper at ICLR 2023
Correlation Functions in MFT The dynamical mean mean-field theory (DMFT) for the nODE follows the logic presented in many previous works, see e.g. Crisanti & Sompolinsky (2018): Helias & Dahmen (2020). Proceeding via the Martin-Sigge Siggia- Rose statistical field theory, in the saddle-pond point approximation, valid for large N, h described by Gaussian process with zero mean and covari- ance determined by the self-consistent DMFT equation
= (0: + 1) + I)Ca(1))) Cr((((),
(13)
where we have chosen the convention to represent correlation functions
1 Cn((,)= - E (ii(t)),tt))) Cx(t,t)= = 'N EE(O))(()))
1
(14)
and the averages are taken over the random parameters.
In order find self-consistent solution, need to express Cp as a function of C.. This can be accomplished by appealing to well-known known results in the literature on the neural network Gaussian process (NNGP) kernel for the MLP defined by Fo (see e.g. [Wiliams] (1996): Lee et al. (2017)) To get the desired correlation function, or kernel, we define a hidden layer kernel function
N₂
(t,tt)
1
((())(())
(15)
which satisfies the recurrence relation
(1(11)
+
(16)
1+a1 m
(()(),))) = 1+ a2+1
(K(K((( + +
(17)
(t,) = 144+11
(t,t))
(18)
Cr(t,t'= = (t,tt)
(19)
where
R((L))- (
(<()))) k((),)) )
k²(t,t) k((t,t)
(20)
Here, we have defined the correlators
(x) = f dxx
1x" R 100 (x1()((()2))
(21)
2n det K
(K) = / dx
22
2K (xx))
(22)
2
Asympotic Stability Let us consider the divergence of trajectories. The usual trick is take two replicas with different initial conditions but identical weights (Derrida & Pomeau 1986 Schuecker
al. 2018). This will change the DMFT in the following way
(o. + 1) + 1)Cx = C+b(t,t')
(23)
with a, 6 1,2. Here, the RHS is obtained from the recurrence relations
Kl.ab 1 a1 (1,'), + of o²
(24)
kl+l,abll,1)
o
Co (1,)) + +o?
(25)
1+ a+41
(tb(t,t') L.eb (t,
(26)
15

---

Under review as a conference paper at ICLR 2023
We assume steady state which is time- translaannlttion invariance, so the correlation functions depend only on the difference T = |t-t'|. Then, expanding around the replica symmetric solution (T) = (T( + = will give the eigenvalue equation for Q
1)2 02) Q X
(27)
where we have used
8CF
Cor (KL-11)
Cø x
w
(28)
(T) 1 aL
1 02
1 a1
= (T)
(29)
Here, have defined the susceptibility xe((, which satisfies its recurrence relation (sup- pressing the time arguments)
- Co (Rt) X6, XI I+a1
Xe+1
(30)
1 cx+1
The susceptibility at unequal times is typically not studied in the FNN setting (Schoenholz et al,, 2017 Doshi al., 2021). Usually, the equal-time susceptibility x(0) is sufficient, since it charac- terizes the behavior of gradients. The object which appears here X(T) is tantamount to studying the overlaps of the gradient of the FNN output for two different inputs. However, if we are instead interested in fixed points, we have quite simply
11+1)==
(31)
The susceptibility which appears here x0) is precisely the object typically studied for FNN. So, we use the intuition from feedforward networks and initialize at criticality, we will find a marginally stable fixed point in the nODE.
Fixed-Point Jacobian Radius Proceeding, we wish to determine the edge of stability for fixed- points. To do so, must first use the MFT to find fixed points according to the self-consistent equation
C₁ = KL.
(32)
In the large N limit, the spectral of the Jacobian D depends only on the distribution of h, and thus on Cr. Furthermore, since it is uniformly shifted by the identity, the spectral radius of J. which we denote p(JJ, is enough to determine stability. One can show that the squared spectral radius p(T) given by
(
1 N
x(J)²
(33)
(L-1 CKw (=1
1
21.
=x.(0))
(34)
1 1+ar ae
Since the correlation functions that appear depend only on the distribution of h, and thus only on Ch once the MFT fixed-point equation is solved, the solution can be plugged into this expression for the spectral radius to determine stability.
Note also that the squared spectral radius equal to the static susceptibility defined above, as must. A common set up will have N₂ = N2 N, while all hidden layers have the same dimension N1 = ... Nr.1 H. Then defining a (H/N, 0) and B (1,0) for (Glorot, Kaiming), we get
II (cCK)))
L-1
p(J)²=
(35)
(1 3) (1 xx
In Figure 4 we compute the critical curve in the ow og plane along which p(J) 1. We show how this curve changes with increasing depth. For concreteness, we choose Kaiming scaling and ((x) (a))( activation.
16

---

Under review as a conference paper at ICLR 2023
With biases exactly zero, the zero fixed point typically determines the edge of chaos. The spectral radius for the zero FP is
Te Le Ø' 2'())((L-1 (1+ 21 a)
0
(36)
Stable
3
0
L
2-
L=6 6
1-
Chaotic
0-
1.25 1.50
1.75
2.00
2.25
2.50
2.75
3.00
3.25
2
Figure 4: To obtain these curves, we used ((x) tanh(x). This shows the critical curve separating stability from chaos as a function of bias and weight variances. The different curves correspond to MLP functions of differing depth. We used Kaiming scaling such that 3 0. For fixed depth L, the region below the plotted curve is chaotic, whereas the region above the plotted curve is stable.
Explicit Solutions for ReLU Networks (6 ReLU) If the MLP utilizes only the ReLU activa- tion, there does not appear to be a chaotic phase. When tanh is applied as the final nonlinearity for Fe. the system has chaotic regime. The suggested initializations Equations [1 52) are valid for both Fg with and without the final nonlinearity tanh.
We will make use the integral identities for one-point functions
Co(K)=
( Dæ ( (VKa)+)
1 = K, Car Car(K) /
-
dh
(37)
2
2
and for two-point functions, setting x = x2), and assuming a iime-translation invariant kernel
Ko Kr
Ko ) K, S Ko
K =
(38)
K
have
Ck(
00
dx dæjdxz
#112E R--=''x R-xx
(39)
Jo
2T V det K
1
Kr
K,
+ tan
2T V
K3 K?
(40)
4
VK3 -??
Co.
1 00
|
dx dx
e
4x" R-1x
1
1
2
tan
Kr
(41)
2n V det K
1
VK3 K?
(42)
Fixed-Points We begin by analyzing the time-independent fixed points. The fixed-point can be determined exactly using the recurrence relations. Define the coefficients
o2
o
a1 2(1+a)
a2
2(1 + 3)
b ==0
(43)
17

---

Under review as a conference paper at ICLR 2023
Then compute the kernel for the ReLU MLP via the recurrence relations
K" = 201 Ch + Oz
(44)
kt+l =az b, 8 = 1.,, 2-2
(45)
K". 0² 1 k +
(46)
1
1-2
=maaz (2a1 Cn +z + - 022
b b.
(47)
02
The dynamical fixed- fixed-point of the nODE is determined by C₁ K which implies
L-2
C₁ 2a7a2
0102 2000 d1
L-2
012
(48)
02
Therefore, fixed point exists for
2a L-2 < 1.
(49)
Note that since the LHS here is precisely equal to the squared spectral radius, if a fixed point exists, then it must also be stable.
Criticality will correspond to the spectral radius of the input-output Jacobian being precisely equal to unity. The resulting equation can solved for o and yields
de + a) 21-11/ /2+1)////1111110 8)1-2/1,
Critical init.
(50)
Specifying for the two popular initialization schemes discussed above gives
0e = 221-1/1, Kaiming scaling
(51)
0 33/21 (1 a1+ Glorot scaling
(52)
Comparing these the traditional choices for these initializations, find that Kaiming initializa- tion with d= = /2 will place the network in the unstable regime. Conversely, Glorot initialization with o= = 2 will initialize the network in the stable regime.
trivial corollary our analysis thus far is that randomly initialized nODE without a leak term always unstable, since the condition for stability in this setting is p(J) which implies critical
B COMMON FEATURES OF GATING ACROSS ARCHITECTURES
Following Krishnamurthy et al. 2022). we did an analysis on the empirical Jacobian spectrum of LEM (Rusch al.] (021) with gating and without gating, and compared them those of mGRU (Ln, = 1, = 1) and gnODE (Ln, 3, Lg 1) (Figure 5). To generate the plots in Figure 5 we d. tanh, initialized Wi, according Equation and similarly initialized W Wi,i with:
W 5.y9 ~ ~ (0,
(53)
where = 1s, and N = 1000. We discretized the network dynamics with the forward Euler method with At 1s for mGRU and gnODE, and discretized LEM with the forward-backwant Euler method with At 1s (following exactly Equation (3) in Rusch al. (2021)). To ensure that the dynamics reached steady-state, we ran the solvers up until 1000s, and evaluated the eigenvalues of the numerical Jacobian of the approximate stead We found that the spectrums of the networks we get are roughly similar when we discretize the dynamics with the Tsitouras 5/4 Runge- Kutta method, except for the spectrums of the LEM, which had shapes similar to those of the mGRU.
18

---

Under review as a conference paper at ICLR 2023
mGRU
gnoDE
LEM
1
1
No Gating
Im(A)
0
0
-1
-1
-1
-2
-1
-2
-1
2
-
0
Strong Gating
1
1
1
Im(A)
0
-1
-1
-1
-2
-1
-2 -1
-2 -1
Re(A)
Re(A)
Re(A)
Figure 5: All non-gating weights (W1) Kaiming normal with overall scale σu 1.5. All gating weights (W") set to Kaiming normal with overall scale a₂ = 0 (i.e., no gating, top row) and o, 5 (bottom row).
When the LEM does not have gating (Figure 5 top right), see that, compared to mGRU gnODE without gating (Figure 5 top left and middle), the special anti-diagonal block structure of W lets the LEM stay close criticality. This may partially be due the fact that the LEM without gates can be mapped to a Hamiltonian dynamical system. However, when we add gating to LEM presented in Rusch et al. (202T); Figure 5 bottom right), it nullifies the effect from the special anti-diagonal block structure, and we see a robust "pinching" of the Jacobian spectrum leading to eigenvalues clustering near zero and thus long timescales/stable gradients, which is ubiquitous for the gated networks (gnODE, mGRU, LEM; Figure []] bottom row). This pinching results in long- lived modes, contributing to all of these gated networks' ability to learn long time dependencies.
C GARDNER VOLUME FOR TRAJECTORY FITTING CAPACITY
In this section, we derive the capacity of spherical perceptron to store a random time series by mapping the problem to Gardner's original calculation (Gardner 1988) This result also appears in Bauer & Krey (1991] Taylor (1991); Bressloff & Taylor (1992), which studied storage capacity for time-delay RNNs. Previous work has also studied storage capacity for temporal sequences in RNNs with Hebb rule structured connectivity (Sompolinsky & Kanter 1986 Nadal 1988).
We start by setting up the problem in more generality. In the main text, pursued definition of expressivity that involved fitting a random time series. The ability to fit such noise is intimately connected storage capacity of perceptron.
Consider a discrete-time nODE (or a generalized RNN)
h Fohh))
(54)
with which we want fit random time series
(Eo, E E2,
(55)
where & are i.i.d random variables. A perfect fit will require set of parameters 0 that satisfy the set T' equations
ε Fo(Ee), 0,2, T-1
(56)
We will try to find the volume in parameter space which satisfy this equation. A similar question was asked in Brunel (2016), which was interested the structure of solutions which store the optimal length sequence.
19

---

Under review as a conference paper at ICLR 2023
We allow for error in the fit, and want to find all 8 which satisfy these constraints. There are different formulations depending on the activation functions. In general, for smooth activation functions, we can define an indicator function
T x(E) = II + (5+1 Fe(5)) +c) + Fo(G) +e).
(57)
the weights are such that the trajectory of the nODE follows & within some margin e, then X = 1; otherwise, O. is also necessary to insert some sort of regularizer, so that the volume in 0 space does not explode. This will have the effect replacing the measure d0 - dug with a regularized measure that converges, and which assume is normalized f ddu 1. With these ingredients, the volume in the space of parameters is given by
/ due x(E)
V v=
(58)
Specifying this setup to the spherical perceptron considered by Gardner, we use F(E) sign (N -1/__E)) with parameters J. and binary patterns E E {-1,+13. This is the set-up an- alyzed in Bauer & Krey [1991]; Taylor (1991): Bressloff & Taylor 1992), where strated that the calculation ends up being identical to the Gardner calculation. For convenience, we show here how the temporal sequence storage problem be mapped to the storage of fixed-point
also demon-
storage.
Due the threshold activation, the indicator function be written
T-1
x(E) = II
1/2 e
Σ
e++
Jo
é
(59)
t=0
The total volume will be given by
1 V = 7 / dp))(E),
duu
II
dJiz8
Σ S3 P )
Z =
-//Il dps
(60)
Alight
Alisti
After expressing the Heaviside step function using its Fourier representation, the expression for the volume can be seen to factorize into a product
V
(61)
where the volume V; is calculated over all entries in fixed row of the connectivity matrix J:
V.
dus / duc
1
/
I[
dr exp t=0
ixa (xc -
N-1/2 1/2 Σ st
J +
(xxe)) ))
(62)
In order to calculate the disorder (pattern) average of log V,, it necessary to introduce replicas and calculate (V.*) and subsequently take - 0. The replicated volume is written
/
e++ Jo +c +x))
du duddxfdX exp
ix
X - 1-2
(63)
Z
aiiti
Averaging over random patterns will introduce into the integral the term proportional to
II 1=0
II
cos
1/2
N
1/2
Σ
2 Jo
(64)
20

---

Under review as a conference paper at ICLR 2023
This is the point where make the mapping directly onto Gardner's calculation. Notice that after disorder averaging, the integrand factorizes into a product of terms at different times. This is identical to the factorization for different fixed-point patterns in Gardner 19888. This demonstrates that the equivalence between the volumes for fixed-point storage and temporal sequence storage is non-perturbative, and valid for any N. Technically, [Talor (199T); Bressloff & Taylor (1992) demonstrate the equivalence in the large N setting. Thus, the calculation proceeds as in the original work, but with the the total trajectory length I + 1 replacing the number of patterns p. This yields the critical capacity as at T/NN 2. In other words, the maximal length of trajectory scales as 2N.
D EXPERIMENT DETAILS D. CODE
All of the networks presented in this work (vanilla RNN, mGRU, GRU, LSTM. LEM, nODE and gnODE) are implemented with our Julia (Bezanson et al. 2017) package, RNNTools.jl. This pack- age is based Flux.jl Differentialliquationn,, and Flux. (Innes 2018 Rackauckas & Nie 2017
Rackauckas et al. 2020).
D.2 GATING ARCHITECTURE
We let the gating function G. (h, x)) to be a ( + U,æ + b9) (i.e., b, = 1) for Sections 5.2 5.4.2 in the main text. We default to this architecture unless otherwise noted. For Section 5.4.3 we assumed that Lg = 2. Anecdotally, architectures with Lg- = 2 appears to perform better than architectures with L = 1.
D.3 CHOICE DISCRETIZATION
In our experiments, we choose to discretize our networks (vanilla RNN, mGRU, GRU, LSTM, nODE and gnODE) using the canonical forward Euler method, and the LEM with the forward- backward Euler method in |Rusch et al. (202T) (we also present results for LEM discretized with the forward Euler method in the corresponding Sections in the Appendix). While the optimal choice of discretization method may depend on the problem, find that the simple Euler solver achieve strong performance while taking less training time than an adaptive solver in our experiments. Of- ten, the number of function evaluations (NFEs) in nODE training for adaptive schemes, and several regularization methods have been introduced to reduce NFEs (Kelly al.) 2020) Ghosh et al. 2020 Finlay et al. 2020 Pal et al., 2021). On the other hand, we can control the NFEs explicitly by changing the timestep At in a fixed-timestep solver, such the Euler method. While the Euler method does not have guarantees on the growth of error, it may in fact allow representing more functions compared to adaptive methods that provide such guarantees, precisely because of the errors from the discretization (Dupont et al., 2019). We do not lose the benefit of being able to train nODEs on irregularly-sapple Fuler solver. For the n-bit flip-flop task in Section 5.2 changing the Fuler method (used for pre- senting results in the main text) to the Tsitouras 5/4 Runge-Kurt Kutta method did not make significant qualitative difference. For fitting our networks to the OU trajectory in Section 5.3, having an explicit control over the NFEs is crucial for a fair comparison, and the Euler solver was the natural choice. We also see that Euler discretization was sufficient to achieve good performances on the tasks in Section 5.4.1 and Section (5.22 which involve seaguulrly-saaplee trajectories. For Section [.4.1] it is interesting to see that our Euler-discretized mGRU and GRU show accuracies that are higher than the accuracies of GRU-ODEs (De Brouwer al.) 2019 Jordan al. (02T) (which use a modern adaptive solver) reported in Kidger et al. (2020). This suggests that the Euler discretization (which does not necessarily assume At can be a fast, practical alternative to adaptive methods.
become extremely large during
time series when we use the
D.4 CHOICE OF ADIOINT
For all networks consider, we backpropagate through the operations of the solver-thai is, use the cilssrrtizze-heenoptiiiie approach, as is standard in training a RNN, instead of using the "optimize-theenddiseeiiii approach used in Chen et al. (2018) to train nODEs. A few studies show
21

---

Under review as a conference paper at ICLR 2023
that the former produces more accurate gradients than the latter and can yield better performances
Gholami et al. 2019 Onken & Ruthotto 2020).
E BIBIT FLIP--I.OP TASK
For all versions of the n-bit flip-flop task in this section, the total length of each trial was 1s, binned into 10ms bins. Thus each trial had 100 time-bins. The width of each pulse was set to be 20ms. 600 trials were generated total, where 500 trials were used for training and the remaining 100 trials were used for validation.
Networks considered in this section (vanilla RNN, mGRU, GRU, nODE and gnODE) were initial- ized with Glorot uniform initialization (Glorot & Bengio 2010) with zero bias. T = 0.01s in all of the networks. We used A (Loshchilov & Hutter 2019) for training.
E.1 FIXED- 3-B B-BIT FLIP-FLOP TASK
This is the version of the task that was originally introduced in Sussillo & Barak (2013). For this task, determined the total number of pulses (summed across n channels) on each trial by sampling number k from the Poisson distribution with mean 12. We then randomly chose l: indices from to 100 without replacement. These k indices were the indices at which the pulses occur. For each of the k indices. we randomly chose which one of the channels the pulse will occur. Then for the channel where the pulse appears, we chose either +1 or -1 randomly as the value to be taken by the pulse (Figure(AA)
We trained our networks on 500 trials of this task for 200 epochs. The initial states of the networks were not learned, and were initialized with hg ~ (0, E), where Σ = N+1
2 I was the variance.
For vanilla RNN, mGRU, GRU, we varied the phase-space dimension N = {6. 12, 18), where for each N, we used the learning rate 17 = 10-2 2, rate of weight decay A, = 10 and the batch size B 100. For nODE and gnODE, we similarly varied N {6, 12, 18). and used 17 = 10 w = 10 and B = 100. For nODE and gnODE, Fg had 3 hidden layers with 100 units each layer (i.e., L 4 and H 100). We logged the validation MSE traces of mGRU, GRU and gnODE of N = 6 over 200 epochs (or 200 x (500/100) = 1000 iterations), and found that mGRU, GRU and gnODE all achieved validation MSEs < 0.01 at least at some point over the 200 epochs. Similarly, we logged the validation MSE traces of vanilla RNN and nODE of N = 6. These networks reached minimum validation MSEs of 0.033 and 0.016. respectively. over the 200 epochs. All networks reached minimum validation MSES < 0.01 during 200 epochs when N For further analyses of the trained networks (e-g. performing PCA over the trajectories taken by the networks, and finding the fixed points of the networks), we used the set of parameters that achieved the minimum validation MSES over the 200 epochs. All networks, when the reached minimum validation MSE was 0.01, used similar strategies for this task the networks created 8 stable fixed points to solve the task, where each of the stable fixed points represented each output that the networks should take (Figure 6C for vanilla RNN; other networks not shown). For details how the fixed points were found, see Section E.3
(12, 18} (Figure (B).
E.2 CRIIBBE--AMPITURE 3-BIT FLIP-FLOP TASK
We determined when the pulses occur and in what channel the pulses occur in the same way as Section E. Then for the channel where the pulse appears, we drew sample m from U[-1, and let vn be the value to be taken by the pulse Trigure[A in main text).
To ensure fair comparisons across different networks (vanilla RNN, mGRU, GRU, nODE and gn- ODE), for each network, we ran 3 x 3 x 3 27 different configurations of (1), w. B), where 7 = {10-4, ,10-3, 10 10-2), J).00 w = {10-3, 10 10-1} and B = {10, 50, 100). For each network and each configuration, we trained for 600 epochs, and determined the set of parameters that gives the minimum validation MSE over the 600 epochs. Each circle in Figure I is the minimum validation MSE achieved over 600 epochs for single configuration, with total of 27 circles for each network. For nODE and gnODE, Fg had 3 hidden layers with 100 units each layer (i.c., = and H 100). We let N be either 6 (Figure
[ TA-C) or 100 (Figure RD.
22

---

Under review as a conference paper at ICLR 2023
A
Example traces
M
hc
B
dim(h)
dim(h) 12
dim(h) 18
1.2
RNN
1.0
mGRU
Validation MSE
GRU
0.8
nODE
gnODE
0.4
0.0
1000
0
1000
# of iterations
# iterations
iterations
Figure 6: Networks performing the original exed-amplitude 3-bit flip-flop task (Sussillo & Barak
2013). (A) An example validation trial with inputs in each channel shown in black, and the trained vanilla RNN traces maintaining the previous pulse value shown in colors. (B) Validation loss traces function of the number of iterations. (C) The first 3 principal components of the vanilla RNN
(N = 18) trajectories and fixed points. Cyan indicates stable fixed point. Magenta indicates unstable
fixed point.
E.3 PIIND-POINT FINDER
The finder should find some which satisfies h To find such h, we define some function f such that h f(h). In the case of nODE, for example, f(h (1/r) (-h h (2 =+))), where assume that az(1) (t) = O, and 0 is the set of trained parameters of the nODE. We used Newton's 's method (implemented in Julia's NLsolve.j] package; MMgensen & Riseth (2018)
find the root of the nonlinear function f (ft). From starting point, we ran the method for 100 iterations, and terminated whenever (h) ||| 0.01. Choosing what starting point to use can be important, especially when N is large. Following (ussillo & Barak (2013), we used points in the trajectories taken by the network the validation trials as the starting points the finder.
detected fixed points by running the finder 10, 000 times, each with 10, 000 different starting points. Once we detect some h that satisfies checked whether each element h; for all {1, 2, ..., N\ satisfy 2q < h, < 2r to
||hl| < 0.01, ensure that the detected fixed (or slow) point is not too far from the trajectories taken by the network. Here, q min, (min, (y:)) where y is of the points in the trajectories taken by the network in the validation trials. Similarly, r = max, (axx,yymxx, (max, (y:)).
We also explored different criterion to identify fixed points that near the latent trajectories whenever the identified fixed point less than 1 in Euclidean distance from any of the points actually traversed by the networks, include the fixed point in the plot. Even for this criterion, still saw a result similar to what was presented in Figure
E.4 STABILITY OF FIXED POINTS
Figure 1 suggests that the vanilla RNN (N 100) may be reaching its solution using a com- bination of marginally-stabl and stable fixed points, while the gated networks (mGRU, GRU and gnODE) mostly rely marginally-stabe fixed points.
We further projected the 100-dimensional fixed points of the vanilla RNN to the 3-dimensional PC space and found that the unstable fixed points scattered around the stable fixed points, suggest-
23

---

Under review as a conference paper at ICLR 2023
40
maxRe(A(O1)
mGRU
gnODE
Figure 7: Networks assuming N 6 performing the variable- amplitude 3-bit flip-flop task. Each circle is the spectral abscissa of the Jacobian evaluated at a detected fixed point. Bold horizontal lines indicate medians.
ing that the unstable fixed points may be facilitating the network to fall into one of the stable or marginally stable fixed points.
We did similar analysis for networks assuming N and find similar results (Figure 7 The medians and quartiles of the plotted circles in Figure T of the main text and those of Figure are provided in Table
Table We compute the spectral abscissa max
max(Re (Re [A(D))) at a numerically-detectel fixed point. We provide below the medians and quartiles of the distribution of max(Re (A(D))).
(a) N 100
(b) N
Model
Median
Quartiles
Model
Median
Quartiles
RNN
2.587
1.999 6.360]
RNN
2.750
[-0.768, 4.703]
mGRU
-0.472
[0.473, -0.470]
mGRU
0.004
[0.002, 0.009]
GRU
-0.033
-0.474, 0.070]
GRU
0.005
[0.001, 0.023]
nODE 4.988
-1.627,9.45
nODE
2.626
[-2.030, 13.151]
gnODE
0.842
-0.461, 1.698)
gnODE
0.002
[0.002, 0.003]
E.5
THE FAMILY OF 2-B1T FLIP-FLOP TASKS
E.5.1 4 STABLE FIXED POINTS
We determined when the pulses occur and in what channel the pulses occur in the same way as Section E. except that now n = 2. We trained the gnODE for 200 epochs with 27 different hyperparameter configurations, similar to Section E.2 The initial state of the gnODE was learned the initial state was assumed to be an affine transformation of the input the first time-bin. The gnODE's Fg had 2 hidden layers with 316 units each layer (i.e., L 3 and = 316). For Figure A used the gnODE that reached the lowest validation MSE (3.203e-8) among the 27 different runs.
E.5.2 SQUARE ATTRACTOR
We determined when the pulses occur and in what channel the pulses occur in the same way as Section E.2 except that now n = (Figure [A). We trained all networks (vanilla RNN, mGRU, GRU, nODE and gnODE) for 200 epochs, each with 27 different hyperparameter configurations, similar to Section
(Figure BB) 3 The initial states of the networks were learned - the initial state assumed to be affine transformation of the input at the first time-bin. For nODE and gnODE, F, had 2 hidden layers with 316 units each layer (i.e., = 3 and H = 316). For Figure we used the gnODE that reached the lowest validation MSE (0.008) among the 27 different configrrations.
24

---

Under review as a conference paper at ICLR 2023
A
B
0.25
1
Validation MSB
Channel
Channel
0.15
0.05
grÔDE
Figure 8: Networks assuming = performing the square 2-bit flip-flop task (Section E.5.2 (A) An example validation trial with inputs in each channel shown in black, and the trained gn- ODE traces maintaining the previous pulse value shown in colors. (B) For each network, we tried 27 different hyperparameter configurations. Each circle represents the minimum validation MSE achieved during 200 epochs of training. Circles with black edges represent the minimum out of the 27 configurations.
Charcll
Channel
Figure 9: The gnODE assuming N performing the rectangle and the disk 2-bit flip-flop task. (A) An example validation trial on the rectangle task with inputs in each channel shown in black, and the trained gnODE traces maintaining the previous pulse value shown in colors. (B) Same as (A) but for the disk task.
E.5.3 RECTANGLE ATTRACTOR
We determined when the pulses occur and in what channel the pulses occur in the same way as Section E.5.2 except that the pulse value for Channel was drawn from 2 2]. while the pulse value for Channel 2 was drawn from U[- (Figure 9A). For Figure we used the gnODE that reached the lowest validation MSE (0.0137) among the different configurations. Architecture used for gnODE for this task was the same the one used in Section E.5.2
E.5.4 DISK ATTRACTOR
We determined the total number of pulses (summed across n channels) on each trial by sampling number i from the Poisson distribution with mean 6. We then randomly chose indices from 1 to 100 without replacement. These k indices were the indices at which the pulses occur. For each of the & indices, where c; is the pulse value in Channel i. Thus the input pulses in n channels appear at the same time. For Figure
drew random samples C1. C2. ... Cn which satisfy 1 V c +... + cf 2,
used the gnODE that reached the lowest validation MSE (1.910e-4) among the 27 different configurations. Architecture used for gnODE for this task was the same as the one used in Section E.5.2
E.5.5 RING ATTRACTOR
We determined when the pulses occur in way similar to Section E.5.5 However, k was drawn from the Poisson distribution with mean 12, not 6. Also, the constraint that C1. C2. Cn had to satisfy was / +
c 2.
We trained the gnODE for 200 epochs with 27 different hyperparameter configurations this task. Architecture used for gnODE for this task was the same as the one used in Section [.3.2] Figure I0 shows the flow fields of gnODE trained this task. Figure 10A shows the flow field of gnODE
25

---

Under review as a conference paper at ICLR 2023
when we let the initial state that the network should take (in the output space) be (0,0) (validation MSE: 8.345e-5). We see a structure that is similar to what we see when we train the gnODE on the disk 2-bit flip-flop task (Section E.5.4). When we instead let the initial state that the network should take be 0), we see a more ring-like structure near (2. 0) (Figure 103: validation MSE: 0.007). We also tried changing the initial state of gnODE from (0,0) to (1.5, for the gnODE trained on
the disk 2-bit flip-flop task (Section E.5.4), but did not see a significant qualitative difference.
A
B
2-
Channel 2
1-
Channel
0
-1
-1
-2
-1 0 Channel 1
-2
0
i
2
Channel
Figure 10: Flow fields of gnODE with N = 2 performing the ring 2-bit flip-flop task. (A) The initial state of gnODE is set to be (0,0), 0). (B) The initial state of gnODE is set to be (2, 0).
E.6 How ARE CONTINUOUS ATTRACTORS GENERATED IN THESE TASKS?
When we evaluated the gate output G.p x) (with Lg = 1) and - h + Fo x) on several points inside the attractor, the norm of -h Fo a)) generally closer than the norm of the gate output. Thus, after training, it may be possible for gnODE to do the tasks without the gate. However, this does not mean that gates do not help with training/peffommncce For the points evaluated, we found that each component of the gate had values ~ 0.1, and this helps - + Fo x) become even closer zero.
F FITTING AN ORNSTEIN-UHLENDECC TRAJECTORY
The task that the network has to perform is to perfectly fit finite number of samples from an -aaeebn-Uhlenneck (OU) process,
Toudz = dou + ædt + doudw,
(65)
where w is a Wiener process. In our analysis. we set dim(z) = 30, TOU = 1s, Aou -11, a()) 1. and oou and generate single trajectory from this process for 100s using the SOSRI method and sample at every 1s of this trajectory. This gives us total of 100 samples from this OU process. We trained our networks on a single trajectory of these 100 samples in single batch, using AdamW [Loshchilov & Hutter 2019). We trained vanilla RNN, mGRU, GRU, nODE and gnODE on this task. The network weights were initialized with Glorot Normal initialization (Glorot & Bengio, 2010], and biases were initialized with zero-mean Gaussian with variance 10-o. The networks received st)) = as their inputs. Note that (i)(x) = dim(z) = 30. The initial states of the networks were learned the initial state was assumed to be an affine transformation of the input at the first tine-bin. For the vanilla RNN, mGRU and GRU, we systematically varied the phase- space dimension N and 7 of the model. For the nODE and gnODE, number of hidden layers = I,, Fg and the number of units N₂ in each hidden layer of Fe. along with N and r. We assumed that the number of units N₂ is the same across the hidden layers (i.e., H N₁ =.. - N
set Lg 1, and varied the
To ensure we using the appropriate learning rate and the rate of weight decay, trained each network with 9 different combinations of the learning rate and the rate of weight decay, and picked the best model out of the that reached the lowest training loss anytime during the 2000 epochs of training. The training MSE values plotted in Figure 3 of the main text consider only the lowest training losses out of the 9. With the learning rates and rates of weight decay determined, we ran
Thii the default SDE solver in the DifferentialFquationnsjj. package in Julia (Rackauckas Nie 2017
26

---

Under review as a conference paper at ICLR 2023
the experiment times with different random seeds. We chose the learning rate from one of 10 10 and 10 10-2, and the rate of weight decay from 10 3 10 and 10
We that for all networks, when the model is closer to yoU 1, generally achieve lower training MSEs (Figure 3D and Figure TTp. This confirms our intuition that networks perform best when their timescales match correlation time of the data.
RNN
B
MGRU
C
GRU
D
nODE
0.4
0.4
0.4
0.3
108
03
Taining 100
0.3
0.3
Vainnn
Talnno
02
02
02
0.2
0.1
10s
0.1
0.1
0.1
0.0
T-
0.0
0.0
0.0
10
10
10
10a
10
10
30
10a
Phase-space dimension
Phase-apazz dimersson
Phase-spece dimersion
Prave-ipace dimansion
10s, 30s, 100s} fitting samples from the OU traje- Figure 11: (A) Vanilla RNN assuming {1s, 10s,
tory. (B) mGRU. (C) GRU. (D) nODE with 1 hidden layer (H N₁ 316).
We also see that, for nODE and gnODE, when increase the number of units H in each hidden layer, the networks become more expressive (Figure 3C and Figure 12]. When model T = 1s, nODEs tend to be more expressive than vanilla RNNs when the phase-space dimension is low, but see that as we increase the phase-space dimension, RNNs become more expressive (Figure -B). increase model т, nODEs become consistently more expressive than RNNs when H [Figure[2[cDD T2C-D and Figure A-).. We find that gnODEs are consistently more expressive than GRUs across different model T' E [1s, 10s, 30s, 100s) }) and different numbers of phase-space dimensions N (Figure BC and Figure I2E and Figure 13 This confirms our intuition that increasing networks have to solve, and that, because h evolves very slowly for very large т, this places greater burden on Fe.
However, is sufficiently large (Figure
equivalent effectively increasing the difficulty of the task that the
A
B
RMN
=
0.4
RODE.
Training MSE
0.3
Traiming MSE
0.3
n
NODE, 316
RODE, - 1000
0.2
GRU
0.1
== 1s
0.1
r=
15
gnODE, # 10
gnODE,
0.0
L=
0.0
gnODE.N= 100
gnODE. 316
Phase-space dimension
Phase Phase-space dimension
gnODE,
D
0.45
0.35
Training MSE
0.40
Training MSF
0.350
0.375
Training MSE
0.25
0.30
0.35
0.325
0.20
0.30
0.15
0.10
T= 30s
0.250
==
0.05
1s
=
0.225
L=
0.00
L
10'
102
10a
10" Phase-space dimension
10'
10'
Phase-space dimension
Phase-space dimension
Figure 12: Increasing the number of units H each hidden layer of F9 increases practical expres- sivity of networks.
(A) Training MSEs of nODE with 1 hidden layer (i.e.. L = 2). assuming 1s. (B) Training MSEs of nODE with 2 hidden layers (i.e., L 3), assuming
1s. (C) Training
MSEs of nODE with I. 2 and 30s. (D) Training MSEs of nODE with I. = 3 and 30s. (E) Training MSEs of gnODE with and T 1s.
27

---

Under review as a conference paper at ICLR 2023
A
0.44
B
C
0.4
RNN
0.40
MGRU
0.42
GRU
0.35
0.3
Training MSE
0.40
Itaining MSE
Training MSE
NODE
gnODE
0.36
0.20
0.34
0.15
0.32
N = 10
N = 32
0.0
N 100
0.10
10a
10
10
10a
10
10 109
10a
10a
Figure 13: (A) N 10 across all networks. nODE and goODE has Fa with 1 hidden layer, where H N1 1000. (B) N 32 across all networks. nODE and gnODE has Fg with 1 hidden layer, where = N2 = 1000. (C) N = 100 across all networks. nODE and gnODE has F6 with 1 hidden layer, where N1 1000.
Lastly, we did not observe that increasing the number of hidden layers in Fg significantly increases expressivity of nODE and gnODE (Figure [4]. This is perhaps related the observation that for regression problems, width matters much more than depth (Radhakrishnan et al. 2022).
t= 15
10s
100s
A
B
D
4.26
0.31
0.30
4.30
noDE
02
22
2033
0,27
0.32
5.20
0.30
6:35
12
107
10
1IF
Phase-pece dimimioo
E
F
G
0,375
H 0.375
0.550
0.35
0.32
0.325
gRODE
025
0.300
8 0,30
R0275
0.75
0.22
2 0239
0.25
6:5
0.12
0.80
0000
der
10
Phose SpRC0 deversion
1P
12
space dinersion
19
Figure 14: Increasing the number of hidden layers in Fg does not significantly increase practical 1000, assuming T 1s. Mint indicates nODE with 1 hidden layer (H Nj 1000), and pink indicates nODE with 2 hidden layers (H N1
expressivity of networks. (A) nODE with
N2 = 1000). (B) nODE with H = 1000, assuming т = 3s. (C) nODE with И = 1000, assuming
10s. (D) nODE with H 1000, assuming 100s. (E) gnODE with H 1000, assuming
т = 1s. (F) gnODE with H = 1000, assuming = 3s. (G) gnODE with = 1000, assuming
10s. (H) gnODE with H 1000, assuming 100s.
G REAL-WORLDS TASKS
G.1 NETWORK INITIALIZATIONS
For RNNs, mGRUs, GRUs, LSTMs and LEMs, we applied either the Kaiming or the Glorot Normal initialization. Both of these initializations should already give criticality. For nODEs and gnODEs, the standard Kaiming or Glorot initialization do not give criticality. Whenever we applied either the Glorot or Kaiming Normal initialization to these networks, the final nonlinearity dn. of F& was set to be Z. Whenever we applied the critical initialization in Appendix A dn was set to be tanh. We found that, for all the real-world datasets we consider (Sections 5.4. 5.4.3). having no tanh degrades performance, consistent with the observation in (Kidger et al. 2020). We suspect that this is due to the fact that a nODE which does not have does not tanh as its final nonlinearity does not have chaotic phase (Appendix A).
28

---

Under review as a conference paper at ICLR 2023
G.2 LATIN ALPHABET CHARACTER TRAJECTORY CLASSIFICATION
We trained mGRU, GRU, LSTM. LEM, nODE and gnODE on the task of classifying 20 different Latin alphahets based on the pen-tip trajectories and forces applied to the tip. This dataset had 2858 trials total (2000 trials for training, 429 trials for validation and 429 trials for testing), with each trial being a time-series with 182 time-bins. The time-series was 4-dimensional, with the first dimension being the time stamp, the second being the position of the pen, the third being the y position, and the fourth being the pen tip force. We took the preprocessed data as is from the Neural CDE repository. Further details on the dataset and the preprocessing step can be found in the repository and Kidger et al. (2020). We used the "30% dropped" dataset (30%, because 30% of the samples in the trajectories were randomly dropped to make the time-series irregularyysaampledd to determine the best set of hyperparameters for each network (using grid search), and used the same set of hyperparameters for the "50% dropped" and "70% dropped" datasets.
We assumed each time-bin in the data is Is-long (thus each trial is 182s long), and set 7 At 1s. We trained our networks for the total of 1300 epochs, where we first trained only the first 14 time- bins for 100 epochs, and then the first 28 time-bins for the next 100 epochs, until we reached 182 time-bins. This method of "iteratively growing the fit" is sometimes used to train a RNN (Hafner) 2017) or nODE (Rackauckas et al. 2020). The initial states of the networks were also learned - the initial state was assumed to be an affine transformation of the input at the first time-bin. We determined the set of network parameters that achieves the lowest validation loss over the 1300 epochs, and used this validation loss (the cross entropy loss for this classi- fication task) as the measure of performance nODE and gnODE-v1, to determine the best set of hyperparameters, we performed a grid search over the learning rate 17 E {10-4, 10 10 10-2). rate of weight decay Aw E {10-3, 10 10 10-1), ") initialization scheme (Glorot Normal or Kaiming Normal; biases were always initialized with zero-mean Gaussian with variance 10-9), phase-space dimension N E (32, 100), and the num- ber of units H {100, 316, 1000} in each hidden layer of Fg. For gnODE-v2. we used the same hyperparameter-scunn Equation 51 of Appendix A For nODE, gnODE-v1 and -v.. 2 hidden layers for F6 (i.e.. L, = 3). For gnODE-v1 and gnODE-v2, we only considered L2 = 1. For mGRU and GRU, we performed grid search over the learning rate 1) {10-4, 10-3, 10 10-3), rate of weight decay Av E {10-3, 10 10-2, 10-1), initialization scheme (Glorot Normal or Kaiming Normal; biases were always initialized with a zero-mean Gaussian with variance 10-6) and the phase-space space dimension N E {100, 316, 1000}. The batch size B was set to be 32 for all networks, following the suggestion in Kidger et al. (2020). With the hyperparameters that achieved the min- imum validation loss, we trained the networks 5 times with different random seeds, and evaluated the test losses from those 5 runs. Table 3 shows the best set of hyperparameters found by the grid search, and Table shows the means and standard deviations of each networks' test accuracies for the "30% dropped". "50% dropped" and "70% dropped" datasets with the hyperparameters in Ta- ble 3 Table A of the main text shows results for the "70% dropped" dataset with the best set of hyperparameters found by the subset of the search space where N is fixed to be 100. The gnODE performance in Table [TA is from gnODE-v2 in Table 4
determine the best set of hyperparameterss. For
space as that of gnODE-v1, except that we applied the critical initialization in only considered FNNs with
We see that the gated networks (mGRU, GRU, LSTM, LEM, gnODE) achieve accuracies similar to that of nCDE reported in Kidger et al. 2020 The gated networks' accuracies are also higher compared to the non-gated network (i.e., nODE). Lastly, we see that mGRU, with fewer parameters, performed similarly to a GRU, confirming that the functional roles of the update and reset gates are similar and the reset gate can be taken out (Krishnamurthy et al. 2022). Similar observations have been made in Ravanelli al. (2018).
G.3 WALKER2D KINEMATIC SIMULATION PREDICTION
In this experiment, the networks (mGRU, GRU, LSTM, LEM, nODE and gnODE) were given the task of predicting what the future dynamics should be given data samples up until the current time- point in time series generated from the MuJoCo physics engine kinematic simulations (Todoroy et al. 2012).
This dataset had 12, 893 trials total (9684 trials for training, 1272 trials for validation and 1937 trials for testing), with each trial being time-series with 84 time-bins. The time-series was 17-
29

---

Under review as a conference paper at ICLR 2023
Table 3: The set of hyperparameters that achieved best performance on the classification of Latin alphabet character trajectories. Along with the best number of parameters identified through the hyperparameter-sent
the minimum and maximum number of parameters possible in the search space are also shown in the table. G is Glorot Normal and K is Kaiming Normal initialization. C' is the critical initialization proposed in Appendix Euler.
FE
Forward Euler, FBE- Forward-Backwe -aakward
Model
# of parameters (min/max)
Am, N, Ne,, init)
mGRU
23, 520 (23, 520/2, 035, 020)
1103,10-1 100, N/A, G)
GRU
312, 228 (34, 020/3, 040, 020)
(10-3,10-3,310 N/A, G)
LSTM
4, 0,050, 020 (45, 020/4, 050, 020)
1003,,103,1000 N/A, G)
LEM (FE)
415,24 244 (45, 020/4, 050, 020)
(10-3,10-3,311 N/A, K)
LEM (FBE)
4, 050, 020 (45, 020/4, 050, 020)
(10-3,1033, 1000, N/A, G)
nODE
33, 220 (17, 852/1, 620)
(10-A.10-3, 100, 100, K)
gnODE-v1
178, 072 (19, 036/1, 120)
(10-3,10-3 316, K)
gnODE-v2
1,219, 120 (19, 036/1, 219, 120)
(10-3,10-3 100, 1000, C)
Table 4: Test accuracy for classification of Latin alphabet character trajectories (mean + std, error bars computed from training the networks 5 times with different random seeds) with the hyperpa- rameters in Table 3
Model
Test Accuracy
Name
30% dropped
50% dropped
70% dropped
mGRU
0.987 + 0.005
0.987 + 0.001
0.983 + 0.002
GRU
0.990 = 0.001
0.990 + 0.004
0.987 + 0.003
LSTM
0.990 = 0.002
0.990 = 0.004
0.990 = 0.002
LEM (FE)
0.986 + 0.006
0.990 + 0.002
0.987 + 0.005
LEM (FBE)
0.990 + 0.004
0.991 0.004
0.987 0.001
nODE
0.924 = 0.095
0.807
0.898 + 0.089
gnODE-vl
0.984 + 0.005
0.981 0.003
0.986 + 0.003
gnODE-v2
0.987 = 0.002
0.987 0.004
0.986 + 0.005
dimensional. We took the preprocessed data as is from the ODE-LSTM repository, where 10% of the samples were dropped along the trajectories, and 1 1% of all actions were overwritten by random actions (Lechner & Hasani 2020). Further details on the dataset and the preprocessing step can be found in the repository and Lechner & Hasani (2020).
We assumed each time-bin bin in the data is 1s-long (thus each trial is 84s long), and set 7 = At 1s. Our networks received the 17-dimensional time-series, along with an extra dimension specifying the time stamp, as input, and emitted their predictions of what the 17-dimensional state will be on the very next time-step. We trained networks for the total of 700 epochs, where first trained only the first 14 time-bins for 100 epochs, and then the first 28 time-bins for the next 100 epochs, until reach 84 time-bins. When train the full 84 time-bins, we train for 200 epochs instead of 100 epochs, thus making the total 700 epochs. The initial states of the networks were learned - the initial state was assumed be an affine transformation of the input the first time-bin. We determined the set of network parameters that achieves the lowest validation loss over the 700 epochs, and used this validation loss (the MSE loss for this prediction task) as the measure of performance to determine the best set of hyperparameters. We performed the same hyperparameter grid search as in Section G.2 The batch size was set to be 256 for all networks, following the suggestion in Lechner & Hasani (2020). With the hyperparameters that achieved the minimum validation loss, we trained the networks times with different random seeds, and evaluated the test losses from those runs. Table 5 the means and standard deviations of each networks' test MSEs along with the hyperparameters found by grid search. The gnODE performance in Table B is from gnODE-v2 in Table 5
30

---

Under review as a conference paper at ICLR 2023
Table 5: Test MSE for Walker2D (mean + std, error bars computed from training the networks times with different random seeds).
Model
# of parameters (min/max)
(1,, Aw, N1, Ne, init)
Test MSE
mGRU
223, 113 (27, 417/2, 074, 017)
(10-2,10-3, N/A, G)
1.074 + 0.070
GRU
328, 973 (39, 317/3, 093,017
(10-2,10=9,316 N/A, K)
0.772 + 0.028
LSTM
53, 117 (53, 117/117/,,4 137,
(10-2,10~3,00 N/A, K)
0.865 + 0.009
LEM (FE)
440,83 837 (53, 117/4,131,011
(10-2,10-3,0 3,,,,,,,,,,,,, G)
1.030 0.022
LEM (FBE)
440, 837 (53, 131,0117
(10:2,1033,316.VV. G)
0.699 = 0.010
nODE
1, 222,, 717 (19, 201/1,,23,77
(10-3,10=3,100, 1000, G)
0.707 = 0.023
gnODE-v1
1, 086, 833 (21, 233/1, 235, 617)
(10-8,10-1,,3, 1000, K)
0.552 = 0.019
gnODE-v2
1, 2,2, 617 (21, 233/1, 238/1,235,601 617)
1000,,003,100,1000,,, C)
0.588 = 0.003
G.4 SPEECH COMMANDS CLASSIFICATION
We trained mGRU, GRU, LSTM, LEM, nODE and gnODE on the task of classifying ten spoken words, such as " "Stop" and "Go", " based on one-second audio recordings of these words. The dataset is originally from Warden 2018 and preprocessed using the pipeline in the Neural CDE repository (Kidger et al., 2020 There are total 34, 975 time series (70% training, 15% validation, and 15% test data), where each time series has 20 channels of 161
data points. Further details on the dataset and the preprocessing step can be found in the repository and Kidger et al. (2020).
Section 5.3 showed that networks perform best when the timescales assumed by the net- works match correlation time of the data. To further probe this effect, we varied T from (0.006s, 0.062s, 0.621s). We trained our networks for the total of 300 epochs the entire time series. The initial states of the networks were also learned - the initial state was assumed to be an affine transformation of the input at the first time-bin. We determined the set of network pa- rameters that achieves the lowest validation loss over the 300 epochs, and used this validation loss (the cross entropy loss for this classification task) as the measure of performance to determine the best set of hyperparameters. For nODE and gnODE, we tested two different architectures one without the final tanh nonlinearity (nODE-vl and gnODE-v1 in Tables ( and one with the tanh (nODE-vl and gnODE-v2 in Tables 8 and gnODE-v3 in Table 6). To determine the best set of hyperparameters, we performed a grid search over the learning rate 17 E {10- 4 10 ..102 2), rate of weight decay X {10 3 10 10 100'),, }.. phase-space dimension N e {100, 316). and the number of units Ne. E {1000, 2000, 40003 in each hidden layer of Fa. We only considered FNNs with 1 hidden layer where I,, 2. For nODE-vl and gnODE-v1, we applied both the Glorot Normal and Kaiming Normal to determine the best initialization scheme. For nODE-v2, gnODE-v2 and gnODE-v3, we applied the critical initialization (Appendix (A). Biases were always initialized with a zero-mean Gaussian with variance 10-5). For gBODE-v3, the number of units N₂, in the hidden layer of G. was set to be 1000. For mGRU and GRU, we performed a grid search over the learning rate v) E (10-1, 10-3 10-2), rate of weight decay A E (10-3 10-2 10-1), initialization scheme (Glorot Normal or Kaiming Normal; biases were always initialized with a zero-mean Gaussian with variance 10 (0-5), and the phase-space dimension N E {100, 316, 10003. The batch size B was set to be 256 for all networks. With the hyperparameters that achieved the minimum validation loss, we trained the networks times with different random seeds, and evaluated the test losses from those runs. Tables 6 show each networks' test accuracies for 7 {0.006s. 0.062s, 0.621s} with the hy- perparameters found by grid-search. The hyperparameter-searre
space of Tables 7 was the same
as that of 6 6 The results for nODE and gnODE in Table IC are from nODE-v2 in Table and gnODE-v3 with 7 = 0.001, A = 0.001, N = 100, Nr, = 3500 and Ne.- = 1000.
Consistent with Figure 3D, we observe that the increased complexity in F& of nODE/gnODE be- comes more useful as T is increased. When 7 is decreased, the increased complexity in F& does not appear to help. Across all nODE, and makes nODE more robust to changes in T.
{0.006s, 0.062s, 0.621s), gnODE shows an improvement over
When 7 is small (0.006s or 0.062s), nODE/gnODE with the final tanh (i.e., nODE-2//mODD that is critically initialized (following Equation [1 in Appendix (A) showed improvement in perfor- mance compared to nODE/gnODE without the final tanh (i.e., nODE-v1/gnODE.w/), If there is no
31

---

Under review as a conference paper at ICLR 2023
Table 6: Test accuracy (mean = std, error bars computed from training the networks 5 times with different random seeds) for classification of Speech Commands with 7 0.006s.
Model
# of parameters (min/max)
Aw, NN, No., init)
Test Accuracy
mGRU
27, 610 (27, 610/2, 076, 010)
(10-2, 100, N/A, K)
0.809 + 0.018
GRU
3, 098, 010 (39, 810/3, 098, 010)
110-9.103,1000 N/A, K)
0.855 + 0.001
LSTM
306,30 (54, 210/4, 142, 010)
(10-3,10-3, N/A, G)
0.807 + 0.004
LEM (FBE)
4,142, 010 210/4, 142, 010)
(10-3,1033,1000 N/A, K)
0.836 = 0.005
nODE-vl
1,318, 438 (225, 310/2, 626, 438)
(10-4,10-9,316, 2000, G)
0.110 + 0.010
nODE-v2
225,310 310/2, 626, 438)
(10-3.10-3,100, 100,,
0.140 + 0.012
gnODE-v1
903,510
(237,510/2, 733, 246)
(10-3:10-3,100,4000 4000,
0.795 + 0.005
gnODE-v2
771, 246 (237,510) 733, 246)
(10-3,10-3.316,
1000, C)
0.815 + 0.004
gnODE-v3
1,972, 754 (447, 410/3, 280, 754)
(10-3,10-) 316,, 1000, C)
0.844 + 0.002
Table 7: Test accuracy for classification of Speech Commands with r 0.062s.
Model
# of parameters
(n, Aw, N, Ne,,, init)
Test Accuracy
mGRU
223, 738
N/A, K
0.796 + 0.002
GRU
330, 546
(10-2,10-3,316 N/A, K)
0.809 = 0.006
LEM (FBE)
444, 306
(10-2,10-3,316, N/A, K)
0.785 = 0.005
nODE-v1
225,31 310
(10-3,10-3,100, 1000, K)
0.246 + 0.050
nODE-v2
1,318, 438
(10-4,10-3,316, 2000, C)
0.169 + 0.015
gnODE-vl
1,425. 246
(10-3,10-3,31) 2000, G)
0.790 0.018
gnODE-v2
237,510 510
(10-2,10-8,100 100, 1000, C)
0.769 + 0.072
tanh at the end, critical initialization does not appear to help. This may be related to the observation in Appendix A that there is no chaotic regime for networks that only have ReLU activations. This consistent with the observation in Kidger et al. (2020) that applying tanh at the end improves performance of a nCDE.
For gnODE, using Gpp. with Lg gave better performance than using Go with La (see gnODE-v2 and -v3 in Table 6
G.5 LEM PERFORMANCE ON WORLD TASKS
Our results with LEM the real-world tasks suggest that the performance of an LEM is similar to an mGRU or a GRU, especially when LEM is discretized with the forward Euler method, instead of the forward-backward Euler method suggested in Rusch al. (202T) (see Appendix B for
possible explanation and see Appendix G.2 G.3 for LEM results with the forward Euler method). In Rusch al. (202 l) At in LEM treated as a hyperparameter, while At for other models is taken to be equal to I. For fairer comparisons, At could have been equal across all models. Our
experiment in Section 5.3] and particularly F makes exactly the point that if we change (which is effectively the same as changing At in Rusch et al.) (2021) we are effectively giving the models different problems to solve, with different timescales. our experiments, we set 7 to be the same across all models compared (see also Appendix G.r discussion of how changing 7 affects network performances on a real-world task). Rusch Mishra (2020) discusses training At, and it would be interesting if making trainable in models leads improvements in performance, though this would be beyond the scope of this work.
H DEFINITION OF CONTINUOUS ATTRACTOR
We use the terminology "continuous attractor" in the main text, which is very common in neuro- science, but possibly less known in the broader dynamical systems and machine learning commu- nities. In this Appendix, we give a precise definition and attempt to establish a connection between continuous attractors and center manifold theory (Carr 1981).
32

---

Under review as a conference paper at ICLR 2023
Table 8: Test accuracy for classification of Speech Commands with 0.621s.
Model
of parameters
(1,, Am, N, Ne,, init)
Test Accuracy
mGRU
2, 076, 010
(10-2,10-) 1,000,0000,,,,,,,,,,111110 N/A, G)
0.733 0.005
GRU
3, 3,098, 010
(10-2,10-3, 1000, N/A, G)
0.743 + 0.008
LEM (FBE)
4,142,010
(10-2,10-3, 1000, N/A, K)
0.713 + 0.003
nODE-v1
447,310
(10-2,10=3, 2000, K)
0.710 0.020
nODE-v2
664,
(10-2,10-3,316 1000, C)
0.725 + 0.009
gnODE-v1
771,246
(10-2.10-3,316, 1000, G)
0.720 + 0.039
gnODE-v2
771,246
100-210-3,3000 1000, C)
0.762 + 0.005
By a continuous attractor, we mean a connected manifold of fixed points. More precisely, first order ODE à = f(z) for x E R", is said to have continuous attractor $ if the following conditions hold:
1. S C IR is a didimnnsion manifold (usually with boundary of dimension d 1) embedded in the full phase space, d S n.
2. Vz S. f(x) = O.
3. Defining the Jacobian D(z) = (af (x), and the spectral abscissa (or largest real part of the spectrum) v ())())). then for S, the spectral abscissa ()())) 0.
Unlike limit cycles or chaotic attractors, the dynamics stationary on the continuous attractor S, since by definition & f(x)
by Item 2. Another almost trivial consequence of the items above are that 5 is an invariant manifold of the dynamics, since for any initial condition x(0) E S. axt) E for all O. Indeed, xxt) z0)) Item ensures that perturbations off the manifold S will decay back toward the manifold, implying it is an attractive manifold.
We now want argue that given Items above, S is also center manifold. Let us now consider the tangent space around a point zu E S. This will be spanned by the d zero mode eigenvectors tj of the Jacobian:
D 0, k l,...,d. d.
(66)
In other words, we have that f(x + eth) = for & = 1,...d d. Let us consider a decomposition of the displacements from an:
J To + Σ
ugth
Σ kmd++1
(67)
Here we use the fact that nonzero modes * will be normal to the manifold. Now we have new global coordinates which align with the tangent space (ux) and transverse space (yk). However, in these coordinates, the constraint for the attractor manifold S becomes
v()) 0.
(68)
We now seek to determine y(0) and y0) By construction, y(0) = 0. Taking derivatives of the implicit equation for S gives
0f, y(u)) Dup
Oyw n = 0. Ouk
Σ Dis (0x,/Oux) = Σ Dij 19
Σ
(69)
k=d+1
Since * 0, this implies
1
Oyk
Σ
Eiin
=
(70)
Ouk
#=d+1
33

---

Under review as a conference paper at ICLR 2023
Since Dr * 0 by construction, must have that Ou Vu 0. which is what wanted to show. Therefore, the attractor manifold is an invariant manifold that is parameterized by a function (u) which satisfies y(0) Dy(0) = 0. According to (Carr 1981), this means S is also center manifold.
I COMPUTING INFRASTRUCTURE
We trained our networks on a Dell computer cluster with 320 NVIDIA P100 GPUs across 80 Broad- well nodes. Our networks took up of these nodes, using GPU that has 16GB of memory.
34


Reviews:
[{"review_id": "ShKtGflWeth", "reviewer": "Reviewer_5GnR", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The authors introduce a gated neural ODE, an architecture that extends the neural ODE with a recurrent gating variable. They demonstrate that this architecture can learn continuous attractors, provide somewhat interpretable solutions and perform well for some select benchmarks.", "summary_of_the_review": "Overall, the gnODE is a very interesting architecture that might have the potential to be an interpretable performant architecture. But the paper only uses a narrow set of benchmarks and compares with very few other relevant architectures, so the generality and scalability of the model proposed here is not clear.", "strength_and_weaknesses": "## Strengths:\n\n- The method is a simple extension of neural ODEs that seem to have many useful properties and benefits.\n- The ability to learn continuous attractors in an interpretable way is interesting.\n- Performs well on considered benchmarks.\n- The expressivity measure seems useful and relevant to the community.\n\n## Weaknesses:\n\n- Very limited set of benchmarks are considered.\n- No comparison with architectures such as LSTMs, LMU, S4 etc.\n- The specific motivation of constructing this architecture is not clear -- is it to provide a better performing architecture? If so comparison on lots more benchmarks and architectures are needed.\n- gnODE seems to be very parameter inefficient compared to other methods here.\n\n## Other questions:\n\n- Fig. 1: why does gnODE have much higher variance than other RNN models?", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "JQSotMamyt", "reviewer": "Reviewer_XxBD", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The manuscript aims to build an accurate, explainable digital twin for the biological neural networks. The neural ODEs have a capability to address complex behaviour. The Authors suggest a gated version to ensure also long term memory effects. The gated systems, the authors find are leading to simpler dynamics, all in increase the trainability, expressivity and interpretability. The dynamics of the gated neural ODE is discussed quite lengthly relating to the other already existing alternatives. \n\nExperimental trials are made for N-bit flip flop task, fitting a finite number of samples from an Ornstein-Uhlenbeck (OU) process, and a 2d walker kinematic simulation predictor among others. The experiments shows that the gated version can utilise lower dimensional phase spaces. The paper provides an extensive Appendix where the performance of the gated system is studied from multiple point of views.\n", "summary_of_the_review": "Excellent, experimental review the performance of gated neural ODEs with some mathematical analysis in the appendices.  A good addition to the toolbox for creating models for multi-scale signals.  The capability to the architecture is demonstrated with a wide set of examples. However, I am waiting for the real process with importance that can be \"understood\" and \"interpreted\" by using this architecture in the fit. For example, a discovery of type of biological neural system that can be characterised and interpreted with this novel tool.", "strength_and_weaknesses": "The paper is comprehensive description of the gated neural ODE and provides  basic intuition of its capabilities. There was one practical example on its capability to fit continuous processes. The strength of the paper is that is describes a novel model that could \"explain\" or \"interpret \"different data sets and provide a model that has the ability to shed light  to the fundaments in the process described by the data, and, for example, show a fixed point structure in the particular data.\n\nThe weak point is that this remains quite an exercise of experimental trainings on selected examples without practical use cases. The question that is hinted at the introduction remains open. Could this be a tool that can reveal some of the algorithms biological neural networks are running. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "ESjfeIniAb", "reviewer": "Reviewer_oDgk", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper advances the applicability and interpretability of neural ODE by injecting a regulatory mechanism to adjust the conditional time scale of the vector field. The authors show improvement on a variety of tasks respect to previous methods, and provide insights about the structure of the fixed points by inspecting the maximum eigenvalue of the Jacobian. ", "summary_of_the_review": "Intriguing paper that I enjoy reading. It shows better performance in a variety of tasks especially when the nODE needs to learn complex random trajectories. Yet, many questions remain in terms of what's the underlying dynamical structure generated in the autonomous regime.", "strength_and_weaknesses": "What the authors call the principle of expressivity is a very interesting phenomena of doing more with less. That reveals itself best fitting random trajectories indicates that gnODEs can potentially thread through unstable saddles as shown in Fig. 3\n\nThe simple proposal of using the dissipative -h term to achieve stability for many initial conditions might be perceived as trivial for people working in dynamical systems. And it's surprising that other have not used before.\n\nIt’d be informative to also plot the time scaling of the vector field in each of the tasks. Is gnODE freezing the dynamics when learning the variable amplitude flip-flop?\n\nCan the authors elaborate on the structure of the marginally-stable fixed points for a given autonomous dynamical system realization? The presence of unstable saddle is interesting because those are the main ingredients of heteroclinic orbits which are structurally unstable but can create sequences of quasi-stable states that are input dependent. This type of dynamics is referred in the literature as winnerless-competition principle.\n \nFor a given size N, gNN and non-time dependent input, how many stable and saddle fixed points are detected? GRU and mGRU appear to learn only stable fixed points which is not an interesting dynamics. gnODE, on the other hand, appear to have a combination of saddles and stable FPs, but I’d like to know if the system has a complex structure of fixed points for non-explicit time dependence on the inputs (autonomous dynamical system). Then, tracking the separatrices of the saddles would enrich our understanding of what gnODE can learn.\n\nCan the authors explain what they mean by continuous attractors? There most common attractors one finds in these types of dissipative dynamical systems are limit cycles or even chaotic attractors. Are the authors referring to center manifolds?\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "4sMKKO1--P", "reviewer": "Reviewer_fH84", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes a Gated Neural ODE (gnODE) model and empirically investigates the trainability, expressivity, and interoperability of the proposed model.", "summary_of_the_review": "This is a half-baked draft with significant issues with clarity. ", "strength_and_weaknesses": "## Strength\n\n1) This paper proposes a new model with an interesting combination between the gating mechanism and Neural ODEs.\n\n\n## Weaknesses\n\n1) The organization of this paper is a bit messy. The paper investigates three largely independent aspects of Neural ODEs with empirical experiments. It is unclear about the synergy between these three aspects. It is also unclear how these aspects relate to \"understanding how the dynamics in biological and artificial neural networks implement the computations\". \n\n2) Several arguments of the paper are not well-supported. For example, it is uncommon to claim that the model is interpretable just because it can have a relatively lower hidden dimension for achieving the same performance compared to other models, especially when this result is obtained on a single synthetic dataset. \n\n3) The empirical comparison is insufficient. The authors compare the proposed method with very closely relevant baselines such as ODE-LSTM by the reported test MSE in the literature. However, the vanilla Neural ODE and GRU models experimented by this paper already have better test MSE than the reported test MSE of ODE-LSTM in the literature, which casts doubt if the comparison is fair.\n\n4) The paper seems to be a half-baked draft. For example, there is no legend in Figure 3B. The template Author Contributions and Acknowledgements are not removed from the draft. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject