GRSN: Gated Recurrent Spiking Neurons for POMDPs and MARL

The development of SRL can be mainly divided into three periods: (1) the basic research of SNNs and RL Williams (1992); Bohté et al. (2002); (2) the synaptic-plasticity-based SRL algorithm Seung (2003); Urbanczik and Senn (2009); Frémaux et al. (2010); and (3) the combination of deep RL and SNNs Zhang et al. (2021); Liu et al. (2022); Qin et al. (2023b). Early-stage works often study the combination of synaptic plasticity and RL theory, which aims to reveal how reward mechanisms in the brain correlate and combine with RL algorithms. In the later stage, with the development of DRL, the SRL algorithm focuses on applying SNNs to DRL, aiming for better performance.

Due to the inherent temporal nature of spiking neurons, there have been many attempts to use recurrent spiking neural networks for temporal data processing. Some work Lotfi-Rezaabad and Vishwanath (2020); Rao et al. (2022) has designed spike-based long short-term memory networks and tested them on audio or text datasets. Other work Yin et al. (2020); Ponghiran and Roy (2022); Liu et al. (2023) is to explore the long-term dependence of spiking neurons and complete structure optimization or time series forecasting tasks.

For all variables, superscripts refer to the number of layers in the neural network, and subscripts refer to their time steps in the time window. For example, $\bm{u}^{l}_{t}$ denotes the membrane potential in layer $l$ at the $t$-th time step. We follow the conventions representing vectors and matrix with bold italic letters and bold capital letters, respectively, such as $\bm{o}$ and $\bm{W}$. We use $\odot$ to signify element-wise multiplication for two vectors.

Different from traditional neural networks, SNNs use binary spikes as information carriers. In order to compensate for the lack of binary spikes in information expression, the time dimension, or latency, is introduced into SNNs. SNNs accept event streams as input, and the forward propagation of SNNs is repeated for $T$ time steps in the temporal dimension to calculate the output spike trains. Therefore, the basic computing unit of SNNs, spiking neurons, has unique spatio-temporal dynamics. Here we introduce the widely used Leaky Integrate-and-Fire (LIF) neuron Wu et al. (2018) as the benchmark in this work. In each time step $t$, the membrane potential $\bm{u}^{l}_{t}$ of spiking neurons at the $l$-th layer will integrate the input current $\bm{c}^{l}_{t}$ and the decay voltage $\beta\bm{\hat{u}}^{l}_{t-1}$:

Membrane potential will be reset to $u_{r}$ (default set to $0$) when a spike is emitted at the last time step:

The neurons will emit output spikes $\bm{o}^{l}_{t}$ whenever the membrane potential $\bm{u}^{l}_{t}$ exceeds the threshold $\vartheta$:

where $\Theta(x)$ is the Heaviside function:

Due to the non-differentiability of the binary firing function $\Theta(x)$, SNNs often employ surrogate gradients Neftci et al. (2019) during back-propagation (BP). We adopt the arc-tangent function to replace the derivative of the binary firing function:

where $\alpha$ is a hyper-parameter that controls the width of the surrogate function. With the surrogate gradients for binary firing functions, we can directly use the BP method to train SNNs.

Reinforcement learning tasks can be modeled as Markov decision processes (MDP) ($\mathcal{S},\mathcal{A},\mathcal{R},p,\gamma$), with state space $\mathcal{S}$, action space $\mathcal{A}$, scalar reward function $\mathcal{R}$, transition dynamics $p$, and discount factor $\gamma$ Sutton and Barto (2018). RL agents are controlled by policy $\pi$, with the goal of maximizing the expected discounted return $\mathbb{E}_{\pi}[\sum_{t=0}^{\infty}\gamma^{t}r_{t+1}]$. According to the Bellman equation Bellman (1966), we can obtain the iteration form of the value function:

When the policy is determined, Eq 7 can be simplified as:

Observing Eq 8 and Eq 1, we can find that the state transitions in MDPs exhibit a striking resemblance to the state changes of spiking neurons. Therefore, the key to solving the temporal mismatch problem is to utilize this similarity.

We propose a novel SRL paradigm according to this similarity: within the framework of RL algorithms, the application of SNNs in a manner similar to RNNs (Figure 2A), aligning each step of spiking neuron updates $\bm{u}_{t}\rightarrow\bm{u}_{t+1}$ with each state transition $s_{t}\rightarrow s_{t+1}$ in an MDP. The advantage of this approach lies in the natural utilization of the temporal structure of SNNs while significantly reducing the number of time steps required for the entire task. Figure 2A and  2B show the proposed TAP and the detailed dynamics of spiking neurons.

In Section 3.2, we introduce the basic computational unit of SNNs: LIF neurons. It uses a constant leaky factor $\beta$ to play the roles of forget gate and input gate. The use of this constant factor results in weaker temporal correlations in the neuron’s internal state ($\bm{u}_{t}$), making it challenging to handle data with long-term dependencies. In order to enhance the memory capacity of SNNs in the temporal domain, we need to analyze and optimize the LIF neuron.

According to Eq 1-5, it can be observed that the LIF neuron has two internal states: membrane potential $\bm{u}^{l}_{t}$ and input current $\bm{c}^{l}_{t}$. To compare the importance of these two internal states in backpropagation, we need to calculate the gradient of the loss function for connection weight $\bm{W}^{l}$. Suppose $\mathcal{L}$ is the loss function that we would like to minimize. According to Eq 1-5 and the chain rule, its gradient is:

$T$ is the total number of time steps. The first term $\frac{\partial\mathcal{L}}{\partial\bm{o}^{l}_{t}}$ in Eq 9 is determined by the decoder and loss function and does not affect the flow of the gradient. According to Eq 4 and Eq 6, the second term could be derived as:

Then we calculate the third term $\frac{\partial\bm{u}^{l}_{t}}{\partial\bm{W}^{l}}$:

Bring Eq 1 into Eq 11:

Obviously, $\frac{\partial\hat{\bm{u}}^{l}_{t-1}}{\partial\bm{c}^{l}_{t}}=0$. Let

Then Eq 12 can be simplified as:

Continuing substitution $\frac{\partial\bm{u}^{l}_{t}}{\partial\bm{W}^{l}}$ leads to a closed-form expression:

When the variable subscript $t=0$, its value is the artificially set initial value.

To analyze the main carriers of gradient, we propose the following theorem:

To further verify Theorem 1, we conducted experiments in CartPole-V (see Section 5) environment and visualized the parameter distribution during the training process (Figure 3). From figure 3, the maximum absolute value of the term $\beta\delta_{t}$ is less than $1$, and most of the membrane potential values are distributed around $0$. Hence, $\prod_{i=1}^{T}\beta\delta_{t}$ will tend to $0$ when $T$ is large enough. This result is consistent with Theorem 1.

According to Theorem 1, the gradient of the loss function $\mathcal{L}$ is mainly transmitted through the input current $\bm{c}$. Therefore, adding gated recurrent connections to the input current $\bm{c}$ can enhance the memory capability of spiking neurons more effectively than membrane potential $\bm{u}$.

Optimizing the dynamics of neurons and improving memory ability is mainly divided into three steps. The first step is to change the decay factor $\beta$ to a learnable parameter, which can be jointly optimized with the entire network. The learnable decay factor can dynamically adjust the forgetting ability of neurons during the training process and adaptively find the optimal value Rathi and Roy (2023). Then, we use a soft reset instead of a hard reset (Theorem 1 still holds in the soft reset setting; see Figure 3 and appendix for details). Compared to the hard reset method, soft reset can reduce temporal information loss, and it has better performance when the network is shallow Ledinauskas et al. (2020). Finally, we added recurrent connections and adopted gating functions to enhance the temporal correlation of the neuron’s internal states $\bm{c}$.

We named the proposed neuron model gated recurrent spiking neurons (GRSN). Its membrane potential update rules, firing rules, and definitions of Heaviside function are consistent with those of LIF neurons (see Eq 1, Eq 4 and Eq 5). Different from LIF neurons (Eq. 3), the reset rule of GRSN is soft reset:

GRSN also adds recurrent connections at the current level:

where $\mathscr{F}(\cdot)$ is the forgetting gate that regulates the proportion of forget and input, and $\mathscr{I}(\cdot)$ is the input gate that processes input signals. Both $\mathscr{F}(\cdot)$ and $\mathscr{I}(\cdot)$ are composed of a fully connected (FC) layer and an activation function. Their detailed definitions are shown as follows:

where $\sigma$ is the sigmoid function.

Figure 2C shows the dynamics of GRSN in the temporal alignment paradigm. Compared with Figure 2B, red lines indicate the recurrent connections. The pseudocode of GRSN with the temporal alignment paradigm is shown in Algorithm 1.

In Section 4, we propose the TAP that can solve the temporal mismatch problem and greatly reduce time steps, and the GRSN that can enhance the temporal association of spiking neurons and improve performance. To demonstrate the impact of these two on the training speed and algorithm performance more intuitively, we conducted multiple groups of ablation experiments. We conducted experiments in both classical control tasks and SMAC environments. Similar to the previous setting, we conducted five independent experiments at each of the different settings and averaged them as the final results (Table 3).

SNNs use discrete spikes for information transmission, so their energy consumption is estimated differently than that of ANNs. We follow the convention of the neuromorphic computing community Wang et al. (2023); Qin et al. (2023b); Wu et al. (2023) by counting the total synaptic operations (SOP) to estimate the energy consumption of SNN models. Specifically, the SOP for SNNs correlates with the neurons’ firing rate, fan-out $f_{out}$, and time window size $T$:

where $L$ is the total number of layers and $N^{l}$ is the number of neurons in layer $l$, $f_{out}$ is the number of outgoing connections, and $\bm{o}_{t}^{l}[j]$ denotes the output spike of the $j$-th neuron in layer $l$ at the $t$-th time step.

For ANNs, SOP is only related to network structure. It is defined as:

where $f_{in}^{l}$ represents the number of incoming connections to each neuron in layer $l$, $L$ and $N^{l}$ are consistent with their definitions in SNNs.

SNNs use efficient addition operations, while ANNs use more expensive multiply-accumulate operations. According to Han et al. (2015), we measure 32-bit floating-point addition operations by $0.9pJ$ per operation and 32-bit floating-point multiply-accumulate operations by $4.6pJ$ per operation. Here, we randomly select 1024 samples to estimate the average SOP and energy consumption for SNNs. Table 4 shows the results of different networks in different environments, and it can be seen that SNN can reduce energy consumption by up to about 3800 times compared to RNNs. The total energy consumption is also reduced by about 50%, which effectively reduces the agent’s resource requirements.