Volume Transmission Implements Context Factorization to Target Online Credit Assignment and Enable Compositional Generalization

We thank reviewer n2dL for the insightful and thorough review. We are glad you found the work "original... and very thorough in its empirical analyses." We have worked to address your important concerns about clarity, justification, and the connection to experimental neuroscience.

Major Concern 1: How does temporal credit assignment work?

This is a crucial point that we abbreviated in the main text. Our learning rule is a temporally truncated approximation to the full BPTT gradient. For a parameter like $T_{ijk}$, the update is proportional to: $\Delta T_{ijk} \propto n_k(t) \cdot \epsilon_i(t) \phi'(A_i(t)) r_j(t)$ This rule is "context-localized" by $n_k(t)$ and "temporally-localized" by using the instantaneous error signal $\epsilon_i(t)$. While this truncates the full gradient through time, we argue it is a powerful and plausible approximation for two reasons. First, similar to algorithms like e-prop, the recurrent state $\vec{r}(t)$ and, critically, the slow neuromodulatory state $\vec{n}(t)$ act as eligibility traces, carrying the necessary temporal context from the past. Second, our empirical results, especially on the RL and multitasking benchmarks which require effective credit assignment over intermediate timescales, demonstrate that this approximation is potent enough to solve complex tasks. The network learns to encode the necessary long-term information into its slow context variables, allowing the local update rule to be effective.

Major Concern 2: Why neuromodulators? Low-dimensionality and testable predictions.

This is an excellent and deep question.

• Why NMs & Low-Dimensionality: We chose NMs as our inspiration because their biological function—globally and diffusely reconfiguring circuits—is an exemplar physical substrate for the computational concept of "context." We agree that a handful of NMs are insufficient to capture the full complexity of brain-wide state control. We frame our model as a proof-of-principle showing how even a very low-dimensional modulatory system can powerfully gate high-dimensional dynamics. We agree with your suggestion that this system likely works in concert with other mechanisms (like thalamocortical gating) to form a richer context space. Our goal is to show how volume transmission provides a powerful, efficient, and interpretable means of broadcasting context but, like all models of neural dynamics, this is at best only part of the story.
• Testable Predictions: Our model makes several concrete, testable predictions. 1) As shown in Fig 5C, it predicts that during multitask learning, ambient NM concentrations in a relevant cortical area should become decodable for task context, a hypothesis that is becoming directly testable with new imaging techniques (e.g., GRAB sensors). 2) The model predicts a relationship between the learned timescales of different NM systems and the statistical regularities of the environment. 3) In RL settings, it predicts that some component of the neuromodulatory state space should explicitly encode RPEs, consistent with known roles for dopamine and norepinephrine.

Minor Concerns:

• NM Release/Dale's Law: We agree our phrasing should be improved. In a revision, we would state that our model is a functional abstraction which can be interpreted in neurobiological terms. A unit "releasing" an NM means its activity drives the release of that NM systemically. We will rephrase "dopamine neuron" to "a unit whose activity promotes dopamine release as indicated by the non-zero entry on the $R$ matrix". Indeed, the one-to-one mapping of nm0 to "dopamine" is speculative as an evocative example, but we acknowledge that it may do more harm than good in this context. If given the opportunity to revise, we would remove it in favor of more precise language.

• Baseline for Figure 5: This is a critical point we have now addressed. We trained a parameter-matched standard RNN on the multitasking benchmark which have matched previous reported results in Yang 2019. The results show that our spatially-embedded e-nmRNN achieves comparable performance across the 20 tasks when similar effort is put into hyperparameter tuning (using a Bayesian Search with a Tree Structured Parzen estimator). This demonstrates that the inductive biases of our architecture provide tangible value on this complex benchmark, by yielding emergent, interpretable structures (modularity, timescale separation) that a standard RNN does not as readily admit. Indeed, we hypothesize that there is value to more closely matching the inductive biases inherent to neurobiological implementation for both understanding and application.

We believe these clarifications and new results substantially improve the paper and directly address your thoughtful concerns. Thank you once again for your thoughtful review of our work and your contributions to its improvement.

