Adaptive behavior with stable synapses

Consider a generic learning task, in which the response $y$ to a state $x$ and is influenced by a set of parameters $w$. The latter are adjusted in function of an internal error $e$. This can be rewritten in a generic formulation as follows:

where $f$ is feedback from the environment (e.g. the reward, the target behavior, …). Alternatively to the standard interpretation of $w$ as synaptic weights of a neuronal network, we consider them as the state variables of a dynamical system coupled to the dynamics of $y$, expressed by the activity of an auxiliary network. For this reason, we also refer to them as “virtual” weights.

We consider the task of learning a target temporal trajectory $y^{targ}(t)$, We define $y(t)=\bm{w}\cdot\bm{x(t)}$ as the current estimation of $y^{targ}(t)$, where $\bm{w}$ and $\bm{x(t)}$ are virtual weights and input vectors respectively, since we extended to formulation to a multidimensional dataset. In this case, there is not a stationary projection of the training set, but the current value of the target signal itself is projected at every time stem $t$. We do not separate the target estimation on a test set and on a training set, as a consequence, $y(t)=y^{train}(t)$. In this case, learning can be formulated as follows:

where we removed the dependence on time $t$ for simplicity. The operations required are nonlinear, and usually are naturally implemented by the plasticity rule and by the multiplication between presynaptic activity and synaptic weights. However, here $\bm{w}$ are not actual weights but rather dynamical variables.

We propose, as possible implementation of this, that such non-linear functions are computed by two neural networks networks $W_{\Theta^{w}}$, and $Y_{\Theta^{y}}$:

In particular, we used two RNNs (see details below in the following Method subsections) that receive $\bm{x},\bm{w},e$ as inputs and provide the proper output thanks to a suited training of their readout weights ($\Theta^{y}$ and $\Theta^{w}$, following reservoir computing paradigm). In addition, we introduce a novel concept known as a gain-modulated network. This concept is inspired by the remarkable ability observed in biological neurons, particularly L5 neurons as discussed in previous studies (e.g., [22, 21]). These neurons exhibit the capacity to non-linearly integrate segregated inputs, a process critical for various cognitive functions. More details in the section "Gain Modulated Reservoir Computing". We empirically demonstrate that this network architecture outperforms standard RNNs in approximating and generalising the required virtual update rules described above, suggesting that gain modulation might be an important requirement for the adaptive behaviour observed in biological agents. The formulation shown above can only be used to tackle linear problems, by learning linear relationships between $\bm{x}$ and $y$. However, if we consider that $\bm{x}$ is the activity of another RNN that operates as a reservoir computer extracting nonlinear features of the input sequence, a wider class of tasks can be tackled without changing the constraint of a linear readout $y=wx$.

In the context of reinforcement learning, our framework outlines a systematic approach for modelling agents that can dynamically adapt their behaviour across diverse environments. We start by defining a policy network, denoted as $\bm{\pi}=\mathrm{softmax}(\bm{y})$ which implements a policy mapping the agent state encoded by the vector $\bm{x}$ to a probability distribution over actions. For the sake of simplicity, We assume a linear agent such that $\bm{y}=\bm{w}\cdot\bm{x}$. This assumption is done without loss of generality since this could be easily extended by resorting to a reservoir computer as an intermediate layer as described above. The policy depends on the virtual weights $\bm{w}$, determined by the activity of an additional RNN. This auxiliary network adjusts its internal activity based on the rewards received, effectively implementing policy gradient updates of the virtual weights and thereby modulating the agent behaviour in real time. This can be formalized in a formulation that is very similar to the one used above, by changing the definition of $e(t)$ as follows:

This is the dynamics obtained by evaluating the policy gradient with respect to the virtual weights $\bm{w}$. $a(t)$ is an integer value indicating the action at time $t$ among $\mathrm{D}$ possible ones. $\mathds{1}_{a(t)}$ represents the ’one-hot encoded’ action (as defined in [23, 24]) at time $t$. It is a $\mathrm{D}$-element vector where the $a(t)$-th element is one, and all other elements are zero.

This formulation holds for a null discount factor [23, 24]), we refer to the Appendix section C for the description of the general case. The above equation can be rewritten in order to refer only to two networks, one for estimating the gradients, and one for the scalar product.

To demonstrate the ability of this architecture to learn in context, we train networks Y and W using an Algorithmic Distillation protocol, as defined in [25]. Namely, we consider a family of environments $\mathcal{A}$ with different reward distributions ¹¹1A family of environments is a set of possible environments with the same action-state transitions but different reward distributions.. We then train the aforementioned networks using $N_{s}$ learning histories obtained through policy gradient methods from a subset of these environments, designated as $\mathcal{A}_{\text{ID}}$. The network performance is then tested out of distribution, i.e., on the rest of the environments $\mathcal{A}_{\text{OOD}}:=\mathcal{A}\setminus\mathcal{A}_{\text{ID}}$

