Paper: Emergent mechanisms for long timescales depend on training curriculum and affect performance in memory tasks (ICLR-2024)

Reviewer comment (Evaluation):
I think this paper could and should be valuable to neuroscientists, but at the moment the link between the paper results and biology is not so clear. For example, do the authors surmise that optimised, heteregenous instrinsic neuronal timescales should not be employed in the brain? Would there be any experimental predictions one could infer from this work?

Author response:
Thank you for recognizing the neuroscientific interest in our results; we are primarily motivated by insights from and for neuroscience. To make this more clear, we added a new section 5 on the neuroscientific implications of our findings. 

- Single neuron vs network timescale:

We agree with the referee; this is a fascinating point, and we should expand on it and highlight the possible biological implications. Our results indicate that adapting single-neuron timescales to the task requirements is less robust than learning based on adjustment of synaptic weights. This is greatly in line with how the neural networks in the brain primarily train: although there are possibilities to adjust single-neuron timescales, primary learning is implemented in adjusting the connectivity.  However, the heterogeneity in single-neuron timescales might be very important in the brain and our networks. We changed the duration of each stimulus presentation and found that the individual timescales track the input timescale in single and multi-head networks (Appendix C). This proposes one possible mechanism for how single-neuron timescales (equivalent to the membrane constant in our formalism) might be important for computations in reality, where the input timescales could vary a lot, e.g., different input timescales might be tracked by neurons with different taus.

Reviewer comment (Experiments):
The tasks selected are relatively simple and it is not clear if these results generalise to harder temporal tasks, though this is addressed as a limitation of the study. More generally, I would caution the authors' with the implication that their findings can generalise to "networks on sets of related tasks instead of a single task". The authors do demonstrate their case well for the setting of learning different N simultaneously during curriculum learning, but whether these results apply when, e.g. language processing in two different languages, is not clear to me.

Author response:
- Task choice and generalization:

We agree with the referee and weakened our statements and highlighted this in limitations.

Reviewer comment (Writing):
The motivation for two working memory tasks: N-parity and N-DMS, is not very well presented. What are the differences between the tasks. The authors note that Fig 6b reflects "distinct computational requirements for each task" but this is not elaborated upon.

Author response:
- N-DMS and parity:

The difference between these two tasks is that N-DMS only requires memory of input, whereas N-parity also needs computation (though quite simple). We add this briefly to the text (section 3.1).

Reviewer comment (Theory):
In the 4th paragraph of section 2 it is written tau_{net} <= tau, should it not be the other way around?

Author response:
- Typo in section 2: Thank you, we corrected it.

Reviewer comment (Presentation):
For fig 3a, which task is that? N-parity? The learning curves on this plot also appear to be still going up; it is not also clear whether these N will eventually be learned.

Author response:
- Results for N-DMS task (Fig.3a ,4a) and Learning curves:

Figure 3a. is for the N-Parity task (we now added to the caption). We have also made an equivalent figure for the N-DMS task (Appendix E, Fig. S13). These figures are meant to demonstrate how the rate of learning decreases as the task complexity (N) increases. In the case of N-DMS, we can explicitly see it fails to solve the task for N > 20 (Fig. S13) and for N-parity, this is also true for slightly larger N, though we only included up to N=20 in Fig. 3a. Furthermore, this figure illustrates how slow the learning in this case is, particularly compared to a multi-head network that reaches N=50 at 50 epochs. 

An equivalent figure for Fig 4a. is added for the N-DMS task (Fig. S10) for values of tau fixed at 1 and 2. We observe similar patterns as with the N-Parity, but with smaller differences, most likely due to a simpler task.

Reviewer comment (Writing):
Section 3.1: "N -parity: The network has to output the binary sum (XOR) of the last N digits, indicating an odd or even number of non-zero digits". I don't understand this; is it that the model should output whether the sum of the last N digits is odd or even?

Author response:
- Description of the tasks, particularly N-parity:

 Thanks for pointing out the problem with our description. Our aim with that description was to clarify that since the computations are online (the network needs to produce an output every time step), it needs somehow to keep a memory of the last N digits, but, indeed, we do not know the exact mechanism for solving the task. We have edited the text accordingly.