We thank reviewer 8PJy for your thorough and critical review. We sincerely apologize that a lack of clarity in our initial submission led to a frustrating reading experience. We acknowledge, and we appreciate your efforts to identify these specific issues with such helpful granularity. We believe the concerns raised can be fully addressed through clarification, and your feedback will be invaluable in improving the manuscript. We have worked to address your points below and are confident that a revision would be substantially clearer and stronger.

On Clarity (Weaknesses 1, 2, 3):

We are grateful for the detailed list of clarity issues. In a revision, we would ensure every variable is defined upon its first appearance and restructure the appendix.

We address your specific examples:

1. Undefined variables: You are correct. In a revision, we will define:

• LIF: Leaky Integrate-and-Fire neuron model, the starting point of our derivation.
• $S_j(t)$: The spike train of the $j$-th neuron in the biophysical model.
• $\phi, \theta$: The nonlinear activation functions for the rate ($\phi$) and neuromodulator ($\theta$) dynamics. We use functions like tanh or ReLU that enforce positive-only firing rates and saturation, consistent with the closed-form solution of the steady-state firing rate for a LIF neuron with a finite refractory period.
• $D$: The linear readout matrix that maps neural activity $\vec{r}(t)$ to the output $y(t)$.
• $A_i(t)$, $\epsilon_i(t)$: The pre-activation input or "synaptic drive" ($A_i(t)$) to neuron $i$'s nonlinearity, and the instantaneous error signal ($\epsilon_i(t) = y_i(t) - y_i^*(t)$).

2. Appendix Organization: We agree the appendix should be restructured for a clearer narrative flow. We take your comments here seriously.

Task Details: We will expand all task descriptions.

• Multitask (Fig 5): The 20 tasks are from the well-established benchmark by Yang et al. (2019) and cover a range of cognitive functions. We will explicitly state this and cite the source more clearly. In the supplement, we will add a section dedicated to building the reader's intuition about the task set.

• Bandit Task (Fig 4): Your observation about the periodic switching is valuable. To address this potential confound, we ran a new experiment where the high-reward arm switched at stochastic, unpredictable intervals. The agent adapted robustly, confirming it learns a genuine model of the task state and is not merely entraining to a fixed frequency.

On Biological Plausibility (Weakness 4):

Our goal is a computationally powerful model that is inspired by biology, not a perfectly literal biophysical replica. We view our choices as principled abstractions designed to isolate core computational ideas.

• "Well-Mixed" Assumption: We agree this is a simplification. We interpret it as a timescale separation assumption, where diffuse neuromodulator concentrations establish a stable "context" for fast neural dynamics ($\tau_n \gg \tau_r$). While diffusion itself is slow, the establishment of a new ambient concentration across a local circuit can occur on a timescale relevant for trial blocks in cognitive tasks. We feel this is supported by two observations: 1) The integrative nature of the neuromodulator concentration acts as a low-pass filter on the population dynamics, which in turn modifies the mean population dynamics. This is a tool for achieving rich multi-timescale dynamics. 2) The neurobiological observation that single neuromodulator-releasing neurons project widely across the cortex suggests that biology could accelerate the approach to a well-mixed state by increasing the density of release loci.

• Interaction Term ($Z$ matrix): This term models effective interactions, not literal chemical reactions. It allows for one neuromodulatory system to influence another (e.g., release of one NM promoting or inhibiting the release/clearance of another), a known biological phenomenon (Morón JA, 2002, J Neurosci.). In practice, we often find this matrix is learned as approximately diagonal, simplifying to self-decay terms.

• Nonlinearities: These are standard in rate-based models to ensure stability and represent the saturating, positive firing rates of neurons. These nonlinearities also provide important numerical stability during training across a much wider range of hyperparameters.

Specific Questions:

• Define "Context Factorization": We define this as the decomposition of a network's effective connectivity into a fast-varying neural activity and a slow-varying modulatory state (the 'context'). Mathematically, the effective synaptic weight from neuron $j$ to $i$ is factorized as: $W_{ij}^{\text{eff}}(t) = W_{ij}^{0} + \sum_{k=1}^{M} \underbrace{T_{ijk}}_{\text{"dynamics"}} \underbrace{n_k(t)}_{\text{"context"}}$ This allows the low-dimensional context vector $\vec{n}(t)$ to dynamically reconfigure the high-dimensional effective connectivity of the network. We use the word "factorization" to evoke the mathematical method of decomposing a function into a product of components—here, the components are "context" (encoded in neuromodulator concentrations) and "dynamics" (encoded in the resulting synaptic connectivity patterns).

