Structured flexibility in recurrent neural networks via neuromodulation

Thank you very much for your feedback. We have addressed some of your comments in the general rebuttal, in particular the second point under Weaknesses (regarding performance comparison to LSTMs and vanilla RNNs). We would also like to respond to the additional weaknesses noted.

While we currently have no comparison with neural data, we believe this model can offer testable predictions of how diseases involving neuromodulator deficiencies and dysfunction impact movement and timing. As RNNs have been used to generate testable predictions of neural function, we believe our model offers a biologically meaningful extension by scaling synaptic weights as biological neuromodulation has been shown to do.

We did not compare our work with the prior methods mentioned in the Related Work section because of a few key differences in motivation and implementation. The work of Tsuda et al., while similar, requires prespecification of (1) which neurons are neuromodulated and (2) the constant value that the neuromodulation takes. Our paradigm allows both of these parameters to be learned by the model. It was unclear how to set the aforementioned parameters in the Tsuda model in order to make a comparison. While Liu et al. also studied the impact of neuromodulation in RNN models, their work focused primarily on the impact of neuromodulation on training and not during task performance. Due to the difference in motivation, we did not compare to their model.

We agree that training our model on more complex tasks will be an important step forward. Since vanilla RNNs are used to model more complex tasks and combinations of tasks [Driscoll], we believe our model should also be amenable to training in these situations. We appreciate the suggestion for studying reaching/control perturbation tasks and agree that training the NM-RNN in this task setup could help provide valuable insights into the relationship between neuromodulation and adaptation.

To answer your questions:

Have the authors tried training the neuromodulatory subnetwork alone while treating the task subnetwork as a reservoir? This could potentially provide for a parameter-efficient multi-task learning framework (loosely related, see the recent paper by Williams et al., 2024 [3]).

We like this idea and think it would be an interesting extension for the multitask section of the paper. While we did not try random reservoir dynamics as in Williams et al., we did analyze the effect of retraining only the neuromodulatory signal while leaving the low-rank dynamical components fixed when learning a new task (paper section 5). These results could be viewed as a proof-of-concept that the NM-RNN can effectively reuse existing dynamics when learning new tasks.

Could the authors clarify what the three different NM-RNN curves are in Fig. 5C? If they are different seeds or configurations, this should be specified.

We apologize for the confusion—these are different hyperparameter settings for the NM-RNN, parameter-matched to the N=10 LSTM. We will update the caption of this figure in the final version to indicate the exact hyperparameters used in each curve. For reference the settings are (M=5, N=18, R=8), (M=5, N=13, R=12), and (M=10, N=12, R=7). Our overall goal in this figure was to show that LSTMs and NM-RNNs are able to solve the task, while vanilla RNNs cannot train on it.

Have the authors tried using other activation functions in calculating z(t), i.e., not restricting the neuromodulatory signal to values between 0 and 1 (and allowing negative values)? Could this lead to improved performance?

While this may lead to improved performance, we believe it would negatively impact the biological plausibility of the modulatory signal. We capped the neuromodulatory signal at 1 to model a saturating effect of neuromodulators on downstream synaptic strength—the assumption being that synaptic strength cannot increase indefinitely. Similarly, we did not allow negative saturation, instead setting the floor value of the neuromodulator to 0 to indicate complete silencing of synaptic connections and preserve the sign of the connection.

Thank you very much for your feedback. We have addressed some of your comments from the Weaknesses section in the general rebuttal (in particular, concerns about Evaluation and Baselines). We would also like to clarify why we did not initially provide an LSTM baseline for the Measure-Wait-Go and multitask settings. The goal of this project is to offer a new model relevant to the field of computational neuroscience. Low-rank RNNs are commonly used tools in computational neuroscience due to their low-dimensional internal dynamics. We augmented the low-rank RNN with an interpretable modulatory parameter, and showed the comparative effectiveness of low-rank RNNs and NM-RNNs on tasks related to neuroscience. We also showed that NM-RNNs help to recover some of the performance gains of LSTMs. However, LSTMs are not commonly used as “brain-like” models when studying neural population dynamics. One of our contributions was to show the similarities between our more “brain-like” NM-RNN and the popular LSTM. Like LSTMs, transformers are not typically used to model neural population dynamics, so we did not include comparisons to these models.

To address your questions:

Why choose rank-1 for modulatory subnetwork?

To clarify, the modulatory subnetwork is full-rank and the output-generating subnetwork is low-rank. In our study of the Measure-Wait-Go task we chose to present results with rank-1 and rank-3 NM-RNNs. We analyzed rank-1 networks due to their theoretical tractability, and rank-3 networks in order to compare to existing work which used rank-3 low-rank RNNs on this task [1]. For the multitask setting, we chose rank-3 NM-RNNs by sweeping over ranks and determining the rank required to consistently train on the first three tasks. For the element-finder setting we tried many combinations of rank-size parameters to compare to the LSTM.

