Low Tensor Rank Learning of Neural Dynamics

Summary

Proposes a low tensor rank recurrent neural network (ltrRNN) architecture, in which the tensor constructed by stacking the RNN weight matrices of different trials is constrained to have low tensor rank. Empirically shows that ltrRNNs can fit neural recordings during a motor learning task, achieving lower unexplained variance than baseline methods. Then, demonstrates that an RNN trained to perform the same motor learning task yields dynamics that can also be fit well with ltrRNNs (i.e. low tensor rank dynamics). Lastly, theoretically analyzes how gradient-based optimization can lead to low rank, establishing upper bounds on matrix and tensor ranks for RNNs. Discloser: 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.

Strengths

1. Reads relatively well. 2. 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. 3. 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. [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). [2] Razin, Noam, Asaf Maman, and Nadav Cohen. "Implicit regularization in tensor factorization." International Conference on Machine Learning. PMLR, 2021. [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. [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. [5] Timor, Nadav, Gal Vardi, and Ohad Shamir. "Implicit regularization towards rank minimization in relu networks." International Conference on Algorithmic Learning Theory. PMLR, 2023.

Weaknesses

I found the empirical evidence to be somewhat unsatisfactory since it only includes two datasets. Experiments on further tasks/datasets may greatly solidify the viability of ltrRNN. Such experiments can reveal whether the low tensor rank dynamics are specific to the type of neural recording data examined here or it is more general (and thus significant). An 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”.

Questions

Have you applied the ltrRNN to other datasets? Or is the purpose of the architecture specific to the kind of data reported in the paper?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

The authors have adequately addressed possible limitations of their work.

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.

Rating

4: Borderline reject: Technically solid paper where reasons to reject, e.g., limited evaluation, outweigh reasons to accept, e.g., good evaluation. Please use sparingly.

Confidence

1: Your assessment is an educated guess. The submission is not in your area or the submission was difficult to understand. Math/other details were not carefully checked.

Soundness

2 fair

Presentation

2 fair

Contribution

2 fair

Limitations

N/A

Summary

The work "Low Tensor Rank Learning of Neural Dynamics" investigate the low-rankness of RNNs with application to neural data, i.e. neural signals of a test subject performing a motor task. The 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. In their theory section, the authors propose two theorems that show the boundedness of the singular values of RNN weight matrices.

Strengths

This is a solid paper. Particularly interesting is: - Investigation of low-ranked RNNS as a model for real-world data - in this case neural signals for a motor task. - Analysis of the gradient dynamics, i.e. the singular values of RNN gradients, - Extensive numerical tests to validate the propositions, supplemented by code examples. - Comprehensive Related work section, which is important for such interdisciplinary work.

Weaknesses

- 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. This is the major drawback on this papers presentation, in my opinion. - The computational cost of training and used hardware should be described.

Questions

- Line 97: Please clarify: Do you ask if the weight tensor itself is low-rank or if the updates are low-rank? - 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. - Line 96: please explain the meaning of x and u in the context of your application. - 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. - Sec4: The authors need to specify * a) which low-rank tensor format is chosen to save the weight tensors * b) which update method / low-rank integrator or optimizer is used to compute the weight updates - 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? - 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: 1. 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? 2. How comes, that the neural data seem to be somewhat chaotic? I am by no means an expert in this application field, but find this intriguing. - 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? - Line 256: One should also mention [2] as one of the fundamental works of adjoint-based automatic differentiation. [2] Griewank, Andreas, Walther, Andrea; Introduction to Automatic Differentiation; PAMM; https://doi.org/10.1002/pamm.200310012

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

The authors have described various limitations of their work, and proposed concepts on how to deal with them in future work.

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.

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

2 fair

Contribution

3 good

Limitations

Yes, the authors have adequately addressed the limitations and potential negative societal impact of their work.