We consider the task of autonomously predicting a temporal trajectory $\{y^{targ}(t)\}_{t}$. Following the reservoir computing paradigm, we employ an RNN (Fig. 1A, pink), whose dynamics follows Equation (6), to extract temporal features. During training, the reservoir receives the target dynamics $y^{targ}(t)$ as input ($x=y^{targ}(t)$ in Eq. (6), while $R$ is Gaussian matrix with zero mean and variance, $\sigma_{R}^{2}$) and is tasked with predicting the subsequent step of the trajectory via a linear readout of its activity (open loop, see Fig. 1B). This setup can be reformulated as a dynamical supervised learning problem, as described in the methods section, with the inputs $\bm{x}$ replaced by the features $\bm{z}(t)$ extracted by the RNN. The readout virtual weights can then be dynamically adjusted, minimizing the error between the target and the current prediction $y(t)$. This results in the following dynamics:

Here, the network $W_{\Theta^{w}}$$({\color[rgb]{1,0.58984375,0.70703125}\bm{z}},e)=\textrm{GMNN}_{\Theta^{w}}($$\bm{z}$$,e|\gamma_{e}e)$ is dedicated to estimating the gradient of virtual weights, while $Y_{\Theta^{y}}$$($$\bm{z}$$,w)=\textrm{GMNN}_{\Theta^{y}}($$\bm{z}$$,w|\gamma w)$ is tasked with estimating the predicted $y(t)$ as a function of the new virtual weights (Fig.1A green and orange respectively). Here, the parameters $\gamma,\gamma_{e}$ define the strength of the gain modulation, as explained in the Methods section.

These gain-modulated architectures are pre-trained to replicate, respectively, gradient descent updates and scalar products obtained on a set of $N_{train}$ training sequences $\{\{y_{\alpha}^{targ}(t)\}_{t}:\ \alpha\in[N_{train}]\}$. More specifically, the target sequences are $5$ sinusoidal functions with different frequencies (see vertical red lines in Fig. 1C.)

In the closed-loop phase, the features $\bm{z}$ are obtained by directly feeding the network estimation $y(t)$ as the input to the RNN ($x=y(t)$ in Eq. (6)). This results in an autonomous dynamical system that reproduces the target trajectory. (see Fig. 1B). In Fig.1B we report an example of successful dynamical adaptation of our model. Errors, represented by the blue trajectory, are fed back to the initial network to assess necessary updates to virtual weights, shown in the green trajectory. These updates are then transmitted to the second network, modifying the decoding of reservoir dynamics indicated by the orange lines. Initially, the reservoir receives the target trajectory as input in an open-loop fashion before the red vertical line, which is later replaced by the estimated trajectory itself in a closed-loop configuration after the red vertical line. The networks were pre-trained on five target frequencies marked by red vertical lines in Fig.1C and then tested across a range of frequencies. Evaluation was performed by assessing the mean squared error (MSE) between the target and estimated trajectories in closed-loop scenarios (see Fig.1C, solid and dashed lines represent respectively, median and 20th/80th percentile range. Statistics is evaluated over 10 realizations of the experiment).

Dynamical adaptation is now investigated within the framework of reinforcement learning. We provide robust evidence supporting the hypothesis that gain modulation represents a fundamental component in implementing in-context learning in biological agents. We first explore stateless environments $\mathcal{A}^{bandits}$, which are represented by Bernoulli K-armed bandits [30]. Within each environment $\alpha\in\mathcal{A}^{bandits}$, there exists a subset $P_{\alpha}\subset[K]$ of arms that yield a reward with high probability ($p=0.95$). During the training phase, the in-distribution environments give a high reward probability to even-numbered arms, whereas during testing, the out-of-distribution environments assign a high reward probability to odd-numbered arms.

In this simplified bandit scenario, where state information is absent, the virtual weights $\bm{w}$ directly parameterize the policy probabilities: $\bm{\pi}=\mathrm{softmax}{(\bm{w})}$. Consequently, our focus lies primarily on analyzing the behavior of the network $W_{\Theta^{w}}$$(\cdot)$ acting on the virtual parameters (Fig. 2A, green). This serves as an ideal test bed for evaluating the network’s ability to learn the policy gradient update rule and generalize it to out-of-distribution scenarios.

To parameterize $W_{\Theta^{w}}$$(\cdot)$ we employ a gain-modulated architecture $W_{\Theta^{w}}$ $(\bm{\pi},a,r)=\mathrm{GMNN}_{\Theta^{w}}(\bm{\pi},a,r|\gamma r)$, where $r$ is the reward and $a$ is the action. We train this network to approximate the policy gradient update rule. As training data, we use policy gradient estimates over one single learning trajectory (1000 rounds, learning rate $lr=0.1$) in an in-distribution environment. We then test the distilled policy gradient networks with a higher learning rate ($lr=1.0$) models in out-of-distribution (OOD) environments. We train the network to approximate the policy gradient update rule using policy gradient estimates from a single learning trajectory (1000 rounds, learning rate $lr=0.1$) in an in-distribution environment. We then test the distilled policy gradient networks with a higher learning rate ($lr=1.0$) in out-of-distribution (OOD) environments.

