Sharpness Aware Surrogate Training for Spiking Neural Networks

Spiking neural networks are trained with smooth approximations but deployed with discontinuous spikes. SAST closes that gap by seeking flat loss minima, here is the full theory, the intuition behind it, and a concrete worked example.

Spiking neural networks

The LIF neuron

The workhorse model is the Leaky Integrate-and-Fire (LIF) neuron. At each discrete time step $t$ a neuron accumulates charge into its membrane potential $u[t]$, leaks a fraction of it away, and fires a binary spike $s[t] \in \{0,1\}$ whenever that potential crosses a threshold $\vartheta$. After firing the potential resets.

Concretely, for a layer with weight matrix $W$ and input spikes $s^{\ell-1}[t]$ from the layer below:

where $\tau \in (0,1)$ is the membrane decay constant and $H(\cdot)$ is the Heaviside step function. The term $(1-s^\ell[t-1])$ implements the reset: once the neuron fires, the potential is zeroed before it integrates new input.

Surrogate gradients

Training requires backpropagating through $H$, which has zero gradient almost everywhere and an undefined gradient at zero. The standard fix is to replace $H'(x)$ with a smooth surrogate $\sigma'_\beta(x)$ during the backward pass while keeping the forward spike binary. A common choice is the arctangent surrogate:

SAST takes a different route: rather than keeping the forward pass hard and only smoothing the backward pass, it trains a surrogate-forward network where the spike function itself is replaced by $\sigma_\beta$ everywhere. This means the gradient is exact for the model being trained — no straight-through approximation anywhere in the computation graph.

The transfer gap

After surrogate training the learned weights are evaluated under two regimes. In surrogate-forward mode the smooth $\sigma_\beta$ is used at inference — this is what training optimised for. In hard-spike mode the smooth activations are replaced with $H$ — this is what a real neuromorphic deployment requires.

The transfer gap is simply:

For a standard baseline on N-MNIST, $\Delta = 30.3$ percentage points. SAST reduces this to $2.5$ pp. Understanding why requires thinking about the shape of the loss landscape.

Why flat minima help

A concrete example

Consider a single hidden neuron whose membrane potential at some time step is $u = 0.999$ and whose threshold is $\vartheta = 1.0$.

The solution is to make the network stop relying on neurons that sit on the fence. If every neuron whose surrogate output is near 0.5 were nudged away from threshold — so that $|u - \vartheta| \gg 0$ — then replacing $\sigma_\beta$ with $H$ would barely change anything, because both functions agree far from zero.

SAST achieves this indirectly, without explicitly penalising fence-sitter neurons. It does so by seeking flat regions of the parameter space.

The loss-landscape view

Imagine the surrogate training loss as a mountain range and the parameters $w$ as a location in that landscape. Standard gradient descent finds a valley — a local minimum — but says nothing about how steep the valley walls are. A narrow, sharp valley means that tiny perturbations to $w$ send the loss shooting upward. A wide, flat valley means the loss stays low across an entire neighbourhood.

Switching from surrogate to hard spikes is exactly such a perturbation. Every neuron near threshold gets its activation discretised — a jump that cannot be predicted from the gradient alone. If the network sits in a flat valley, this perturbation stays within the low-loss region. If it sits in a sharp valley, it likely escapes it.

Sharpness-aware minimization finds flat valleys by construction. That is the geometric reason SAST reduces the transfer gap.

The SAST method

SAM recap

Standard training minimises the empirical risk $\mathcal{L}(w)$. Sharpness-Aware Minimization (Foret et al., 2021) instead minimises the worst-case loss in a ball around $w$:

The inner maximisation is intractable in general but admits a first-order approximation. Expanding around $w$:

The outer update then follows the gradient at the perturbed point $w + \varepsilon^*(w)$, not at $w$ itself. This pushes $w$ toward regions where the loss is simultaneously low and locally flat.

The SAST update

SAST applies SAM to the surrogate-forward loss $\mathcal{L}_{\text{surr}}(w)$ — the cross-entropy of the SNN when evaluated with $\sigma_\beta$ everywhere. Each training step proceeds in two phases:

The independence of $\mathcal{B}$ and $\mathcal{B}'$ is not just an implementation detail — it is the condition required by the convergence theorem (Theorem 3 below). In code:

for each step:
  g  = grad surrogate_loss(w; B)           # phase 1 batch
  eps = rho * g / (norm(g) + delta)
  w -= lr * grad surrogate_loss(w + eps; B')  # phase 2 batch

Theory

The theoretical results are scoped to the smooth auxiliary model — the surrogate-forward SNN that SAST actually optimises. They do not claim to govern the hard-spike network directly; that connection is empirical.

State stability

The first result establishes that the surrogate state sequence is well-behaved. Under the assumption that the spectral norm of each weight matrix is bounded:

the surrogate state map is a contraction in the $\ell_2$ norm across time steps. Concretely, for two input sequences $x$ and $x'$ that differ only in one time step:

where $\kappa = \tau\gamma < 1$. This fading-memory property means early perturbations wash out over time — a necessary precondition for the Lipschitz and smoothness results that follow.

Lipschitz bound and smoothness

Under the same spectral-norm bounds and a bound on the surrogate derivative $\|\sigma'_\beta\|_\infty \leq C_\sigma$, the network output is Lipschitz in its input with constant:

The empirical surrogate loss is then $\beta$-smooth in the parameters — meaning its Hessian is bounded:

Smoothness is the key property SAM theory requires. Without it, the first-order perturbation $\hat{\varepsilon}(w)$ would not reliably approximate the true worst-case direction. Because the surrogate-forward model is smooth by construction (every activation is differentiable), this holds cleanly. It would not hold for the hard-spike model.

Convergence

Given $\beta$-smoothness and independent minibatches, SAST satisfies a standard nonconvex convergence bound. After $T$ steps with learning rate $\eta = O(1/\sqrt{T})$:

This guarantees convergence to a first-order stationary point of the surrogate loss at the standard $O(1/\sqrt{T})$ rate for nonconvex stochastic optimisation. The independent-minibatch requirement is what keeps the gradient estimate unbiased at the perturbed point — if the same batch were used for both phases, the estimate would be biased and the bound would not hold.

The paper also includes a local mechanism proposition: under a bound on the per-sample Jacobian $\|J_w s_i\|_F \leq C_J$, the input gradient norm satisfies:

Smaller parameter gradients (as SAM encourages) imply smaller $C_J$, which in turn implies smaller input sensitivity. This is the local theoretical connection between sharpness and the transfer gap — but it is presented as a proposition under local conditions, not a universal theorem.

Results

The empirical story is about the transfer gap, not raw surrogate accuracy. Surrogate accuracy changes modestly; hard-spike accuracy changes dramatically.

On N-MNIST, surrogate accuracy improves from 96.1% to 97.2% — a 1.1 pp gain. But hard-spike accuracy improves from 65.7% to 94.7% — a 29 pp gain. The transfer gap collapses from 30.3 pp to 2.5 pp. On DVS Gesture the pattern repeats: the surrogate gain is modest (1.9 pp) while the hard-spike gain is large (31.5 pp).

The overhead is real. Each SAST step requires two forward-backward passes — one to compute $\hat{\varepsilon}$ and one to compute the update gradient — giving a roughly 2× per-epoch cost. Whether this is acceptable depends on the deployment context.

Limitations

References