Title: Low Tensor Rank Learning of Neural Dynamics
Venue: NeurIPS 2023
Accepted: yes
Keywords: Recurrent Neural Networks

Computational Neuroscience

Neural Data Analysis

Tensor

Learning
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

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 largescale 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.

1 Introduction

Populations of neurons perform tasks through their dynamics, and these computations can be understood through the lens of recurrent neural networks (RNNs) [1, 2]. Recent work has shown that RNNs trained on idealized versions of behavioural tasks from neuroscience experiments can be reverse-engineered to 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, an 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 changes in neural connectivity [12, 13]. For example, it has been observed that gradient descent in task-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, we consider how RNN dynamics evolve as a result of structured changes in connectivity over learning. Specifically, we consider changes in RNN connectivity over multiple

Angus Chadwick ∗

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 us 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 a constraint on the covariance of the trial factors we are also able to separate smooth changes over learning from condition-specific 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-rank RNNs (ltrRNNs) 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 in 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 ∈ R N × N , where each element W ij represents the connectivity from neuron j to i . Here, we consider a set of K discrete samples of the weight matrix over learning: W ( k +1) = W ( k ) +∆ W ( k ) for k ∈ { 0 , ..., K -1 } . In neuroscience, ∆ W ( k ) could represent plasticity-induced changes in connectivity strength following each trial in the experiment (e.g., a perceptual decision or motor action), whereas in machine learning they could represent weight updates arising from the application of a given learning rule to a set of training data. Learning can then be summarised by the weight tensor W ∈ R N × N × K formed by stacking together the weight matrices W ( k ) on all trials 1 , i.e., W ijk = W ( k ) ij .

Here, we investigate the multilinear structure of the weight tensor W in order to gain insight into the relationship between learning, recurrent weights, and network dynamics (Fig. 1a). In particular, given recent studies suggesting that trained RNNs are typically low (matrix) rank [14], we ask whether the weight tensor is low tensor rank , i.e., if the weight tensor can be written as a sum of R rank-1 components:

where a r , b r ∈ R N and c r ∈ R K for r ∈ { 1 , . . . , R } (Fig. 1b). A conse-

a

RNNweighttensor

b

Trial

Trial

1.0

Trial

0.0

Neuron

-1.0

Neuron

Neuron

Neuron

Neuron

Neuron

C

d

Learning

X1

X1

Condition

Condition

ndu

Trial

Time

Time

Time

aj

Xn

a1

Xn

Condition: 1,2

X2

X3

a2

X2

X3

a2

quence of the low tensor rank assumption is that the weights vec ( W ( k ) ) must evolve within an R -dimensional subspace of R N 2 spanned by the set of rank-1 matrices { vec ( a r ⊗ b r ) } R r =1 .

Low tensor rank RNNs. To investigate the structure of weight tensors formed over learning in recurrent networks, we consider a continuous-time RNN (Fig. 1c). The RNN dynamics on trial k are given by:

for B ∈ R N × N input , u ( k ) ( t ) ∈ R N input . If the low tensor rank hypothesis is satisfied, the system can be written as:

2

(k)

a

Neuron 2

0

2

-2

0

2-2

0

2

Neuron 1

Neuron 1

2-

Trial factor

Bifurcation

Learning

0

0

20

40

60

80

100

which reduces to a low-rank RNN on every trial [15]. As a result, x ( k ) ( t ) is constrained to an ( R + N input ) -dimensional subspace of R N spanned by { a r } R r =1 ∪ { B i } N input i =1 , and this subspace is fixed across trials (Fig. 1d).

k (Trial)

The weight matrix evolves across trials via a simple rescaling of the rank-1 components by the trial factors c ( k ) r . 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 r ⊗ b r components of the weights. In fact, non-trivial changes in the flow field (e.g., bifurcations) can occur due to the nonlinear activation ϕ ( x ) (Fig. 2; Supplementary Material A). Moreover, even in the case of a linear ltrRNN ( ϕ = id), a rich repertoire of dynamics is possible. First, the system can switch between different vector fields by allowing the corresponding trial factors to transition to zero (or nonzero) values. More generally, the eigenvalues are non-trivially related to the c r 's. The real eigenvalues can be shown to satisfy the following equality (Supplementary Material A):

λ r = 1 2 cos( θ ˜ v , v rr ) -1 ∑ r ′ c r ′ [cos( θ ˜ v , a rr ′ + θ v , b rr ′ ) + cos( θ ˜ v , a rr ′ -θ v , b rr ′ )] , where ˜ v r and v r are the r th left and right eigenvectors and θ x , y rr ′ is the angle between the r th x and r ′ th 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 r ′ } R r ′ =1 and { b r ′ } R r ′ =1 as well as the angle θ ˜ v , v rr which quantifies the non-normality of the resulting weight matrix.

Fitting ltrRNNs to data. These results suggest that ltrRNNs 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 ltrRNNs could be fit to a dataset (for example, the activations of an RNN over the course of training, or recordings from a population of neurons as an 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 ∈ R N data × T × K , comprising the activity of N data neurons measured on K trials, each of duration T time points, with single trial activations y ( k ) ( t ) . To fit the ltrRNN to such a tensor, we minimise the loss function L ( a , b , c , B, M ) = ∑ T,K t,k ∥ Mϕ ( x ( k ) ( t )) -y ( k ) ( t ) ∥ 2 , where M ∈ R N data × N is a readout matrix mapping the activation of the ltrRNN 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 .

Singleneuron

LtrRNN Data

activity

Trial 3

Trial 2

Trial 1

Time→

Learning

Above we considered how the low tensor rank assumption constrains the possible changes in RNN dynamics for the case in which the c r can change arbitrarily across trials. However, here we are interested specifically in understanding how these dynamics change over the course of learning which

typically imposes additional structure on the c r '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 c r = ( L + σI ) ¯ c r where LL T = A is the Cholesky decomposition of the smooth covariance matrix A (assuming ¯ c r ∼ N (0 , I ) , see Supplementary Material A). In contrast, we constrain the inputs to the RNN u ( k ) ( 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-independent inter-trial variability in the data comes from learning-induced 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 ( k ) ( t ) , as well as the weights vec ( W ( k ) ) , evolve in low-dimensional subspaces of R N and R N 2 , 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

LtrRNNs integrate a 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 ltrRNNs 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-induced 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 to 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 low-tensor-rank 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 a 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, 7], 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 infinite-dimensional 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 finite-dimensional 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 × time × 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 LtrRNN dynamics capture neural activity during motor learning

