Policy optimization emerges from noisy representation learning

Thank you for your valuable and detailed feedback! We have provided our responses below.

The goal of this paper is to address the question: Can policy optimization be implemented with plasticity rules designed for representation learning? This question is biologically relevant given the mounting evidence that most plasticity in the brain is designed for representation learning. NETO is a new biologically plausible learning algorithm. But more importantly, it introduces the novel concept that the brain could exploit its intrinsic noise to implement policy optimization from a system otherwise designed to extract useful information from the environment.

Neuroscience (part 1)

[Response to comments on the value of the NETO noise hypothesis and the role of activity noise]

1. On the roles of activity and weight noise: When we write that “NETO treats the brain's intrinsic noise as a feature rather than a bug,” we mean both the weight and activity noise. We choose to set the activity noise to zero in this work for simplicity, as the activity noise has two conceptually distinct effects on the network. As mentioned in section 3.2.2, the activity noise makes the agent’s policy stochastic because the network’s responses to stimuli are noisy. However, the activity noise also acts as an effective weight noise; the network weights evolve according to the activities, so if there is noise in the activities, this noise will carry through to the weight updates. Therefore, in principle, the brain could take advantage of the weight and activity noise to learn optimal policies (hence the statement in quotes). We focused on the former because the activity noise has the additional effect of creating a stochastic policy, which in principle could be useful but adds unnecessary complexity to the discussions in this work. We appreciate this important point and have added a new section (Appendix A.3) explaining this nuance. We direct the reader to this section when introducing weight noise in line 202.
2. On the value of the NETO noise hypothesis: The hypothesis has value on three fronts. First, it proposes a normative role for noise in the brain by showing that a noisy representation learner can also learn to optimize a policy if its noise is modulated. This follows a similar line of reasoning to recent work that makes hypotheses based on normative computational principles, see [1-3], which is a powerful method for understanding neural computation. Second, there is evidence for noise modulation in the brain. Dopamine is known to control the signal-to-noise ratio in (at least) medial prefrontal cortex [4-6], as noted in Appendix A.3 and A.4. Though this typically refers to the activity signal to noise ratio, as previously noted, activity noise is an effective weight noise. Third, the noise hypothesis has the potential to reproduce several experimentally observed phenomena, as discussed in section 5.1. We thank you for this question and added a note on line 207 pointing the reader to Appendix A.4 for a brief discussion on the biological plausibility of noise modulation.

[Response to comments regarding the comparison to other biological models]

1. We agree that if more complex NETO agents could not solve Atari games, this would be a major strike against the framework. However, this work does not argue that NETO is an interesting model of learning because it is powerful (though it very possibly is, and future work will explore how far the framework can be taken in terms of task complexity with larger networks that use more powerful representation learning objectives). Rather, it argues that NETO is exciting because it presents a new normative role for noise, and reconciles the empirical evidence that much of the brain’s plasticity is designed for representation learning with the fact that organisms are excellent policy optimizers. Moreover, it achieves this reconciliation without proposing disjoint network modules with different plasticity rules (like Capone et al. 2024 do).
2. To your comment about experimental validation of plasticity rules: The SNNs in Capone et al. use two different plasticity rules in two different network modules. In contrast, NETO uses one plasticity rule across the entire network. In the example that we gave, this was the Hebbian rule (derived from the SM objective), which is experimentally validated. In general, NETO can use any number of plasticity rules proposed through computational principles that have seen success in reproducing experimental results (for instance, those in [7]). The specific form of the modulation (9) is an arbitrary choice taken to demonstrate the framework, as stated below the equation. We do not intend to claim that it has been validated experimentally.

(We placed the citations at the end of our multi-part response.)

[Response to final question]

