Elucidating the theoretical underpinnings of surrogate gradient learning in spiking neural networks

TL;DR

The relation of surrogate gradients to two theoretically well-founded approaches are investigated and it is found that smoothed probabilistic models give surrogate gradients a theoretical basis in stochastic spiking neural networks, where the surrogate derivative matches the derivative of the neuronal escape noise function.

Abstract

Training spiking neural networks to approximate universal functions is essential for studying information processing in the brain and for neuromorphic computing. Yet the binary nature of spikes poses a challenge for direct gradient-based training. Surrogate gradients have been empirically successful in circumventing this problem, but their theoretical foundation remains elusive. Here, we investigate the relation of surrogate gradients to two theoretically well-founded approaches. On the one hand, we consider smoothed probabilistic models, which, due to the lack of support for automatic differentiation, are impractical for training multi-layer spiking neural networks but provide derivatives equivalent to surrogate gradients for single neurons. On the other hand, we investigate stochastic automatic differentiation, which is compatible with discrete randomness but has not yet been used to train spiking neural networks. We find that the latter gives surrogate gradients a theoretical basis in stochastic spiking neural networks, where the surrogate derivative matches the derivative of the neuronal escape noise function. This finding supports the effectiveness of surrogate gradients in practice and suggests their suitability for stochastic spiking neural networks. However, surrogate gradients are generally not gradients of a surrogate loss despite their relation to stochastic automatic differentiation. Nevertheless, we empirically confirm the effectiveness of surrogate gradients in stochastic multi-layer spiking neural networks and discuss their relation to deterministic networks as a special case. Our work gives theoretical support to surrogate gradients and the choice of a suitable surrogate derivative in stochastic spiking neural networks.

Figures (10)

• Figure 1: SD are equivalent to derivatives of expected outputs in SPM and smoothed stochastic derivatives in binary Perceptrons.(A) Membrane potential dynamics of an LIF neuron (maroon) in comparison to the Perceptron. When input spikes (from input neurons $n_{\mathrm{in}_1}$ and $n_{\mathrm{in}_2}$) are received, they excite the LIF neuron which causes the membrane potential to increase. Once it reaches the threshold, output spikes $n_\mathrm{out}$ are emitted. In the limit of a large simulation time step ($dt \gg \tau_{mem}$) and appropriate scaling of the input currents, the LIF neuron approximates a Perceptron receiving time-locked input (gray line). (B) Left: The simplified computational graph of a deterministic (blue) and a stochastic (yellow) Perceptron. Right: Forward pass in a deterministic (top) or stochastic (bottom) Perceptron. The colored arrows indicate that both use the derivative of $\sigma_{\beta_\mathrm{SG}}(\cdot)$ on the backward pass, which is the derivative of the expected output of the stochastic neuron in case $\beta_\mathrm{SG} = \beta_\mathrm{N}$. In the case of the deterministic neuron, this constitutes the SG used instead of the non-existing derivative of the step function. (C) Left: Network output over multiple trials in the deterministic Perceptron. Right: The SD (blue) is the derivative of a sigmoid (gray), which is used to approximate the non-existing derivative of the step function (black). (D) Same as (C) but for the stochastic Perceptron. Left: Escape noise leads to variability in the spike trains over trials. Right: The expected output follows a sigmoid, and we can compute the derivative (yellow) of the expected output.
• Figure 2: Derivative computation in MLP. Schematic of an example network for which (surrogate) derivatives are computed according to different methods. The colored arrows indicate where partial derivatives are calculated. (A):SG descent relies on the chain rule for efficient gradient computation in a deterministic MLP. Thus, the derivative of the output with respect to a given weight is factorized into its primitives, which are indicated by the colored arrows. (B)SPM approach the problem of non-differentiable spike trains by adding noise and then smoothing the output based on its expected value. Since this method does not allow the use of the chain rule, the derivative for each weight must be computed directly. (C) The derivative and the expected value are not interchangeable, which makes this option mathematically invalid. Furthermore, it is not possible to achieve the necessary smoothing using the expected value after such an interchange. (D) Smoothed stochastic derivatives in stochAD use the expected value of each node to compute the derivative. However, the method relies on expectation values conditioned on the activity of a specific forward pass.
• Figure 3: SG correspond to smoothed stochastic derivatives in stochastic SNN. The tree illustrates the discrete decisions associated with the binary spike generation process at different points over different layers of a stochastic MLP. A single forward pass in the network corresponds to a specific path through the tree which yields a specific set of spike trains. Another forward pass will result in a different path and spike patterns. Computing gradients using finite differences requires randomly sampling paths from the network and evaluating their averaged loss before and after a given weight perturbation. Although this approach is unbiased for small perturbations, the random path selection results in high variance. Furthermore, that approach is not scalable to large networks. Stochastic SG descent is equivalent to smoothed stochastic derivatives in the stochAD framework. To compute the gradient, we roll out the network once and sample a random path in the tree which we now keep fixed (yellow). At each node, we then compute the expected output given the fixed activation of the previous layer $\mathbb{E}[h_i|h_{i-1}]$, which yields a low-variance estimate (see inset: spike raster, selected trial shown in yellow, spike trains of other trials in gray, expectation shown as shaded overlay). By choosing a surrogate function that matches the escape noise process, both methods give the same derivative for a spike with respect to the membrane potential. Deterministic SG descent can be seen as a special case in which the random sampling of the path is replaced by a point estimate given by the deterministic roll-out (blue).
• Figure 4: SG deviate from actual gradients in differentiable MLP.(A) Schematic of a differentiable network (left) with sigmoid activations (right) for which we compute an SD using the derivative of a flatter sigmoid (yellow) in contrast to the actual activation (black). (B) Top row: Network output (solid gray), smoothed network output (dashed), and integrated SD (yellow) as a function of $w$. The triangles on the x-axis indicate the minimum of the corresponding curves. Bottom row: Derivatives of the top row. Left and right correspond to a flatter ($\beta_{\mathrm{SG}}=15$) and a steeper ($\beta_\mathrm{SG}=25$) SD, see Table \ref{['tab:sign_flip_params']} for network parameters. Note that the actual derivative and the surrogate can have opposite signs. (C) Heatmap of the optimization landscape along $v_1$ and $v_2$ for different $\beta_\mathrm{SG}$ values (top to bottom). While the actual gradient can be asymptotically zero (see yellow dot, bottom), the SD provides a descent direction (yellow arrow), thereby enabling learning (top and middle).
• Figure 5: SG are not gradients.(A) Heat map of the network output while moving along two random directions in the parameter space of the example network with step activation (top) and sigmoid activation (bottom) (see Fig. \ref{['fig:bias_in_sigmoid_networks']} A). The circles indicate a closed integration path through parameter space, starting at the arrowhead. (B) Integral values of SD as a function of the angle along the closed circular path shown in (A). Different shades of yellow correspond to different values of $\beta_\mathrm{SG}$. The blue lines correspond to the actual output of the network with step activation function (solid) or sigmoid activation (dashed). The integrated actual derivative of the network with sigmoid activation matches the output (dashed line) and is thus not visible in the plot. (C) Absolute difference between actual loss value and integrated SD as function of the number of integration steps. The numerical integrals converge to finite values. Thus the observed difference is not an artifact of the numerical integration.