Why the training dynamics (MSE loss) for all models are noisy and fluctuated, any optimization techniques (regularization, normalization, scheduler might help with the training dynamics), how to guarantee the convergence of the model?

We agree that the training dynamics for the Element-Finder task (EFT) are quite noisy, and would likely become smoother with the use of regularization techniques. However, our main point with this figure was to show general trends, i.e. vanilla RNNs struggle to solve EFT while some LSTMs and NM-RNNs are able to learn it consistently. In addition, we did not show the training dynamics for any of the other tasks, but can share empirically that they were much smoother than in the EFT, suggesting that the noisy learning may also be due to the relative complexity of this task.

Any computational efficiency that NM-RNN might bring in?

All models are run sequentially, so while there is a slight difference in computational efficiency between each individual step, the scaling laws likely don’t change. In general, you raise an interesting point that the matrix multiplication cost for running a step of a low-rank RNN is smaller than that of a vanilla RNN. So, low-rank RNNs and NM-RNNs will be faster to evaluate compared to a vanilla RNN with similar neuron count. We will add a discussion of computational implications to the final version.

Citations [1] Manuel Beiran, Nicolas Meirhaeghe, Hansem Sohn, Mehrdad Jazayeri, and Srdjan Ostojic. Parametric control of flexible timing through low-dimensional neural manifolds. Neuron, 111(5):739–753, 2023.

Thanks for the authors' response. While I am not convinced that LSTM and transformers are not "brain-like", while RNNs are. I view them all more as predictive models instead of mechanistic models (i.e. hodgkin huxley model). And a lot of recent works have started to use transformers model neural population dynamics [1][2][3], and outperforms RNNs. I think adding more comparisons with baselines, and improving the optimization will be helpful. I decided to maintain my original score.

[1] Representation learning for neural population activity with Neural Data Transformers.

[2] STNDT: Modeling Neural Population Activity with a Spatiotemporal Transformer.

[3] Neuroformer: Multimodal and Multitask Generative Pretraining for Brain Data.

I just wanted to step in and provide my perspective here as another reviewer, because I believe that the papers this reviewer mentions are orthogonal to what I believe is the goal of this paper.

The goal here is to propose a neuromodulatory mechanism inspired by synaptic gain scaling, as a circuit-level mechanistic model of how modulatory signals could reconfigure the dynamics of a recurrent network to perform different tasks. I view it as a mechanistic model rather than a predictive one, i.e., it describes how a mechanism that is computationally similar to what is observed in the brain enables multi-task flexibility, rather than serving as a better predictive model or advancing performance from a deep learning perspective. In short, it is trying to computationally analyze the effect of synaptic gain scaling on multi-task performance.

Several works have used recurrent neural networks to describe how the brain could be performing certain task-related computations, including influential work by Yang et al. [1] and Driscoll et al. [2]. While I'm not arguing here that somehow vanilla RNNs are more bio-plausible or better brain models than LSTMs, I believe there are different levels of modeling and biological realism in the computational neuroscience literature. Circuit-level computational models such as this work eschew synapse-level biological plausibility and aim to study computations through the lens of dynamics [3], task performance, or other hallmarks of neural computation like flexibility and generalization.

I would thus argue that RNNs are good and simple/minimal models to use here – we are aware of recurrent circuits in various brain regions [4,5,6,7], but there is no similar support for transformer-like architectures or self-attention, to my knowledge (although on the other hand, these models are, recently, used at the cognitive level of modeling [8,9]). And given that the kind of tasks being performed here are related to working memory, which relies on the prefrontal cortex [10], which is further known to contain several recurrent microcircuits [6,7], I'd argue that RNNs are better computational models to use here than transformers if the goal is to model neural computation (which I strongly believe to be the case here).

Finally, and perhaps most importantly, the papers that the reviewer refers to here are not trying to model how the brain could compute, or how mechanisms in the brain could enable multi-task flexibility. Works such as [11,12,13], using transformers to model neural population activity, are harnessing advances in deep learning to serve as better representation learning methods or predictive models of behavior from neural activity, for e.g., towards the goal of building better decoders from brain-computer interfaces – they are not meant to serve as models of how the brain could compute.

Overall, I just wanted to provide my views here and I hope to engage in further discussion with this reviewer and/or the authors on this.

References:

1. Yang, Guangyu Robert et al. “Task representations in neural networks trained to perform many cognitive tasks.” Nature neuroscience vol. 22,2 (2019): 297-306.
2. Driscoll, Laura N et al. “Flexible multitask computation in recurrent networks utilizes shared dynamical motifs.” Nature neuroscience vol. 27,7 (2024): 1349-1363.
3. Driscoll, Laura N et al. “Computation through Cortical Dynamics.” Neuron vol. 98,5 (2018): 873-875.
4. Douglas, Rodney J, and Kevan A C Martin. “Recurrent neuronal circuits in the neocortex.” Current biology : CB vol. 17,13 (2007): R496-500.
5. Wang, Xiao-Jing. “Decision making in recurrent neuronal circuits.” Neuron vol. 60,2 (2008): 215-34.
6. Mante, Valerio et al. “Context-dependent computation by recurrent dynamics in prefrontal cortex.” Nature vol. 503,7474 (2013): 78-84.
7. Fuster, J M. “Memory networks in the prefrontal cortex.” Progress in brain research vol. 122 (2000): 309-16.
8. Didolkar, Aniket, et al. "Metacognitive Capabilities of LLMs: An Exploration in Mathematical Problem Solving." arXiv preprint arXiv:2405.12205 (2024).
9. Webb, Taylor, et al. "A Prefrontal Cortex-inspired Architecture for Planning in Large Language Models." arXiv preprint arXiv:2310.00194 (2023).
10. Funahashi, Shintaro. “Working Memory in the Prefrontal Cortex.” Brain sciences vol. 7,5 49 (2017).
11. Ye, Joel, and Chethan Pandarinath. "Representation learning for neural population activity with Neural Data Transformers." arXiv preprint arXiv:2108.01210 (2021).
12. Le, Trung, and Eli Shlizerman. "Stndt: Modeling neural population activity with spatiotemporal transformers." Advances in Neural Information Processing Systems 35 (2022): 17926-17939.
13. Antoniades, Antonis, et al. "Neuroformer: Multimodal and multitask generative pretraining for brain data." arXiv preprint arXiv:2311.00136 (2023).

Thank you very much for your helpful comments. We have addressed some shared concerns regarding performance comparisons to LSTMs and vanilla RNNs in the general rebuttal. To respond to your comments regarding weaknesses of the paper:

In general, we don’t believe the NM-RNN will outperform the LSTM. As you mentioned, there is a plateau in performance despite similarities in structure. While our NM-RNN captures some aspects of the LSTM via its synaptic weight scaling, other key features of the LSTM are not implemented. In particular, the input and output weights of the NM-RNN are held constant. Perhaps modulating these weights would help recover more LSTM-like performance. However, our main goal was to offer this model as an alternative to commonly-used low-rank RNNs, as a tool to model biological task completion with neuromodulation. In contrast, LSTMs are not commonly used as “brain-like” models of neural dynamics.

We give example outputs of each model on the Measure-Wait-Go task to highlight specific performance (paper fig. 3D and rebuttal fig. B). For the multitask setting and EFT, example outputs are noisier and are less visibly distinguishable between model types, so we did not visualize them, instead choosing to summarize performance via % correct and loss metrics.

We chose to have neuromodulation differentially impact low-rank components to mimic the selective way neuromodulators act on particular synapses [1], and to situate our model alongside the widely used low-rank RNN. However, this choice is certainly not the only way to model neuromodulation in an RNN—one alternative might be to implement neuromodulatory weights on a mixture of sparse weight matrices, instead of low-rank components. As another example, in Tsuda et al. the authors choose particular subsets of weights on which to apply neuromodulation. However, these subsets are prespecified instead of learned during training.

To answer your questions:

The architecture seems to require a prior specification of the maximum rank of the connectivity matrix. From an application standpoint, how would this be determined?

The rank of the connectivity matrix is indeed a hyperparameter that must be determined before training. In general, we performed a parameter sweep over various numbers of ranks to determine the lowest rank that achieved satisfactory performance. For the MWG task, we chose rank-1 networks due to their theoretical tractability (and because even at this lowest possible rank, the NM-RNN can solve the task). However, to compare to the low-rank RNN we chose rank-3 NM-RNNs, to match prior work which used rank-3 RNNs to model this task [2].

Is the proposed method amenable to batch training? Can the authors characterize the trainability/computational implementation considerations of the proposed architecture, relative to say an LSTM? Certainly, there would seem to be more overhead relative to a vanilla RNN.

This method is amenable to batch training, which we used in the multitask and EFT settings. For the MWG task, the space of possible samples was small enough to do full batch training. In terms of computational implementation, all models must be run sequentially, so while there are slight differences in training time the overall scaling rules likely do not change. We will add a discussion of computational implications to the final version.

Citations

[1] Peter Dayan. Twenty-five lessons from computational neuromodulation. Neuron, 76(1):240–256, 2012.

[2] Manuel Beiran, Nicolas Meirhaeghe, Hansem Sohn, Mehrdad Jazayeri, and Srdjan Ostojic. Parametric control of flexible timing through low-dimensional neural manifolds. Neuron, 111(5):739–753, 2023.

