Exact Gradients for Stochastic Spiking Neural Networks Driven by Rough Signals

Summary

Spiking neural networks face the problem of non-differentiability of loss function due to the Heaviside activation function, which works as the spiking function. To back-propagate the loss through the network, Heaviside is replaced with surrogate functions as a workaround. The paper establishes the theoretical existence of the exact gradient for stochastic spiking neural networks, where stochasticity is introduced in spike timing, and the gradient of spike timing is computed with respect to the model parameters. It uses the existing theories of stochastic differential equations (rough path theory) to handle discontinuous spiking events and introduces a novel kernel (Marcus signature kernel) to handle such discontinuities. They further provide an algorithm that can perform parameter estimation of a two-layer network on toy examples.

Strengths

The paper uses stochastic leaky-intergrate-and-fire neurons, which have randomness in the evolution of membrane potential as well as spike timings. It uses the mathematical framework of rough path theory to obtain the gradient of the spike timings with respect to network parameters and formally prove the existence of such a gradient in this framework. The paper is well-written and provides intuitive descriptions along with rigorous mathematical proofs. It brings a novel theoretical insight into spiking neural network training.

Weaknesses

The paper demonstrates it implementation through a toy parameter estimation problem on an one-hidden layer feed forward network of dimension $4 \times 16 \times 2$, where the network parameters are estimated. Such a network, having a tiny input dimension, is not readily applicable to realistic datasets. Further, rather than implementing the algorithm as a parameter search (as done in standard SNN training), it is implemented as a parameter estimation problem. The algorithmic aspects of the paper are somewhat open; the paper discusses several possible avenues, such as "discretize-then-optimize" vs. "optimize-then-discretize" frameworks. The authors must discuss the computation complexity of the proposed EventSDESolve algorithm, which uses the discretize-then-optimize approach.

Questions

**Doubt:** Assumptions 3.3, 3.4 and 3.5 require the event and transition function to be continuously differentiable [line 194]. How do such functions model discontinuous spikes? In contrast, line 229 mentions that trajectories and spike times are almost surely differentiable. Please clarify. **Measure of Randomness:** It is shown that the zeroth-order derivative of the Heaviside function approximates a surrogate function depending upon the noise distribution [1]. How, in the present case, does the variance of the introduced noise (in the Brownian motion) affect the exact gradient? Does it recover the gradient of Heaviside in a limiting case? Which of the two noises (Brownian motion vs. transition noise [line 166]) used in the SLIF modelling is essential to obtain the exact gradient -- are either of them a mathematical convenience? **Computational Feasibility:** Please discuss the computational complexity of EventSDESolve. How does the computational complexity of the gradient scale with the number of parameters? Can the algorithm handle real datasets? Why did the authors choose to solve a parameter estimation problem (through estimation of input current) instead of standard training of the network parameters? [1] Mukhoty, Bhaskar, et al. "Direct training of snn using local zeroth order method." Advances in Neural Information Processing Systems 36 (2024).

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

Yes, authors do discuss some of the limitations of their work.

Summary

This work develops a mathematical framwork to compute the gradients of stochastic spiking neural networks, or more generally Event SDEs. The proposed framework is an alternative to existing surrogate gradient frameworks and an extension of prior adjoint based work (e.g. EventProp) in the presence of stochasticity. Jax-based code is provided with the submission.

Strengths

The manuscript offers an attractive solution to the problem of solving SNNs, based on rigorous demonstrations. The derivations and description makes too much use of very specific jargon, making it difficult to evaluate with respect to the existing methods. At this point, if there is indeed practical merit in the proposed approach, it is unlikely to have an impact to the field as is. The lack of testing on standard SNN benchmarks doesn't help this problem. I believe these problems could be fixed by making more parallels with existing work and at least ont benchmark that demonstrates the effectiveness of the proposed method. The yin yang dataset in [55] would be a good candidate. - Mathematically rigorous demonstration of stochastic SNN gradients is novel and relevant. Strickly speaking, existing methods already demonstrated exact gradients in SNNs in the absence of gradients and when events are known [55], but the extension to stochastic inputs and dynamics is new to my knowledge. Note that in this sense, the statement L260 is already no longer popular belief. - Connection to three-factor rules in the feed-forward case extends prior work in the field to stochastic models

Weaknesses

