Gradient Descent for Spiking Neural Networks

TL;DR

This paper introduces a gradient descent method for training spiking neural networks by making them differentiable, enabling efficient supervised learning for dynamic, spike-based computations.

Contribution

It presents the first exact gradient calculation method for optimizing spiking neural networks, facilitating supervised learning in these models.

Findings

Successfully trained recurrent spiking networks on dynamic tasks

Optimized spike timing and network behavior at multiple time scales

Demonstrated improved learning capabilities for spike-based computation

Abstract

Much of studies on neural computation are based on network models of static neurons that produce analog output, despite the fact that information processing in the brain is predominantly carried out by dynamic neurons that produce discrete pulses called spikes. Research in spike-based computation has been impeded by the lack of efficient supervised learning algorithm for spiking networks. Here, we present a gradient descent method for optimizing spiking network models by introducing a differentiable formulation of spiking networks and deriving the exact gradient calculation. For demonstration, we trained recurrent spiking networks on two dynamic tasks: one that requires optimizing fast (~millisecond) spike-based interactions for efficient encoding of information, and a delayed memory XOR task over extended duration (~second). The results show that our method indeed optimizes the spiking network dynamics on the time scale of individual spikes as well as behavioral time scales. In conclusion, our result offers a general purpose supervised learning algorithm for spiking neural networks, thus advancing further investigations on spike-based computation.

Equations

$$\tau \dot{s}$$

$$\int sdt=\int g\dot{v}dt=\int gdv=1.$$

$$\dot{v}$$

$$\vec{I}=W\vec{s}+U\vec{i}+{\vec{I}}_o,$$

$$\vec{o}(t)=O\vec{s}(t),$$

$$-{\dot{p}}_v$$

$$-\tau {\dot{p}}_s$$

$$\vec{\xi }=W^{⊺}\overrightarrow{(f_I{p}_v)}+{\vec{\partial }}_{s}l,$$

$${\nabla }_{W}C$$

$${\nabla }_{U}C$$

$${\nabla }_{I_o}C$$

$${\nabla }_{O}C$$

$$l=\frac{\Vert \vec{o}-{\vec{o}}_d{\Vert }^2+\lambda \Vert \vec{s}{\Vert }^2}{2},$$

$$f(v,I)=I.$$

$$\tau {\dot{\vec{o}}}_d=-{\vec{o}}_d+\vec{i},$$

$$\tau {\dot{\vec{o}}}_d=-{\vec{o}}_d+A{\vec{o}}_d+\vec{i},$$

$$f(v,I)=(1+\cos (2\pi v))/{\tau }_v+(1-\cos (2\pi v))I,$$

$$\mathcal{H}$$

$$=\sum _i({\bar{p}}_{v_i}+g(v_i){\bar{p}}_{s_i}/\tau )f(v_i,I_i)-({\bar{p}}_{s_i}/\tau )s_i+l(\vec{s}),$$

$$-{\dot{\bar{p}}}_{v_i}$$

$$-{\dot{\bar{p}}}_{s_i}$$

$$\mathcal{H}$$

$$-{\dot{p}}_{v_i}$$

$$-\tau {\dot{p}}_{s_i}$$

Full text

1 Introduction

The brain executes highly dynamic computation over multiple time-scales: Individual neurons exhibit spiking dynamics of millisecond resolution and the recurrent network connections produce slower dynamics on the order of seconds to minutes. How the brain organizes the spiking neuron dynamics to form the basis for computation is a central problem in neuroscience. Nonetheless, most analysis and modeling of neural computation assume static neuron models that produce analog output, called rate-based neurons [1, 2]. These simplified models are compatible with the advanced tools from the deep learning field that can efficiently optimize large scale network models to perform complex computational tasks [3]. However, such rate-based neural network models fail to describe the fast dynamics of spike-based computation in the brain.

The main difficulty in optimizing spiking neural networks stems from the discrete, all-or-none nature of spikes: A spiking neuron generates a brief spike output when its membrane voltage crosses the threshold, and silent at other times. This non-differentiable behavior is incompatible with the standard, gradient-based supervised learning methods. Consequently, previous learning methods for spiking neural networks explored various ways to circumvent the non-differentiability problem.

