Low Tensor Rank Learning of Neural Dynamics
Authors: Arthur Pellegrino
N Alex Cayco Gajic
Angus Chadwick
Venue: NeurIPS 2023 poster (2023)
Keywords: Recurrent Neural Networks
Computational Neuroscience
Neural Data Analysis
Tensor
Learning
TL;DR: Learning in both biological and artificial recurrent neural networks is low tensor rank, enabling the inference of learning-driven changes in population dynamics from large-scale neural recordings.

Abstract: Learning relies on coordinated synaptic changes in recurrently connected populations of neurons. Therefore, understanding the collective evolution of synaptic connectivity over learning is a key challenge in neuroscience and machine learning. In particular, recent work has shown that the weight matrices of task-trained RNNs are typically low rank, but how this low rank structure unfolds over learning is unknown. To address this, we investigate the rank of the 3-tensor formed by the weight matrices throughout learning. By fitting RNNs of varying rank to large-scale neural recordings during a motor learning task, we find that the inferred weights are low-tensor-rank and therefore evolve over a fixed low-dimensional subspace throughout the entire course of learning. We next validate the observation of low-tensor-rank learning on an RNN trained to solve the same task. Finally, we present a set of mathematical results bounding the matrix and tensor ranks of gradient descent learning dynamics which show that low-tensor-rank weights emerge naturally in RNNs trained to solve low-dimensional tasks. Taken together, our findings provide insight on the evolution of population connectivity over learning in both biological and artificial neural networks, and enable reverse engineering of learning-induced changes in recurrent dynamics from large-scale neural recordings.

Low Tensor Rank Learning of Neural Dynamics
Arthur Pellegrino
School Informatics University of Edinburgh pel pellegrino arthurred ac. uk
N Alex Cayco-Gajic* Département d'Etudes Cognitives Ecole Normale Supérieure natasha. cayco. www.iaaaji.aaa.e. fr
Angus Chadwick" School of Informatics University Edinburgh
angus chadwick@ed, ac uk
Abstract
Learning relies coordinated synaptic changes in recurrently connected popu- lations of neurons. Therefore, understanding the collective evolution of synaptic connectivity over learning is key challenge in neuroscience and machine learning. In particular, recent work has shown that the weight matrices of task-trained RNNs are typically low rank, but how this low rank structure unfolds over learning is unknown. To address this, we investigate the rank of the 3-tensor formed by the weight matrices throughout learning. By fitting RNNs of varying rank large- scale neural recordings during a motor learning task, we find that the inferred weights are Now-tensor-rank and therefore evolve over fixed low-dimensional subspace throughout the entire course of learning. We next validate the observation of Now-tensor-rank learning on an RNN trained to solve the same task. Finally, we present set of mathematical results bounding the matrix and tensor ranks of gradi- ent descent learning dynamics which show that low-tensor-rank weights emerge naturally in RNNs trained to solve low- www-dmensional tasks. Taken together, our findings provide insight on the evolution population connectivity over learning in both biological and artificial neural networks, and enable reverse engineering of dearning-induced changes in recurrent dynamics from large-scale neural recordings.
1 Introduction
Populations of neurons perform tasks through their dynamics, and these computations can be un- derstood through the lens of recurrent neural networks (RNNs) [1, 2]. Recent work has shown that RNNs trained on idealized versions behavioural tasks from neuroscience experiments can be reverse-engineered
better understand the dynamical principles by which they perform tasks [3, 4, 5, 6. 7]. RNNs can also be fitted to neural data to infer the latent dynamics that drive neural activity in specific tasks [8, 9, 10, 11]. However,
understanding of task learning in the brain requires methods to understand how these latent dynamics evolve over trials, and to map these computational changes to learning-induced induced changes in neural connectivity [12, 13]. For example, it has been observed that gradient descent fask-trained RNNs tends to drive low-rank weight updates [14]. Yet the structure of learning dynamics itself remains largely unknown, both in the context of gradient-based optimization of neural networks and biological learning in neural systems.
To address this question, consider how RNN dynamics evolve result of structured changes in connectivity over learning. Specifically, we consider changes in RNN connectivity over multiple
These last authors contributed equally.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