To test whether neural population dynamics during learning are consistent with a low tensor rank framework, we fit ltrRNNs 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 u ( k ) ( t ) with a neural ODE (see Supplementary Material B) with solution ˜ u ( i ) on a 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 ltrRNN model. For comparison we also fit a full tensor RNN (i.e. all N 2 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 learning-induced changes [16]. However, of this remaining variability, a ltrRNN with only 5 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 ltrRNN 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 ϕ = id, the activity of the ltrRNN is constrained to an R -dimensional subspace, therefore the MSE (without cross-validation) is lower-bounded by that of the rank R SVD by the Eckart-Young theorem. This suggests that the the higher performance of ltrRNN compared to PCA is due to the nonlinear mapping from the membrane potential space to the firing rate space.

Since after t = 0 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 a 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 ltrRNNs reduce to a 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 r [15]. The membrane potential can therefore be visualized via a projection onto each a r to observe the low-dimensional activity of the network [10]. We find that, compared to applying PCA directly on the neural data, ltrRNNs yield more interpretable visualizations of the condition and learning-related variability in the neural (Fig. 3d,e). Additionally, since ltrRNNs 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 to target-related dynamics, while those which are persistently changed are the temporal variability factors (Fig. 4i).

Overall, we find that learning-related 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.

a

Baseline

Perturbed

Washout

Hand movement

Early

Late

Early

Late

PCD

PMd

CS

M1

AS

cm

1 cm

b

C

P

i

2.8 -

Cross-validated

unexplained variance

0.80 -

Time

口

Perturbed

Washout

Baseline

Target angle

1PerturbedO

{BaselineO

10

10

Trial

0.70-

Train I

Test

0

1.8 *

0.65-

3.2

-10

0.60-

12345

Rank

10

PC2

PC1

-10

PC2

0

PC1

-10

0

LtrRNN, n=50

→PARAFAC

→→—Low-rank RNN

LtrRNN, n=200

→SVD

RNN

e

Tensor RNN

f

Perturbed

Baseline

3.4 *

1.8 *

Target angle

{BaselineO

PerturbedO

15

al2

2.0

10*

2.50 *

10.0

Latent 3

0

0.0

g

h

-15

Im

1.0

0.75

2.50 *

Perturbation

Go cue

-15

Latent2

0.5

Re

3.0

0.0

5.0 *

Mov onset -...

Latent1

20

0

0.5

2.2

0

2.0

2.75

15

20

○Go cue+100ms

1.0

C1

2.6

2.00

C2

1.00-

0

500

1000

00L-

200400

600

Trial

(sw) aw

5 Low tensor rank learning in task-trained RNNs

We next decided to test the low-tensor-rank learning hypothesis in a model in which we had direct access to the ground truth weight tensor. Towards this end, we trained an 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 to the hand:

where κ ( k ) ( t ) ∈ R 2 is an Ornstein-Uhlenbeck process representing noise in the execution of movement [36]. To create a 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-training the RNN to move the hand to the goal position at the go 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 to the velocity of the hand for 100 trials (Fig. 5b):

where R ∈ R 2 × 2 is a 90 ◦ clockwise rotation matrix, and c ∈ R 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.07-

Perturbation

W*-Wo

W-W*β1

MSE hand to target

1.0 -

1.0 -

Target

0.05-

0.8-

0.9-

-60

0.03-

0.8 -

Hand

0.6-

0.5 -

0.7 -

0.01-

0.4 -

50

150

250

2

4

8

10

2

Iteration

Matrix rank

Tensorrank

P

Baseline

Perturbed-early

Perturbed-late

Washout-early

Washout-late

e

f

1.0-

Im

g

Task-trained RNN

Baseline

Perturbation

h

Example trial factor

Cross-validated

unexplained variance

0.4-

0.5-

Ground-truth

bound

Re

2

PC3

1.4 -

1.0.

Task-trained RNN

PC2

1.2 -

1

0.1-

1.0 -

0.0

Minimum.

2

3

4

5

Re1

LtrRNN

0.8-

Perturbation

Tensor rank

0

50

150

250

1

LtrRNN

Trial

We next analyzed how the structure of the weight matrix changed over the baseline, perturbation, and washout blocks. The change in weights as a result of the pre-training W ∗ -W 0 , where W 0 is 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 3 rank1 terms from pre-training will be visible. However, running PARAFAC on the 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 is different than those resulting from perturbation learning (Supplementary Material C).

LtrRNN application . We next ask whether ltrRNN can uncover this low-tensor-rank structure in the weights, especially the small changes due to perturbation learning. Using our cross-validation procedure, we find that the variance of the activity of the task-trained RNN that is unexplained by the fitted ltrRNN plateaus at tensor rank 5 (Fig. 5e), consistent with our analysis of the ground truth weights. Furthermore, the weights of the ltrRNN share similar spectral properties, and a 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 ltrRNN 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, ltrRNN was able to uncover this structure in 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 a set of mathematical results regarding the tensor rank of gradient-based learning in RNNs. Towards this end, we use the method of the adjoint [37, 38] to derive a dynamical system whose solution is the gradient of a loss functional with respect to the RNN weights [39].

The adjoint. For a dynamical system ˙ x = f ( x , θ ) , the state adjoint of the loss functional L : R n → R at a particular time point t is defined as a x ( t ) = dL ( x ( T )) d x ( t ) where T is the time of loss evaluation (but see Supplementary Material E for the case of a loss functional that is integrated over time). From the dynamics of x , the state adjoint dynamics can be derived:

However, to understand gradient-based learning, we require the parameter adjoint: a θ ( t ) = dL ( x ( T )) d θ ( t ) . This can be accomplished by concatenating the original dynamical system by its parameters z ( t ) = [ x ( t ) , θ ( t )] (noting that ˙ θ ( t ) = 0 ) to define the the augmented adjoint a z ( t ) = [ a x ( t ) , a θ ( t )] , whose dynamics can be shown to follow

with terminal condition a z ( T ) = [ dL ( x ( T )) d x ( T ) , 0 ] . The gradient of the loss with respect to the parameters can then be found by integrating the parameter adjoint dynamics backwards through time to get a θ (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], it 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 ˙ x = Wϕ ( x ) -x + B u ( t ) , where x ( t ) ∈ R n , u ( t ) ∈ R m , W ∈ R n × n , B ∈ R n × m , and ϕ : R n → R n an element-wise nonlinearity, and define a loss functional of the linearly decoded RNN state, L ( Dϕ ( x ( T )); y ( T )) for y ∈ R d . Furthermore, let W = ∑ R r α r ⊗ β r . The adjoint dynamics are then given by,

with terminal conditions

a

x

(

T

) =

ϕ

(

x

(

T

))

⊙

∑

d

r

′

L

(

D

r

′

·

ϕ

(

x

(

T

));

y

r

′

)

D

′

′

′

and

a

(

T

) = 0

.

2

r

W

From the adjoint dynamics, it can be seen that the singular values of the gradient of the loss with respect to the weights of the RNN ( ∇ W L = a W (0) ) depend on the dimensionality of the subspaces over which ϕ ( x ( t )) and a x 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:

where σ v r denotes the r th singular value of v .

This theorem demonstrates that there is a natural bound to the numerical rank [41] of gradient-based learning in RNNs. Numerically, we find that the singular value spectrum of both ϕ ( x ) and a x tend to decay exponentially, even in the case of a chaotic RNN (Fig. 6). In Supplementary Material E we 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 to the fastest decaying singular value spectrum of the adjoint.

So far, we haven't made any assumption on the architecture of the RNN such as constraints on W or ϕ . 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 1 with ϕ = id and W (0) of rank R . Furthermore suppose x ( k ) (0) is constrained to the subspace spanned by the columns of W ( k ) and B . Consider the weight tensor W = [ W (0) , W (1) , ... ] where W ( k +1) = W (0) + α ∑ k j =1 ∇ W L ( D x ( j ) ( T ); y ( j ) ) , with y ( j ) ∈ R d where x ( j ) denotes the activity of the RNN after the j th weight update. Then,

The rank of the gradient at the first step is at most rank( ∇ W L (0) ) ≤ R +1 .

The rank of the trial slices of the weight tensor is at most rank( W ( k ) ) ≤ 2 R + m + d .

The tensor rank of the weight tensor is at most rank( W ) ≤ (2 R + m + d ) 2 .

Wenote 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 n ) ⊗ 3 tensor is n 2 -n -1 [42]. In particular, in the case of W (0) = 0 and m = d = 1 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 , m = 2 in Fig. 6. We find that the true ranks fall within these bounds - strictly below due to limited machine precision - and that the numerical ranks are 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 to lie in a low-dimensional subspace of the state space. This, in turn, further pushes weight updates to lie in a low-dimensional subspace of the weight space.

The proofs of Lemma 1 and Theorems 1 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 as 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, as 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 a principled choice of low-dimensional

a

b

C

p

e

1.4-

1 -

1.00-

(s) otbol anjeA lejnus bo7

(1s) 01601

2-

0-

1.00--*

1.2-

0-

0-

0-

10101010101

05:

1.0-

-1

-2.5 0.0 2.5

0.00-

2-

010-

50

000100

value

2-

10 -

0

50

100

X

中

4-

enus

4

Wo

a

0.4-

o

（0）7

0.2-

0.0-

25

50

75

100

-1.00-

2

-6-0-

-8

(x(0)

a(0)

10

1o1

a(100)

Φ(x(100)

20

Logs

-6

8

1

101

W100

min（o(，p)

Wo:100

10

a

20

0.80-

2

4

6

2.5-

2.5-

1.0-

0.5-

singular value log1o(si)

2-

R+1

2R+m+d

(1s) 01601

2-

R+1

2R+m+d

1.00...*

2.0-

(x)Φ

0.0-

0.0

0-

0-

a

1.5-

2.5-

2-

singularvalue

2-

MSE

1.0-

2.5

0.0

X

2.5

4-

-6-

o1

(o)x

-6-

4 -

(0)7

Wo

a

a

0.5-

-00

1.0-

Log

0.0

-8·101

a(00)

x(100)

a(0)

Log

101-0

18-101

101

W100

Wo:100

min（oo²,00)

a

0.94-

0

25

50

75

100

00

0.5

1.0

1.5

2.0

T

10

T

10

2

4

6

Iteration

Time

Singular value index i

Singular value index i

Tensor rank

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 a set of unconnected neurons. In our application to neural data, we 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-specific changes, neural recordings show substantial variability across consecutive trials, which are thought to reflect a combination of i) unmeasured covariates such as 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 is able to explain many features of motor learning in neural data [17, 18], as well as by more general support for an optimization-based framework to understand neural learning [43]. The tensor rank of learning in RNNs with local synaptic rules (e.g., Hebbian) remains an open question. Towards this end, theoretical work has established links between Hebbian and gradient-based learning [44, 45], opening the possibility of an 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.

Acknowledgments

We thank Francesca Mastrogiuseppe, Friedrich Schuessler, and Sina Tootoonian for feedback on the manuscript, and Matthias Hennig, Matt Nolan, and Srdjan Ostojic for helpful discussions. We are particularly grateful to Matthew Perich for sharing his data. This work was supported by the Agence Nationale de la Recherche (ANR-20-CE37-0004; ANR-17-EURE-0017).

References

Niru Maheswaranathan, Alex Williams, Matthew Golub, Surya Ganguli, and David Sussillo. Universality and individuality in neural dynamics across large populations of recurrent networks. Advances in neural information processing systems , 32, 2019.

Saurabh Vyas, Matthew D Golub, David Sussillo, and Krishna V Shenoy. Computation through neural population dynamics. Annual review of neuroscience , 43:249-275, 2020.

Valerio Mante, David Sussillo, Krishna V Shenoy, and William T Newsome. Context-dependent computation by recurrent dynamics in prefrontal cortex. nature , 503(7474):78-84, 2013.

Jing Wang, Devika Narain, Eghbal A Hosseini, and Mehrdad Jazayeri. Flexible timing by temporal scaling of cortical responses. Nature neuroscience , 21:102-110, 2018.

Christopher J 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:1912.10189 , 2019.

Abigail A Russo, Ramin Khajeh, 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-758, 2020.

Alexis Dubreuil, Adrian Valente, Manuel Beiran, Francesca Mastrogiuseppe, and Srdjan Ostojic. The role of population structure in computations through neural dynamics. Nature neuroscience , 25(6):783-794, 2022.

Matthew G Perich, Charlotte Arlt, Sofia Soares, Megan E Young, Clayton P Mosher, Juri Minxha, Eugene Carter, Ueli Rutishauser, Peter H Rudebeck, Christopher D Harvey, et al. Inferring brain-wide interactions using data-constrained recurrent neural network models. BioRxiv , pages 2020-12, 2020.

Feng Zhu, Andrew Sedler, Harrison A Grier, Nauman Ahad, Mark Davenport, Matthew Kaufman, Andrea Giovannucci, and Chethan Pandarinath. Deep inference of latent dynamics with spatio-temporal superresolution using selective backpropagation through time. Advances in Neural Information Processing Systems , 34:2331-2345, 2021.

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-24086, 2022.

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 arXiv:2202.06159 , 2022.

Lea Duncker, Laura Driscoll, Krishna V Shenoy, Maneesh Sahani, and David Sussillo. Organizing recurrent network dynamics by task-computation to enable continual learning. Advances in neural information processing systems , 33:14387-14397, 2020.

Harsha Gurnani and N Alex Cayco-Gajic. Signatures of task learning in neural representations. Current Opinion in Neurobiology , 83:102759, 2023.

Friedrich Schuessler, Francesca Mastrogiuseppe, Alexis Dubreuil, Srdjan Ostojic, and Omri Barak. The interplay between randomness and structure during learning in rnns. Advances in neural information processing systems , 33:13352-13362, 2020.

Francesca Mastrogiuseppe and Srdjan Ostojic. Linking connectivity, dynamics, and computations in low-rank recurrent neural networks. Neuron , 99(3):609-623, 2018.

Matthew G Perich, Juan A Gallego, and Lee E Miller. A neural population mechanism for rapid learning. Neuron , 100(4):964-976, 2018.

Barbara Feulner, Matthew G Perich, Raeed H Chowdhury, Lee E Miller, Juan A Gallego, and Claudia Clopath. Small, correlated changes in synaptic connectivity may facilitate rapid motor learning. Nature communications , 13(1):5163, 2022.

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.

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-137, 2018.

Lea Duncker and Maneesh Sahani. Dynamics on the manifold: Identifying computational dynamical activity from neural population recordings. Current opinion in neurobiology , 70:163-170, 2021.

Adil G Khan, Jasper Poort, Angus Chadwick, Antonin Blot, Maneesh Sahani, Thomas D Mrsic-Flogel, and Sonja B Hofer. Distinct learning-induced changes in stimulus selectivity and interactions of gabaergic interneuron classes in visual cortex. Nature neuroscience , 21(6):851-859, 2018.

Angus Chadwick, Adil G Khan, Jasper Poort, Antonin Blot, Sonja B Hofer, Thomas D Mrsic-Flogel, and Maneesh Sahani. Learning shapes cortical dynamics to enhance integration of relevant sensory input. Neuron , 111(1):106-120, 2023.

Sukbin Lim, Jillian L McKee, Luke Woloszyn, Yali Amit, David J Freedman, David L Sheinberg, and Nicolas Brunel. Inferring learning rules from distributions of firing rates in cortical neurons. Nature neuroscience , 18(12):1804-1810, 2015.

Ian Stevenson and Konrad Koerding. Inferring spike-timing-dependent plasticity from spike train data. In J. Shawe-Taylor, R. Zemel, P. Bartlett, F. Pereira, and K.Q. Weinberger, editors, Advances in Neural Information Processing Systems , volume 24. Curran Associates, Inc., 2011.

Scott Linderman, Christopher H Stock, and Ryan P 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.

Andrew M Saxe, James L McClelland, and Surya Ganguli. Exact solutions to the nonlinear dynamics of learning in deep linear neural networks. arXiv preprint arXiv:1312.6120 , 2013.

Tamara G Kolda and Brett W Bader. Tensor decompositions and applications. SIAM review , 51(3):455-500, 2009.

Alex H Williams, Tony Hyun Kim, Forea Wang, Saurabh Vyas, Stephen I Ryu, Krishna V 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 , 98(6):1099-1115, 2018.

Hugo Soulat, Sepiedeh Keshavarzi, Troy Margrie, and Maneesh Sahani. Probabilistic tensor decomposition of neural population spiking activity. Advances in Neural Information Processing Systems , 34:1596915980, 2021.

Arthur Pellegrino, Heike Stein, and N Alex Cayco-Gajic. Disentangling mixed classes of covariability in large-scale neural data. bioRxiv , pages 2023-03, 2023.

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-1520, 2019.

Richard J 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):123-128, 2022.

Manuel Beiran, Nicolas Meirhaeghe, Hansem Sohn, Mehrdad Jazayeri, and Srdjan Ostojic. Parametric control of flexible timing through low-dimensional neural manifolds. Neuron , 2023.

Xulu Sun, Daniel J O'Shea, 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):274279, 2022.

Darby M Losey, Jay A Hennig, Emily R Oby, Matthew D Golub, Patrick T Sadlter, Kristin M Quick, Stephen I Ryu, Elizabeth C Tyler-Kabara, Aaron P Batista, Byron M Yu, et al. Learning alters neural activity to simultaneously support memory and action. bioRxiv , pages 2022-07, 2022.

Robert J Van Beers, Patrick Haggard, and Daniel M Wolpert. The role of execution noise in movement variability. Journal of neurophysiology , 91(2):1050-1063, 2004.

L Bittner. Ls pontryagin, vg boltyanskii, rv gamkrelidze, ef mishechenko, the mathematical theory of optimal processes. viii+ 360 s. new york/london 1962. john wiley & sons. preis 90/-, 1963.

Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. Advances in neural information processing systems , 31, 2018.

Balasubrahmanyan Srinivasan, Upadrasta Ravikirana Prasad, and Nalam Jaganmohan Rao. Back propagation through adjoints for the identification of nonlinear dynamic systems using recurrent neural models. IEEE Transactions on Neural Networks , 5(2):213-228, 1994.

Xuechen Li, Ting-Kam Leonard Wong, Ricky T. Q. Chen, and David Duvenaud. Scalable gradients for stochastic differential equations. International Conference on Artificial Intelligence and Statistics , 2020.

Per Christian Hansen. Rank-deficient and discrete ill-posed problems: numerical aspects of linear inversion . SIAM, 1998.

Jarosław Buczy´ nski and Joseph M Landsberg. Ranks of tensors and a generalization of secant varieties. Linear Algebra and its Applications , 438(2):668-689, 2013.

Blake A Richards, Timothy P Lillicrap, Philippe Beaudoin, Yoshua Bengio, Rafal Bogacz, Amelia Christensen, Claudia Clopath, Rui Ponte Costa, Archy de Berker, Surya Ganguli, et al. A deep learning framework for neuroscience. Nature neuroscience , 22(11):1761-1770, 2019.

Cengiz Pehlevan, Tao Hu, and Dmitri B Chklovskii. A hebbian/anti-hebbian neural network for linear subspace learning: A derivation from multidimensional scaling of streaming data. Neural computation , 27(7):1461-1495, 2015.

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.

Steffen Schotthöfer, Emanuele Zangrando, Jonas Kusch, Gianluca Ceruti, and Francesco Tudisco. Lowrank lottery tickets: finding efficient low-rank neural networks via matrix differential equations. arXiv preprint arXiv:2205.13571 , 2022.

Charles H Martin and Michael W Mahoney. Implicit self-regularization in deep neural networks: Evidence from random matrix theory and implications for learning. The Journal of Machine Learning Research , 22(1):7479-7551, 2021.

Sanjeev Arora, Nadav Cohen, Wei Hu, and Yuping Luo. Implicit regularization in deep matrix factorization. Advances in Neural Information Processing Systems , 32, 2019.

Othmar Koch and Christian Lubich. Dynamical low-rank approximation. SIAM Journal on Matrix Analysis and Applications , 29(2):434-454, 2007.

John A Rhodes. A concise proof of kruskal's theorem on tensor decomposition. Linear Algebra and its Applications , 432(7):1818-1824, 2010.

Tamara G Kolda. Orthogonal tensor decompositions. SIAM Journal on Matrix Analysis and Applications , 23(1):243-255, 2001.

Grigorios A Pavliotis. Stochastic processes and applications: diffusion processes, the Fokker-Planck and Langevin equations , volume 60. Springer, 2014.

Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980 , 2014.

David Zoltowski, Jonathan Pillow, and Scott Linderman. A general recurrent state space framework for modeling neural dynamics during decision-making. In International Conference on Machine Learning , pages 11680-11691. PMLR, 2020.

Walter Rudin. Functional analysis . McGraw-Hill, 1973.

Roger A Horn and Charles R Johnson. Topics in matrix analysis cambridge university press. Cambridge, UK , 1991.

Rudolf Emil Kalman. Mathematical description of linear dynamical systems. Journal of the Society for Industrial and Applied Mathematics, Series A: Control , 1(2):152-192, 1963.

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 ∈ R N × N × K ,

In particular, at trial k , the dynamics of the RNN can be described as,

for B ∈ R N × N input , u ( k ) ( t ) ∈ R N input , so that the RNN is a low rank RNN [15]. When training, we initialize a r ∼ N ( 0 , I ) , b r = a r . Furthermore, the weights are parameterized such that || a r || = || b r || = 1 , so that the magnitude of a component is captured by c r .

Smoothness constraint over trials . We further constrain the initial covariance in trial of the trial factors by parameterizing them as c r = ( L + σI ) ¯ c r where LL T = A is the Cholesky decomposition of the smooth covariance matrix A , and ¯ c r is initialized as ¯ c r ∼ N ( 0 , I ) . In particular, we use a rational quadratic kernel s 2 (1+(2 l ) -1 ( k i -k j ) 2 ) -1 , where k i is the i th trial index. This is equivalent to performing a 3 -mode matrix-tensor product on the weight tensor itself ( L + σI ) × 3 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). In particular, given that ( L + σI ) is invertible, any possible c r can in theory be obtained upon optimization of the ¯ c r . By additionally setting a regularization on ¯ c r , we penalize the smoothness of c r as non-smooth solutions have diverging ¯ c r . In this way, we bias the optimization process towards smoother c r '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 a failure of the animal to perform a given trial does not imply that a trial wasn't informative, or that learning did not occur.