- Use of jargon and mathematical assumptions which are sometimes not sufficiently connected with the studied problem. This is a missed opportunity IMO to connect with the relevant communities. e.g. I have little clue what this means: "train NSDEs non-adversarially using a class of maximum mean discrepancies (MMD) endowed with signature kernels [51] indexed on spaces of continuous paths as discriminators". Why adversarial? what is a continuous signature kernel? Maybe I'm the wrong reader for this paper, but then this would likely apply to most readers interested in training SNNs. - Lack of a benchmark and experimental comparison to prior work based on eventprop or surrogate gradients. The framework is written in jax, and should in principle allow a scalable simulation, why not run on a simple SNN benchmark? The provided experiments are overly simple and leaving this open as a vague limitation statement is not enough.

Questions

- Please explain some ambiguous sentences and less known facts: - L211: "If the diffusion term is absent we recover the gradients in [6]" from which expression to we recover the gradients in [6]? - Give at least a one sentence description of cadlag (right continuous, left limit). Why is it relevant to the modeling of SNNs? - What is a continuous signature kernel? Can you give an intuition without having to read another technical paper? - Definition 3.1and its eq (3-5) are to my understanding a general version of eq (2). Is it really necessary to define these two models separately? I would recommend using the same notation if possible, or to the very least to connect the expressions across the two definitions. - Can you solve at least one practical SNN training problem? - How does the training effort (wall time, memory) compare to existing methods? Is this an impediment to test on a benchmark?

Rating

6

Confidence

2

Soundness

4

Presentation

2

Contribution

3

Limitations

Some limitations are discussed in the conclusions, but leaving a simple SNN benchmark to future work is not a limitation, but a shortcoming of the manuscript in my opinion.

Summary

The paper introduces a mathematical framework using rough path theory to model stochastic spiking neural networks as stochastic differential equations with event discontinuities, driven by càdlàg rough paths. This framework accommodates potential jumps in both solution trajectories and driving noise. Furthermore, the authors introduce Marcus signature kernels to extend continuous signature kernels to càdlàg rough paths and using them to define a general-purpose loss function for training stochastic SNNs as generative models.

Strengths

1. The paper is generally well-written. Despite its technical nature and numerous technical definitions, it is clear that the authors have made efforts to present the material in an accessible way. 2. The framework in the paper identifies sufficient conditions for the existence of pathwise gradients of solution trajectories and event times with respect to network parameters, and deriving a recursive relation for these gradients.

Weaknesses

I am not an expert in rough path theory and I don't see any major weaknesses.

Questions

1. How does the performance of the proposed gradient-based training compare to other established methods in the literature, such as those in [1,2,3]? More exhaustive empirical evaluations would be valuable and this would help to understand the benefits and limitations of this approach in practice. 2. Can the framework be extended to the case of single spike scenario (for instance in [1,2,3])? [1] Göltz, J., Kriener, L., Baumbach, A. et al. Fast and energy-efficient neuromorphic deep learning with first-spike times. Nat Mach Intell 3, 823–835 (2021). [2] I. M. Comsa, K. Potempa, L. Versari, T. Fischbacher, A. Gesmundo and J. Alakuijala, "Temporal Coding in Spiking Neural Networks with Alpha Synaptic Function," ICASSP 2020 - 2020 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), Barcelona, Spain, 2020, pp. 8529-8533, doi: 10.1109/ICASSP40776.2020.9053856. [3] Stanojević, Ana et al. “An Exact Mapping From ReLU Networks to Spiking Neural Networks.” Neural networks : the official journal of the International Neural Network Society 168 (2022): 74-88.

Rating

6

Confidence

2

Soundness

3

Presentation

3

Contribution

3

Limitations

There is no potential negative societal impact.

Summary

This paper introduces a general mathematical framework to model stochastic spiking neural networks (SSNN) as stochastic differential equations with event discontinuities, and identifies sufficient conditions ensuring the existence of gradients. With a newly defined loss function, SSNNs can be trained as generative models with an end-to-end autodifferentiable solver. Some empirical verifications demonstrate the effectiveness of the proposed method, and also, discussions on the connection with bio-plausible learning algorithms are provided.

Strengths

1. This paper proposes a mathematically rigorous framework to model SSNNs as SDEs with event discontinuities and driven by cadlag rough paths. The paper also identifies sufficient conditions for the existence of gradients, which strictly generalizes previous results only considering ODEs. 2. This paper provides the first gradient-based training of a large class of SSNNs with noise processes in both the spike timing and the network’s dynamics. 3. This paper discusses how the results lead to bio-plausible learning algorithms.

Weaknesses

1. There are several assumptions in the analysis. It would be better to have more discussions on whether these assumptions (e.g., Assumptions 3.1 and 3.2) can hold in common conditions. 2. This paper only conducts experiments on two toy settings. It would be better to discuss how to apply the method to more application conditions.

Questions

What is the computational complexity of SDESolveStep and RootFind?

Rating

6

Confidence

2

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors discussed limitations in Section 5.