Reviewer comment (Theory):
Same N-Parity section: "the value of the digit presented at t − N needs to be subtracted from the running sum". I am not such a fan of this description as the authors' are suggesting the underlying mechanism for how the RNN solves the problem is known. Perhaps the RNN is incorporating another solution; e.g. perhaps N running sums are simultaneously encoded in the RNN and are reset every N timesteps.

Author response:
- Description of the tasks, particularly N-parity:

 Thanks for pointing out the problem with our description. Our aim with that description was to clarify that since the computations are online (the network needs to produce an output every time step), it needs somehow to keep a memory of the last N digits, but, indeed, we do not know the exact mechanism for solving the task. We have edited the text accordingly.

Reviewer comment (Experiments):
For Figs 3a, 4a, do these results also apply for the N-DMS task? Why are only

Author response:
- Results for N-DMS task (Fig.3a ,4a) and Learning curves:

Figure 3a. is for the N-Parity task (we now added to the caption). We have also made an equivalent figure for the N-DMS task (Appendix E, Fig. S13). These figures are meant to demonstrate how the rate of learning decreases as the task complexity (N) increases. In the case of N-DMS, we can explicitly see it fails to solve the task for N > 20 (Fig. S13) and for N-parity, this is also true for slightly larger N, though we only included up to N=20 in Fig. 3a. Furthermore, this figure illustrates how slow the learning in this case is, particularly compared to a multi-head network that reaches N=50 at 50 epochs. 

An equivalent figure for Fig 4a. is added for the N-DMS task (Fig. S10) for values of tau fixed at 1 and 2. We observe similar patterns as with the N-Parity, but with smaller differences, most likely due to a simpler task.

Reviewer comment (Experiments):
Is k > 1 ever considered? For the multi-head case the authors claim that tau -> k, and this seems true for k = 1; was this also tested for higher k?

Author response:
- 4a: Results for N-parity and N-DMS:

Apologies for the misleading k-notation; and added a more detailed explanation about k interpretation in new section 5 and Appendix C.

Reviewer comment (Reproducibility):
For Fig 3 was tau fixed? So only recurrent connectivity optimised? This should be made clear.

Author response:
- Fig. 3: fixed timescales:

The timescales are fixed only in Fig. 4a (and corresponding Fig. S10).

Reviewer comment (Presentation):
Is Fig 5a for the single-headed network?

Author response:
- Figure 5. Meaning of stalling the timescales:

Panel a is for single-head networks (now added to the caption). Our interpretation for stabilization of tau_net is that, for large N, networks need to find solutions other than increasing the timescales of local outputs (measured by tau_net). They are possibly employing combinations of neurons that can have different timescales on the population level than individual neurons if they are well-coordinated. Since we cannot assess all the combinations of local outputs' contributions to the readout, we designed a proxy. We took the whole population and summed up the activity of all neurons, obtaining the timescale of the population activity tau_pop. For all tasks and training curricula, tau_pop grows with N (Fig. S4).

Reviewer comment (Evaluation):
For Fig 5d in each task the network timescale tau_net increases up to about N=30 then shrinks. Is this shrink surprising? It is also unclear to me how these multi-headed network can perform the task so well, when their tau_net is overall considerably smaller (about half the size) of the worse-performing single-head networks.

Author response:
- Figure 5. Meaning of stalling the timescales:

Panel a is for single-head networks (now added to the caption). Our interpretation for stabilization of tau_net is that, for large N, networks need to find solutions other than increasing the timescales of local outputs (measured by tau_net). They are possibly employing combinations of neurons that can have different timescales on the population level than individual neurons if they are well-coordinated. Since we cannot assess all the combinations of local outputs' contributions to the readout, we designed a proxy. We took the whole population and summed up the activity of all neurons, obtaining the timescale of the population activity tau_pop. For all tasks and training curricula, tau_pop grows with N (Fig. S4).

Reviewer comment (Evaluation):
Fig 6b: I was surprised at how high dimensional these network activities are. In my experience trained RNNs produce quite low dimensional dynamics (< 10 or often < 3). Is this high dimensionality surprising to the authors?