Condition-wise 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 a deep neural network (DNN)

where v ( i ) ( t ) ∈ R N NODE and D ∈ R N input × N NODE . Throughout this work, we used a fully-connected 3-layer DNN with layers of size 150 and ReLU nonlinearities. This provides inputs whose dynamics are considerably less constrained than those generated by low-rank RNNs. Thus, we 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 to 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, or 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 to learning.

Loss . In sections 4 and 5 we focus on optimizing for the mean squared error (MSE)

where y are firing rate estimates data and M ∈ R N × N data . In section B we present supplementary results on optimizing the Poisson log-likelihood with respect to spike data,

where y are binned spike data.

Pseudocode. The following pseudocode summarizes the steps of fitting an ltrRNN to neural data. For the sake of clarity, we present the simplest case: an autonomous ltrRNN with fixed initial state. 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 to change over conditions and trials, the evaluation of the dynamical system can be done over {-T 0 ∆ t, ..., 0 , ..., T ∆ t } ,where x 0 is now the state at -T 0 ∆ t , and the fit to data is still done on the non-negative time states.

Algorithm 1 Low tensor rank recurrent neural network fit to data

inputs:

initializations:

randomly initialize a r , b r ∈ R N for r = 1 , . . . , R randomly initialize ¯ c r ∈ R N for r = 1 , . . . , R randomly initialize q ∈ R R , M ∈ R N data × N