To systematically investigate the impact of gain modulation ($\gamma=1$) on out-of-distribution performance, we compare it with networks where the reward does not modulate the network gain ($\gamma=0$). For each case, we select the optimal hyperparameters through a grid-based search (additional details in the Appendix).

Our experimental findings are presented in Fig.2. First, comparing models of different sizes (number $N$ of hidden units), we find that models with gain modulation require significantly fewer neurons to achieve maximum scores in OOD environments (Fig.2C).

In Fig.2B, we report the regret per round distribution at the 100th round and compare it with the distribution obtained by policy gradient with the same learning rate and iterations. A gain-modulated network achieves a regret distribution comparable to the policy gradient in both ID and OOD environments, with a lower median. In contrast, networks without gain modulation show significantly higher regret. Analyzing the regret curves in Fig.2D, we observe that a model with gain modulation often learns a more data-efficient algorithm than its source, even in OOD environments. Conversely, a model without gain modulation fails to generalize to OOD environments and does not converge to zero regret.

In summary, gain modulation enables the network to consistently and efficiently distill the correct gradient update rule and generalize it to unseen environments, predicting the correct virtual weight update in regions of the input space far from the training data.

We examine a reaching task, known as the dark room task (a simple instance of the water maze task [31]), set in a 2D maze within the domain $(-1,1)\times(-1,1)$ with a grid size of $0.1$. Within this grid-like environment, the agent navigates by selecting one of four actions: up, down, left, or right, thereby determining its subsequent position. The primary goal is to locate a concealed object within the maze, with the agent having sole awareness of its own position. Feedback is provided via rewards, where the agent receives a reward of $10$ if the distance to the object is $0.1$, $15$ if the distance is $0$, and $0$ otherwise. Through iterative exploration, the agent develops a strategy to efficiently traverse the maze and pinpoint the object despite the limited visibility. The position of the agent is encoded separately using 25 input units $\bm{x}$ each, employing Gaussian activation functions distributed on a 5x5 grid in the maze and with a width of $0.2$.

To accomplish this task we consider a network architecture composed of two networks. One GM-network, $W_{\Theta^{w}}$$(r,a,\bm{\pi})=\textrm{GMNN}_{\Theta^{w}}(r,a,\bm{\pi}|\gamma_{r}r)$ is responsible for estimating the necessary update of virtual weights through policy gradient (Fig. 3A, green network), while another GM-network $Y_{\Theta^{y}}$$(\bm{x},\bm{w})=\textrm{GMNN}_{\Theta^{y}}(\bm{x},\bm{w}|\gamma\bm{w})$ estimates policy reconfiguration (Fig. 3A, shown in orange) resulting from changes in virtual weights. Here, the parameters $\gamma,\gamma_{r}$ define the strength of the gain modulation, as explained in the Methods section. Their dynamics can be described by the following equation Eq.(5). Indeed, we compare the case with (Fig. 3A) and without (Fig. 3B) gain modulation. We refer to Supporting Information for further details on the training procedure.

Firstly, we assessed the performance of the policy gradient algorithm using a set of 8 food locations (refer to Fig. 3C, indicated by red crosses), where the total reward averaged over 200 trials was observed against the number of trials (Fig. 3E, depicted by blue lines, thin lines for individual positions and thick lines for the average across all positions). After multiple trials, the agent successfully achieved precise targeting of the circles. Data collected from these experiments were utilized to train our networks to estimate gradients and scalar products, as defined in Eq. (5).

To validate that our trained model is capable of dynamically implementing policy gradient itself, we tested it on both the training set locations (ID, Fig. 3C, left panel) and new test positions (OOD, Fig. 3C, right panel). For each food position (coded with different colors), we illustrate a sample trajectory executed by the agent (in corresponding colors) to reach the target at the end of the training. The agent’s precision closely matches that of the plastic policy gradient learning rule. We present the reward plotted against the number of trials for both training (Fig. 3E, right panel, pink line) and test (Fig. 3C, right panel, pink line) food locations.

We compared these performances against an architecture lacking gain modulation ($\gamma=0$), observing worst performances (see Fig. 3B, D, F). Notably, while performances for ID food locations are acceptable (Fig. 3E, left panel, pink lines), those for OOD cases are extremely poor (Fig. 3E, right panel, pink lines). This observation, coupled with the results from the preceding section, suggests that gain modulation is a crucial component in facilitating the generalization of adaptive behavioral capabilities in recurrent networks.

We demonstrate, that our architecture is capable to perform the temporal computation required to learn delayed action-reward temporal relations (see section C in Appendix), requiring evaluating temporal credit assignment.