Author response:
- High dimensionality of the outputs compared to other RNN studies:

 It is indeed surprising when we compare the dimensionality to other neuro studies with RNNs. Part of the reason for this can be that the common neuro tasks are by design low-D. For instance, tasks usually have a fixed trial structure, and the low-D trajectories often follow the events within a trial. For example, a DMS task is often designed to start with one stimulus, followed by the delay and then the 2nd stimulus. The only variable is the delay period. However, in our case, the computations have to be performed at every time step, and very different input patterns (the number patterns grow with N) should be mapped to the same outputs (0,1). Nevertheless, it is a point for further investigation.

Reviewer comment (Presentation):
Add (AC) after autocorrelation in the Fig 1 caption.

Author response:
- Typos and suggestions:

Thanks for noting these. We corrected the typos and edited the text.

Reviewer comment (Writing):
In the 4th paragraph of section 2: tnet -> tau_{net}.

Author response:
- Typos and suggestions:

Thanks for noting these. We corrected the typos and edited the text.

Reviewer comment (Theory):
I would appreciate a more formal, mathematical formulation of how tau_{net} was obtained for each neuron in the appendix.

Author response:
Moreover, since the definition of tau_net is an important part of our paper, we edited the definition in the main text to provide more detailed explanations.

Reviewer comment (Writing):
End of first paragraph of section 4.1, remove comma after K.

Author response:
- Typos and suggestions:

Thanks for noting these. We corrected the typos and edited the text.

Reviewer comment (Presentation):
I'd recommend the related work prior to the results.

Author response:
We did not change the location of the Related Work section since it includes a comparison to our findings; instead, we edited it to better reflect these comparisons.

Reviewer comment (Presentation):
I'd recommend all supplementary figures together at the very end. Their appearance amongst the different appendix sections seems quite random/distracting.

Author response:
- Typos and suggestions:

Thanks for noting these. We corrected the typos and edited the text. 

We moved all supplementary figures to the end.

Reviewer comment (Evaluation):
Claims in conclusion seem overstated. This is more a significance issue than a quality issue.

Author response:
We edited the paper and made our statements more grounded and concrete (e.g., Relevance to Neuroscience, Related Work sections) and extended the limitations section

Reviewer comment (Experiments):
I would say more experiments are needed, but it's not a matter of specific additional experiments - for what it is, this seems to be a careful and informative set of experiments. It's more that the paper needs considerable expansion, which may also be a significance issue.

Author response:
While we could not generate an entirely new set of experiments in a short time, we propose the general set of requirements these experiments need to satisfy (in the limitation section). Moreover, we described an additional task that would fulfill these requirements (Appendix D).

Reviewer comment (Evaluation):
Figure 8: the curves do show a difference but not a drastic one, making it hard to judge how much more robust multi-head is. Significance testing, grounding according to some baselines (including visually on the figure), or comparison to similar results in other contexts may help.

Author response:
We run the significance test for different perturbation strengths and retraining N (Appendix M). For all perturbations, when the perturbation size is very small, there is almost no difference; for larger perturbations, there is a significant difference, and then performance for both networks degrades to the level where they are indistinguishable again.

Reviewer comment (Experiments):
As stated in the limitations, there is only one testbed in this paper and it is very toy. Unfortunately, the claims of this paper seem to be about improving RNN performance, which means that some kind of scale or breadth is requisite.

Author response:
We added to the limitation section a requirement for the task. Moreover, we added Appendix D, discussing a possible task for future research: pattern generation, which is a common task. We adapt this continuous-time task and turn it into a curriculum for generating the sum of sine waves.

Reviewer comment (Evaluation):
The motivation connects to biological neural networks, but that stops at the intro - the paper never connects the investigated mechanisms back to biological intelligence in any way after the prose in the intro, let alone through experimentation. It's therefore unclear how significant these results are, though they are intriguing.

Author response:
We see our work as strongly motivated and interpretable in the neuroscience context. We added a new section 5 to better discuss this relevance. In brief, our results can be extended to continuous-time settings relevant to modeling brain networks. Moreover, we discuss that using recurrent connectivity to learn the tasks and develop long timescales aligns well with previous experimental findings.

Reviewer comment (Experiments):
What other experimental settings/testbeds might be easily accessible for your setup?