1. We view the SM objective as only one possible choice of representation learning objective. Crucially, the framework is intended to be more general and incorporate objectives that can build much more powerful representations. The reason that we chose to introduce NETO with the SM objective in great depth (as opposed to testing it on multiple objectives) is that this gives good intuition for why and how NETO agents learn optimal policies: they diffusively search near the solution space of their representation learning objective. With this intuition, our opinion is that it is clear the framework will work with other representation learning objectives. The one caveat, alluded to in section 4.2, is that NETO agents are most likely to find reasonable policies if these policies exist near the solution space of their representation learning objective. (This is perfectly reasonable biologically, since nervous systems must extract “useful” information from the environment in order for organisms to survive.) As long as the objective extracts information suitable for the task, our analysis will generalize. The experiment demonstrating our claims about NETO agents using the NNSM objective will be included in the camera-ready version and will also further support this claim.

Citations:

[1] Lipshutz, David, et al. "Normative framework for deriving neural networks with multicompartmental neurons and non-Hebbian plasticity." PRX Life 1.1 (2023): 013008.

[2] Duong, Lyndon, et al. "Adaptive whitening with fast gain modulation and slow synaptic plasticity." Advances in Neural Information Processing Systems 36 (2024).

[3] Pehlevan, Cengiz, Tao Hu, and Dmitri B. Chklovskii. "A hebbian/anti-hebbian neural network for linear subspace learning: A derivation from multidimensional scaling of streaming data." Neural computation 27.7 (2015): 1461-1495.

[4] Kroener, Sven, et al. "Dopamine modulates persistent synaptic activity and enhances the signal-to-noise ratio in the prefrontal cortex." PloS one 4.8 (2009): e6507

[5] Winterer, Georg, and Daniel R. Weinberger. "Genes, dopamine and cortical signal-to-noise ratio in schizophrenia." Trends in neurosciences 27.11 (2004): 683-690.

[6] Vander Weele, Caitlin M., et al. "Dopamine enhances signal-to-noise ratio in cortical-brainstem encoding of aversive stimuli." Nature 563.7731 (2018): 397-401.

[7] Halvagal, Manu Srinath, and Friedemann Zenke. "The combination of Hebbian and predictive plasticity learns invariant object representations in deep sensory networks." Nature neuroscience 26.11 (2023): 1906-1915.

Neuroscience (Part 2)

[Response to figure comment]

1. Thank you for raising this point. We have edited the figure to reflect the structure of the mammalian brain. However, we would like to comment that this figure is purely schematic. In principle, NETO could be relevant to any organism with a neuromodulatory system. The position of the elements in the mammalian brain is intended only to visually clarify the setup.

[Response to comments on experimental connections]

1. We greatly appreciate these comments. To start, NETO is not a theory for representation learning but a theory for how policy optimization can emerge from a representation learner. As such, its predictions do not necessarily compete with the predictions of representation learning theories but rather add to them. The specific formulation of NETO that we discuss was designed to give intuition for the framework and highlight how it learns. It is abstract from a biological perspective because it has two neurons and operates under the SM objective in simple tasks. As such, it is outside of the scope of this work to make specific falsifiable experimental claims regarding representation drift. The purpose of including experimental connections here was to highlight that NETO can qualitatively reproduce several experimentally observed phenomena, and therefore shows promise for a biological theory of learning. The purpose was not to give conclusive evidence or specific predictions, though this will be the subject of future work.
2. Regarding your comment about NETO operating under the NNSM objective: this is also a good point. We hesitated to elaborate on the connections to experiment further because we were already short on space. We are working on an experiment to demonstrate our claims empirically. We will incorporate it into the camera-ready version.

Machine Learning

[Response to comment on architectures considered]

1. The choice of including recurrent architectures in Figure 1 was intentional but nothing more than suggestive. NETO agents learn rewarding policies because their reward-dependent noise induces a distribution in and around the solution space of a representation learning objective. If a representation learning objective were implemented in a recurrent architecture, mathematically speaking, this same phenomenon would occur, so NETO should still “work.” However, testing a variety of architectures is outside the scope of this work, as the primary purpose is to introduce the NETO framework and associated concepts.

[Response to comment on comparisons with other ML algorithms]

1. Since the primary aim of the paper is to introduce NETO as a framework for learning in biological systems, the relevant comparative baselines are the components and assumptions of other biological models, not the learning efficiencies of other algorithms in these simple tasks. NETO is not to be understood as a competitor to classical TD-learning-style algorithms in terms of efficiency.