SpikeProp [4] considered the spike times of neurons as state variables and used the differentiable relationship between the input and the output spike times to minimize the difference between the actual and the desired output spike times. However, the creation and deletion of spikes are non-differentiable, so the number of output spikes must be pre-specified. Memmesheimer et al [5] considered the problem of generating spikes at desired times and remaining silent at other times as a binary classification problem and applied the perceptron learning rule (See also [6]). Pfister et al [7] considered stochastic spiking neurons and maximized the smooth likelihood function of the neurons spiking at desired times. More biologically inspired methods based on spike-time-dependent-plasticity (STDP) have also been proposed [8, 9, 10, 11]. All of these methods, however, require target spiking activity of individual neurons at desired times, which significantly limit their range of applicability.

Alternative methods have also been proposed: Instead of directly optimizing a spiking network, these methods optimize a network of static analog neurons and replicate the optimized solution with a spiking network. Hunsberger and Eliasmith [12] used analog units that closely approximated the firing rate of individual spiking neurons. Instead of replicating individual neuron’s dynamics, Abbott et al [13] proposed replicating the entire network dynamics using recently developed methods from predictive coding [14]. Although these approaches are applicable to a wider range of problems, the replicated spiking networks can only mimic the solutions of the rate-based networks, rather than exploring the larger space of spike-time based solutions.

In this paper, we introduce a novel, differentiable formulation of spiking neural networks and derive the exact gradient calculation for gradient based optimization. This method optimizes the recurrent network dynamics on the time scale of individual spikes for general supervised learning problems.

2 Methods

2.1 Differentiable synapse model

In spiking networks, transmission of neural activity is mediated by synaptic current. Most models describe the synaptic current dynamics as a linear filter process which instantly activates when the presynaptic membrane voltage $v$ crosses a threshold: e.g.,

[TABLE]

where $\delta(\cdot)$ is the Dirac-delta function, and $t_{k}$ denotes the time of threshold-crossing. Such threshold-triggered dynamics generates discrete, all-or-none responses of synaptic current, which is non-differentiable.

Here, we replace the threshold with a gate function $g(v)$: a non-negative ($g\geq 0$), unit integral ($\int g~{}dv=1$) function with narrow support111Support of a function $g:X\to\mathbb{R}$ is the subset of the domain $X$ where $g(x)$ is non-zero., which we call the active zone. This allows the synaptic current to be activated in a gradual manner throughout the active zone. The corresponding synaptic current dynamics is

[TABLE]

where $\dot{v}$ is the time derivative of the pre-synaptic membrane voltage. The $\dot{v}$ term is required for the dimensional consistency between eq (1) and (2): The $g\dot{v}$ term has the same $[\text{time}]^{-1}$ dimension as the Dirac-delta impulses of eq (1), since the gate function has the dimension $[\text{voltage}]^{-1}$ and $\dot{v}$ has the dimension $[\text{voltage}][\text{time}]^{-1}$. Hence, the time integral of synaptic current, i.e. charge, is a dimensionless quantity. Consequently, a depolarization event beyond the active zone induces a constant amount of total charge regardless of the time scale of depolarization, since

[TABLE]

Therefore, eq (2) generalizes the threshold-triggered synapse model while preserving the fundamental property of spiking neurons: i.e. all supra-threshold depolarizations induce the same amount of synaptic responses regardless of the depolarization rate (Figure 1A,B). Depolarizations below the active zone induce no synaptic responses (Figure 1E), and depolarizations within the active zone induce graded responses (Figure 1C,D). This contrasts with the threshold-triggered synaptic dynamics, which causes abrupt, non-differentiable change of response at the threshold (Figure 1, dashed lines).

Note that the $g\dot{v}$ term reduces to the Dirac-delta impulses in the zero-width limit of the active zone, which reduces eq (2) back to the threshold-triggered synapse model eq (1).

The gate function, without the $\dot{v}$ term, was previous used as a differentiable model of synaptic connection [15]. In such a model, however, a spike event delivers varying amount of charge depending on the depolarization rate: the slower the presynaptic depolarization, the greater the amount of charge delivered to the post-synaptic targets.

2.2 Network model

To complete the input-output dynamics of a spiking neuron, the synaptic current dynamics must be coupled with the presynaptic neuron’s internal state dynamics. For simplicity, we consider differentiable neural dynamics that depend only on the the membrane voltage and the input current:

[TABLE]

The dynamics of an interconnected network of neurons can then be constructed by linking the dynamics of individual neurons and synapses eq (2,3) through the input current vector:

[TABLE]

where $W$ is the recurrent connectivity weight matrix, $U$ is the input weight matrix, $\vec{i}(t)$ is the input signal for the network, and $\vec{I}_{o}$ is the tonic current. Note that this formulation describes general, fully connected networks; specific network structures can be imposed by constraining the connectivity: e.g. triangular matrix structure $W$ for multi-layer feedforward networks.

Lastly, we define the output of the network as the linear readout of the synaptic current:

[TABLE]

where $O$ is the readout matrix.

The network parameters $W$, $U$, $O$, $\vec{I}_{o}$ will be optimized to minimize the total cost, $C\equiv\int l(t)~{}dt$, where $l$ is the cost function that evaluates the performance of network output for given task.

2.3 Gradient calculation

The above spiking neural network model can be optimized via gradient descent. In general, the exact gradient of a dynamical system can be calculated using either Pontryagin’s minimum principle [16], also known as backpropagation through time, or real-time recurrent learning, which yield identical results. We present the former approach here, which scales better with network size, $\mathcal{O}(N^{2})$ instead of $\mathcal{O}(N^{3})$, but the latter approach can also be straightforwardly implemented.

Backpropagation through time for the spiking dynamics eq (2,3) utilizes the following backpropagating dynamics of adjoint state variables (See Supplementary Materials):

[TABLE]

where $f_{v}\equiv\partial f/\partial v$, and $\xi$ is called the error current. For the recurrently connected network eq (4), the error current vector has the following expression

[TABLE]

which links the backpropagating dynamics eq (5,6) of individual neurons. Here, $f_{I}\equiv\partial f/\partial I$, $({f}_{I}{p}_{v})_{j}\equiv f_{I_{j}}{p}_{v_{j}}$, and $(\partial_{s}l)_{j}\equiv\partial l/\partial s_{j}$.

Interestingly, the coupling term of the backpropagating dynamics, $g\dot{p}_{s}$, has the same form as the coupling term $g\dot{v}$ of the forward-propagating dynamics. Thus, the same gating mechanism that mediates the spiked-based communication of signals also controls the propagation of error in the same sparse, compressed manner.

Given the adjoint state vectors that satisfy eq (5,6,7), the gradient of the total cost with respect to the network parameters can be calculated as

[TABLE]

where $(\partial_{o}l)_{j}\equiv\partial l/\partial o_{j}$. Note that the gradient calculation procedure involves multiplication between the presynaptic input source and the postsynaptic adjoint state $p_{v}$, which is driven by the $g\dot{p}_{s}$ term: i.e. the product of postsynaptic spike activity and temporal difference of error. This is analogous to reward-modulated spike-time dependent plasticity (STDP) [17].

3 Results

We demonstrate our method by training spiking networks on dynamic tasks that require information processing over time. Tasks are defined by the relationship between time-varying input-output signals, which are used as training examples. We draw mini-batches of $\approx 50$ training examples from the signal distribution, calculate the gradient of the average total cost, and use stochastic gradient descent [18] for optimization.

Here, we use a cost function $l$ that penalizes the readout error and the overall synaptic activity:

[TABLE]

where $\vec{o}_{d}(t)$ is the desired output, and $\lambda$ is a regularization parameter.

3.1 Predictive Coding Task

We first consider predictive coding tasks [14, 19], which optimize spike-based representations to accurately reproduce the input-ouput behavior of a linear dynamical system of full-rank input and output matrices. Analytical solutions for this class of problems can be obtained in the form of non-leaky integrate and fire (NIF) neural networks, although insignificant amount of leak current is often added. The solutions also require the networks to be equipped with a set of instantaneously fast synapses, in addition to the regular synapses of finite time constant [19].

The membrane voltage dynamics of a NIF neuron is given by

[TABLE]

To ensure that the membrane voltage stays within a finite range, we impose two thresholds at $v_{\theta+}=1$ and $v_{\theta-}=-1$, and the reset voltage at $v_{\text{reset}}=0$.222The reset process, which may seem non-differentiable, does not influence the gradient calculation. For simplicity, we allow the $v_{\theta-}$ threshold to trigger negative synaptic responses, which can be turned off if desired.