definitions:

while the loss l hasn't converged do

end for

SGDUpdate ( a r , b r , ¯ c r , M, q ) end while

▷ The initial state ▷ For all trials (parallelizable) ▷ Matrix of activity of the RNN during trial k

Code availability . The ltrRNN implementation can be found at https://github.com/ arthur-pe/LtrRNN .

A.2 The dynamics of ltrRNNs

Rich changes of dynamics through oblique columns . Unlike for a matrix rank decomposition, a tensor rank decomposition can be unique even for non-orthogonal factors. A sufficient condition for uniqueness is that r a + r b + r c ≤ R -2 where, without loss of generality, r a denotes the maximum number of linearly dependent columns of a [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 case where the a j 's are not orthogonal, non-trivial qualitative changes in the vector fields can occur. Since for any given trial k , an ltrRNN is simply a low rank RNN [15],

▷ Neural data tensor of shape neuron × time × trial ▷ Time points of evaluation of the dynamical system

so that the dynamics of x ( k ) are constrained to span { a j } ∪ { B j } . Unlike low rank RNNs, ltrRNN are not necessarily invariant under changes of bases of a j 's. Nevertheless, we can introduce an orthonormal basis { ˜ a j } so that,

In particular, notice that the dynamics along all ˜ a i could potentially be affected by varying c ( k ) j .

Conversely, the a i 's and b i 's being respectively orthonormal - e.g. as in a singular value decomposition - is not a sufficient condition for uniqueness of the tensor rank decomposition [51]. Nevertheless, constraining the a i to be orthogonal and ensuring that the Kruskal constraint is satisfied, the vector fields are then orthogonal. In that case, varying c ( k ) i for some i corresponds to rescaling the vector field along a i . Nevertheless, as is illustrated below with a single component, the leak term allows the system to display typical properties of linear and nonlinear dynamical systems.

Bifurcation in a tensor rank-one RNN . A classical example of bifurcation in a two-neuron system is that of the pitchfork supercritical bifurcation. Here, we show that it is essentially a tensor rank one RNN. We also illustrate how the corresponding linear RNN bifurcates. Let a = b = [ -1 / √ 2 , 1 / √ 2] T so that,

That is x 1 = -x 2 . We consider two cases, ϕ = tanh and ϕ = id, both odd functions. Introducing the two in the previous equation, -ϕ ( x 2 ) = ϕ ( x 1 ) . So that x i = cϕ ( x i ) . Now considering each activation function separately,

tanh: i) c > 1 has two solutions. ii) c ≤ 1 one solution (the origin). At c = 1 the origin is a non-hyperbolic fixed point.