trials by assuming that the weight tensor, i.e. the 3-tensor formed by stacking the weight matrix over all trials, has low tensor rank. This allows to identify how distinct components of the weight matrix vary over trials while simultaneously restricting the parameter complexity and benefiting from the interpretable framework of low (matrix) rank RNNs at each trial [15]. Furthermore, by imposing constraint on the covariance of the trial factors we are also able to separate smooth changes over learning from condition-spel specifie variability on individual trials. In contrast to classic tensor decomposition methods (such as PARAFAC), which instead constrain the neural activity itself to be low rank, low-tensor-ranee RNNs (lrRNNs) capture high-dimensional neural activity resulting from nonlinear neural dynamics, while preserving the interpretability of these linear methods through having interpretable low-dimensional factors.
Main contributions. We first apply this method to neural data during a motor learning task [16] to show that the resulting neural activity can be captured with surprisingly low-tensor-rank weight updates. Next, we aim to find intuition for the observation that learning can be low tensor rank by turning to gradient- based optimization, which has recently been shown to be able to explain changes in neural activity patterns over motor learning [17, 18]. Towards this end, we show numerically that an RNN that is trained to perform the same task also results low-tensor-rank learning dynamics. Finally, we provide analytical intuition in the form of upper bounds on both the matrix rank and the tensor rank of gradient- based optimization in RNNs. Ultimately these results provide evidence for low tensor rank learning structure of neural dynamics both in the brain and in RNNs.
2 Low tensor rank learning
The weight tensor. Learning in recurrent neural networks, both biological and artificial, involves the continual update of a weight matrix W E TRNN,, where each element Wa, represents the connectivity from neuron to i. Here, we consider set of K discrete samples of the weight matrix over learning: w(K+1) (Kk) + (k for k E { K - 13. In neuroscience, (W could represent plasticity-indus induced changes in connectivity strength following each trial in the experiment (e.g., perceptual decision or motor action), whereas in machine learning they could represent weight updates arising from the application of given learning rule to set of training data. Learning can then be summarised by the weight tensor W E I Nxxx formed by stacking together the weight w((k) (k) 1j
matrices W on all trials', i.e., W Wysa
Here, we investigate the multilinear structure of the weight tensor W in der to gain insight into the relationship between learning, recurrent weights, and network dynamics (Fig. 1a). In particular, given recent studies suggest- ing that trained RNNs are typically low (matrix) rank [14], we ask whether the weight tensor is low tensor rank, i.c., if the weight tensor can be written as a sum of R rank-1 1 components:
RNN weight tonsor
Trisl
Traal
Trisk
10
Do
Neurrn
Neuron
Neuron
Nauron
Neuron
+...
Neuron
d
Leariing
Imput
Iria
Coosdtion
Gooddtion
Time
Time
Condition:
1.2
Figure 1: Low tensor rank recurrent neural networks. a. We consider the 3-tensor formed by stacking the weights over learning. b. Constraining the weight tensor to be low rank allows for interpretable analysis of the evolution of the weights over learning. c. LtrRNNs partition neural variability into task condition-specific inputs and learning- related weight changes over trials. d. LirRNNs capture changes in dynamies constrained to a fow-dimensional sub- space, which reshape neural representations over learning.
W
& by @ Cr
(N(M) 11
a by
=ar,bbrgge (N)
where a,, by E RN and c, E R for {I...., R) (Fig. 1b). A conse- quence of the low tensor rank assumption is that the weights vec(W(()) must evolve within an R-dimensional subspace of R spanned by the set of rank-1 matrices (vec(a, @ b
For consistency with our analysis of neural data, we use the terminology trial throughout to describe the index of the weight samples k denoted in superscript (cf. neuron and time indices denoted in subscript).

---

Low tensor rank RNNs. To investigate the structure of weight tensors formed over learning in recurrent networks, we consider a continuous-time RNN (Fig. Ic). The RNN dynamics on trial & given by:
[R] ((x)(((())
x(R)
(1)
for B E I Nx Nupe, (t) E R Napu, If the low tensor rank hypothesis is satisfied, the system can be written as:
R
xx(
h (ar @0 b,))(x(()) - -x x(1) + Bu(h(()
which reduces to a low-rank RNN on every trial [15]. As result, (t) is constrained to an (R + Ninpur) dimensional subspace of INN spanned by {ar U {B.) and this subspace is fixed across trials (Fig. Id).
Nouronz
The weight matrix evolves across trials via simple rescaling of the rank-1 components by the trial factors (()) This may at first appear to restrict the possible changes in dynamics over trials to simple scalings of the flow field along the directions determined by the corresponding a, & b, components of the weights. In fact, non-trivis triviri changes in the flow field (e.g., bifurea- tions) can occur due to the nonlinear activation 6(x) (Fig. 2: Supplementary Material A). Moreover, even in the case of a linear luRNNN (Ø =id), a rich repertoire of dynamics is possible. First, the system can switch be- tween different vector fields by allowing the correspond- ing trial factors to transition to zero (or nonzero) values. More generally, the eigenvalues are non-trivially related to the c, c, The real eigenvalues can be shown to sat- isfy the following equality (Supplementary Material A); X = cos (0%:) Cp (00 + OM/b) cos 00,= @rrb),, ). where v, and v, are the rth left and right eigenvectors and 6%.5 is the angle between the rth x and y. This non-trivial relationship between the eigendecomposition and tensor rank decomposition endows considerable flexibility in terms of the possible changes in dynamics over trials, depending on the relative orientations of (a, 1R and (b, ] as well as the angle 0 which quantifies the non-normality of the resulting weight matrix.
Neuron
RB
Leaming
60
8
Figure 2: Bifurcation in the flow
field of an ItrRNN. Here a = b
1-1/V2, 1/v2,". so that the weights any given trial k are, (kk) [1 1.. A supercritical pitchfork when
2
bifurcation occurs at ckk) = tanh; for & =id (Sup. Fig. 1) a line attractor emerges at c(k)_ 1.
Fitting ItrRNNs to data. These results suggest that lrRNNS comprise a highly flexible and expressive class of neural networks for relating the changes in recurrent dynamics to low-dimensional structure of the weight updates. We next considered how ItRNNNs could be to a dataset (for example, the activations of an RNN over the course of training, or recordings from population of neurons animal learns to perform a task) in order to elicit a low tensor rank description of the weight updates governing the evolution of the dynamics over trials.
We consider a dataset in the form of a tensor Y E R.NaxxXxX comprising the activity of Naai neurons mea- sured on trials, each of duration T time points, with
Single neuron activly
LIITINN Data Loss
Trial
Trial 2
Trial
single trial activations y (A)) (t). To fit the ItrNNN to such tensor, we minimise the loss function L(a, b, c. B, B) T.K (M() (f) ()))1 where M E Na readout matrix mapping the activation of the ltRNN units onto the neural data. The loss function is minimised directly via gradient descent with respect to the parameters of a rank R tensor decomposition of W.
Time
Leaming
Figure 3: LtrRNN fit to neural data. LirRNNs can be fitted to neural popula- tion data (here shown for single neu- ron and for three trials), thus capturing smooth changes in activity over learning.
Above we considered how the low tensor rank assumption constrains the possible changes in RNN dynamies for the case which the c, can change arbitrarily across trials. However, here we are interested specifically in understanding how these dynamics change over the course of learning which
3

---

typically imposes additional structure on the 's. Here, we add an inductive bias that the components determining the weights change smoothly over trials, which we impose by parameterizing the trial factors as = (I++T)) TT), E where I.I," A is the Cholesky decomposition of the smooth covariance matrix A (assuming C ~ [0, r), see Supplementary Material A). In contrast, we constrain the inputs to the RNN u (8) (t) to depend only on task condition (i.e., the type of trial, such as the stimulus presented or set of instructions given), thus assuming that most of the condition-indeeendent inter-trial variability in the data comes from learning-induce
changes in the dynamics of the system rather than changes in its inputs. By partitioning the variability in this way, we ensure that the activations x() (t), well as the weights vee(W(()), evolve in low-dimensional subspaces of R and RN= respectively, throughout the course of learning. In the following sections we ask how accurate this low-dimensional view of learning is in the context of both biological and gradient-based learning.
3 Related work
LirRNNs integrate broad range of concepts in neuroscience and machine learning. including low (matrix) rank RNNs, tensor decompositions, fitting RNNs to neural data, and analytical studies of gradient descent dynamics. Here we review the relationship between ItrRNNs and previous results across these domains.
Inferring weights from neural activity. A diverse range of methods have historically been used to infer synaptic connectivity from neural recordings [19]. Our work specifically builds upon a recent line of investigation attempting to infer the recurrently generated dynamics x = f(x, W) of the network from samples of its trajectories [20]. [21] and [22] inferred learning- induuedd changes in population dynamics from neural data, but this approach was limited to pre-post comparisons at two timepoints in learning. Others have developed methods infer the learning rules governing weight updates from post-learning neural activity [23], or spike train recordings [24, 25]. However, no previous study has directly inferred the evolution of weights over the course of learning in neural data. Our wwwteensor-rnk approach enables this inference by constraining the parameter complexity, which allows more efficient application to neural data. Moreover, in contrast to previous approaches [21, 22] our method ensures that the dynamics at different phases of learning remain jointly interpretable due to the existence of stable latent space in which the network activity unfolds.
Low rank RNNs. Previous work has introduced low rank RNNs as a powerful framework to uncover the low-dimensional dynamics underlying task performance [15, 71, and which can be inferred from neural population data [10]. In the context of learning, one could naively fit a separate low rank RNN to data recorded on each trial, but there is no guarantee that the resulting O(NK) parameters could be related to each other. By instead assuming the weights have low tensor rank structure, we ensure that smooth changes in the dynamics can be mapped over trials with O(N + K) parameters, while benefiting from the low rank RNN framework on each trial.
The dynamics of learning. Recent work has investigated the dynamics of learning via gradient descent in artificial neural networks [26, 14]. [26] showed that deep linear networks progressively learn the leading singular values of the input-output covariance matrix, which naturally leads to low rank weight matrices when the input-output mapping is itself low rank. [14] extended these findings to show that the weight updates of an RNN are low rank. However, the proof assumed an dimenti-ad linear dynamical system at steady-state and with Gaussian- distributed weights. Using adjoint sensitivity analysis, our work extends this result by deriving bounds on the matrix rank on the weight updates of dimensianai nonlinear networks away from steady-state. Furthermore, to investigate the evolution of weights over learning, we provide bounds on the tensor rank of learning.
Tensor decomposition methods. Our framework can be used to infer low dimensional latent structure in neural data. An alternative approach is to directly fit a low rank approximation of the neuron x time x trial data tensor [27, 28, 29, 30]. Such linear methods have become increasingly popular because of their interpretability but have two key shortcomings: first, neural representations are generally nonlinearly embedded due to their dynamics and task structure [31, 32, 33]. Second, these methods do not provide insight into the dynamics underlying neural computations. Here, we address both issues by parameterizing the dynamics themselves as being low tensor rank, which allows changes in the dynamics to be mapped directly while also enabling RNN activity over trials to be visualized in a fixed subspace spanned by the weight tensor column factors.

---

4 LirRNN dynamics capture neural activity during motor learning
To test whether neural population dynamics during learning are consistent with low tensor rank framework, we fit ItrRNNs of varying rank to recordings from the motor and premotor cortex of the macaque during a motor learning task in which the subject must adapt to a force field perturbation in order to reach the target endpoint [16] (Fig. 4a). Following evidence that motor cortical initial states are set by upstream regions during the preparatory period [16, 34], we parameterize (t) with neural ODE (see Supplementary Material B) with solution ( on trial of condition i during motor preparation, after which the dynamics evolved autonomously. As preprocessing, we first Gaussian filtered the spiking activity of each neuron (std = 40 ms) to obtain smooth estimates of the instantaneous firing rates. To compare across models, blocks of consecutive entries in time and trials were held-out for cross validation, and the remaining entries of the data tensor were used to fit the parameters of each lurNNN model. For comparison we also fit full tensor RNN (i.e. all N K entries of W were fitted), as well as full and low matrix rank RNNs whose weights stayed constant over all trials. We quantified model performance as the unexplained variance on the interiors of the held out blocks while discarding borders to reduce temporal correlations between the train and test set (Fig. 4b inset, Supplementary Material B).
Interestingly, when we compared the performance of the full tensor RNN to a static RNN, we found only a 5% difference in the variance explained, indicating that much of the variability in the data is determined by task condition and dynamics, with the difference in performance attributable to marming-inducce changes [16]. However, of this remaining variability, a ltRRNN with only components was able to achieve similar performance to the full tensor RNN (Fig. 4b), supporting the hypothesis that learning dynamics are low rank. Using this cross-validation procedure also allowed us to compare the performance of the ItrRNN model with low-rank matrix and tensor decompositions which do not fit the underlying nonlinear neural dynamics. We found that for low ranks, our method outperformed truncated SVD (applied to the trial-concatenated data) and PARAFAC (Fig. 4b). Interestingly, in the case that &a =id, the activity of the ItrRNN is constrained to an R-dimensional subspace, therefore the MSE (without cross- (oossvalidattion is lower-bounded by that of the rank R SVD by the Eckart Young theorem. This suggests that the the higher performance of ltRNNN compared to PCA is due to the nonlinear mapping from the membrane potential space to the firing rate space.
Since after f the network evolved autonomously, all of the information regarding its trajectory was contained in the the recurrent dynamics and initial state. We find that the initial states inferred by the model reflect the topology of the task variables (Fig. 4c.). In comparison, in the perturbation block of trials there was only small change to the initial states, consistent with the finding that the majority of the variability in the data was due to changes in task condition rather than learning (Fig. 3b, Supplementary Material B).
Since ItrRNNs reduce low-rank RNN on any given trial, their membrane potential x(t) is constrained to lie in the space spanned by the column vectors a, [15]. The membrane potential can therefore be visualized via a projection onto each a, to observe the low-dimensional activity of the network [10]. We find that, compared to applying PCA directly on the neural data, ItrRNNs yield more interpretable visualizations of the condition and learning- related variability in the neural (Fig. 3d.e). Additionally, since ItrRNNs parameterize changes in dynamics, we can visualize the vector field in the subspace spanned by the column vectors. Consistent with our task-trained RNN results, and those of the literature [17], small changes in the vector field are sufficient to account for learning the perturbation (Fig. 4f). These changes in the dynamics of the network can be easily interpreted through the trial factors. We find that the dynamics along certain directions in the membrane potential space change during learning (Fig. 4g). Furthermore, consistent with the hand kinematics, and recent experimental evidence [35], some of these changes in the trial factors do not simply revert back to baseline during the washout period (in which the force field perturbation is removed; Fig. 4h,i). Dynamics along some columns capture target variability (Fig. 4j. top 2 components), whereas others capture mainly temporal variability (Fig. 4j. bottom 3 components). Interestingly, the trial factors which revert to baseline during the washout period are those corresponding target-related dynamics, while those which are persistently changed are the temporal variability factors (Fig. 4i).
Overall, we find that relaeed variability can be accounted for by low-tensor-rank changes in the recurrent dynamics of neural populations. LtRRNN allows uncovering these changes in large-scale neural data and vizualizing their effect in an interpretable fashion.
5

---

Euscino
Perlibed
Washout
Eedy
Aanu
d
ad reptence
0.755
Pertureed
Wissoot
Buodino
Tiege wede
Buadre
Thme
0.0-
-10
PARARAD
Lo7thi ro200
RV3
Tonsor PRN
Basolin
Target and angle
Basclinc
Parturbed
L
h
10
Laser
Latent
00
Tmm
20
Go que 100ma
-.0
2.00 T1e
Timo (90
Figure 4: LtrRNN separates learning- from condition-related variability in motor neural data. We apply ItRNN to recordings of the motor and premotor cortex during learning of perturbed reach task. (Schematic adapted from [16]). b. We hold out for testing blocks of trials and time points (100 ms by 50 trials). This reduces temporal correlations between the train and test compared to holding out individual tensor entries. e. State of the ItrRNN (R = 5) 100 ms after the go cue. LtrRNN captures the topology of the task. Furthermore, different initial states seem to emerge during the perturbation period. d. PCA directly applied to the neural data. e. Activity of the ItrRNN projected the column vectors of the first three components. Vector field along pairs of column vectors. We constrain the activity to the columns of the tensor such that x a-a,- + ar'a, for gr, 9r' [- 2] and we compute the vector field (ő(x)) b,) + a, (Ø(x)) by ). g. Eigenspectrum of (A),, Saturation gradient indicates trial. h. The trial factors e, can be seen as latent variable of learning, such that they describe the evolution of the weights in -inddimnsioonal subspace of the weight space spanned by a, & b,. Trial factors. The black line computed using a j. LtrRNN activity projected onto each column factor (ordered as in i). In particular the activity along the top two components seem be more specific to reach angle.
as
Low tensor rank learning in task-trained RNNs
We next decided to test the low-tensor-rank learning hypothesis in model in which we had direct access to the ground truth weight tensor. Towards this end, we trained RNN with unconstrained weight structure to perform the same motor task.
Task-trained RNN model. At any given trial k, the RNN linearly controls the force applied the hand:
(k)) Do (+) + K (*)
where N: (+) R22 is an Ornstein-Uhlenbedd process representing noise in the execution of move- ment [36]. To create purely ballistic model, we provide the RNN the target position and a hold cue. For the objective we use the integrated hand to target position so that the RNN simply has to push the
hand as fast as possible towards the target (Fig. 5a).

---

After first pre pre-training the RNN to move the hand to the goal position at the onset, we then probe motor learning following the same protocol as the neural data. Specifically, the RNN first performs 50 trials in the baseline condition, after which we introduce a force perturbation orthogonal and proportional the velocity of the hand for 100 trials (Fig. 5b):
gx(h) (0(xx)())))))) cRt cRty
(t) (t) +x (t)
where R E R2x2 is a 90º clockwise rotation matrix, and E is the coefficient of perturbation. Finally, the RNN performs the original unperturbed task for another 100 trials (washout). Throughout the perturbation and washout periods, the weights are updated following SGD.
a
b
0.77
Ranurtation
MSE hand to target
Target
Vartanee oxplined
0.05
07-
0.C3-
Hand
08-
05
0.01
04-
0
30
150
260
laration
Matrix
Tensor ranks
d
Esse ine
Peetureed - early
Perturbed lite
Wishout early
Wase out lste
g
Bcseline
Forturationn
h
Tine and Rive
Excmpie
0.5-
Crouncl-tuth bound
Croes calaadddddd
0:4-
1.4
0.8
Takkera ned RNN
PG2
0.2
0.1-
1.0-
0.8
Tensic 1971
0 60
150
250
LOANN
Tral
Figure 5: Learning-induced weight changes in task-trained RNNs are low tensor rank. a. RNN
model. b. Average MSE between hand and target positions integrated throughout the trial. c. Left: Variance explained of the weights resulting from pre-training W The original task training results rank-3 RNN. Right: Variance explained of the tensor of updates W W* due retraining. The retraining procedure results in rank-2 tensor. We further found that the subspaces spanned by the pre-training columns (resp. rows) and retraining columns (resp. rows) were different, suggesting tensor of rank at most 5. d. Hand movements during various periods of learning, where "early" and "late" describe respectively the first and last trial during which the perturbation introduced or removed. Using the same cross-validation as in Fig. 4 uncovers the low tensor rank structure. f. Eigenspectrum of the weights (K) over learning. Top: Ground truth weights of the task- trainerained RNN. Bottom: Weights of the ltrNNN fit the task-trained RNN activity. Imaginary eigenvalues emerge to counter the rotational effect of the perturbation, and are uncovered by ItRNN.. g. Activity projected on PCs. The rotational activity during perturbed trials is uncovered by ItRNNN, h. Example of trial factor uncovered by ItRNN which correlates with learning at the level of the behaviour.
We next analyzed how the structure of the weight matrix changed over the baseline, perturbation, and washout blocks. The change in weights result of the pre-training W Wg. where Wo the random initialization, and W the weights after pre-training, is of matrix rank 3 (Fig. 5c, left). Since the weight updates resulting from learning the perturbation are small (consistent with the literature [17]). we cannot simply apply PARAFAC on the weights recorded during perturbation learning, as only the rank-1 terms from pre-training will be visible. However, running PARAFAC on updates W - W* @ 1 (where 1 is the vector of ones), reveals tensor rank 2 updates (Fig. 5c, right). Furthermore, the subspace spanned by the columns (or rows) resulting from pretraining different than those resulting from perturbation learning (Supplementary Material C).
7

---

LirRNN application. We next ask whether IuRNN can uncover this low-tensor rank structure in the weights, especially the small changes due to perturbation learning. Using our cross- vols-valiiation procedure, we find that the variance of the activity of the task-trained RNN that is unexplained by the fitted ItrRNN plateaus tensor rank (Fig. 5e), consistent with our analysis of the ground truth weights. Furthermore, the weights of the ItRRNN share similar spectral properties, and similar evolution over learning, compared to ground truth (Fig. 5f). In particular, in both cases, imaginary eigenvalues, corresponding to rotational activity, emerge over learning, consistent with the behaviour being rotational post-learning to counter the force field (5d), The emergence of rotational activity is also visible at the level of the neural activity (Fig. 5g).
Therefore, consistent with the results found through luRNN fit to neural data, we found that the weight updates in an RNN trained on a perturbed ballistic reach task had low-tensor-rank structure. Furthermore, ItrRNN was able to uncover this structure the weights of the task-trained RNN, and its evolution over learning.
6 Gradient-based learning constrains the tensor rank of weight updates
We have so far observed that learning leads to low-tensor-rank weight changes in both biological data and task-trained RNNs. To gain deeper insight into why this is the case, we next present set of mathematical results regarding the tensor rank of gradient- based learning in RNNs. Towards this end, use the method of the adjoint [37, 38] derive dynamical system whose solution is the gradient of a loss functional with respect to the RNN weights [39].
The adjoint. For dynamical system x = f(x, 0). the state adjoint of the loss functional L: R - R a particular time point defined as ax (t) dL(x(T)) where 1' is the time of loss evaluation (but see Supplementary Material E for the case a loss functional that is integrated over time). From the dynamics of x, the state adjoint dynamics can be derived:
dx((t)
df (x(()))) T
åx (t) =
ax
dx(t)
However, understand gradient-based learning, we require the parameter adjoint: at dL(x(T)
d0(2)
This can be accomplished by concatenating the original dynamical system by its parameters z(1) = (x(t), (t)] (noting that ((t) = 0) to define the the augmented adjoint az (t) = [ax((), ao (a) (t)), whose dynamics can be shown to follow
df(x(t)) x(t)
(d((x(()) do
(1)
ax ax()),
(x)
with terminal condition a, (T) (L())))
0] The gradient of the loss with respect to the parame- ters can then be found by integrating the parameter adjoint dynamics backwards through time to get
dx((T)
ag (0). In Supplementary Material E we provide a short derivation of the adjoint adapted from (38).
Main results. While the adjoint method is extensively used as a numerical tool in autodifferentiation packages [38, 40], can also provide analytical insight into the gradient dynamics of dynamical systems. Towards this end, we now return to the case of an RNN, with the parameter of interest being the weight matrix. We demonstrate several mathematical results with the aid of the adjoint.
Lemma 1. Consider the RNN & W Wx(x) x + Bu(t), where x(t) R", u(t) RT, W R=x", B E IR and ò : R" + RR" an element-wise nonlinearity, and define a loss functional of the linearly decoded RNN state, ((Do(((T)) for E Rd. Furthermore, let W = = ar 8 Br The adjoint dynamics are then given by,
åx La. 8 (3, o (()00))0(((()))
ax ax.
åw ax 2 ((x(1))
with terminal conditions (T) = d' o E L'DD, (((D)/)(((())))))))())) ((x(T)): Drr and aw (T) = 0.
Where C denotes the Hadamard product and o' the element-wise derivative of 6
8

---

From the adjoint dynamics, it can be seen that the singular values of the gradient of the loss with respect the weights of the RNN (V aw (0)) depend on the dimensionality of the subspaces over which (()( and ax evolve. This can be used to show that the gradient's singular values are bounded by the singular values of both the activity and the adjoint, which we formalize in the following theorem.
Theorem 1. Consider an RNN be defined as in Lemma 1. Then the singular values of its gradient can be bounded as:
max
{ (10x)
0x on the of/00] S or vwLk s min o of ( }
o(x) ax of
of
where of denotes the rth singular value vf
This theorem demonstrates that there is natural bound the numerical rank [41] of gradient-based learning in RNNs. Numerically, we find that the singular value spectrum of both ((x) and ax tend to decay exponentially, even in the case of a chaotic RNN (Fig. 6). In Supplementary Material E repeat the same analysis with various activation functions, initial weight variances and ranks: overall, we find that smooth activation functions (such as the most commonly used tanh) tend to lead the fastest decaying singular value spectrum of the adjoint.
So far, haven't 't made any assumption the architecture of the RNN such as constraints on W or 6. Under such constraints, stronger and more explicit bounds on both the matrix and tensor ranks can be obtained.
Theorem 2. Consider an RNN defined as in Lemma with á =id and ( of rank R. Furthermore suppose (k()) (0) is constrained to the subspace spanned by the columns of (k)) and B. Consider the weight tensor W = [ where ((++)) = =W(") + a Eiz V (T)x with y(i E Rdd where x(i) denotes the activity of the RNN after the jth weight update. Then,
W()),...)
1. The rank of the gradient at the first step is at most rank
(olll(Www w 11 S R
2. The rank of the trial slices of the weight tensor is at most rank W (+) - 2R + m + d.
3. The tensor rank of the weight tensor at most rank (W S (2R + m + d)2
We note that these bounds are tight, in the sense that there exists networks for which they are equalities; therefore, they cannot be improved upon without restricting the set of architectures considered (e.g. to normal weight matrices). We also point out that they are non-trivial as the current best upper bound on the max rank of a (R tensor is n² n 1 [42]. In particular, in the case of (() = 0 and m considered by [14], the tensor rank is at most 4. For arbitrary weight initializations, the matrix and tensor mathematical ranks can be high. Nevertheless, in most cases the matrix and tensor numerical ranks will fall vastly below these bounds due to Theorem 1.
We illustrate this result with R - 3, d = 2. = 2 in Fig. 6. We find that the true ranks fall within these bounds- strictly below due limited machine precision and that the numerical ranks extremely low. Intuitively, decoding the firing rate of the RNN into a low-dimensional space produces weight updates that push the RNN and adjoint activity lie in a low-dimensional subspace of the state space. This, in turn, further pushes weight updates to lie in low-dimensional subspace of the weight space.
The proofs of Lemma and Theorems and 2 are provided in Supplementary Material E. We also discuss two additional cases, namely that of a loss integrated over some period of time, and the gradient of the loss w.r.t. to other parameters of the RNN. Finally, we show that momentum- based optimization methods such ADAM have the same property. Together, these analytical results provide insight as to why gradient dynamics bound the matrix rank of the weight updates in generic RNNs, well as the tensor rank of learning dynamics in the linear case.
7 Discussion
Summary. In the present work, we explored the tensor rank of learning in artificial and biological neural networks. We showed that learning leads to low-tensor-rank weight updates, which can be exploited to uncover smooth changes in dynamics along with principled choice of low-dimensional

---

14-
Log Log sa.grgarrr 1og (S)
12.
Log eng to
+
00.
0.57
00
25
75
100
25.
2.--
100
20
00
25
80
1.
We
afl
T
-12
mirretee apet)
0
25 75
150
105 2.0
1cra
Tinc
Singulr WiLL
Singular
Figure 6: The singular value spectrum of the gradient of random RNNs with random objective. Top: Full rank RNN with tanh activation function in a chaotic regime. Bottom: Low-rank (R 3) linear RNN. In both cases m = d = 2. The networks are trained to map a time-varying input (parameterized as an LDS) to given target output. Error shade represents the standard deviation over 5 random initialization. a. Loss over training. Inset: activity of two neurons over learning, red is post-learning. c. Singular value spectrum of the activity and of the adjoint. d. Singular value spectrum of the weights, gradient, and bound we derive. d. Tensor rank of the weight tensor minus initial weights. Variance computed per trial slice. factors over which weights and activity evolve throughout the entire course of learning. Finally, derived upper bounds on the singular values of gradient dynamics of nonlinear RNNs, and on the matrix and tensor ranks in the linear case.
b. Example of network activity at one iteration.
Modelling limitations and future work. Inferring connectivity from neural recordings is in general an ill-posed problem [19]: given any observed pattern of activity, it is always possible that the data could be explained entirely in terms of a external input to set of unconnected neurons. In our application to neural data, resolved such ambiguities by assuming that inputs were fixed on each trial of a given condition, forcing the changes in neural activity across trials of the same condition to be captured by (smoothly varying) changes in weights. We note that, in principle, our framework and code could also implement residual trial-to-trial variability in the inputs.
In addition to learning- and condition-speciti changes, neural recordings show substantial variability across consecutive trials, which are thought to reflect a combination of i) unmeasured covariates such behavioural or intrinsic state, ii) stochasticity in the neural system itself, and iii) changes in the initial state at trial onset. Future work could incorporate behavioural covariates, allow for trial-specific initial states, and model the data using a stochastic dynamical system within each trial.
Here, we focused on gradient- based learning, motivated by recent evidence that it able to explain many features of motor learning in neural data [17, 18), as well as by more general support for an optimization-base
framework to understand neural learning [43]. The tensor rank of learning in RNNs with local synaptic rules (e.g.. Hebbian) remains open question. Towards this end, theoretical work has established links between Hebbian and gradient-basee learning [44, 45], opening the possibility of extension of our mathematical results to biologically plausible learning rules.
Our analytical results provide intuition for our observation that the rank of learning dynamics is limited by task complexity. This supports previous findings of low (matrix) rank weight changes in RNNs [14] and in deep networks [46, 47, 48], Our results on the numerical matrix rank also has interesting ties to work that uses rank compression for more efficient training in deep networks [46], and for numerical solutions to systems with time-varying dynamics [49].
Broader impact. We introduce a novel framework for understanding learning in the brain and artificial neural networks. Our mathematical results on the rank of RNN gradients have broad relevance to the machine learning community, while our results on motor neural data could drive future applications to brain computer interfaces. Overall, our work makes novel contributions towards understanding the emergence of computations through learning in neural systems.
10

---

Acknowledgments
We thank Francesca Mastrogiuseppe, Friedrich Schuessler, and Sina Tontoonian for feedback on the manuscript, and Matthias Hennig, Matt Nolan, and Srdjan Ostojie for helpful discussions. We are particularly grateful to Matthew Perich for sharing his data. This work was supported by the Agence Nationale de la Recherehe (ANR-20-CE37-000) ANR-17-EURE-0017)
References
Niru Maheswaranathan, Alex Williams, Matthew Golub, Surya Ganguli, and David Sussillo. Universality and individuality neural dynamics across large populations recurrent networks. Advances in neural information processing systems, 32, 2019.
[2] Saurabh Vyas, Matthew D Golub, David Sussillo, and Krishna Shenoy. Computation through neural population dynamics. Annual review of neuroscience, 43:249-275, 2020.
[3] Valerio Mante, David Sussillo, Krishna Shenoy, and William T Newsome, Context-dependent computa- tion by recurrett dynamics in prefrontal cortex. nature. 503(74744:78-844 2013.
[4] Jing Wang, Devika Narain, Eghbal A Hosscini, and Mchrdad Jazayeri. Flexible timing by temporal scaling of cortical responses. Nature neuroxcience, 21:102-110, 2018.
[5] Christopher Cueva, Peter Y Wang, Matthew Chin, and Xue-Xin Wei. Emergence of functional and structural properties of the head direction system by optimization of recurrent neural networks. arXiv preprint arXiv: 10189, 2019.
[6] Abigail A Russo, Ramin Khajehh, Sean R Bittner, Sean M Perkins, John P Cunningham, Laurence F Abbott, and Mark M Churchland. Neural trajectories in the supplementary motor area and motor cortex exhibit distinct geometries, compatible with different classes of computation. Neuron, 107(4):745- 2020.
2020.
[7] Alexis Dubreuil, Adrian Valente, Manuel Beiran, Francesca Mastrogiuseppe, and Srdjan Ostojic. The role of population structure computations through neural dynamics. Nature newroscience, 25(6):783-794, 2022.
[8] Matthew G Perich, Charlotte Arlt, Sofia Soares, Megan E Young, Clayton P Mosher, Juri Minxha, Eugene Carter, Ueli Rutishauser, Peter Rudebeck, Christopher Harvey, al. Inferring brain-wide interactions using data-constrained recurrent neural network models. BioRxiv, pages 2020-12, 2020.
[9] Feng Zhu, Andrew Sedler, Harrison Grier, Nauman Ahad, Mark Davenport, Matthew Kaufman, Andrea Giovannucci, and Chethan Pandarinath. Deep inference of latent dynamics with resolution using selective backpropagation through time. Advances in Neural Information Processing Systems, 34:2331 2345, 2021.
super-
[10] Adrian Valente. Jonathan W Pillow, and Srdjan Ostojic. Extracting computational mechanisms from neural data using low-rank rnns. Advances in Neural Information Processing Systems, 35:24072 15:24072-244066, 2022.
[11] Justin Jude, Matthew G Perich, Lee E Miller, and Matthias H Hennig. Robust alignment of cross-session recordings of neural population activity by behaviour via unsupervised domain adaptation. arXiv preprint 2xXiv:2202.06/99 2022.
[12] Lea Duncker, Laura Driscoll, Krishna V Shenoy, Mancesh Sahani, and David Sussillo. Organizing recurrent network dynamies by task-computation to enable continual learning. Advances in neural information processing systems, 11114488-13399 2020.
[13] Harsha Gurnani and N Alex Cayco-Gajic. Signatures of task learning in neural representations. Current Opinion in Neurobiology, 13:102759, 2023.
[14] Friedrich Schuessler, Francesca Mastrogiuseppe, Alexis Dubreuil, Srdjan Ostojic, and Omri Barak. The interplay between randomness and structure during learning in mns. Advances in neural information processing systems, 33: 3352--1333,, 2020.
[15] Francesca Mastrogiuseppe and Srdjan Ostojic. Linking connectivity, dynamics, and computations in low-rank recurrent neural networks. Newron, 99(3))609--6)) 623, 2018.
[16] Matthew G Perich, Juan Gallego, and Lee E Miller. A neural population mechanism for rapid learning. Newron, 100(4):964 2018.
[17] Barbara Feulner, Matthew G Perich, Raeed H Chowdhury, Lee Miller, Juan A Gallego, and Claudia Clopath. Small, correlated changes in synaptic connectivity may facilitate rapid motor learning. Nature communications, 13(1):5163, 2022.
[18] Peter C Humphreys. Kayvon Daie, Karel Svoboda, Matthew Botvinick, and Timothy P Lillicrap. Bci learning phenomena can be explained by gradient-based optimization. bioRxiv, pages 2022- 12, 2022.
11

---

[19] Ildefons Magrans de Abril, Junichiro Yoshimoto, and Kenji Doya. Connectivity inference from neural recording data: Challenges, mathematical bases and research directions. Neural Networks, 102:120-337. 2018.
[20] Lea Duncker and Maneesh Sahani. Dynamies the manifold: Identifying computational dynamical activity from neural population recordings. Current opinion in neurobiology. 70:163-170, 2021.
[21] Adil G Khan, Jasper Poort, Angus Chadwick, Antonin Blot, Maneesh Sahani, Thomas D Mrsic-Flogel, and Sonja B Hofer. Distinct learning-induce changes in stimulus selectivity and interactions of gabaergic interneuron classes in visual cortex. Nature newroscience, 21(6):851-83 2018.
[22] Angus Chadwick, Adil G Khan, Jasper Poort, Antonin Blot, Sonja B Hofer, Thomas D Mrsic- Maneesh Sahani. Learning shapes cortical dynamics to enhance integration of relevant sensory input. Neuron, 111(T):06-120,, 2023.
Mrsic-Flogel, Flogel, and
[23] Sukbin Lim, Jillian McKee, Luke Woloszyn, Yali Amit, David Freedman, David L. Sheinberg, and Nicolas Brunel. Inferring learning rules from distributions of firing rates cortical neurons. Nature neuroscience, 18 (8(12):1804-11 1810, 2015.
[24] lan Stevenson and Konrad Koerding. Inferring spike-timing-deeendenn plasticity from spike train data. In J. Shawe-Taylor, R. Zemel, P. Bartlett, F. Pereira, and K.Q. Weinherger, editors, Advances in Neural Information Processing Systems, volume 24, Curran Associates, Inc., 2011.
[25] Scott Linderman, Christopher H Stock, and Ryan Adams. A framework for studying synaptic plasticity with neural spike train data. In Z. Ghahramani, M. Welling, C. Cortes, N. Lawrence, and K.Q. Weinberger, editors, Advances in Neural Information Processing Systems, volume 27, Curran Associates, Inc., 2014.
[26] Andrew M Saxe, James I. McClelland, and Surya Ganguli. Exact solutions to the nonlinear dynamics of learning in deep linear neural networks. arXiv preprint
2013.
[27] Tamara G Kolda and Brett Bader. Tensor decompositions and applications. SIAM review, 51(3):555-500 2009.
[28] Alex H Williams, Tony Hyun Kim, Forea Wang, Saurabh Vyas, Stephen Ryu, Krishna Shenoy, Mark Schnitzer, Tamara G Kolda, and Surya Ganguli. Unsupervised discovery of demixed, low-dimensional neural dynamics across multiple timescales through tensor component analysis. Neuron, 19991111-1 2018.
[29] Hugo Soulat, Sepiedeh Keshavarzi, Troy Margrie, and Maneesh Sahani. Probabilistic tensor decomposition of neural population spiking activity. Advances in Neural Information Processing Systems, 34:15969 15980, 2021.
[30] Arthur Pellegrino, Heike Stein, and N Alex Cayco-Gajic. Disentangling mixed classes of covariability in large-scale neural data. bioRxir, pages 2023 00, 2023.
[31] Rishidev Chaudhuri. Berk Gerçek. Biraj Pandey, Adrien Peyrache, and Ila Fiete. The intrinsic attractor manifold and population dynamics of a canonical cognitive circuit across waking and sleep. Nature neuroscience, 22(9):1512-1222 2019.
[32] Richard Gardner, Erik Hermansen, Marius Pachitariu, Yoram Burak, Nils A Baas, Benjamin A Dunn, May-Britt Moser, and Edvard I Moser. Toroidal topology of population activity in grid cells. Nature, 602(7895): 128, 2022.
[33] Manuel Beiran, Nicolas Meirhaeghe. Hansem Sohn, Mehrdad Jazayeri. and Srdjan Ostojic, Parametric control of flexible timing through low-dimensional neural manifolds. Neuron, 2023.
[34] Xulu Sun, Daniel J O'Shca, Matthew D Golub, Eric M Trautmann, Saurabh Vyas, Stephen I Ryu, and Krishna V Shenoy. Cortical preparatory activity indexes learned motor memories. Nature, 602(7896)2274- 279, 2022.
[35] Darby M Losey, Jay A Hennig, Emily R Oby, Matthew D Golub, Patrick T Sadlter, Kristin M Quick, Stephen I Ryu, Elizabeth C Tylle-Kabbar,, Aaron P Batista, Byron M Yu, et al. Learning alters neural activity simultancously support memory and action. bioRxiv, pages 2022-07, 2022.
[36] Robert I Van Beers, Patrick Haggard, and Daniel M Wolpert. The role of execution noise in movement variability. Journal of neurophysiology. 91(2):050-1163 2004.
[37] L Bittner. Ls pontryagin, vg boltyanskii, rv gamkrelidze, ef mishechenko, the mathematical theory of optimal processes. viii+ 360 new york/london I 1962. john wiley & sons. prcis 90/-, 1963.
[38] Ricky TQ Chea, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. Advances in neural information processing systems, 31, 2018.
[39] Balasubrahmanyan Srinivasan, Upadrasta Ravikirana Prasad, and Nalam Jaganmohan Rao. Back propaga- tion through adjoints for the identification of nonlinear dynamic systems using recurrent neural models. IEEE Transactions on Neural Networks, 5(2):213-22) 1994.
12

---

[40] Xuechen Li, Ting-Kan Kam Leonard Wong, Ricky T. Q. Chen, and David Duvenaud. Scalable gradients for Stochastic differential equations. International Conference on Artificial Intelligence and Statistics, 2020.
[41] Per Christian Hansen, Rank-deficient and discrete ill-posed problems: numerical aspects of linear inversion. SIAM, 1998.
[42] Jarostaw Buczyáski and Joseph M Landsberg. Ranks tensors and generalization secant varieties.
Linear Algebra and its Applications, 438(():668-66)) 669))696666688666 2013.
[43] Blake A Richards, Timothy P Lillierap, Philippe Beaudoin, Yoshua Bengio, Rafal Bogacz, Amelia Christensen, Claudia Clopath, Rui Ponte Costa, Archy de Berker, Surya Ganguli, et al. deep learning framework for neuroscience. Nature neuroscience, 22(11):1761-1771 2019.
[44] Cengiz Pehlevan, Tao Hu. and Dmitri Chklovskii. A iabbian/anti-hebbune neural network for linear suhspace learning: A derivation from multidimensional scaling of streaming data. Neural computation, 27(7):1461-1495, 2015.
[45] Shanshan Qin, Shiva Farashahi, David Lipshutz, Anirvan M Sengupta, Dmitri B Chklovskii, and Cengiz Pehlevan. Coordinated drift of receptive fields in hebbian'anti-hebbian network models during noisy representation learning. Nature Neuroscience, pages 1-11, 2023.
[46] Steffen Schotthöfer, Emanuele Zangrando, Jonas Kusch, Gianluca Ceruti, and Francesco Tudisco. Low- rank lottery tickets: finding efficient low-rank neural networks via matrix differential equations. arXiv preprint arXiv:2205. 13571, 2022.
[47] Charles H Martin and Michael W Mahoney. Implicit self reuulreaulaiizaaiinn in deep neural networks: Evidence from random matrix theory and implications for learning. The Journal of Machine Learning Research, 22(1):7479- 1551, 2021.
[48] Sanjeev Arora, Nadav Cohen, Wei Hu, and Yuping Luo. Implicit regularization in deep matrix factorization. Advances in Neural Information Processing Systems, 32, 2019.
[49] Othmar Koch and Christian Lubich. Dynamical low-rank approximation. SIAM Journal on Matrix Analysis and Applications, 29(2)433-4454, 2007.
[50] John Rhodes. concise proof kruskal's theorem tensor dccomposition. Linear Algebra and its Applications, 432(():1818-1822 2010.
[51] Tamara G Kolda. Orthogonal tensor decompositions. SIAM Journal Matrix Analysis and Applications, 23(1):243-255, 2001.
[52] Grigorios A Pavliotis. Stochastic processes and applications: diffusion processes, the Tokker-Plannc and Langevin equations, volume 60. Springer, 2014.
[53] Diederik Kingma and Jimmy Ba. Adam: method stochastic optimization. arXiv preprint
1rXiv:1412.699 2014.
[54] David Zoltowski, Jonathan Pillow, and Scott Linderman. general recurrent state space framework for modeling neural dynamics during becision-making In International Conference on Machine Learning. pages 11680-11691. PMLR, 2020.
[55] Walter Rudin. Functional analysis. McGraw-Hin Hill, 1973.
[56] Roger A Horn and Charles R Johnson. Topics in matrix analysis cambridge university press. Cambridge, UK, 1991.
[57] Rudolf Emil Kalman. Mathematical description of linear dynamical systems. Journal of the Society for Industrial and Applied Mathematics, Series A: Control, 1(2):122-192, 1963.
13

---

Supplementary Material
A Low tensor rank recurrent neural networks
A.1 Architecture
Low tensor rank weights. In order to probe for the tensor rank of learning in neural data, we introduce an RNN architecture that captures the evolution of neural activity over slow timescales. We first recall the description of the architecture of the main text, W E NNNKK
R
Σ a, & b, & Cy-
In particular, at trial k the dynamics of the RNN can be described as,
R
T ((()) 6(x) x Bu (2)
(u
a, @ b,)o(x) x Bu (t)
for B E NxNua u (1) E RR so that the RNN is a low rank RNN [15]. When training, we initialize a, NO, T), b, ar. Furthermore, the weights are parameterized such that ||a,|| = |b, - 1, so so that the magnitude of a component is captured by c,,
Smoothness constraint over trials. We further constrain the initial covariance in trial of the trial factors by parameterizing them as C, = (L + of)c, where LL" A is the Cholesky decomposition of the smooth covariance matrix A, and C, is initialized as C, ~ (O Z). In particular, we use rational quadratic kernel s1 +1++ (2))-((k; (11; )))))_1, where k; is the ith trial index. This is equivalent to performing 3-mode matrix-tensoo product the weight tensor itself (L oI) X3 W so that its entries over trials are linear combinations of smooth functions up to observation noise.
This parameterization is similar to that of Gaussian process regression, except no probabilistic objective is set (kernel regression). particular, given that (I+ aT) invertible, any possible c, can in theory be obtained upon optimization of the c,.. By additionally setting a regularization on č we penalize the smoothness of c, as non-smooth solutions have diverging č,. In this way, we bias the optimization process towards smoother C, 's. As illustrated in Fig. 4b, the cross-validated loss remains similar to the non-smooth, full-rank, case.
A key advantage of having smooth trial factors is that missing trials can be easily accounted for. Indeed, in most large-scale neural datasets, such as the one explored in the present work, potentially many trials may have been discarded, for example due to behavioural performance being outside the range set by the experimentalist. The assumption being made here is that the across-trial covariance is preserved when such trials occur. That is, we assume that failure of the animal to perform a given trial does not imply that a trial wasn't 't informative, or that learning did not occur.
Condition wiis inputs. We parameterize the condition-wise inputs to the network with neural ordinary differential equations [38], i.e. as a dynamical system whose right hand side is parameterized by deep neural network (DNN)
(()
(N) u ((Dv
where v(0(() E Nonne and D E I Nuju X XNNon Throughout this work, we used a fully-connected 3-layer DNN with layers of size 150 and ReLU nonlinearities. This provides inputs whose dynamics considerably less constrained than those generated by low-rank RNNs. Thus, do not make any assumption on the activity of upstream brain regions which drive the activity of the brain region being recorded. While here we chose model the inputs using autonomous neural ODEs, one might imagine feeding the neural ODE with behavioural or task covariates to relax the condition specificity assumption, fitting residual inputs to capture trial-by-trial variability arising from unmeasured variables [21]. This could account for some of the variability that can neither be explained by the task condition, nor by changes in the dynamics due learning.
Loss. In sections 4 and 5 we focus on optimizing for the mean squared error (MSE)
K T
b,b,,, = Σ1
M (1)
(((0)))((((()))))) y()((t)))
(A.1.I)
14

---

where y are firing rate estimates data and M E IR N In section B we present supplementary results on optimizing the Poisson likeliiddoood with respect spike data,
L(a, b, B,
Σ Σ log( (o (t))) (tz))))
(A.1.2)
where y are binned spike data.
The following pseudocode summarizes the steps of fitting an ItrRNN to neural data. For the sake of clarity, present the simplest The neural ODE-driven case can be achieved by coupling f with a deep neural net. Additionally, to allow for the initial state (w.r.t. the data) of the RNN change over conditions and trials, the evaluation of the dynamical system can be done over { ..., O, ..... T&t), where Xo is now the At, and the fit data is still done on the non-negative time states.
Pseudocode.
autonomous ItrRNN with fixed initial state.
Algorithm Low tensor rank recurrent neural network fit to data
inputs: Y R MNuxxxx
Neural data tensor of shape neuron time trial
t = At, TAt]
Time points of evaluation of the dynamical system
initializations:
randomly initialize ar, by E RR- for r = I....,R R
randomly initialize C, E. for I...., randomly initialize q E R" M E R N.
definitions:
fw (x) = Wo(x) xx
RNN dynamics for a given weight matrix Trial covariance matrix defined by smooth kernel & Cholesky decomposition
A R R such that Aij t- x(i,))
I. <- Cholesky(A (A + o²1)
while the loss hasn't converged do
0
+ Le
W + E === ar
b, @ Cr
X, + Σ Σx=1 ar
The initial state
for k = 1, K do
b> For all trials (parallelizable) Matrix of activity of the RNN during trial
(k) + ODESolve(fwwa) Xo, t)
1 + 1+ ||M - ( (*)
Fir to data
Regularization
end for SGDUpdatc (a,, b, cr, M, q)
D E.g. a, f- a, 1) 1)xX aw da,
dx dW
end while
Code availability. Arthur-pe/LtrRie
The ItrRNN implementation can be found at https://github. com/
A.2 The dynamics of ItrRNNs
Rich changes of dynamics through oblique columns. Unlike for matrix rank decomposition, tensor rank decomposition
be unique even for non-orthogonal factors. A sufficient condition for where, without loss of generality, r' denotes the maximum
uniqueness is that ra + Tb + re < number of linearly dependent columns of [50]. In other words, fitting the changes of dynamics over trials as opposed to a single low rank RNN shared over all trials gives additional information regarding the columns and rows. In the where the a, not orthogonal, non-trivial qualitative changes in the vector fields can occur. Since for any given trial k an ItrRNN is simply low rank RNN [15].
(A) Σ ag(C ((x(11)) (() + (k)
(A.2.1)
15

---

Table I: Hyperparameters of the ItrRNN models. Bold indicates values specific to section 4 and 5, indicates cross-validated hyperparameters, other hyperparameters were tuned by hand.
Neural data (S4)
Simulated data (S5)
LtrRNN
R
5
5*
T
200"
200"
Ø
tanh
tanh
Smoothness
50
13
0.1
0.1
0.1
0.1
Neural ODE
Layers
150 ReLU
150 x 150
N/A N/A
Regularization
a
0.01
Cross-validation
Train blocks Test blocks 1 x x 10
1 x 10 x 20
x 10 x 10
5 x
neuron time x trial
Table 2: Approximate training time of ItrRNNs. Here on the neural data of section 4. The variables which impact the most training time are the trial and neuron dimensions of the ltRNN (not the rank or data time steps), well as the neural ODE architecture. Hardware desktop with RTX 3090 Nvidia GPU and 17-12700K Intel CPU. indicates the one used section 4.
Neurons
200
400
Trials
370
12min
22min
740
16min
40min
so that the dynamics of (R) are constrained to span(a,) U (BB}. Unlike low rank RNNs, ItRRNN are not necessarily invariant under changes bases of aj's Nevertheless, we can introduce an orthonormal basis (ā,) so that,
W(A) (((())
a. D a. (b)(b) (x)))c
(A.2.2)
In particular, notice that the dynamics along all â; could potentially be affected by varying cj
Conversely, the a;'s and b, 's being respectively orthonormal - c.g. as in a singular value de- composition - is not sufficient condition for uniqueness of the tensor rank decomposition [51]. Nevertheless, constraining the to be orthogonal and ensuring that the Kruskal constraint is satisfied,
(k)
the vector fields are then orthogonal. In that case, varying c for some corresponds rescaling
the vector field along a, Nevertheless, as is illustrated below with a single component, the leak term
allows the system display typical properties of linear and nonlinear dynamical systems.
two-neuron system is that of the pitchfork supercritical bifurcation. Here, show that it is essentially tensor rank one RNN. We also illustrate how the corresponding linear RNN bifurcates. Let a = = [-1/v2, 1 1/v2T so that,
Bifurcation in a tensor rank-one RNN. A classical example bifurcation
(h))
1
-1
x=0
X1
c ((xi) ((x2) 2 -xxx)) + o(x2)
(A.2.3)
2
-1
X2
That is -X2 We consider two cases, Φ = tanh and Φ id, both odd functions. Introducing the two the previous equation, -x(x2) ((x1). So that x; cœ((x). Now considering each activation function separately,
16

---

tanh: i) e > has two solutions. ii) S one solution (the origin). At 1 the origin is non-hyperbolic fixed point.
id: i) c 1 for any x The non-zero eigenvalue of the Jacobian of the system is negative, therefore a line attractor. ii) x for any c. Then the Jacobian of the system at the origin has both positive and negative eigenvalues for e > 1 and only negative for S 1.
2-2.0
C-1.0
6-0.5
Nautonn
RVV
LDS Sdeenii
INsuroo 2
Supplementary Figure 1: Bifurcation in a tensor rank one RNN.
More general changes in vector fields. At the cost of increasing the rank, an ItrRNN can possibly transition between two arbitrary vector fields of ranks R₃ and R2.. For example, let c₁- [1, 0.9, .... 0] for j E {1,.... R1} and c [0, ..., 0.9, 1] for j E {Rı + 1..,. R1 + R2). There might however be multiple bifurcations between the first and last trial. More generally, given that any tensor has (possibly high rank) tensor decomposition, any weight tensor can in theory be captured by an ltrRNN.
This further illustrates the relevance of the result we found that ItRNN of very low ranks fit data as well as full tensor RNNs.
A.3 Relationship between rank decomposition and eigendecomposition
In this section, we investigate the relationship between an arbitrary low rank decomposition of a matrix and the eigendecomposition. We assume that, on a given trial k; (()) has the rank decomposition
nf'
((k) = adiag(c (k) b²²
ab""..
(A.3.1)
where a E N b E CR' e()) E R and a,, b; are the rows/colums of a, b respectively. One such low rank decomposition is the eigendecomposition"
(4) VAV-1 Vdiag(A)
Σ A.v.T.
By the rank-nullity theorem, R is the rank of W(+), so that R' 2 R. with R R' when Equation (A.3.1) is minimal rank decomposition. Equating the two decompositions gives general expression for the eigenvalues:
A V- adiag(c)b" = A, aV= = Σ c ay)(v; by).
Asssuiin (k) diagonalisable. A similar argument using the Jordan normal form holds for defective matrices.
'Here, write the left eigenvectors (rows of V-1, transposed into column vectors) as V.. We assume without loss of generality that the right eigenvectors v; normalised to unit length, in which case the orthonormality of left and right eigenvectors gives V. v, dy = || | cos(0 1/ cos(0% ). where 0%v is the angle between the ith left and right eigenvector.
= ||
17

---

If A4, a, b, c(0) E R. this gives rise to the result stated in the main text:
Σ
e (
cos "0) [b cos 0%)/)) |y
(A.3.2)
||رل|| cos 00 (1 cos hy
c
(A.3.3)
cos(0
||
ا cos 0j cos "j
Σ g
(A.3.4)
cos (D*
||aر ||اا (cos(0) + 07 + cos(O o2) = 2 cos(
E
(A.3.5)
Note that the above derivation makes no assumptions about the form of the low rank decomposition, other than that it is real-valued. Low rank decompositions commonly set ||a.]| = ||b,l| = 1, and often enforce orthogonality on the a, and/or b,, thereby introducing additional constraints on the relationship between the c's and X's. Normal matrices have cos (6);;) = 6ij and v; = vs. in which case further simplifications can be made.
B Motor learning
Pre-processing. Motor (n - 72 for Fig. 4, n 70 for Sup. Fig. 4) and premotor (n 231 for Fig. 4. 137 for Sup. Fig. 4) cortical neurons were used. The data were Gaussian filtered with a standard deviation of 40 ms (4 time bins). It was then centered by its baseline activity through subtracting neuron-wis the mean activity from around target onset to go-cue, and rescaled by dividing neuron-wise by the standard deviation of the execution period. Namely, the activity of a neuron (t) was given by
y.(t) (t)eso00 3:(1) ((y.(1)
(B.0.1)
where 0 is the go cue. Example of activity upon this pre-processing is given in Supp. Fig. 2.
Modeling assumptions. We assumed that motor and premotor cortex was driven into an initial state by imputs from upstream regions during the preparatory period, after which the input shuts off that the resulting activity during the reach evolves autonomously via the recurrent dynamics dynamics from that initial state. We therefore set u(!) = for > 100 where = 0 is the time of the go cue. Where the 100 ms account for sensory delay.
Cross-validation procedure. We cross-validated the optimal rank and number of neurons of the 1tRNNN. Low matrix or tensor rank models can be cross-validated by holding out specific entries of the matrix or tensor for training, and then used for testing. However, neural data has temporal correlation, such that the entry of the time- time-by-triale b----b-t--b----eeeeeee and 1,jk. For example, assuming the data are continuous, a simple average of these entries will give an optimal estimate Tijk in the limit of small time bins. Thus, the test set can be trivially inferred from the train set. We validated this intuition by performing the same cross- vols-vvaiittionn as section but with 1 1 x blocks, and found the test loss was similar as using the train loss over the whole dataset (Sup. Fig. 3).
data tensor Tijæ is strongly correlated with Ti Ti-15k
To counter this effect, sets of contiguous entries [Tijk:-.. Tetn.jk] can be held out of training, and the interior of these blocks [Ti+g.kk...... ..., Ti+n-q.ik) used for testing [30]. Here, given that we are interested in uncovering smooth changes in neural activity over slow timescales, we hold out n--y-m matrices [[T],,,... ..., [t], [m, ..., Timn.i+m.h (Fig. 4b. inset).
As mentioned in Supplementary Material A, our method allows inferring the dynamics of held-out trials. We found that using the cross-validation procedure from [30] infers similar ranks holding out entire trials (Sup. Fig. 3). We nevertheless applied this procedure for the sake of being able to compare different classes of models.
Poisson log likelihood. Another method for fitting firing-rate models to spike data to use the negative Poisson log likelihood loss [11, 10]. We fitted a luRNN (R 5, 200 as in the MSE
18

---

loss case) with softplus activation using negative Poisson log likelihood loss. We found similar but overall noisier results (Supp. Fig. 5). This may be due the presence of high firing rate neurons which are normalized by the preprocessing procedure but not likelihood fitting.
a
Daselin
Perturbed early
Perturteed
Washous eatl
Washout lato
7.-
00
-2
6
6
8
komi
tom
lom
Baselhne
Pertubedd
Wast
1.-
Eangrin addria
0+
0.0
Eiing
OuO
250
500
500
750
Mow crest
750
100
Washout
Doseline to
Purtub
Per Lubedd
Washou. Bally
Wachout
6.
26
60 (mm
ds
2s
do X(bng
2
8.
oo kimm
24
28
o.b
2
0. 1omi
25
50
kom
Baseline
Pertubedd
West out
Fingiaa astrana
n.-
a.--
Fooains
Rertuee
Washour
360
250 Time mml
500
750
260
260 Trmse Oret
sdo
750
250
Times Sred
260
500
750
-100
-A5
PC1
Supplementary Figure 2: Persistent effects of motor learning. Top: session used in the main text (Fig. 4). Bottom session used supplementary material (Sup. Fig. 4). Hand movement during the first and last 80 trials of perturbation learning and washout. The hand trajectories of some reach directions do not revert back post- washout (e.g. light green for top: dark green for bottom) b. Single neuron activity averaged within each condition. c. State go cue - 100ms. Larger full color marker are median within a condition. For some reach directions, the washout tends be more similar the perturbed state than the baseline.
C Task-trained RNN model of motor learning
Task design. A trial is split into preparatory period I. E [0, To.) and execution period (Tg. Tena] Here we set Tgo 2, Tend To model a ballistic reach, the RNN receives input the target information and hold cue during the preparatory period. During the execution period it evolves
19

---

Entire trel
tricts 100
Unexplained varisnce
0.75
5 Single entriee
0.70
0.65+
0.86
2.
Pank
Supplementary Figure 3: Comparison of cross-validation procedures. Applied to the neural data of the session used in the main text.
0.78
0.98
Crosd
0.00
12828
Park
Haselin
Porfuttation
10-
0.5
0.0
0.5
20
46
1.50
Leoind
2.50
Laticn 2
200
Supplementary Figure 4: LtrRNN applied to additional recording session. Cross-validated loss with held out blocks of size 100ms by 20 trials. b. PCA on preprocessed data. c. State of the IurNNN go cue +100ms. f. Projection on first three ay. Projection of the vector field along ay (see main text). f. Eigenspectrum of (k(k) over trials. g. First three cy. h. Trial factors cj. i Projection of x (*)) on the corresponding aj-
autonomously.
(k) dx [Wo(x( (1) + L Brarget Uarget + LIsT Bhotau hold dt + odW (C.0.1)
(k)
where Brarget E 1Rax2 iarget = [cos( (0(K)) sin (0(K)) is static vector representing the position of the target, Bhold E Rax1, u hold = a cue indicating to hold movement, and dW the infinitesimal increments of Wiener process [52]. The dynamics of the hand are given section of the main text. The loss is taken to minimize the distance between the hand (*)(t) (t) and the target (k) throughout the execution period, while keeping the hand still during the preparatory period, that is L(W, Buargets Bhold, D)
K
Too 1y (1) + Tend Tgo TTp
1
Tend (
(1)( v(0))pd))
(C.0.2)
K
k=
go
J0
20

---

PC7
PC1
2.5
8.
1.1a
Lawent
0.5
Latent
00
03
10
13
4200
-10
-15
Supplementary Figure 5: Poisson log-likelihood fitting. a. Projection on first three ay. b. State of the ltrRNN at go cue +100ms. First three cj. d. Eigenspectrum of (2() over trials. e. Trial factors cy. Projection of (x+)(1) (1) on the corresponding a,.
In particular, the speed of the reach is only constrained by the noise of the RNN and the hand. The dynamical system as whole is evaluated with differentiable adaptive step SDE solver [40] and trained with ADAM [53] during initial training, and SGD during motor perturbation learning.
Analysis of the weights. We found that, consistent with the literature [17], the changes in weights resulting from the initial training were much larger than those of motor perturbation learning. PARAFAC on the full tensor of updates W Wo & 1 captured the weight tensor in 3 components (not shown), whose columns and rows were essentially those of performing SVD on W Fitting additional PARAFAC components revealed that residual variability in the updates of pretraining was larger than the motor perturbation learning variability. Furthermore, unlike SVD, there is no guarantee that the components of fitting rank + PARAFAC model will be related to those of fitting a rank i model. Nevertheless, motor perturbation learning had significant change on the eigenvalues and activity the RNN (Fig. 5f.g)
To uncover an upper bound on rank of the weight tensor, we split the analysis into the pre-training and motor perturbation learning. We first performed SVD the change in weights matrix W Wo. where W are the weights the network post-training, but pre-motor perturbation learning (Fig. 5b.). We found that the changes in weights were well captured by rank-3 decomposition. Then, we performed PARAFAC on the change in weights tensor W - W @ 1 of the motor perturbation learning. We found that this change of weight tensor was well approximated by tensor rank 2 decomposition. The combination of these results upper-bounds the tensor rank of the overall changes weights 5. Finally, we compared the subspace spanned by the columns of the SVD and PARAFAC decomposition by projecting the weight tensor W - W' & 1 on the first three column and row singular vectors of W Wo and found approximately a remaining 0.2 unexplained variance, suggesting that the columns of Wa W and W were not orthogonal, but did not span the same subspace. Combined, these results suggest that the numerical tensor rank least 4 and at most consistent with the results uncovered by ItRNN from the RNN activity (Fig. 5e).
21

---

D Task-trained RNN models of additional neuroscience tasks
To investigate the generality of the low tensor rank framework, we additionally trained RNNs on three non-motor tasks commonly used in neuroscience.
Sensory evidence accumulation task [54] (Fig. 6i). The RNN receives one-dimensional Ornstein- Uhlenbeck (OU) process input whose expected steady-state either positive negative. The target output is either +1 or -1 if the mean of the input is positive or negative, respectively.
diccenniinniionll input. The first two inputs are independent OU processes as the previous task. The third input is binary and constant, and determines which of the stochastic inputs must be integrated. The target output is +1 (or-1) 11) if the mean of the input of the OU process indicated by the contextual input is positive (or negative).
Contextual decision making task [10] (Fig. 6ii). The RNN receives
Working memory task [14] (Fig. 6ii). The RNN receives a I-dimensional input consisting two stimuli of different amplitudes separated by delay in time. The target output is the identity of the stimulus (1 or 2) which had larger amplitude.
Note that, in contrast to the motor adaptation task. these tasks do not contain any baseline period. Therefore, as we could not compare weight updates to pre-trained solution, we simply analyzed the tensor rank of the weights starting from random initialization. That is, we first trained an RNN on each of these tasks using gradient descent, then used PARAFAC on W W @ 1 to determine the tensor rank of the resulting neuron neuron iteration tensor of weights over training. In each of these tasks we found that the variance explained indeed saturated at low tensor ranks (at R = 1, and 3; Fig. 6e).
0.90-
07'
12"
ea-
D0
0.--
0.00
04-
04-
4.75
0.90-
0.-
00-
011
11
1.0-
1.0
1.25
0.0
0.6-
2.8
1.00
2.5
0.6
0.50
02
2.5
1.0
0.25
2
100
800
1.0"
0.00
16-
022-
02-
0.6-
45-
016-
0.2
02-
810
-1.
200
T1mm
Time
hasionn
Dur
Supplementary Figure 6: Validation of low-tensor-rank framework on RNNs trained on addi- tional tasks. Each row indicates a different task. Row Sensory evidence accumulation. a. Example inputs. Each line represents different trial; red/blue indicate target identity. b. Example outputs after training (for the inputs shown in a). Loss curve over iterations. d. Eigenspectrum of W*A)) W()), Color gradient indicates training iterations (with à = 1 yellow). e. After training, we perform PARAFAC on W (1)) @ estimate the rank of learning dynamics. Row ii. Contextual decision making. Row iii. Working memory task.
22

---

E
Mathematical results
E.1
Adjoint derivation
In this section, we present a derivation of the adjoint mainly following [38]. Then, we derive the adjoint of recurrent neural networks.
E.1.1 State adjoint
Consider the dynamical system = f(x, 0) Rd where R& is set of parameters. Furthermore, let the functional L _ + IR such that L(x(T)) is our loss. First, define the state adjoint,
dI((x(T))
a(t) dx(()
(E.1.I)
Notice that since x(t + At) is function of x(t),
dL(x(T)) dx(t At) a(l AI) dx(t + At) dx(() By Taylor expanding x(t + At) = x(t) + Atf(x(()) + O(AP), we get,
a(I) dx(t + At) dx((t)
(E.1.2)
a(t) a(t AT) dx((t) ((x(1)) + +((A))
(E.1.3)
where I the identity matrix. By rearranging. a(t At) a(t)
a(l At) df(x(t)) O(A)). dx(())
(E.1.4)
Al
Taking the limit as AL 0,
da(()
a(t) dx((())
(E.1.5)
dt
dx(t)
We now have the dynamics of the adjoint; all that remains is that find an initial (or rather terminal) condition. For this, notice that
a(T) dL(x(T))
(E.1.6)
dx(T))
the usual gradient of L w.r.t. to its argument.
E.1.2 Parameter adjoint and gradient
In the above we derived the state adjoint a. However, for our purposes we also require the parameter adjoint dxL(x(T)) /dO) For this, it suffices to augment the original dynamical system with its parameters & [f 0), 0] and initial (later terminal) condition z(0) = [x(0), 0J. Defining the loss L(z(T)) ((x(T)), the adjoint of this augmented system is,
dT)z [di(z(T)) dl(z(T)) az(t) dz((t
[x((((T) dL(x(T)) dx(()
dL(x(T))
ax (t),
dx(t)
d0
d0
d0
(E.1.7)
which contains the desired term dL(x de. It now remains describe the dynamics of the augmented system. By the same argument as for the state adjoint (E.1.5),
dag (t)
ag (t) dē
(E.1.8)
dt
dz
By (E.1.7) and by unconcatenating
ax (t). d[()) d0
df(x(())
dx(!) dxc())
de
(E.1.9)
d0
dx(()
a(t) df (x))))
df (x(()))
a(t)
(E.1.10)
x(t)
d0
'Here [.]]] denotes concatenation.
23

---

Hence the following dynamical system can be evaluated, x(t), att, dL(x(T))
(x(())
df(x(())
= f(x((). a),
x(t) a(t)
(E.1.11)
dt
d0
d0
with terminal condition
dT((x(T))] xT)),
-T() dl((((T)) dx(TT)
x x(),, (T).
.0
(E.1.12)
d0
Notice that obtain the terminal condition, since it depends x(T), the original dynamical system must be evaluated forward once.
E.1.3 RNN adjoint
Let & f(x, W) WØ(x) + Bu(t) and L(x(T)) [Io((x(T)) y/1 for y R.. Furthermore, as it will be convenient, let W = = a, @ 3,. We will derive one by one the terms needed to characterize the parameter adjoint. First the Jacobian is
df(x(1),) dx((t)
R.
Σ
a, s (3, の (x(())) -T
(E.1.13)
where @ denotes the element-wise product. I the identity matrix, and ø' the derivative of é. Next, df (x(t), 0):
0
$ $
(E.1.14)
dW,k jk
((xx)
i ==
that is
df (x(x))) 0) dW
63 ((x).
(E.1.15)
Finally the terminal condition,
dl(x(T)) (D, o(x) 0z)² = Σ (Dj o(x) - (;)(Di(''()))
(E.1.16)
dx((T); dr,
that is,
dL(x(T)) Σ(D, 6(x) oxxx) d'(x) 3 Σ (D, 6(x) -1)), (E.1.17)
dxx(T)
Or more explicitly, åa
åx Σ a ao00 @ o
ax(T) (, D3)
åw ax 2 ((x),
(E.1.18)
E.2
Rank of the gradient
E.2.1
The gradient as composition of operators
In this section, we prove Theorems 1. In order to derive bounds on the singular values of V V = aw (0), we shall now consider ax and o(x) linear operators with integration. Namely,
axy = [' (t))()(t)
(E.2.1)
where y E for some suitable Hilbert space, such as 2 for our case.
More formally, let ax: ((x) E Bu.o, where Ba, is the Banach space of i) compact ii) bounded operators from to R", such that ax, x(x) H. - R.. In particular, we note that compactness follows from the image of ax, that is a (H), to be a vector subspace of IR" and therefore be of finite rank. By the same argument, Φ(x) is compact, and therefore so ((x)* " by Schauder's theorem [55], where is adjunction. Furthermore, notice that Φ(x), ax are solutions of dynamical systems with differentiable right hand side and therefore bounded ") if they are evaluated for finite time, and therefore bounded when seen operators. The following result now be applied:
24

---

Lemma 2 ([55)). Let E Bu,v, then T admits a singular value decomposition. Furthermore, this singular value decomposition of finite rank.
We can now prove the main theorem.
Proof of Theorem I. Notice that the composition of the two operators is: ax o x(x)* Furthermore, the singular values of V L are,
VWL.
Vw min max USR YEU, dim yy||=
Ly
= min max USR yeu, dim llyl==
[aax(x)*>)]
(E.2.2)
which can be bounded as,
min max |axo(x))))l UCR" yEt,
min max | ||ax]llo(x)) UCR" yet!
(E.2.3)
dim U= [y|==
dim t= ||||=1 n-
min max [l UCR yeU, dim U== ||y
(E.2.4)
12
that is,
(E.2.5)
LwL S of o((x)* 0
(E.2.6)
Now notice that, for any operator, akin to the matrix case, Tú, T2, the adjoint of their composition is (2)" T. Furthermore,
(w
min max (1(a,
max U UC YEU, dim U= ||||==1
min
||o(x)a[yl
(E.2.7)
UCR" yEU.
dim U= ||y|==1
S 0o ((x)
min max ||ayl||
(F.2.8)
UCR" yeU,
dim U= ||||==1
0
(x) 22
0,
(E.2.9)
Noticing that 0, a2 :. and 0: TwL of (VwL)"
L & 0 0,th
o(x)
(E.2.10)
Combining E.2.6 and E. we obtain the sought upper bound of Theorem 1,
VWL < min
{ o of
o(x)
o(x)
(E.2.11)
Similar steps can be used to derive the lower bound of Theorem 1, using instead the identity 22t |Tay|| < [|TyTyy|| where T7,Th T2 By, and of is the smallest non-zero singular value of Ti
We however note that, unlike the upper bound, the lower bound we provide does not have any numerical use, as the smallest singular value of the adjoint or of the firing rate is practically 0 (for example, well bellow machine precision).
We further point out that a more explicit characterization the singular values of the gradient can be obtained. Let U,*hx, 0,((x) be the right singular vectors of respectively ax and ø(x), so that Via* (t), vo((x) (t) are the left singular vectors and . o(x) the singular values.
o d(x)) /- T U.- (17)))00000000)) (t)dt
V Σ (vex @ V 0rt0 J=
(E.2.12)
25

---

The characterization of the rank of the gradient has thus shifted to the firing rate and adjoint spaces. Furthermore, the integral just the inner product between their right singular vectors and therefore of magnitude bounded by 1. Hence, for the dynamics of the weights to be large in given direction in weight space V @ V, all three of: the singular values of the firing rate, the state adjoint, as well as their cofluctuation in time, must not be small.
E.2.2 Weight gradient for time discretizations
Finally, we mention that our detour through functional analysis was for the sake of mathematical rigour, and that in practical applications. the RNN and its adjoint are evaluated at discrete time steps 0, At... gát. gåt. In that case, the gradient can be estimated as simple matrix-matrix product. Let A (0), .... ax (q and B (0) ((x) (gat)). Then L = AtAB", and the bounds E.2.6 and E.2.10 10 follow from classic matrix-matrix product bounds [56]. In particular, given that (xx) and ax are smooth, without loss of generality,
w
lim (1+
o
AB
(E.2.13)
This simple matrix-matrix product opens up the possibility of fast RNN adjoint implementations as, 4f unlike computing ax aw in general adjoint solvers, which requires (ggn²) time and On?²) memory complexity for an n-dimensional RNN and time steps, here the time complexity drops to Ognr²).
E.2.3 Rank of linear RNNs
In this section we prove Theorem 2.
Proof of claim 1. I >n a a, then by Lemma 1, the dynamics of the adjoint are constrained to the span of the rows of W, namely, åx E (pan( Therefore, ax E span(B, (3,) U (7)), which is most R+1 dimensional subspace. If ax is constrained & x is also constrained to at most R + 1 dimensional subspace. Since
ft
at most R + dimensional
subspace, then åw aw (T)) aw constrained to the same subspace its dynamics, and in particular, rank aw (0)
rank w L (0) < R+11.
Proof of claim 2. Suppose (R)) = - 2R+m+d (k) (3(0) } U
@ 3/a0 where a; " (a E saan (0)) U
(k)
a; == V. In particular, notice that V is only dependent on the initial weight. Then by similar argument as above, ax (k) E V, which implies aw (0)(k) Vw [()) V.
(k+1) E V. That is Therefore ((++1) (()) + +VW = E, -28+m+d (k+1) & g with a rank W (k- < 2R + m d.
a;
Proof of claim 3. Mutatis mutandis (lk+1) V. that is V. Therefore, ( E 28+m+ (k)
(k),s, ij a; @ aj for some Gj 's where a E V. Or equivalently, W = Σij aj 8 Cij. That is, rank W < (2R m d)2.
22-m+
E.3 Extensions of our results
Loss integrated over time. Commonly, the loss considered might be integrated,
L(T) L(T):= - S." ((x(t)
(E.3.1)
The parameter adjoint dependent only linearly on the state adjoint, we may therefore integrate the state adjoint for all initial conditions.
C(T) C(T)= ([aa(t))
ax(() + S. (z(()))))) dt
(E.3.2)
26

---

The term inside the first integral is just the solution of time-varying autonomous LDS, therefore,
S- (t) ( (0,
xL(x(t),y(t))
dt
(E.3.3)
dx(t)
Where + is the linear dynamical system state transition matrix [57]. But notice that this is the solution to the controlled LDS,
of (x(())
xL(x(t),v(t) dx(1)
åx ax dx(l)
áx')")
(E.3.4)
In the specific case of an RNN,
T
åx
Σ aj 0
ax ax + d'(((t)) O EDD o(x(t) - ,2) (x = 0
(F.3.5)
Therefore, Theorem remains unchanged for loss integrated over time. For Theorem 2, Claims 2-3 remain unchanged, while Claim becomes rank L < R + min(d, m].
Gradient with respect to other parameters. So far we have focused on the gradient of the weights of the RNN. As we have seen, the space over which the state adjoint (x evolves as well as the trajectories of the system itself x(t) determine the space over which the gradient evolves. But those are respectively dependent on the decoder and encoder B of the system. all parameters D, B, W of the system are optimized simultaneously, is most often the case, we may wonder how our bounds hold.
First, for B. notice that by a similar argument as for W, df(x)) = I 2 u(t), so that åø ax @ u(t). Therefore, following essentially the same derivation as that of Theorem the following bound can be derived,
dB
of of
(F.3.6)
By similar argument as for the derivation of Theorem 2, B(ah) E spanfar) (8:) ( (D.).
Second, for D, since df = 0, that is ån = O, 2D
dL(T) ap - ap = v = dD
(E.3.7)
In other words, the gradient of the loss with respect the decoder weights have zero dynamics.
The gradient of a functional w.r.t. a given parameter is independent of the gradient of that functional with respect to another parameter if these two parameters do not depend one another. Therefore these results hold regardless of which combination of W, or optimized.
Batched updates. Most often, the weights are updated in batches. That is AW(k) 0''2, W L (k.g) where q is the index over the batched dimension. Since the column and row spaces of Vw ((ka)) remain unchanged, Theorem 2-3. remain unchanged, while 1. be- comes rank A AW (0)
< R. + minfd, m). For Theorem 1, the common singular value identities [56] 0 (A+ + B) < (A) + (B) and oz(cA) = (A) for E R can be used. Then, TE (AW,) <Q-EE oi, (Vw (ika))
Momentum-basee optimizers. Momentum-based optimizers such as Adam [53] are commonly used to train RNNs on behavioural tasks. Here we focus on the first moment, a similar derivation can be undertaken for higher moments. In that case, momentum variable is introduced. which is updated as (k+1) (M) + V) [k),, where B determines the speed of the exponential decay. The weights are then updated as W(++)) = ((k) (k+1,, where a is the learning rate. Which implies, (A) ((++1) (K) a Σ₁ B) V L6), Using the same identities
for batched updates, oI 4, ( <E -B)' 0, (Ww)
27

---

E,4 Numerical simulations
Similarly [14], we illustrate our mathematical results on random RNNs. Since constant imputs are one dimensional (Bu = Σ B,u),, we instead use varying inputs parameterized with LDS:
i = Mu - u n(0) 11
(E.4.1)
where u()) IR" M3 (O, 1 lmm), uo (0. 1/2). The target outputs are set y E R! Y: 1)) The loss is defined as,
L(W)= = [x) (T) - y1
(F.4.2)
where x is the solution of an RNN as considered thus far.
In Sup. Fig. 7 we show the effect of varying é, the rank R of the initial weights well as the standard deviation (or strength) g of the initial weights such that W sj (0,,g"".
28

---

b
11:0-
35
2
20
06-
3.
100g
-10
100
0.10-
6C5
0.00
21
2.5
26-
0.
-2
00
20
No0
100
100
0.87
-00
11
02-
Log Matp//1
0.
0..0
-25
2.5
04
11o
3-
02-
-04-
0..0.
1v
2.5
1.2-
Logg arpaaii mall locuo (63
26-
00
1.2
Log arguu S
0.00-
0.08
75
15
25
50
150
60
7.5
0.5-
O 0.800
0.4
25
00-
0.60
0.86-
0.050
Bank
1900
Torsor
Cokem
200
Supplementary Figure 7: Singular values of adjoints and gradients. a. Loss over training. Inset: activation function. b. Activity during the last trial (black). Inset: activity of two example neurons over training. e. Singular values of the firing rate and adjoint. Inset: additional singular values. Singular values of the gradient, weights, and the bound of the tensor decomposition of weight tensor W (0)) & 1. We additionally plot the variance explained over performing matrix decomposition on all possible unfoldings of the weight tensor. Architectures. Unless noted, the initial weight std g = 1.5, the input and output dimensions m 2. Rectified tanh. We found that non-smooth activition functions seem to give the slowest decay of the singular values of the firing rate and adjoint. ii-ii. Softplus and tanh, which are the most common activation functions in neuroscience, have exponentially decaying firing rate and adjoint singular values, and therefore (by Theorem 1) exponentially decaying gradient singular values. ív. Low rank (R = 3) linear RNN. As per Theorem 2, the first gradient step is of rank R + = 4. v. Tanh in chaotic regime (g 2.1). Despite being in chaotic regime, the firing rate, the adjoint, and therefore (by Theorem 1) the gradient have exponentially decaying singular values.
derive. Variance explained per rank
29


Reviews:
[{"review_id": "RPtlOy9Jg2", "reviewer": "Reviewer_X8YP", "summary": "In this paper, the authors explore the tensor rank of learning in artificial and biological neural networks. They showed that learning leads to low-tensor-rank weight updates, and derived upper bounds on the singular values of gradient dynamics of nonlinear RNNs, as well as on the matrix and tensor ranks in the linear case. Experimental results effectively support their model's conclusion.", "questions": "None.", "limitations": "Yes, the authors have adequately addressed the limitations and potential negative societal impact of their work.", "rating": 5, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. Previous works have shown weight matrices in well-trained RNNs are low-rank. This paper focuses on whether tensors derived from weight matrices over the process of training are low-tensor-rank, which is meaningful and has a strong motivation.\n\n2. The paper supports its results empirically.\n\n3. The theoretical statements are detailed and valid.", "weaknesses": "1. From the perspective of researchers outside this field, a more structured presentation may be more conducive to understanding the contribution of work and promoting work.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for noting that our paper has a *strong motivation*, and that our claims are supported both by our *empirical* and our *mathematical results*. We agree that providing a clear presentation of these results to researchers unfamiliar with neuroscience is important for the broader impact of our work. We briefly summarize the structural changes to the paper that we will make in response to all reviewers' feedback:\n\n* **Mathematical results.** As pointed out by other reviewers, the bounds we mathematically derive regarding the rank of the gradient of recurrent neural networks may be of interest to the broader machine learning community, but only come late in the paper. To address this, we will summarize them and their associated simulations earlier in the paper.\n* **Terminology.** Some terms such as \"trial\" in our neural data application are not standard in machine learning. To address this, we will define neuroscience-specific terminology more systematically.\n* **Main contributions.** Although the different sections of our paper cohesively support our main hypothesis that learning is low tensor rank, we aknowledge that a subset of them might be more of interest to any particular community at NeurIPS. To address this, we will add a main contributions section at the beginning of the paper highlighting its structure so that readers can navigate to the section that is most relevant to their research.\n* **LtrRNN pseudocode**. In light of multiple reviewers' comments and suggestions, we believe that understanding the fitting procedure of ltrRNN is key to understand the paper as a whole. To address this, we will i) include a short description of the ltrRNN fitting procedure in the main text ii) include pseudocode of ltrRNN fitting procedure.\n* **Supplementary material.** Many of our results rely on Supplementary Material (e.g. all theorems' proofs). To address this, we will more systematically reference specific sections of Supplementary Material.\n\nOverall, we believe that these structural changes will significantly improve the accessibility of our work. Nevertheless, any additional suggestions the reviewer may have to improve readability would be very welcome, as we want to make sure that our work is accessible to readers from across different backgrounds."}, {"review_id": "P708E3WsKO", "reviewer": "Reviewer_T6eW", "summary": "The work \"Low Tensor Rank Learning of Neural Dynamics\" investigate the low-rankness of RNNs with application to neural data, i.e. neural signals of a test subject performing a motor task. \nThe authors describe that RNNs are of low rank in the trial mode when parametrized as a 3-tensor where one dimension represents the different trials. The findings are validated by showing that a low-rank parametrized RNN is able to fit the motor task with similar accuracy to the full-rank network. \nIn their theory section, the authors propose two theorems that show the boundedness of the singular values of RNN weight matrices.", "questions": "\n\n- Line 97: Please clarify: Do you ask if the weight tensor itself is low-rank or if the updates are low-rank? \n\n- Line 118: Add a reference or proof (in the appendix at least) for the statement. You say to yourself that this is a non-trivial statement.\n\n- Line 96: please explain the meaning of x and u in the context of your application. \n\n- Line 168: It is not clear, what you mean by your method. I suppose, that you mean the formulation of an order three tensor of low-rank structure as given in the Eq. of Line 95. If so, it is unclear, how the tensor is represented (Tucker format, Tensor-Trains,...) and how the network is trained. Vanilla on the factors, or with dynamical low-rank methods. Can you comment on this? An algorithm, reference, or some equations to explain the method would be necessary. \n\n- Sec4:  The authors need to specify \n\t* a) which low-rank tensor format is chosen to save the weight tensors\n\t* b) which update method / low-rank integrator or optimizer is used to compute the weight updates\n\n- Line 168: You say that your method outperforms truncated SVD. My question is, how is truncated SVD training applied? Is a full rank update and afterward a truncated SVD performed in each training iteration, or do you do something else? \n\n- Line 182: \"Compared to PCA on neural data, ltrRNNs yield more interpretable visualizations...\" First, (minor comment), there is a typo \"intepretable\". Second, this is an interesting aspect. I think, that training an RNN instead of direct application of PCA means, that you have first a differentiable (and thus smooth) model representation of the neural data, which is then, of course, nicer interpretable, and more visualizable. My questions:\n\t1. Are the neural data smoothed out in some sense, before applying PCA, is the smoothness of Fig.3d a result of the plotting tool or is the neural data smooth, i.e. without large jumps or discontinuities?\n\t2. How comes, that the neural data seem to be somewhat chaotic?\nI am by no means an expert in this application field, but find this intriguing. \n\n\n- Line 225 and the following: You describe how ltrRNN uncovers the low-rank structure. As per Appendix A, it seems that you construct the ltrRNN architecture such that it is low-rank per definition.  Thus the network has now another chance to learn low-rank features. Can you comment on this for clarification?\n\n- Line 256:  One should also mention [2] as one of the fundamental works of adjoint-based automatic differentiation. \n\n\n\n[2] Griewank, Andreas, Walther, Andrea; Introduction to Automatic Differentiation; PAMM; https://doi.org/10.1002/pamm.200310012\n", "limitations": "The authors have described various limitations of their work, and proposed concepts on how to deal with them in future work.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "\nThis is a solid paper. Particularly interesting is:\n\n- Investigation of low-ranked RNNS as a model for real-world data - in this case neural signals for a motor task. \n- Analysis of the gradient dynamics, i.e. the singular values of RNN gradients, \n- Extensive numerical tests to validate the propositions, supplemented by code examples. \n- Comprehensive Related work section, which is important for such interdisciplinary work. ", "weaknesses": "\n- Some method details for the ltrRNN training need clarification (see questions) In particular, an algorithm or some more mathematical details on tensor format and weight updates are required.\nThis is the major drawback on this papers presentation, in my opinion. \n- The computational cost of training and used hardware should be described.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the positive feedback regarding our paper and its potential reach as interdisciplinary work.\n\n**Q1. Clarification of ltrRNN model and training.**\n\n**A.** We agree with the reviewer that having a good grasp of the ltrRNN training procedure is important for the reader to understand the rest of the paper. We had originally provided a detailed definition of the training procedure in Supplementary Material A. In the revised manuscript we will i) provide a short description of the training procedure in the main text, ii) systematically point at the relevant section in Supplementary Material, and iii) provide pseudocode of the ltrRNN fitting procedure.\n\n**Q2. Computational cost of training.**\n\n**A.** We agree that it is useful for the reader to know the computational cost of the method. Supplementary Material A currently includes a table describing the time necessary to fit ltrRNNs of different numbers of neurons (given that this is the main bottleneck) to the neural data application in Figure 1. In the revised manuscript we will point to this table directly when describing the ltrRNN model.\n\n**Q3. Low rank weight tensor or low rank weight updates.**\n\n**A.** We thank the reviewer for pointing out the lack of clarity in line 97 regarding whether the low-rank hypothesis applies to the weight tensor $\\bf W$ or the tensor $\\bf \\Delta W $ that stacks the weight update matrices: $\\Delta W^{(k)} = W^{(k)} - W^{(k-1)}$. We indeed mean that $\\bf W$ itself is low tensor rank. We will correct this wording in the new version to read: \"Here, we ask whether the weights have low tensor rank when they are updated over the course of learning\".\n\n**Q4. Reference for the eigenvalue equation in line 118.**\n\n**A.** We thank the reviewer for noticing this overlook. The derivation is already present in Supplementary Material A.3, but we neglected to mention it in the main text. We will add a pointer to the derivation to the revised version.\n\n**Q5. Meaning of x and u.**\n\n**A.** RNNs of the form in line 96 are commonly used in neuroscience as models of networks of biological neurons [1,2], where the state $\\bf x$ of the RNN models the membrane potential of the neurons, $\\phi(\\mathbf{x})$ models their firing rate (i.e. activity), and $\\bf u$ represents input activity from other brain regions to the neurons. We will incorporate these definitions into the model description of the main text to aid the reader.\n\n**Q6.  Weight tensor representation and training.**\n\n**A.** We agree with the reviewer that the phrasing of line 168 is misleading. By *our method* we mean ltrRNN. We directly parameterize the weights as being low tensor rank $\\mathbf{W} = \\sum_{r=1}^R \\mathbf{a}_r \\otimes \\mathbf{b}_r \\otimes \\mathbf{c}_r$ and perform gradient descent (more precisely, ADAM) on $\\mathbf{a}_r, \\mathbf{b}_r, \\mathbf{c}_r$. For this, we use the framework of neural ODEs which allows backpropagating through solving an ODE [3].\n\nWe believe that the addition of the training procedure pseudocode will help clarify this.\n\n**Q7. Truncated SVD**\n\n**A.** In line 168, we agree that this wording is unclear: we apply SVD and PARAFAC to the neural data itself rather than the weights. When we say SVD and PARAFAC are outperformed by ltrRNN we mean that ltrRNN has lower mean squared error in the test dataset than either SVD or PARAFAC models of the same rank. We included these as a baseline as it is very common in neuroscience to analyze low rank SVD or PARAFAC [4] models fit to neural activity.  \n\n**Q8. Interpretable visualizations with ltrRNN**\n\n**A.** We first thank the reviewer for noting the typo on line 182; we will correct this in the revised paper. In Figure 3, we followed common practice [5] of convolving the spike times with a Gaussian kernel, with $\\sigma=40$ ms (Supplementary Material B for details). However, note that in the Supplementary Material, we also illustrate an application of ltrRNN which is fitted to maximize the log-likelihood of the spike counts assuming a Poisson distribution (Supplementary Figure 5).\n\nThere is a rich literature exploring the computational implications of chaotic dynamics in RNN models in neuroscience (e.g., [1]). However, the amount of variability in neural data (i.e., from trial to trial) precludes any definitive statement regarding whether fluctuations in recorded neural activity is due to the chaotic regime or noise. This question is especially fraught for RNNs due to technical issues in their training procedures that can prevent them from inferring chaotic dynamics from time series [6].\n\n**Q9. Uncovering low-rank structure by definition**\n\n**A.** The ltrRNN is indeed low-rank by definition. We thank the reviewer for pointing out that this sentence is misleading. In the revised paper we will rewrite this sentence as: \"LtrRNNs enables inference of the tensor rank of the weights from the neural activity through crossvalidation.\"\n\n\n**Q10. Reference to adjoint-based automatic differentiation**\n\n**A.** We thank the reviewer for pointing us to this work which we were previously unaware of. We will incorporate the citation into the revised paper.\n\n**References**\n\n[1] Kadmon, Timcheck, and Ganguli. \"Predictive coding in balanced neural networks with noise, chaos and delays.\" *NeurIPS* 2020.\n\n[2] Valente, Pillow, and Ostojic. \"Extracting computational mechanisms from neural data using low-rank RNNs.\" *NeurIPS* 2022.\n\n[3] Chen, Rubanova, Bettencourt, and Duvenaud. \"Neural ordinary differential equations.\" *NeurIPS* 2018.\n\n[4] Soulat, Keshavarzi, Margrie, and Sahani. \"Probabilistic tensor decomposition of neural population spiking activity.\" *NeurIPS* 2022.\n\n[5] Park, Seth, Paiva, Li, and Principe. \"Kernel methods on spike train space for neuroscience: a tutorial.\" *IEEE Signal Processing* 2013.\n\n[6] Mikhaeil, Monfared, and Durstewitz. \"On the difficulty of learning chaotic dynamics with RNNs.\" *NeurIPS* 2022."}, {"review_id": "v9jFJM0ztn", "reviewer": "Reviewer_jwhe", "summary": "The presented work investegated the 3-tensor formed by the weight matrices of RNNs across trials and found it is low-rank. The authors also conducted a mathematical proof that the weights learned by gradient-descent on low-dimensioanl tasks are low-rank.", "questions": "It might be my problem but I found this paper hard to understand. For example, what does \"weights change smoothly over trials\" (line 129) mean? I suppose trials are unordered, so what does smoothness mean here? Also, what does \"smooth covariance matrix\" mean in line 132? $A$ is the covariance matrix of which variable? In the \"slow timescale variability in data\" (line 136), the variability of what? In \"we compared the performance of the full tensor RNN to a static RNN\", what is a static RNN? It would be beneficial if the authors provided clear mathematical definitions for both the model and the task, as I currently feel somewhat out of context.", "limitations": "N/A", "rating": 4, "confidence": 1, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "First I should ackonwledge that I am not an expert in neural science, and I found it hard to fully understand this paper, so I can only provide limited insights.\n\nStrengths:\n- The low-rank property of RNN is interesting and is potentially useful for undertanding the neural dynamics and develop neural network architechures. \n- The mathematical framework could be valuable for comprehendingnature of gradient descent.\n\n------\n\nThe author rebuttal has addressed most of my questions. However, since I'm not familiar with neural science, I decide to keep my current rating to this paper as borderline reject with the lowest confidence value.", "weaknesses": "- The empirical results seems to be limited since the experiments are only performed on motor learning tasks. It's unclear how universal the low-rank property is.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for noting that our framework is valuable for understanding gradient based learning as well as neural dynamics.\n\n**Q1. Generality of low-rank learning dynamics.**\n\n**A.** We have now included three additional task-trained RNN simulations and an additional neural dataset to validate the generality of our results. See Qcae Q4 and General Response to Reviewers for additional details.\n\n**Q2. Lack of clarity for non-specialist readers.**\n\n**A.**  We appreciate the reviewer's feedback that our paper may be unclear to readers without expertise in neuroscience. As our work intersects machine learning and computational neuroscience, we believe it is important that their presentation can be understood by both communities. We will therefore clarify the neural data application section of the manuscript by better defining terminology that is not standard in machine learning. We will also highlight earlier in the manuscript the mathematical results and their associated numerical simulations in order to cue potential readers who may be interested in gradient learning dynamics, and will additionally provide better insight regarding their broader significance in machine learning research.\n\n**Q3. Trial ordering.**\n\n**A.** We thank the reviewer for pointing out the lack of clarity of the smoothness constraint we impose on the weights as well as about its relationship to the concept of a *trial*.\nOne trial corresponds to one repetition of the task of the animal: in Figure 1, this entails moving the cursor from the center of the screen to the target. The different targets in the experiment represent different task conditions. Trials are indeed ordered in the sense that trial $k+1$ occurs after trial $k$ in the original experiment. However each trial is assigned a random task condition.\n\n**Q4. Smooth changes in weights over trials.**\n\n**A.** Over learning, neural activity changes in order to perform the computations necessary for the task. This change in neural activity is generally accepted to be due to changes in the synaptic weights, which evolve over slow timescales due to well-known plasticity mechanisms such as Hebbian learning [1-2]. Because synaptic plasticity is slower than stimulus-driven changes in neural firing, we assume a separation of timescales such that the weight matrix stays constant within a trial but changes from one trial to the next (this separation of timescales has also previously been emphasized within neuroscience [3-4]). This is parametrized by defining the weight matrix on trial $j$ as $W^{(j)} \\in \\mathbb{R}^N$. When we say that the weights change smoothly over trials (e.g., in line 129), we mean that $W^{(j+1)}-W^{(j)}$ should be small.\n\n**Q5. Smooth covariance matrix.**\n\n**A.** As a reminder, on trial $j$ the weight matrix is $W^{(j)} = \\sum_r^R c_r^{(j)} {\\mathbf a}_r \\otimes {\\mathbf b}_r$, where $c_r^{(j)}$ is the $j$th element of $\\mathbf{c}_r$. Therefore the trial factors $\\mathbf{c}_r$ represent how $W^{(j)}$ changes over trials. Our assumption that the $W^{(j)}$ change smoothly over trials thus corresponds to an assumption that the $c_r^{(j)}$ changes smoothly over $j$.\n\nMore practically, we implement smoothness of the weights over trials by constraining the temporal covariance between $W^{(i)}$ and $W^{(j)}$. This can be done by first stacking the trial factors into a matrix $\\mathbf{c}=[\\mathbf{c}_1, ..., \\mathbf{c}_r]$, and by assuming that its covariance matrix of $Cov(\\mathbf{c}) \\in \\mathbb{R}^{K \\times K}$ is given by a smooth kernel. In section 2 of the main text and Supplementary Material A we detail how this can be achieved. We will reword this section of the revised paper to clarify this smoothness constraint for the reader.\n\n**Q6. Slow timescale variability in data.**\n\n**A.** In line 136 we refer to the variability of the neural activity of the recording. By slow timescale, we mean changes from trial to trial (i.e., over minutes or even days), rather than rapid changes (i.e., over milliseconds). Recent work has suggested that this separation of timescales captures different biological processes, with rapid changes representing processing of external stimuli and slower changes representing learning [3]. Here we further assume that the variability of neural activity over trials due to learning can be accounted for by changes in $W^{(j)}$, while the variability of neural activity within a single trial is the result of the ODE defining the network for a fixed $W^{(j)}$ (see also our response to Q4).\n\n**Q7. Static RNN.**\n\n**A.** By \"static RNN\" (line 160) we mean an RNN whose weights remain fixed between trials (i.e., $W^{(i)} = W^{(j)}$ for all $i,j$). That is, we fitted a single set of weights ($N^2$ parameters) for the entire neural recording. We will now refer to this model simply as \"an RNN with fixed weights over trials\".\n\n**Q8. Mathematical definitions of the model.**\n\n**A.** We agree with the reviewer's general point that having a good grasp of the ltrRNN model and training procedure is important for the reader to understand the rest of the paper. We had originally provided a detailed definition of the model and training procedure in Supplementary Material A while the main text only provided a high-level description. In the revised manuscript we will i) provide a short description of the training procedure in the main text, ii) systematically point at the relevant section in supplementary material iii) provide pseudocode of the ltrRNN fitting procedure.\n\n**References**\n\n[1] Confavreux, Basile, et al. \"A meta-learning approach to (re) discover plasticity rules that carve a desired function into a neural network.\" *NeurIPS* 2021.\n\n[2] Stevenson, Ian, and Konrad Koerding. \"Inferring spike-timing-dependent plasticity from spike train data.\" *NeurIPS* 2011.\n\n[3] Soulat, Hugo, et al. \"Probabilistic tensor decomposition of neural population spiking activity.\" *NeurIPS* 2021."}, {"review_id": "4wKAEn9kvk", "reviewer": "Reviewer_Qcae", "summary": "Proposes a low tensor rank recurrent neural network (ltrRNN) architecture, in which the tensor constructed by stacking the RNN weight matrices of different trials is constrained to have low tensor rank. Empirically shows that ltrRNNs can fit neural recordings during a motor learning task, achieving lower unexplained variance than baseline methods. Then, demonstrates that an RNN trained to perform the same motor learning task yields dynamics that can also be fit well with ltrRNNs (i.e. low tensor rank dynamics). Lastly, theoretically analyzes how gradient-based optimization can lead to low rank, establishing upper bounds on matrix and tensor ranks for RNNs. \n\nDiscloser: The research area of the current paper (interplay between machine learning and neuroscience) falls outside my expertise, and so it is difficult for me to assess the novelty and significance of some of the contributions. My review mainly focuses on presentation, soundness, and the theoretical analysis of gradient descent constraining tensor rank of weight updates.\n", "questions": "Have you applied the ltrRNN to other datasets? Or is the purpose of the architecture specific to the kind of data reported in the paper?", "limitations": "The authors have adequately addressed possible limitations of their work.\n", "rating": 7, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Reads relatively well.\n\n2. As far as I am aware, the technique proposed for modeling structure in a task learning process through low tensor rank of weights across different iterations is novel. I found the idea of examining such low rank structure insightful --- existing characterizations of (low rank) structure during gradient-based learning typically focus on weights/representations of a single iteration. I believe this concept may turn out useful for future study of implicit regularization of gradient descent.\n\n3. Analyzing the dynamics of gradient-based optimization is a subject of significant interest in recent years. The bounds on the matrix and tensor ranks of the gradient and weights of a continuous RNN contribute to a line of works suggesting that gradient descent leads to low rank solutions (e.g. in matrix and tensor factorizations as well as non-linear networks [1, 2, 3, 4, 5]). With that said, the form of the RNN considered is unorthodox in terms of its update rule and in being continuous-time, which may limit the impact of these results.\n\n[1] Li, Zhiyuan, Yuping Luo, and Kaifeng Lyu. \"Towards resolving the implicit bias of gradient descent for matrix factorization: Greedy low-rank learning.\" arXiv preprint arXiv:2012.09839 (2020).\n\n[2] Razin, Noam, Asaf Maman, and Nadav Cohen. \"Implicit regularization in tensor factorization.\" International Conference on Machine Learning. PMLR, 2021.\n\n[3] Razin, Noam, Asaf Maman, and Nadav Cohen. \"Implicit regularization in hierarchical tensor factorization and deep convolutional neural networks.\" International Conference on Machine Learning. PMLR, 2022.\n\n[4] Boursier, Etienne, Loucas Pillaud-Vivien, and Nicolas Flammarion. \"Gradient flow dynamics of shallow relu networks for square loss and orthogonal inputs.\" Advances in Neural Information Processing Systems 35 (2022): 20105-20118.\n\n[5] Timor, Nadav, Gal Vardi, and Ohad Shamir. \"Implicit regularization towards rank minimization in relu networks.\" International Conference on Algorithmic Learning Theory. PMLR, 2023.\n", "weaknesses": "I found the empirical evidence to be somewhat unsatisfactory since it only includes two datasets. Experiments on further tasks/datasets may greatly solidify the viability of ltrRNN. Such experiments can reveal whether the low tensor rank dynamics are specific to the type of neural recording data examined here or it is more general (and thus significant).\n\nAn additional (more minor) comment: Some of the terms used are non-standard in the machine learning literature. Since results such as those in Section 6 can be of interest to researchers not familiar with this terminology it may be best to clarify (e.g. in footnotes or an appendix) their meaning. For example, the terms “trial”, “task condition”, “chaotic regime”.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the positive feedback regarding the novelty and utility of our work.\n\n**Q1. Relationship between continuous-time and discrete-time RNNs.**\n\n**A.**  We used a formulation of a continuous-time RNN that is commonly used in neuroscience applications of machine learning research [1-3], as these equations can be seen as approximations to more biologically motivated systems of ODEs describing neural circuit dynamics [4].  More generally, there has been a recent surge of interest in applications of ODEs to machine learning [5], as moving to the continuous domain enables the large body of literature on the mathematical theory behind ODEs and dynamical systems to be leveraged. For example, we exploited the adjoint state method to derive our analytical bounds on the ranks of the weight updates. In broad strokes, the adjoint method works by defining a new (\"adjoint\") dynamical system that is driven by the time-inverted dynamics of the original ODE (see Supplementary Materials D.1.1-2 for a precise mathematical treatment). However, this is not always possible for a discrete-time RNN ($\\mathbf{x}_{t+1} = f(\\mathbf{x}_t,\\theta)$), as a non-invertible $f$ would mean that the dynamics of $\\mathbf{x}_t$ cannot be time-inverted (as opposed to the continuous-time RNN, $\\mathbf{\\dot x} = f(\\mathbf{x},\\theta)$). Therefore, while we expect that an extension of our analyses from continuous  to discrete-time RNNs is possible, it would be a non-trivial exercise.\n\nRegarding the update rule, we note that while our bounds in the main text are based on gradient descent, they can be extended to Adam (line 777 of Supplementary Material D). In the revised version of the manuscript, we will point the reader to the relevant section of the Supplementary Material.\n\nFinally, we thank the reviewer for pointing out these additional citations on implicit regularization, which we had missed. We will cite these in the discussion section.\n\n**Q2. Unfamiliar terminology.**\n\n**A.** We thank the reviewer for pointing out the neuroscience-specific jargon. We will modify the revised version with the following definitions: Typically, in a neuroscience experiment, the animal is required to perform the same task over many repetitions in order to quantify the variance of neural activity. Each of these repetitions is a *trial*. In addition, a typical task often has many \"conditions\" that change the target output of the task (for RNNs different condition corresponds to mapping a different input to another target output). Regarding chaotic dynamics, it has been shown that in RNNs with weights $\\sim \\mathcal{N}(0,\\sigma^2)$, increasing $\\sigma^2$ leads to a transition from non-chaotic dynamics (the system has non-positive Lyapunov exponents and $\\bf x$ settles into a fixed point attractor) to chaotic dynamics (the system has at least one positive Lyapunov exponent and $\\bf x$ displays large fluctuations over time, see e.g. [3]). We will incorporate these definitions in the introduction and in footnotes in the revised version.\n\n**Q3. Application to other datasets.**\n\n**A.** To address the reviewer's question, and to demonstrate that the idea of low-tensor-rank learning is not specific to motor learning, we have now performed additional experiments (see General Response to Reviewers for details) which confirm that other forms of learning in neural data and in task-trained RNNs are often low rank.\n\nIn general, however, we acknowledge that we do not expect *all* changes in brain activity to necessarily have low tensor rank weights. An important counterexample is the bump attractor network [9]. In this model, each neuron has excitatory (positive weight) connections with its adjacent neurons, and inhibitory (negative weight) connections with distal neurons, e.g. $W_{ij}^{(k)}=w_E$ if $|i-j|<r$, else $W_{ij}^{(k)} = -w_I$, where $w_E,w_I>0$ and $r$ is the radius of the excitatory connections. Because $W^{(k)}$ has banded structure along its diagonal, it will have full matrix rank. Therefore in the case of the brain learning a bump attractor network, and since a high matrix rank in one of the slices of the tensor implies a high tensor rank, we expect $\\bf W$ to have high tensor rank. As this is an important counterexample to the low tensor rank learning framework, we will incorporate it into the discussion of our revised paper regarding universality and limitations of the ltrRNN framework.\n\n**References**\n\n[1] Turner, Dabholkar, and Barak. \"Charting and navigating the space of solutions for recurrent neural networks.\" *NeurIPS* 2021.\n\n[2] Schuessler, Mastrogiuseppe, Dubreuil, Ostojic, and Barak. \"The interplay between randomness and structure during learning in RNNs.\" *NeurIPS* 2020.\n\n[3] Kadmon, Timcheck, and Ganguli. \"Predictive coding in balanced neural networks with noise, chaos and delays.\" *NeurIPS* 2020.\n\n[4] Dayan and Abbott. Theoretical Neuroscience. *MIT Press*, 2001.\n\n[5] Chen, Rubanova, Bettencourt, and Duvenaud. \"Neural ordinary differential equations.\" *NeurIPS* 2018.\n\n[6] Pontryagin, Mishchenko, Boltyanskii, and Gamkrelidze. The mathematical theory of optimal processes. *Classics of Soviet mathematics*, 1962.\n\n[7] Miller and Fumarola. \"Mathematical equivalence of two common forms of firing-rate models of neural networks.\" *Neural Computation*, 2012.\n\n[8] Humphreys, Daie, Svoboda, Botvinick, and Lillicrap. \"BCI learning phenomena can be explained by gradient-based optimization.\" *bioRxiv* 2022. https://doi.org/10.1101/2022.12.08.519453\n\n[9] Compte, Brunel, Goldman-Rakic, and Wang. \"Synaptic mechanisms and network dynamics underlying spatial working memory in a cortical network model.\" *Cerebral Cortex* 2000."}]

Author rebuttal:
We thank the reviewers for their helpful and supportive comments. We are pleased to have received positive feedback regarding the novelty and interest of our submission from the reviewers, several of whom are self-described non-experts in neuroscience. We believe this highlights the potential for our work to be relevant to many research areas across the broader machine learning community.

We have made two substantial improvements to the paper in response to the reviewers' constructive criticisms. Results from our new analyses can be found in **Figures R1 and R2** of the associated 1-page PDF.

**New simulations.**
To investigate the generality of our results, we have now expanded our analyses by testing the low-tensor-rank learning hypothesis on three additional RNN models commonly used in neuroscience. These include:

* Sensory evidence accumulation task [1] where an RNN must learn to integrate a noisy instantaneous input (**Figure R1 i**).
* Contextual decision making task [2] where an RNN must decide which of two noisy inputs to integrate (**Figure R1 ii**).
* Working memory task [3] where an RNN must maintain a representation of a past input over time in order to compare it to a later input (**Figure R1 iii**).

Along with the motor adaptation example in the original paper, these four tasks span perceptual, motor and cognitive processes. Moreover, RNNs trained on these tasks have been proposed as models of different areas of the brain involved in those processes (e.g., [4,5]). Thus, these new experiments constitute a more systematic validation of the low-tensor-rank weights hypothesis we put forth.

We followed the same procedure to determine the tensor rank of learning for these tasks as for the task-trained RNN in our original submission (Figure  4). That is, we first trained an RNN on each of these tasks using gradient descent, then used PARAFAC to determine the tensor rank of the resulting neuron $\times$ neuron $\times$ iteration tensor of weights over learning. *In each of these tasks we found that the variance explained indeed saturated at low tensor ranks* (at $R=1,4,$ and $3$; **Figure R1 c**). We will add these results as a supplementary figure to the revised version of the paper.

**New data application.**
In addition, we applied the ltrRNN model to a new dataset consisting of neural recordings in mouse visual cortex during a perceptual learning task [6] (**Figure R2 a**). As a reminder, in the ltrRNN we fit an RNN to neural activity with the explicit constraint that the weight tensor has tensor rank $R$ (Figure 1; Supplementary Material A). We found that the performance of ltrRNN in this new dataset saturated at low ranks around $R = 3$ or $4$ (**Figure R2 b**). The fitted inputs distinguished rewarded from non-rewarded stimuli after the time at which the sensory stimulus was presented (**Figure R2 c**). Some of the trial factors appeared to track slow changes over learning while others remained relatively stable over trials (**Figure R2 d**). *These new results complement our original application to data from monkey motor cortical recordings during an adaptation task, demonstrating that the ltrRNN framework is not specific to a particular kind of learning, brain region, or species.*

***Overall we believe that our new simulations and data analyses, together with our mathematical results which give generic bounds on the ranks of gradient-based learning, validate the breadth of the ltrRNN framework.***

**Clarity and terminology.**
Several reviewers also noted that parts of the manuscript were difficult to understand due to (1) some details or definitions being hidden in the supplementary information and (2) undefined terminology (often, jargon specific to the neuroscience community). We strongly agree that it is important that our manuscript be accessible to the broader machine learning community, as we believe many of our results are of more general interest, beyond applications in neuroscience. To address this, we will make the following changes to the paper:

* To avoid overloading the reader, we had originally hidden some of the model specifications and training in the Supplementary Material. The reviewers fairly pointed out that this made it difficult to understand the paper. We will correct this by 1) incorporating many of these technical details back into the main text, and 2) by including pseudocode of the ltrRNN fitting procedure.
* To aid the non-specialist reader, we will add definitions of neuroscience-specific jargon.
* We will add a *main contributions* section to highlight the two principal components of our paper: 1) ltrRNN as a method for fitting neural data, and 2) the theoretical results regarding gradient learning dynamics. We believe this section will help readers from across different communities of NeurIPS better navigate the paper.

We have provided detailed responses to individual reviewers' questions below.

**References**

[1] Zoltowski, Pillown, and Linderman. "A general recurrent state space framework for modeling neural dynamics during decision-making." *ICML* 2020.

[2] Valente, Pillow, and Ostojic. "Extracting computational mechanisms from neural data using low-rank RNNs." *NeurIPS* 2022.

[3] Schuessler, Mastrogiuseppe, Dubreuil, Ostojic, and Barak. "The interplay between randomness and structure during learning in RNNs." *NeurIPS* 2020.

[4] Feulner, Perich, Chowdhury, Miller, Gallego, and Clopath. "Small, correlated changes in synaptic connectivity may facilitate rapid motor learning." *Nature Communications* 2022.

[5] Mante, Sussillo, Shenoy, and Newsome. "Context-dependent computation by recurrent dynamics in prefrontal cortex. *Nature* 2014.

[6] Khan, Poort, Chadwick, Blot, Sahani, Mrsic-Flogel, and Hofer. "Distinct learning-induced changes in stimulus selectivity and interactions of GABAergic interneuron classes in visual cortex." *Nature Neuroscience* 2018.

Meta-review: The authors in this manuscript study the stages of learning in recurrent neural networks and how low-rank connectivity emerges. The authors use their observation that the low-dimensional space of the low-rank connectivity is fairly consistent over learning to improve learning algorithms. The theory of low-rank learning was followed up by experiments in the domain of motor learning. The reviewers agreed that the theoretical and practical aspects of the paper were interesting and well presented, and thus I recommend this work be accepted.
Decision: Accept (poster)