• Figure 3 Performance: Your observation is correct. The e-nmRNN's higher capacity makes it more prone to overfitting on this smaller dataset compared to a GRU. Our model's core strength is not necessarily to outperform all baselines on all tasks, but to provide a flexible and interpretable context-gating mechanism, the value of which is more clearly demonstrated in the multitasking and RL benchmarks.

• Figure 4 Bandit: In response to your comment, we have trained additional e-nmRNNs on the bandit task with both more stochastic block switching and taking the trained agents (without learning) into new switching periods. In the first set of experiments, we observed that the relative probability of the arms was detectable from the neuromodulator vector though admittedly with a less cyclic embedded trajectory. In the second set of experiments with fixed agents trained on oscillatory environments, we found that the e-nmRNNs where able to adapt to both longer switching periods and shorter switching periods. Given these quick turn around experiments, we expect that the more randomly switched bandits would indeed follow the reviewer's implicit hypothesis that they would more robustly learn online credit assignment that would allow the networks to be more robust to task variations.

• Figure 5 Spatial Plot: Thank you for this excellent suggestion. We performed the analysis you proposed on the trained multitasking network. We found that units specializing in releasing a specific neuromodulator form notable spatial clusters, providing new evidence strengthening our claims of emergent, biologically-reminiscent organization.

We hope these detailed clarifications and new analyses address your concerns and further the value of our work. We are committed to using your thought-provoking feedback to significantly improve the manuscript.

We thank reviewer yq59 for your positive review. We are delighted that you found the paper "clearly written and easy to follow," that our proposed model "presents a novel architectural idea," and that our "experimental results are well-supported by appropriate baselines." We address your insightful questions and comments below.

Weakness: Scalability properties of e-nmRNN are less clear... discuss the computational cost and parallelizability...

This is an important point. As a recurrent architecture, the e-nmRNN is inherently sequential during inference, which makes it less parallelizable than models like Transformers. This represents a trade-off: we gain powerful inductive biases for online, continuous adaptation within a single data stream at the cost of the massive parallelizability that makes Transformers so scalable for large-batch offline processing. We position the e-nmRNN as an architecture well-suited for applications like real-time control, agent-based learning, and brain-computer interfaces, where online processing is paramount. We believe its value also stems from its connection to neurobiology, providing an interpretable model for scientific inquiry even if it is not intended to compete with Transformers on large-scale sequence processing benchmarks.

Q1: The spatially embedded variant outperforms the standard e-nmRNN. Could this improvement be attributed to effective sparsity induced by spatial constraints?

Yes, we agree and believe this is a primary reason for the improved performance. The spatial embedding acts as a strong and biologically-inspired regularizer. It simultaneously imposes sparsity and encourages the formation of local processing motifs by making connection strength decay with distance. This regularization appears to prevent overfitting and guide the learning toward more robust and generalizable solutions, particularly on the complex multitasking benchmark.

Q2: In the multi-task decision-making experiments (section 6), does contextual differentiation emerge autonomously from the model dynamics, or is context explicitly provided as input?

The network is provided a one-hot task ID at the beginning of each trial. However, the crucial finding is how the network learns to use and maintain this information. The architecture's inductive bias encourages this context to be encoded and dynamically maintained within the neuromodulatory state $\vec{n}(t)$. As shown in Figure 5C, a linear decoder trained on these neuromodulator concentrations can reliably identify the current task long after the initial cue has passed, whereas a decoder trained on the neural activity $\vec{r}(t)$ cannot. This demonstrates that the e-nmRNN learns to represent the task context in its endogenous modulatory state, using it to shape the circuit dynamics appropriately for the task at hand.

Q3: What precisely differentiates the modeled "dopamine" and "serotonin" neurons in the architecture?