id: i) c = 1 for any x 1 . The non-zero eigenvalue of the Jacobian of the system is negative, therefore it is a line attractor. ii) x 1 = 0 for any c . Then the Jacobian of the system at the origin has both positive and negative eigenvalues for c > 1 and only negative for c ≤ 1 .

Trial

C=2.0

C=1.0

C=0.5

3

RNN (tanh)

Neuron 2

0

3

3

LDS (identity)

Neuron 2

a

3

-3

0

3-3

0

3-3

0

3

Neuron 1

Neuron 1

Neuron 1

More general changes in vector fields . At the cost of increasing the rank, an ltrRNN can possibly transition between two arbitrary vector fields of ranks R 1 and R 2 . For example, let c j = [1 , 0 . 9 , ..., 0] for j ∈ { 1 , ..., R 1 } and c j = [0 , ..., 0 . 9 , 1] for j ∈ { R 1 +1 , ..., R 1 + R 2 } . There might however be multiple bifurcations between the first and last trial. More generally, given that any tensor has a (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 ltrRNN 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 , W ( k ) has the rank decomposition

where a ∈ C N × R ′ , b ∈ C R ′ × N , c ( k ) ∈ R R ′ , and a i , b i are the rows/colums of a , b respectively. One such low rank decomposition is the eigendecomposition 34

By the rank-nullity theorem, R is the rank of W ( k ) , so that R ′ ≥ R , with R = R ′ when Equation (A.3.1) is a minimal rank decomposition. Equating the two decompositions gives a general expression for the eigenvalues:

If λ i , a , b , c ( k ) ∈ R , this gives rise to the result stated in the main text:

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 i ∥ = ∥ b i ∥ = 1 , and often enforce orthogonality on the a i and/or b i , thereby introducing additional constraints on the relationship between the c 's and λ 's. Normal matrices have cos( θ ˜ v , v ij ) = δ ij and ˜ v i = v i , 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, n = 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-wise 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 ¯ y i ( t ) was given by

where t = 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 inputs from upstream regions during the preparatory period, after which the input shuts off so that the resulting activity during the reach evolves autonomously via the recurrent dynamics dynamics from that initial state. We therefore set u ( t ) = 0 for t > 100 where t = 0 is the time of the go cue. Where the 100 ms account for a sensory delay.

Cross-validation procedure . We cross-validated the optimal rank and number of neurons of the ltrRNN. 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-by-trial-by-neuron data tensor T ijk is strongly correlated with T i -1 ,jk and T i +1 ,jk . For example, assuming the data are continuous, a simple average of these entries will give an optimal estimate T ijk 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-validation as section 4 but with 1 × 1 × 1 blocks, and found the test loss was similar as using the train loss over the whole dataset (Sup. Fig. 3).

To counter this effect, sets of contiguous entries [ T ijk , ..., T i + n,jk ] can be held out of training, and the interior of these blocks [ T i + q,jk , ..., T i + n -q,jk ] used for testing [30]. Here, given that we are interested in uncovering smooth changes in neural activity over slow timescales, we hold out n -bym matrices [[ T ijk , ..., T i + n,jk ] , ..., [ T i,j + m,k , ..., T i + n,j + m,k ]] (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 as 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 is to use the negative Poisson log likelihood loss [11, 10]. We fitted a ltrRNN ( R = 5 , n = 200 as in the MSE

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 to the presence of high firing rate neurons which are normalized by the preprocessing procedure but not likelihood fitting.

a

Baseline

Perturbed -early

Perturbed-late

Washout-early

Washout-late

7.5 -

5.0-

2.5 -

(cm)

0.0

>

2.5 -

5.0 -

7.5 -

-5

0

5

0

5

0

-5

0

5

× (cm)

× (cm)

× (cm)

x(cm)

× (cm)

b

Baseline

Perturbed

4

Washout

C

1.0 -

Firing rate estimate

Firing rate estimate

2

Firing rate estimate

0.5-

0.0-

2

-2-

-2

0.5 -

Baseline

Go cue

Perturbed

-4-

-250

0

250

500

750

-4

-250

0

250

500

750

4 -

-250

0

250

---Mov onset

500

750

-1.0 -

-0.5

0:0

Washout

0.5

()

()

(sw) w

PC1

Baseline

Perturbed -early

Perturbed-late

Washout-early

Washout-late

cm

2

-5.0

0.0

2.5

-2.5

0:0

2.5

-2.50.0、2.5

-2.50.0

2.5

-2.50.0、2.55.0

x(cm)

×(cm)

× (cm)

× (cm)

×(cm)

Baseline

4 -

Perturbed

4 -

Washout

Firing rate estimate

Firing rate estimate

2-

Firing rate estimate

2

0.5-

0-

PC2

0.0-

2-

-2-

-2-

0.5 -

Baseline

Go cue

Perturbed

Mov onset

Washout

-4 -

-250

0

Time (ms)

250

500

750

4 -

-250

0

(s)

250

500

750

4 -

250

0

Time (ms)

250

500

750

-1.0

-0.5

PC1

0.0

0.5

1.0

C Task-trained RNN model of motor learning

Task design . A trial is split into a preparatory period t ∈ [0 , T go ) and execution period [ T go , T end ] . Here we set T go = 2 , T end = 4 . To model a ballistic reach, the RNN receives input the target information and a hold cue during the preparatory period. During the execution period it evolves

0.80-

Entire trials

20 trials × 100 ms x 1 neuron

iance

0.75 -

Singleentries

varia

No cross-validation

0.70-

Unexplained

0.65-

0.60-

0.55-

12345

10

Rank

20

a

unexplained variance

0.75 -

e

h

Cross-validated

0.70-

0.65-

0.60-

1.4 *

0.55 -

3.00 *

1 2 3 4 5

10

Rank

20

b

Perturbation

Baseline

?

Per

BL

-5

5f

Im

wo

1.0

2

5

0.5

PC1

PC1

90g

-5°PC2

0.0

Re

0

1.8 *

P

15

0.5

1.0

2.50 *

Latent 3

g

1.25 *

2.6

15

3.00 *

Mov onset

Go cue

2.0

20

1.50

2.75

C1

C2

-10

Latent 2

0

10-20

Latent

0

2.50

2.00

1.50 -

0

200400600800

FF

10

0

200400600

Trial

(su) ew

autonomously.

where B target ∈ R n × 2 , u ( k ) target = [cos( θ ( k ) ) , sin( θ ( k ) )] is a static vector representing the position of the target, B hold ∈ R n × 1 , u ( k ) hold = 1 a cue indicating to hold movement, and dW the infinitesimal increments of a Wiener process [52]. The dynamics of the hand are given in section 5 of the main text. The loss is taken to minimize the distance between the hand y ( k ) ( t ) and the target v ( k ) throughout the execution period, while keeping the hand still during the preparatory period, that is L ( W,B target , B hold , D ) =

BL

AD

wo

a

PC1

Latent 3

2.2

1.9

2.5

2.1

C13.5

1.>1.8

C2

Go cU6

P

1.5  Im

1.0

0.5

Latent2

0

Re

3

0.0

0.0

0.5

1.0

1.5

2.0

0.5

1.0-

1.5 *

00

400

600

In particular, the speed of the reach is only constrained by the noise of the RNN and the hand. The dynamical system as a whole is evaluated with a 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 -W 0 ⊗ 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 a rank k +1 PARAFAC model will be related to those of fitting a rank k model. Nevertheless, motor perturbation learning had a significant change on the eigenvalues and activity of 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 on the change in weights matrix W ∗ -W 0 , where W ∗ are the weights of the network post-training, but pre-motor perturbation learning (Fig. 5b.). We found that the changes in weights were well captured by a rank3 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 a tensor rank 2 decomposition. The combination of these results upper-bounds the tensor rank of the overall changes in weights to 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 ∗ -W 0 and found approximately a remaining 0 . 2 unexplained variance, suggesting that the columns of W 0 -W ∗ and W ∗ were not orthogonal, but did not span the same subspace. Combined, these results suggest that the numerical tensor rank is at least 4 and at most 5 , consistent with the results uncovered by ltrRNN from the RNN activity (Fig. 5e).

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 a one-dimensional OrnsteinUhlenbeck (OU) process input whose expected steady-state is either positive or negative. The target output is either +1 or -1 if the mean of the input is positive or negative, respectively.

Contextual decision making task [10] (Fig. 6ii). The RNN receives a three-dimensional input. The first two inputs are independent OU processes as in the previous task. The third input is binary and constant, and determines which of the two stochastic inputs must be integrated. The target output is +1 (or -1) if the mean of the input of the OU process indicated by the contextual input is positive (or negative).

Working memory task [14] (Fig. 6iii). The RNN receives a 1-dimensional input consisting of two stimuli of different amplitudes separated by a 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 a 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 (0) ⊗ 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 , 4 , and 3 ; Fig. 6e).

a

b

C

p

e

1.0 -

0.7 -

0.75

0.8-

0.6-

02 - Im

0.25

0.50

0.6-

MSE Loss

0.5-

0.6-

0.8 -

0.00

0.4 -

Re

0.25

0.4 -

0.3-

0.4-

0.50-

0.2

0.2-

0.2-

0.75 -

0.0

0.1

0

50

100

150

200

0:2

-00

10

ii

2.5

1.0

1.0 -

1.25 -

ndul

-2.5

0.0

0.5-

MSELoss

1.00 -

0.8 -

2.5

0.0 -

0.75-

Re

80

0.4

Input 2

0.0

0.5

0.50-

0.2-

2.5

0

1.0

0

2

0.25 -

0

100

200

0.0-

4

6

8

10

ii

1.4 -

1.0 -

0.35

1.0 -

1.2-

1.0 -

0.5 -

0.25-

0.30-

0.5

1.0

explained

0.6-

0.8-

0.8-

0.0-

Re

0.6-

1.0

1.5

2.0

0.4 *

0.4 -

0.5 -

0.15 -

0.5

0.2 -

0.2-

0.10 -

0.0-

0

6

100

200

300

400

0.0-

6

8

10

Time

Iteration

Rank

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 ˙ x = f ( x , θ ) ∈ R n where θ ∈ R k is a set of parameters. Furthermore, let the functional L : R n → R such that L ( x ( T )) is our loss. First, define the state adjoint ,

Notice that since x ( t +∆ t ) is a function of x ( t ) ,

By Taylor expanding x ( t +∆ t ) = x ( t ) + ∆ tf ( x ( t )) + O (∆ t 2 ) , we get,

where I is the identity matrix. By rearranging,

Taking the limit as ∆ t → 0 ,

We now have the dynamics of the adjoint; all that remains is that we find an initial (or rather terminal) condition. For this, notice that

is 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 dL ( x ( T )) /d θ . For this, it suffices to augment the original dynamical system with its parameters 5 ˙ z = [ f ( x , θ ) , 0 ] and initial (later terminal) condition z (0) = [ x (0) , θ ] . Defining the loss ¯ L ( z ( T )) = L ( x ( T )) , the adjoint of this augmented system is,

which contains the desired term dL ( x ( T )) /d θ . It now remains to describe the dynamics of the augmented system. By the same argument as for the state adjoint (E.1.5),

By (E.1.7) and by unconcatenating z ,

Hence the following dynamical system can be evaluated,

with terminal condition

Notice that to obtain the terminal condition, since it depends on x ( T ) , the original dynamical system must be evaluated forward once.

E.1.3 RNN adjoint

Let ˙ x = f ( x , W ) = Wϕ ( x ) -x + B u ( t ) and L ( x ( T )) = || Dϕ ( x ( T )) -y || 2 for y ∈ R d . Furthermore, as it will be convenient, let W = ∑ R i α i ⊗ β i . We will derive one by one the terms needed to characterize the parameter adjoint. First the Jacobian is

where ⊙ denotes the element-wise product, I the identity matrix, and ϕ ′ the derivative of ϕ . Next,

̸

that is

Finally the terminal condition,

that is,

Or more explicitly, ˙ a z =

E.2 Rank of the gradient

E.2.1 The gradient as a composition of operators

In this section, we prove Theorems 1. In order to derive bounds on the singular values of ∇ W L = a W (0) , we shall now consider a x and ϕ ( x ) as linear operators with integration. Namely,

where y ∈ H for H some suitable Hilbert space, such as L 2 for our case.

More formally, let a x , ϕ ( x ) ∈ B 0 , 0 , where B 0 , 0 is the Banach space of i) compact ii) bounded operators from H to R n , such that a x , ϕ ( x ) : H → R n . In particular, we note that compactness follows from the image of a x , that is a x ( H ) , to be a vector subspace of R n and therefore be of finite rank. By the same argument, ϕ ( x ) is compact, and therefore so is ϕ ( x ) ∗ by Schauder's theorem [55], where ∗ is adjunction. Furthermore, notice that ϕ ( x ) , a x are solutions of dynamical systems with differentiable right hand side and therefore bounded (in R n ) if they are evaluated for finite time, and therefore bounded when seen as operators. The following result can now be applied:

Lemma 2 ([55]) . Let T ∈ B 0 , 0 , then T admits a singular value decomposition. Furthermore, this singular value decomposition is of finite rank.

We can now prove the main theorem.

Proof of Theorem 1. Notice that the composition of the two operators is: a x ◦ ϕ ( x ) ∗ = ∇ W L . Furthermore, the singular values of ∇ W L are,

which can be bounded as,

(E.2.5)

Now notice that, for any operator, akin to the matrix case, T 1 , T 2 , the adjoint of their composition is ( T 1 T 2 ) ∗ = T ∗ 2 T ∗ 1 . Furthermore,

that is,

Noticing that σ a ∗ x i = σ a x i and σ ∇ W L i = σ ( ∇ W L ) T i ,

Combining E.2.6 and E.2.10 we obtain the sought upper bound of Theorem 1,

Similar steps can be used to derive the lower bound of Theorem 1, using instead the identity σ T 1 n || T 2 y || ≤ || T 1 T 2 y || where T 1 , T 2 ∈ B 0 , 0 and σ T 1 n is the smallest non-zero singular value of T 1 .

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 of the singular values of the gradient can be obtained. Let U a x i , U ϕ ( x ) i be the right singular vectors of respectively a x and ϕ ( x ) , so that V a x i ( t ) , V ϕ ( x ) i ( t ) are the left singular vectors and σ a x i , σ ϕ ( x ) j the singular values.