[Response to comment on task difficulty]

1. We completely agree. We reiterate the point that the tasks considered in this work are intentionally simple. This choice ensures that the learning process is amenable to mathematical analysis and therefore easy to understand conceptually. This attribute is more important for this work than testing NETO in complex tasks, given that the purpose of the paper is to introduce the framework as a model of learning in biological systems.

[Responses to “Additional weaknesses”]

1. We are grateful that you pointed out the missing citation here. We have corrected this and rephrased our description of the strategy used by the authors. We refrain from being more explicit in their argument because it is quite involved (please see the supplementary information in Qin et al. 2023).
2. Thank you for your concern and insightful suggestion. Empirically, we observe that our argument holds only for certain input normalizations. The main goal of the Cart Pole task is to provide an example where optimal policies do not exist within the solution space of $\mathcal{L}$. As such, we've decided to exclude the normalizations where optimal policies can be found within this space. Instead, we only include the example with no normalization. Additionally, we've included the experiment you suggested in a new section (Appendix A.6) and direct the readers there on line 425. With these adjustments, Cart Pole more clearly serves as an example of a task that the agent can only solve by identifying task-relevant features. It also removes the possible confusion around the meaning of normalizations.

Thank you for your valuable feedback! We have provided our responses below.

[Responses to comments on clarity]

1. There were a range of responses regarding the clarity of the manuscript, from very clear (reviewer DznG) to overly pedantic (reviewer ZG6t) to slightly unclear (reviewer vJYN). We agree that reformatting the paper in terms of theorems and lemmas would highlight the theoretical results. However, this choice would also falsely suggest that the primary contribution of the paper is the results in the contextual bandit task. Rather, the primary contribution is the novel, biologically relevant idea that reward-dependent noise can implement policy optimization in networks otherwise designed for representation learning. This idea is much more general than the two tasks (and even the representation learning objective used), so we hesitate to present results regarding these specific examples through theorems and lemmas. The tasks presented serve to illustrate the main conceptual ideas regarding how NETO agents learn policies.
2. With that being said, we greatly appreciate your suggestion to offer more high-level intuition before delving into the calculations and have added this intuition to lines 229-230 before delving into the math.

[Responses to comments on theoretical contributions]

1. In general, any choice of $\eta_{\theta}(R)$ that monotonically decreases as a function of $R$ is appropriate, consistent with the requirement outlined in Section 3.2.3 that the noise decrease monotonically as a function of $R$. We focused on the example $\eta_{\theta}(R) = e^{-\beta R}$ for clarity, as the purpose of the contextual bandit task is to illustrate the more general concept that modulating noise facilitates policy optimization. In this example, the agent is guaranteed to learn an optimal policy as $\beta \to \infty$ (though in practice, $\beta \approx 5$ works nearly perfectly, as seen in Figure 4). Therefore, the statement on line 376 was referring to the choice of the parameter beta: It should be large. We replaced the phrase “with the appropriate choice” with the phrase “with strong modulation” to emphasize the fact that we are referring to the choice of beta here. However, in general, any monotonically decreasing function of R is appropriate, with more sharply decreasing functions (compared to the reward scale) leading to more bias towards rewarding policies.
2. The specific analysis in Section 4.1 applies only to agents using the similarity matching objective and in the specified contextual bandit task. The NETO framework, however, is completely general in terms of both the representation learning objective used and the task. The examples provided are meant to illustrate how NETO agents learn in scenarios where their learning is analytically tractable and therefore easily understood.

[Responses to additional questions]

