Spikeometric: Linear Non-Linear Cascade Spiking Neural Networks with Pytorch Geometric

Stochastic models can give key insights into the behavior of a network on a systems level without explicitly modeling the subcellular mechanisms of each neuron. They lack some biological plausibility on the neuron level but have been shown to enjoy nice convexity properties, which can be fitted to observed spike data (Paninski, 2004). Traditionally, these are used as encoding models. For example, to study how multi-neuron systems process incoming stimuli (Pillow et al., 2008). Recently, these models have been used as generative models for inverse problems mapping activity to connectivity. As an example, (Das & Fiete, 2020) assessed bias and reconstruction errors in this setting.


Statement of need
Linear non-linear cascade models are much used in computational neuroscience and come in many flavors.Typical examples are SRM, LNP, GLM (Gerstner, 2008;Gerstner et al., 2014;Meyer et al., 2017).What unites these models is that they all model the spike response of a network through the same cascade-like sequence of steps.At each time step, a neuron receives input from its environment and converts it to a firing rate by a nonlinear function.This firing rate parametrizes a probability distribution from which spikes are drawn.This contrasts hard-threshold-based models, in which spikes are emitted whenever a variable representing the membrane potential exceeds a certain value.Although these models share many of the same underlying principles, no unifying framework currently permits easy implementation and direct comparison.
There exists a number of other simulation tools for simulating networks of neurons, notably NEST (Gewaltig & Diesmann, 2007), Neuron (Carnevale & Hines, 2006), Brian2 (Stimberg et al., 2019) and GENESIS (Bower et al., 2013).While all of these are primarily directed at solving systems where the dynamics are given by a system of differential equations, NEST includes a limited selection of pre-built stochastic point-process models similar to the ones found in spikeometric.However, this selection currently consists of only two models, which are not included in the GPU compatible version NEST-GPU (Golosio et al., 2021;Tiddia et al., 2022).There are also some torch-native frameworks like BindsNET (Hazan et al., 2018) and snnTorch (Eshraghian et al., 2023) but these focus on learning rather than simulation and are currently limited to LIF models.In general, spikeometric is more similar to the simulation tools in its focus, but has the automatic parameter tuning and torch-compatibility of the Spiking Neural Network libraries.

Implementation
The spikeometric package is a framework for simulating spiking neural networks using linear non-linear cascade models in Python.It is built on the PyTorch Geometric package and uses its powerful graph neural network modules and efficient graph representation.It is designed to be flexible and easy to use, while also being competitive on speed.Moreover, it's built in a way that accommodates usage or implementation of multiple different models sharing principles from the linear non-linear cascade family of models.
The torch backend makes simulating large networks on a GPU easy, with the extra benefit of having a familiar use pattern, reducing the friction of picking up a new tool.The package relies heavily on PyTorch Geometric (Fey & Lenssen, 2019), with the networks being represented as torch_geometric Data objects and the models extending the MessagePassing base class.The PyTorch Geometric framework is a popular deep learning framework originally designed for Graph Neural Networks (GNNs), a class of neural networks for learning graph-related data (Bronstein et al., 2021).It is the perfect setting for simulating neural networks with tunable parameters, allowing us to formulate the model's equations naturally in terms of vertices and edges, and giving us access to easy automatic tuning of parameters e.g. to match a certain firing rate, provided that the nonlinearity in the model is differentiable.The tuning functionality allows for fitting arbitrary parameters and can provide a starting point for implementing encoding models.
In addition to the models, the package includes dataset classes that can generate random connectivity matrices from a distribution or load pre-constructed connectivity matrices into torch_geometric's Data objects to be passed straight to the model.These objects hold a sparse representation of our connectivity matrices and can be batched together to form isolated subgraphs of a big graph, letting us simulate many networks simultaneously.
Finally, to facilitate the common use pattern of adding an external stimulus to the simulation and recording the resulting activity, we have included various stimulation classes that can be easily added to the model and even tuned to provoke a certain response.