Thank you very much for your thorough comments. Indeed, our motivation was to create a model somewhere between RNNs and biologically-accurate biophysical models. Specifically, we identified neuromodulation as a feature of biological networks that is not often modeled in RNNs. The goal of this paper is to study the impact of synaptic scaling, one specific effect of neuromodulation, on the performance and generalization of RNN models. We have discussed our motivation further in the General Rebuttal and will make this positioning more clear in the final version.

Thank you for suggesting Auzina et al. and Naumann et al., we will add these papers to our related work section to highlight the use of modulatory signals in non-RNN models.

To address your questions:

Conceptual:

By structured flexibility, we mean the introduction of additional parameters (i.e. flexibility) to a model, with these parameters having a clear motivation and interpretation (i.e. structure). In the case of our NM-RNN, the structured flexibility comes from adding a flexible neuromodulatory signal which impacts RNN dynamics in a constrained way (by scaling synaptic weights). In the EFT and multitask settings, we show that this allows the network to distribute different parts of the task across the subnetworks. In the MWG task, we show that ablating each component of the neuromodulatory signal leads to performance deficits (paper fig. 3F). In particular, ablating s_3(t) destroys the ability of the network to produce output ramps of different slopes, implying that the neuromodulatory subnetwork is involved in controlling the output interval length. In all three ablations, the general ramping shape of the output is preserved, implying that this behavior is stored in the dynamics of the output-generating network.

Overall, we learned that synaptic gain scaling can improve generalization capabilities of low-rank RNNs and make them more like LSTMs. More generally, we offer this model as an alternative to commonly-used low-rank RNNs, as a tool to model biological task completion with neuromodulation.

EFT:

Each of our models was first trained in a continual learning regime in which each step of gradient descent was performed using a randomly generated batch of input sequences. Then, at test time, test input sequences were again sampled i.i.d. from the distribution of all valid EFT sequences. (It should be noted that for main paper figs. 5C, D, and F, we randomly generated test inputs while holding the query index and/or target element value fixed.)

As the sequence length T increases while keeping the NM-RNN model size, number of training batches, and training batch size fixed, we predict that model performance will degrade. The set of possible EFT input sequences grows exponentially in T, making it challenging for the model to solve the task within the same number of training iterations; indeed, this is what we observed in our preliminary simulations. However, we expect that model performance would improve if the number of training iterations and/or the model size is suitably increased.

Technical:

Fig. 2A: The corresponding text talks about different components k, but it is not clear what is shown in the figure (a single component across time, or different components? How do changes in s(t) come about?)

This figure is meant as an illustrative example of how the neuromodulatory signal s(t) can impact dynamical timescales. In this example, s(t) is a 1-dimensional signal so k=1. The changes in s(t) are artificially imposed for illustrative purposes, to show how the decay rate of w(t) is modulated by the value of s(t). We will clarify this in the final version.

Could this be a plug-in replacement for GRUs/LSTMs? If so, how would the model perform standard sequence modeling tasks, such as sequence classification?

We are not arguing to replace GRUs/LSTMs in standard ML workflows. However, we show that they do share some capabilities, achieving LSTM-like performance on some tasks.

The authors don't mention the vanishing gradient issue that RNNs severely suffer from. Did the not have this problem?

This issue did not arise for the tasks that we trained on. To speculate, perhaps the NM-RNN does not encounter this issue as often due to its similarities to the LSTM. This could also be a potential reason why the RNNs are failing at the EFT.

Fig 3E: please offer more interpretation for what is shown.

In fig. 3E we show the 3-channel neuromodulatory signal s(t) during the MWG task. On the left, the signals are aligned to when networks receive the measure cue. On the right, the signals are aligned to when networks receive the wait and go cues. The left figure shows how the s(t) responds to the measure cue and evolves throughout the task. The right panel shows how s(t) readies responses for the different interval lengths. In particular, we can see that between the wait and go cues, s_1(t) separates shorter intervals from longer ones, setting up initial dynamics for the go period when the ramp is generated. The third component s_3(t) seems to be involved with ending the output ramp, since it saturates first for the shorter intervals and then the longer ones.

Minor:

ln 104: we will add a citation to Schuessler et al. “The interplay between randomness and structure during learning in RNNs”, NeurIPS 2020

Fig. 4A: we have updated this figure in the PDF attached to the general rebuttal. Please let us know if additional adjustments would help to clarify.

Fig. 5: The squiggly lines are the actual loss curves, while the bold lines show a moving average to indicate overall trends. We will add improved explanations and better-contrast lines to E and F.

Eq. 4: \sigma is the sigmoid nonlinearity, we will clarify this in the final version.

Fig. 3/ln. 208: thank you for catching this, we will amend it in the final version.

ln 267: we have added these plots to the general rebuttal PDF and will include them in the final paper fig. 4D.