Author response:
We added to the limitation section a requirement for the task. Moreover, we added Appendix D, discussing a possible task for future research: pattern generation, which is a common task. We adapt this continuous-time task and turn it into a curriculum for generating the sum of sine waves.

Reviewer comment (Novelty):
How might this connect back to biological intelligence? I realize this isn't necessarily the focus of the paper, but it is discussed extensively in the intro.

Author response:
We see our work as strongly motivated and interpretable in the neuroscience context. We added a new section 5 to better discuss this relevance. In brief, our results can be extended to continuous-time settings relevant to modeling brain networks. Moreover, we discuss that using recurrent connectivity to learn the tasks and develop long timescales aligns well with previous experimental findings.

Reviewer comment (Presentation):
For Figure S12, it would be helpful to readers to point out that the y-axis is on a log scale, so even though visually the linear fit doesn't look good towards larger $t$, the residual errors are much smaller here because this regime has much smaller values than those for earlier $t$.

Author response:
Thank you, we implemented your suggestion (new Fig. S2).

Reviewer comment (Experiments):
In Figure 6, I can infer, but I don't believe it is explicitly stated why N=30 is used as the maximum value for linear regression line fitting.

Author response:
Indeed, we fitted the linear relationship only for small N (N =20). To make it easier to identify, we added the lines as a guide for the eye. The reason for this is to show that the scaling of the dimensionality in the N-DMS task is not linear: for the points with N>>20, the line does not fit, and growth is slower. In the updated version of Fig.6 and new Fig.S20, we present the fits for different ranges.

Reviewer comment (Presentation):
For Figure 8c, I would've liked to see a mean line and N extend to a higher number to see at what relative accuracy each converge and how these compare to one another.

Author response:
We now include the chance-level accuracy in the panel. We retrained the networks for higher N and observed for both single and multi-head curricula, a convergence towards the chance level. Moreover, we calculated the p-values and presented them in Appendix K to demonstrate that there is a significant difference between single- and multi-head networks.

Reviewer comment (Presentation):
Future work is also lacking, both as a result of lacking limitations as well as no suggestions for future work outside of the limitations section.

Author response:
We substantially edited the related work and limitations sections to highlight the significance of our findings in comparison to previous works and its potential extensions for future studies. In brief, although there were theoretical studies analytically deriving timescales for simple networks, those studies were limited to the specific (and simplistic) connectivity and investigated only the relationship between the topology features and local timescales, often without discussing their functional relevance. Here, we extend those findings and relate the dynamical properties directly to learning and computations relevant to working memory tasks, bridging theoretical mechanistic concepts from computational neuroscience to functional concepts in ML. We also show that the multi-head curriculum improves not only the training of RNNs but also of GRUs and  LSTMs.

Reviewer comment (Presentation):
The Limitations section seems like an afterthought that was thrown in to satisfy having "Limitations" of the work mentioned.

Author response:
We substantially edited the related work and limitations sections to highlight the significance of our findings in comparison to previous works and its potential extensions for future studies. In brief, although there were theoretical studies analytically deriving timescales for simple networks, those studies were limited to the specific (and simplistic) connectivity and investigated only the relationship between the topology features and local timescales, often without discussing their functional relevance. Here, we extend those findings and relate the dynamical properties directly to learning and computations relevant to working memory tasks, bridging theoretical mechanistic concepts from computational neuroscience to functional concepts in ML. We also show that the multi-head curriculum improves not only the training of RNNs but also of GRUs and  LSTMs.

Reviewer comment (Novelty):
With the large amount of related work that is close to the proposed work, it seems like the proposed work is primarily an empirical investigation into very simple tasks and limited approaches.

Author response:
We substantially edited the related work and limitations sections to highlight the significance of our findings in comparison to previous works and its potential extensions for future studies. In brief, although there were theoretical studies analytically deriving timescales for simple networks, those studies were limited to the specific (and simplistic) connectivity and investigated only the relationship between the topology features and local timescales, often without discussing their functional relevance. Here, we extend those findings and relate the dynamical properties directly to learning and computations relevant to working memory tasks, bridging theoretical mechanistic concepts from computational neuroscience to functional concepts in ML. We also show that the multi-head curriculum improves not only the training of RNNs but also of GRUs and  LSTMs.