SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Arthur Pellegrino
pellegrino.arthur@ed.ac.uk
N Alex Cayco-Gajic*
Département d'Etudes Cognitives
Ecole Normale Supérieure
natasha.cayco.gajic@ens.fr
Angus Chadwick*
angus.chadwick@ed.ac.uk
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.
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
*These last authors contributed equally.
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 \in \mathbb{R}^{N \times 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)} + \Delta W^{(k)}$  for  $k \in \{0, \dots, K-1\}$ . In neuroscience,  $\Delta 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  $\mathbf{W} \in \mathbb{R}^{N \times N \times K}$  formed by stacking together the weight matrices  $W^{(k)}$  on all trials<sup>1</sup>, i.e.,  $\mathbf{W}_{ijk} = W_{ij}^{(k)}$ .
Here, we investigate the multilinear structure of the weight tensor  $\mathbf{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:
Figure 1: Low tensor rank recurrent neural networks.
$$
\mathbf {W} = \sum_ {r = 1} ^ {R} \mathbf {a} _ {r} \otimes \mathbf {b} _ {r} \otimes \mathbf {c} _ {r}
$$
$$
W_{ij}^{(k)} = \sum_{r = 1}^{R}\underbrace{c_{r}^{(k)}(\mathbf{a}_{r}\otimes\mathbf{b}_{r})_{ij}}_{= a_{r,i}b_{r,j}c_{r}^{(k)}}
$$
where  $\mathbf{a}_r,\mathbf{b}_r\in \mathbb{R}^N$  and  $\mathbf{c}_r\in \mathbb{R}^K$  for  $r\in \{1,\dots ,R\}$  (Fig. 1b). A consequence is that
quence of the low tensor rank assumption is that the weights  $\mathrm{vec}(W^{(k)})$  must evolve within an  $R$ -dimensional subspace of  $\mathbb{R}^{N^2}$  spanned by the set of rank-1 matrices  $\{\mathrm{vec}(\mathbf{a}_r\otimes \mathbf{b}_r)\}_{r = 1}^R$ .
a. We consider the 3-tensor formed by stacking the weights over learning. b. Constraining the weight tensor to be low rank allows for an interpretable analysis of the evolution of the weights over learning. c. LtrRNNs partition neural variability into task condition-specific inputs and learning-related weight changes over trials. d. LtrRNNs capture changes in dynamics constrained to a low-dimensional subspace, which reshape neural representations over learning.
For consistency with our analysis of neural data, we use the terminology trial throughout to describe the index of the weight samples  $k$ , denoted in superscript (cf. neuron and time indices denoted in subscript).
Low tensor rank RNNs. To investigate the structure of weight tensors formed over learning in recurrent networks, we consider a continuous-time RNN (Fig. 1c). The RNN dynamics on trial  $k$  are given by:
$$
\tau \dot {\mathbf {x}} ^ {(k)} = W ^ {(k)} \phi (\mathbf {x} ^ {(k)}) - \mathbf {x} ^ {(k)} + B \mathbf {u} ^ {(k)} (t)
$$
for  $B \in \mathbb{R}^{N \times N_{\mathrm{input}}}$ ,  $\mathbf{u}^{(k)}(t) \in \mathbb{R}^{N_{\mathrm{input}}}$ . If the low tensor rank hypothesis is satisfied, the system can be written as:
$$
\tau \dot {\mathbf {x}} ^ {(k)} = \sum_ {r = 1} ^ {R} c _ {r} ^ {(k)} \left(\mathbf {a} _ {r} \otimes \mathbf {b} _ {r}\right) \phi \left(\mathbf {x} ^ {(k)}\right) - \mathbf {x} ^ {(k)} + B \mathbf {u} ^ {(k)} (t),
$$
Figure 2: Bifurcation in the flow field of an ltrRNN. Here  $\mathbf{a} = \mathbf{b} = [-1 / \sqrt{2}, 1 / \sqrt{2}]^T$ , so that the weights at any given trial  $k$  are,  $W^{(k)} = \frac{c^{(k)}}{2}\left[ \begin{array}{cc}1 & -1\\ -1 & 1 \end{array} \right]$ . A supercritical pitchfork bifurcation occurs at  $c^{(k)} = 1$  when  $\phi = \tanh$ ; for  $\phi = \mathrm{id}$  (Sup. Fig. 1) a line attractor emerges at  $c^{(k)} = 1$ .
which reduces to a low-rank RNN on every trial [15]. As a result,  $\mathbf{x}^{(k)}(t)$  is constrained to an  $(R + N_{\mathrm{input}})$ -dimensional subspace of  $\mathbb{R}^N$  spanned by  $\{\mathbf{a}_r\}_{r=1}^R \cup \{B_i\}_{i=1}^{N_{\mathrm{input}}}$ , and this subspace is fixed across trials (Fig. 1d).
The weight matrix evolves across trials via a simple rescaling of the rank-1 components by the trial factors  $c_{r}^{(k)}$ . 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  $\mathbf{a}_r \otimes \mathbf{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  $\phi(\mathbf{x})$  (Fig. 2; Supplementary Material A). Moreover, even in the case of a linear ltrRNN  $(\phi = \mathrm{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  $\mathbf{c}_r$ 's. The real eigenvalues can be shown to satisfy the following equality (Supplementary Material A):
$\lambda_r = \frac{1}{2}\cos (\theta_{rr'}^{\tilde{\mathbf{v}},\mathbf{v}})^{-1}\sum_{r'}c_{r'}[\cos (\theta_{rr'}^{\tilde{\mathbf{v}},\mathbf{a}} + \theta_{rr'}^{\mathbf{v},\mathbf{b}}) + \cos (\theta_{rr'}^{\tilde{\mathbf{v}},\mathbf{a}} - \theta_{rr'}^{\mathbf{v},\mathbf{b}})]$  where  $\tilde{\mathbf{v}}_r$  and  $\mathbf{V}_r$  are the  $r$ th left and right eigenvectors and  $\theta_{rr'}^{\mathbf{x},\mathbf{y}}$  is the angle between the  $r$ th  $\mathbf{x}$  and  $r^\prime$  th  $\mathbf{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  $\{\mathbf{a}_{r'}\}_{r'=1}^{R}$  and  $\{\mathbf{b}_{r'}\}_{r'=1}^{R}$  as well as the angle  $\theta_{rr'}^{\tilde{\mathbf{v}},\mathbf{v}}$  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  $\mathbf{Y} \in \mathbb{R}^{N_{\mathrm{data}} \times T \times K}$ , comprising the activity of  $N_{\mathrm{data}}$  neurons measured on  $K$  trials, each of duration  $T$  time points, with single trial activations  $\mathbf{y}^{(k)}(t)$ . To fit the ltrRNN to such a tensor, we minimise the loss function  $L(\mathbf{a}, \mathbf{b}, \mathbf{c}, B, M) = \sum_{t,k}^{T,K} \|M\phi(\mathbf{x}^{(k)}(t)) - \mathbf{y}^{(k)}(t)\|^2$ , where  $M \in \mathbb{R}^{N_{\mathrm{data}} \times 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  $\mathbf{W}$ .
Figure 3: LtrRNN fit to neural data. LtrRNNs can be fitted to neural population data (here shown for a single neuron and for three trials), thus capturing smooth changes in activity over learning.
Above we considered how the low tensor rank assumption constrains the possible changes in RNN dynamics for the case in which the  $\mathbf{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  $\mathbf{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  $\mathbf{c}_r = (L + \sigma I)\overline{\mathbf{c}}_r$  where  $LL^T = A$  is the Cholesky decomposition of the smooth covariance matrix  $A$  (assuming  $\overline{\mathbf{c}}_r \sim \mathcal{N}(0, I)$ , see Supplementary Material A). In contrast, we constrain the inputs to the RNN  $\mathbf{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  $\mathbf{x}^{(k)}(t)$ , as well as the weights  $\mathrm{vec}(W^{(k)})$ , evolve in low-dimensional subspaces of  $\mathbb{R}^N$  and  $\mathbb{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.
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  $\dot{\mathbf{x}} = f(\mathbf{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  $\mathcal{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  $\mathcal{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  $\times$  time  $\times$  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  $\tilde{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  $(\mathrm{std} = 40\mathrm{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  $\mathbf{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  $\sim 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  $\phi = \mathrm{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 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  $\mathbf{x}(t)$  is constrained to lie in the space spanned by the column vectors  $\mathbf{a}_r$  [15]. The membrane potential can therefore be visualized via a projection onto each  $\mathbf{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 visualizing their effect in an interpretable fashion.
b
e
d
：