We also introduce the additional fast synaptic current $\vec{s}_{f}$ proposed in [14, 19], which modifies the input current vector to be $\vec{I}=W\vec{s}+W_{f}\vec{s}_{f}+U\vec{i}+\vec{I}_{o},$ where $W_{f}$ is the recurrent weight matrix associated with fast synapses. However, assigning zero time constant to the fast synapses often causes unstable dynamics, because it could lead to one spike immediately triggering more spikes in other neurons. Here, we assign finite time constants for both types of synapses: $\tau_{f}=1$ ms for fast synapses, and $\tau=10$ ms for regular synapses.

Despite its simplicity, the predictive coding framework reproduces important features of biological neural networks, such as the balance of excitatory and inhibitory inputs and efficient coding [14]. Also, its analytical solutions provide a benchmark for assessing results from optimization.

Auto-encoder task

In the auto-encoder task, the desired output signal is a low-pass filtered version of the input signal:

[TABLE]

where $\tau$ is the synaptic time constant [14, 19]. The goal is to accurately represent the analog signals using least number of spikes. We used a network of 30 NIF neurons, and $2$ input and output signals. Randomly generated sum-of-sinusoid signals with period $1200$ ms were used as the input. $\lambda=0.1/N$ ms2. $\Delta=0.1$ was used for training, then set to zero for post-training simulations.

The output of the trained network accurately tracks the desired output (Figure 2A). Analysis of the simulation reveals that the network operates in a tightly balanced regime: The fast recurrent synaptic input, $W_{f}\vec{s}_{f}(t)$, provides opposing current that mostly cancels the input current from the external signal, $U\vec{i}(t)$, such that the neuron generates a greatly reduced number of spike outputs (Figure 2B,C,D). The network structure also shows close agreement to the prediction. The optimal input weight matrix is equal to the transpose of the readout matrix (up to a scale factor), $U\propto O^{\intercal}$, and the optimal fast recurrent weight is approximately the product of the input and readout weights, $W_{f}\approx-UO$ , which are in close agreement with [14, 19, 20]. The regular recurrent connection is not needed for this task and hence $W$ was set to zero. Such network structures have been shown to maintain tight input balance and remove redundant spikes to encode the signals in most efficient manner: The representation error scales as $1/K$, where $K$ is the number of involved spikes, compared to the $1/\sqrt{K}$ error of encoding with independent Poisson spikes.

General predictive coding task

More generally, predictive coding tasks involve linear dynamic relationships between the desired input-output signals of the following form:

[TABLE]

where $A$ is a constant matrix. Here, we trained a spiking network of 30 NIF neurons with $2$ input signals of sums-of-sinusoid and $A=[-0.7,0.36;-2.3,-0.1]$, which strongly modulates the desired output signal dynamics.

Similar to the result shown in Figure 2, the trained network exhibits tightly balanced input current with the network output accurately tracking the desired output. The optimal regular recurrent weight is approximately $W\approx UAO$ (Figure 3), which is also in close agreement with the prediction [14, 19, 20]. The other network structures are similar to the case of auto-encoding task.

These results show that our algorithm accurately optimizes the millisecond time-scale interaction between neurons to find an efficient spike-time-based encoding scheme. Moreover, it also shows that efficient coding can be robustly achieved without introducing instantaneously fast synapses, which were previously considered to be necessary.

3.2 Delayed-memory XOR task

A major challenge for spike-based computation is in bridging the wide divergence between the time-scales of behavior and spikes: How do millisecond spikes perform behaviorally relevant computations on the order of seconds?

Here, we consider a delayed-memory XOR task, which performs the exclusive-or (XOR) operation on the input history stored over extended duration. Specifically, the network receives binary pulse signals, $+$ or $-$, through an input channel and a go-cue through another channel. If the network receives two input pulses since the last go-cue signal, it should generate the XOR output pulse on the next go-cue: i.e. a positive output pulse if the input pulses are of opposite signs ($+-$ or $-+$), and a negative output pulse if the input pulses are of equal signs ($++$ or $--$). Additionally, it should generate a null output if only one input pulse is received since the last go-cue signal. Variable time delays are introduced between the input pulses and the go-cues.

A simpler version of the task was proposed in [13], whose solution involved first training an analog, rate-based neural network and replicating the learned network dynamics with a larger network of spiking neurons ($\approx 3000$), using the method from predictive coding [14]. It also required a dendritic nonlinearity function to match the transfer function of rate neurons.