The characterization of the rank of the gradient has thus shifted to the firing rate and adjoint spaces. Furthermore, the integral is 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 a given direction in weight space V a x i ⊗ V ϕ ( x ) j , 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 , ∆ t..., q ∆ t . In that case, the gradient can be estimated as a simple matrix-matrix product. Let A = [ a x (0) , ..., a x ( q ∆ t )] and B = [ ϕ ( x )(0) , ..., ϕ ( x )( q ∆ t )] . Then ∇ W L = ∆ tAB T , and the bounds E. 2 . 6 and E. 2 . 10 follow from classic matrix-matrix product bounds [56]. In particular, given that ϕ ( x ) and a x are smooth, without loss of generality,

This simple matrix-matrix product opens up the possibility of fast RNN adjoint implementations as, unlike computing a x df dW in general adjoint solvers, which requires O ( qn 3 ) time and O ( n 2 ) memory complexity for an n -dimensional RNN and q time steps, here the time complexity drops to O ( qn 2 ) .

E.2.3 Rank of linear RNNs

In this section we prove Theorem 2.

Proof of claim 1. If W = ∑ R i α i ⊗ β i , then by Lemma 1, the dynamics of the adjoint are constrained to the span of the rows of W , namely, ˙ a x ∈ span { β i } . Therefore, a x ∈ span { β i } ∪ { a x ( T ) } , which is a at most R +1 dimensional subspace. If a x is constrained to a at most R +1 dimensional subspace, then ˙ a W = a x ⊗ x is also constrained to a at most R +1 dimensional subspace. Since a W ( T ) = 0 , a W is constrained to the same subspace as its dynamics, and in particular, rank a W (0) = rank ∇ W L (0) ≤ R +1 .

Proof of claim 2. Suppose W ( k ) = ∑ 2 R + m + d i α ( k ) i ⊗ β ( k ) j , where α ( k ) i , β ( k ) i ∈ span { α (0) i } ∪ { β (0) i } ∪ { B i } ∪ { D i } := V . In particular, notice that V is only dependent on the initial weight. Then by a similar argument as above, a ( k ) x ∈ V , which implies a W (0) ( k ) = ∇ W L ( k ) ∈ V . Therefore W ( k +1) = W ( k ) + γ ∇ W L ( k ) = ∑ 2 R + m + d i α ( k +1) i ⊗ β ( k +1) j with α ( k +1) i ∈ V . That is rank W ( k +1) ≤ 2 R + m + d .

Proof of claim 3. Mutatis mutandis β ( k +1) ∈ V , that is x ∈ V . Therefore, W ( k ) = ∑ 2 R + m + d ij c ( k ) ij α i ⊗ α j for some c ( k ) ij 's, where α i ∈ V . Or equivalently, W = ∑ 2 R + m + d ij α i ⊗ α j ⊗ c ij . That is, rank W ≤ (2 R + m + d ) 2 .

E.3 Extensions of our results