1. We use PDEs that imply continuous time because NETO is meant as a model for learning in biological systems, which live in continuous time. However, to investigate the model's properties computationally, one must use discrete-time, which makes MDPs a natural choice.
2. We included a definition of the similarity matching objective as a footnote in the revised manuscript [line 160].
3. This is also a great question. We can look at the contextual bandit task for some intuition. The matter of whether a diffusive search can escape a local optimum is a question about a specific learning trajectory. As we saw in the contextual bandit task, the most natural way to think about learning in the NETO framework is not to focus on a specific trajectory but to consider the limiting distribution over network parameters. This limiting distribution is defined jointly by the representation learning objective and the noise modulation. If there were isolated locally optimal policies, one would find the distribution peaked at each, with larger peaks corresponding to more rewarding policies. As a result, NETO agents are generally more likely to find global optima than to get stuck in local optima. The specific choice of $\eta_{\theta}(R)$ will determine the relative magnitude of these peaks and will be the subject of future work. However, if one does wish to focus on a specific learning trajectory, the answer is that the agent can escape local optima, though the escape time likely depends nontrivially on the task, the form of noise modulation, and the representation learning objective.
4. Thank you for pointing this out as well. We feel that the distinction between different normalizations was a confusing detail and not necessary to the main point that the agent learns task-relevant features not captured by $\mathcal{L}$ to solve Cart Pole. As such, we removed all normalizations. The agent now receives input from Cart Pole without alteration. We have edited Figure 5 and Section 4.2 to reflect this change.

Thank you for your valuable feedback! We have provided our responses below.

[Responses to comments on task choice]

1. On task simplicity: There is a trade-off between analytical tractability and the complexity of the task. Here, we deliberately consider simple tasks in order to analytically follow and understand the learning dynamics. We make this choice because the primary contribution of this work is to introduce the idea behind NETO: Reward-dependent noise can facilitate policy optimization in networks otherwise designed for representation learning. We agree that in future work it will be important to test more complex tasks.
2. This is also an insightful point. We chose to focus on tasks that do not require online learning because these tasks are more amenable to analysis, which allows us to highlight how NETO agents learn. However, we appreciate the observation and have added a new section to the Appendix (A.2) that extends the contextual bandit by making the context associations switch on some timescale, which means that the agent’s policy must adapt in real time to changing goals (which is only possible through online learning). We see efficient transfer across tasks; the network does not have to relearn to extract the principal subspace of the inputs, it only needs to relearn how to use this information.

[Response to comment on comparative baselines]

1. Since the primary aim of the paper is to introduce NETO as a framework for learning in biological systems, the relevant comparative baselines are the components and assumptions of other biological models, not their learning efficiencies in linear, two-neuron networks. In the Related Work section, we highlight that NETO is more promising than other “baseline” models given the growing consensus that the brain is primarily a network-level representation learning system. To our knowledge, NETO is the only model to date that can solve nontrivial tasks with network plasticity rules that operate primarily on this principle. Moreover, because of its generality, NETO can adopt the previous successes of plasticity rules derived from representation learning objectives in reproducing neural data.

[Response to comment on suitability for ICLR]

1. The aim of this paper is primarily to serve as an introduction to a biologically plausible learning framework. However, the ideas presented in this work could also be of interest for pure machine learning applications in the context of efficient online learning. NETO agents can transfer learned representations across tasks. We demonstrate this idea in Appendix A.2. As tasks change, the agent does not have to relearn to extract the principal subspace of its inputs. Instead, it only needs to relearn how to use this information. While this is a simple example, the concept generalizes to agents with more complex representation learning objectives. For them, this information transfer is crucial because it allows the agents to learn the dynamics of their environment just once, after which they can use this understanding to solve a variety of tasks. Additionally, we note that ICLR has a history of publishing works in theoretical neuroscience. For instance, please see:
   • “A Unified Theory of Early Visual Representations from Retina to Cortex through Anatomically Constrained Deep CNNs” (ICLR 2019)
   • “Relating transformers to models and neural representations of the hippocampal formation” (ICLR 2022)
   • “Disentanglement with Biological Constraints: A Theory of Functional Cell Types” (ICLR 2023)
   • “Hebbian Deep Learning Without Feedback” (ICLR 2023)

[Responses to suggestions]

1. Thank you for this suggestion – we did try log scales but unfortunately it did not make the figures more visually appealing.
2. Thank you for this suggestion too. We agree and have edited Figure 4 to use pi-based labels.