We trained a network of $80$ quadratic integrate and fire (QIF) neurons333NIF networks fail to learn the delayed-memory XOR task: the memory requirement for past input history drives the training toward strong recurrent connections and runaway excitation., whose dynamics is

[TABLE]

also known as Theta neuron model [21], with the threshold and the reset voltage at $v_{\theta}=1$, $v_{\text{reset}}=0$. Time constants of $\tau_{v}=25$, $\tau_{f}=5$, and $\tau=20$ ms were used, whereas the time-scale of the task was $\approx 500$ ms, much longer than the time constants. The intrinsic nonlinearity of the QIF spiking dynamics proves to be sufficient for solving this task without requiring extra dendritic nonlinearity. The trained network successfully solves the delayed-memory XOR task (Figure 4): The spike patterns exhibit time-varying, but sustained activities that maintain the input history, generate the correct outputs when triggered by the go-cue signal, and then return to the background activity. More analysis is needed to understand the exact underlying computational mechanism.

This result shows that out algorithm can indeed optimize spiking networks to perform nonlinear computations over extended time.

4 Discussion

We have presented a novel, differentiable formulation of spiking neural networks and derived the gradient calculation for supervised learning. Unlike previous learning methods, our method optimizes the spiking network dynamics for general supervised tasks on the time scale of individual spikes as well as the behavioral time scales.

Exact gradient-based learning methods inevitably involve discrepancies from biological learning processes. Nonetheless, such methods provide solid theoretical ground for understanding the principles of biological learning rules. For example, our result shows that the gradient update occurs in a sparsely compressed manner near spike times, bearing close resemblance to reward-modulated STDP. Moreover, further analysis may reveal that certain aspects of the gradient calculation can be approximated in a biologically plausible manner without significantly compromising the efficiency of optimization. For example, it was recently shown that the biologically implausible aspects of backpropagation method can be resolved through feedback alignment for rate-based multi-layer feedforward networks [22]. Such approximations could also apply to spiking neural networks.

Here, we coupled the synaptic current model with differentiable single-state spiking neuron models. However, the synapse model can be coupled with any neuron models, including realistic multi-state neuron models with detailed action potential dynamics444Simple modification of the gate function would be required to prevent activation during the falling phase of action potential., including the Hodgkin-Huxley model, the Morris-Lecar model, and the FitzHugh-Nagumo model; and even models with internal adaptation variables. It can also be coupled with neuron models having non-differentiable reset dynamics, such as the leaky integrate and fire model, the exponential integrate and fire model, and the Izhikevich model, although gradient calculation on these models would require additional procedures. This will be examined in the future work.

Acknowledgments

We thank Peter Dayan for helpful discussions.

Supplementary Materials: Gradient calculation for the spiking neural network

Pontryagin’s minimum principle

According to [16], the Hamiltonian for the spiking network dynamics eq (2,3,4) is

[TABLE]

where $\bar{p}_{v_{i}}$ and $\bar{p}_{s_{i}}$ are the adjoint state variables for the membrane voltage $v_{i}$ and the synaptic current $s_{i}$ of neuron $i$, respectively, and $l(\vec{s})$ is the cost function. The back-propagating dynamics of the adjoint state variables are:

[TABLE]

where $f_{v}\equiv\partial f/\partial v$, $f_{I}\equiv\partial f/\partial I$, $g^{\prime}\equiv dg/dv$, and $l_{s_{i}}\equiv\partial l/\partial s_{i}$.

This formulation can be simplified by change of variables, ${p}_{v}\equiv\bar{p}_{v}+g\bar{p}_{s}/\tau$, ${p}_{s}\equiv\bar{p}_{s}/\tau$, which yields

[TABLE]

where we used $\dot{p}_{v_{i}}=\dot{\bar{p}}_{v_{i}}+f_{i}g_{i}^{\prime}\bar{p}_{s_{i}}/\tau+g_{i}\dot{\bar{p}}_{s_{i}}/\tau$.

The gradient of the total cost can be obtained by integrating the partial derivative of the Hamiltonian with respect to the parameter (e.g. $\partial\mathcal{H}/\partial W_{ij}$, $\partial\mathcal{H}/\partial U_{ij}$, $\partial\mathcal{H}/\partial I_{o_{i}}$, $\partial\mathcal{H}/\partial O_{ij}$).