Loss integrated over time . Commonly, the loss considered might be integrated,

The parameter adjoint is dependent only linearly on the state adjoint, we may therefore integrate the state adjoint for all initial conditions.

The term inside the first integral is just the solution of time-varying autonomous LDS, therefore,

Where Φ is the linear dynamical system state transition matrix [57]. But notice that this is the solution to the controlled LDS,

In the specific case of an RNN,

Therefore, Theorem 1 remains unchanged for a loss integrated over time. For Theorem 2, Claims 2-3 remain unchanged, while Claim 1 becomes rank ∇ W L (0) ≤ 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 a x ( t ) 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 D and encoder B of the system. If all parameters D,B,W of the system are optimized simultaneously, as 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 ,B u ( t )) dB = I ⊗ u ( t ) , so that ˙ a B = a x ⊗ u ( t ) . Therefore, following essentially the same derivation as that of Theorem 1, the following bound can be derived,

( k )

By a similar argument as for the derivation of Theorem 2, B i ∈ span { α i } ∪ { β i } ∪ { B i } ∪ { D i } . Second, for D , since df = 0 , that is ˙ a D = 0 ,

In other words, the gradient of the loss with respect to 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 on one another. Therefore these results hold regardless of which combination of W , B or D is optimized.

Batched updates . Most often, the weights are updated in batches. That is ∆ W ( k ) = Q -1 ∑ Q q ∇ W L ( k,q ) where q is the index over the batched dimension. Since the column and row spaces of ∇ W L ( k,q ) remain unchanged, Theorem 2 2-3. remain unchanged, while 1. becomes rank ∆ W (0) ≤ R + min { d, m } . For Theorem 1, the common singular value identities [56] σ i + j -1 ( A + B ) ≤ σ i ( A ) + σ j ( B ) and σ i ( cA ) = cσ i ( A ) for c ∈ R + 0 can be used. Then, σ ∑ i q -Q +1 (∆ W ( k ) ) ≤ Q -1 ∑ Q q σ i q ( ∇ W L ( k,q ) ) .

Momentum-based 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, a momentum variable is introduced, which is updated as M ( k +1) = βM ( k ) +(1 -β ) ∇ W L ( k ) , where β determines the speed of the exponential decay. The weights are then updated as W ( k +1) = W ( k ) -αM ( k +1) , where α is the learning rate. Which implies, ∆ W ( k ) = W ( k +1) -W ( k ) = -α ∑ k j (1 -β ) j ∇ W L ( j ) . Using the same identities are for batched updates, σ ∑ i q -k +1 (∆ W ( k ) ) ≤ ∑ k j (1 -β ) j σ i q ( ∇ W L ( j ) ) .

dD

E.4 Numerical simulations

Similarly to [14], we illustrate our mathematical results on random RNNs. Since constant inputs are one dimensional ( B u = ∑ B i u i ), we instead use time-varying inputs parameterized with LDS:

where u ( t ) ∈ R m , M ij ∼ N (0 , 1 √ m ) , u 0 ∼ N (0 , 1 / 2) . The target outputs are set as y ∈ R l , y i ∼ U ( -1 , 1) . The loss is defined as,

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 as well as the standard deviation (or strength) g of the initial weights such that W (0) ij ∼ N (0 , g 2 ) .

a

b

C

d

e

i

0.30-

1.0-

0.8-

(n) orol anen Jenus o7

2

()06ol ane

2-

1.00 -

0.25-

0.5

-5-

0-

-5

0.98-

0.6-

MSE

0.20-

0.15-

-00

-2.50.0

X

0.0

-2-

50

100

singularv

-2-

-10-

50

0.10-

0.2-

4.

4·

0.92-

0.05-

-00

0.90-

25

50

75

000

0

2

5

10

15

20

5

10

15

20

2

4

6

ii

10-

4

4·

Log singular value log1o()

2.5-

0-

2.5-

1.00-

8-

3-

2

0.0-

0.99-

6-

4.

2.5

82

-2.5-

。.50

-10-

0

X

-5.0-

Varian

2-

1 ·

-7.5-

0.97-

-0

。

0

10.0-

-10.0-

25

50

75

100

0

2

5

10

15

20

5

10

15

20

2

4

6

iii

0.6-

1.

0.5-

1.0-

1:

Log singularvalue log1o (a)

0-

2-

o-

2-

0.95-

1.00 -

-2·

-2·

-1

-2.5 0.0

X

2.5

#

0.0-

5

-6-

-4·

Log singular

-4 -

-6-

50

0.2-

0.5 -

-8-

-8-

0.85-

-00

Po-

-10-

-10-

0

25

50

75

100

2

5

10

15

20

5

10

15

20

2

4

6

iv2.5-

2.0-

1.0-

-2.5-

0.0-

2.5-

-2.5 0.0

2.5

(x)

-0.5-

0.0-

0.5-

1.0-

0.5-

0.5

0.0

g singular value logo(a)

-6-

-4

0-

2

-5-

01

50

Log singular v

-2-

4 ·

6

2-

0°

-5-

0

50

0.98-

0.96

1.00 -

0.99-

0.5-

1.0 -

607

-8-

-8-

0.94-

-00

0

25

50

75

100

1.5 -

0.0

0.5

1.0

1.5

2.0

-10-

2.5

5.0

7.5

10.0

-10-

2.5

5.0

7.5

10.0

0.93-

2

4

6

V

0.8-

1·

1.0-

Log singular value log1o (a)

2

0-

2-

-0

1.000-

0.6-

0.

0.5-

0.

0000

0-

0.975·

-1

2.5 0.0

2.5

#

0.0-

75

-2-

-fo

50

2-

-10-

0

50

X

-4-

Log singular

-4-

0.2-

0.5 -

-6-

-6-

0.875-

-00

0

25

50

75

100

0

2

4

-8-

5

10

15

20

-8-

5

10

15

20

0.850

2

4

6

Iteration

Time

Singular value index i

Singular value index i

Rank

(x10)

(x(100)

。

中(x(01)

a(0)

[01 a(00)

0Φ(x(100)

Wo

0min(o(xo,0(x)

(0)70

Tensor rank

Trial rank

1o1

o

Column rank

Row rank

O△Wo:100

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nProposes 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\nSTRENGTHS:\n1. 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\nWEAKNESSES:\nI 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”.\n\nQUESTIONS:\nHave you applied the ltrRNN to other datasets? Or is the purpose of the architecture specific to the kind of data reported in the paper?\n\nLIMITATIONS:\nThe authors have adequately addressed possible limitations of their work.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-03"}

comments: 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.

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.

Strengths:
- The low-rank property of RNN is interesting and is potentially useful for undertanding the neural dynamics and develop neural network architechures. 
- The mathematical framework could be valuable for comprehendingnature of gradient descent.

------

The 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.

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
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 2
DATE: 2023-07-04

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\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.\n\nWEAKNESSES:\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\nQUESTIONS:\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\nLIMITATIONS:\nThe authors have described various limitations of their work, and proposed concepts on how to deal with them in future work.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

comments: 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.

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.

2. The paper supports its results empirically.

3. 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.

QUESTIONS:
None.

LIMITATIONS:
Yes, the authors have adequately addressed the limitations and potential negative societal impact of their work.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-07