In our model, the differentiation is functional and learned, not hardcoded. We use names of known neuromodulators for illustrative purposes. A unit becomes a "dopamine-releasing neuron" (a term we would clarify to mean "a unit whose activity promotes dopamine release") if the corresponding column in the learned release matrix $R$ becomes selective for that modulator. For instance, if the $j$-th column of $R$, $R_{:j}$, becomes approximately a one-hot vector $[1, 0, 0, ...]^T$, then neuron $j$'s activity $r_j$ primarily drives the dynamics of the first neuromodulator, $n_1$. The functional effects of "dopamine" vs. "serotonin" are then determined by how their respective rows in the learned tensor $T$ gate synaptic transmission throughout the network to solve the given tasks. Here, we primarily use this description as an evocative example of how interpretation could proceed, but intend to reduce the emphasis in future revisions. We thank the reviewer for calling this to our attention.

We sincerely thank Reviewer gAEg for your positive assessment and for recognizing the value of our work. We are particularly grateful for your comments that the paper contains a "large number of carefully-conducted analyses," demonstrates a "strong mix of theory and simulation," and represents a "potentially significant advance for machine learning... and for neuroscience."

Per your suggestion to make our contributions clearer, we offer this bulleted list which we will endeavor to compress into any revised manuscript:

• A Unifying Computational Principle: We propose "context factorization" as a core computational principle implemented by neuromodulation, unifying its dual roles in altering circuit dynamics and gating synaptic plasticity.
• A Bio-Inspired Network Architecture (e-nmRNN): We introduce the endogenously neuromodulated RNN (e-nmRNN), a type of context-factorized hypernetwork derived from biophysical principles, where the network learns to generate its own modulatory context.
• A Mechanism for Online Credit Assignment: We demonstrate that the e-nmRNN architecture naturally enables targeted online credit assignment. Its learning rules are gated by the neuromodulatory state, allowing the network's dynamics to approximate gradient descent and facilitate rapid, in-context meta-learning.
• Strong Compositional Generalization: We show empirically that the e-nmRNN achieves respectable performance on compositional generalization tasks, outperforming and matching several strong baselines.
• Emergent Biologically aligned phenomena: We find that when trained on complex multitasking benchmarks, the model develops emergent properties that mirror neurobiology, such as modularity, cell-type specialization, and distinct functional timescales, enhancing its interpretability.
• Application to Reinforcement Learning: We demonstrate that the e-nmRNN can learn to encode critical learning signals, like Reward Prediction Error (RPE), within its neuromodulator dynamics, providing a mechanism for RPE-gated online adaptation in RL environments.

We now address your specific questions below.

Q1: How does an e-nmRNN differ from an nmRNN (in words and equations)? What is the particular advantage of the "endogenous" part?

This is an excellent question that gets to the heart of our model. The "endogenous" nature of our e-nmRNN is its defining feature: the neuromodulatory state $\vec{n}(t)$ is generated by the network's own recurrent activity $\vec{r}(t)$ through a learned projection matrix $R$. The dynamics are coupled:

This creates a closed-loop, self-modulating system. This contrasts with other models where the modulatory signal might be an external input (e.g., a task-ID vector). The advantage of the endogenous approach is that it allows the network to learn to generate its own internal context states directly from its neural population dynamics. This enables it to autonomously infer the current context from the input stream and adapt accordingly, which is crucial for the meta-learning and RL results we demonstrate, where explicit context signals are not always available.

Q2: Are there situations in which you would expect the algorithm to fail?

Yes, thank you for this insightful question. We see two primary failure modes:

• Data-Limited Overfitting: The model's rich, multiplicative dynamics make it more powerful but also more prone to overfitting than simpler RNNs when training data is scarce. This can be seen in our original Fig. 3 at high task complexity ($K > 15$), where the performance of the e-nmRNN degrades more rapidly than the GRU. In response to your inquiry, we have also performed an analysis of the generalization gap, which shows this phenomenon more clearly as the e-nmRNN’s gap between training and validation loss grows more rapidly than that of the equivalent GRU. A future research topic outside the scope of this work will investigate the effective forms of regularization to avoid this trap.
• Extremely Rapid Context Switching: The neuromodulator dynamics $\vec{n}(t)$ have intrinsic, learned timescales ($\tau_n$). If a task required context switches much faster than these timescales, the modulatory state would be unable to keep pace, leading to performance degradation. This highlights a fundamental trade-off between maintaining a stable context representation and achieving rapid adaptability.

We appreciate the reviewers notes on places where clarity can be improved and take these comments seriously. In future revisions, we will pay special attention to undefined variables, and the use of specialized neuroscience nomenclature without definition (e.g. Dale's law).