Introduction

This is the documentation of sparks-ai, the Python library implementing SPARKS, a Sequential Predictive Autoencoder for the Representation of spiKing Signals. SPARKS learns low-dimensional latent representations of neural recordings, and decodes behaviour or stimuli from them.

Its encoder is built around Hebbian attention: the attention coefficient between two neurons follows spike-timing-dependent plasticity (STDP), so it tracks their relationship as it unfolds, spike by spike. The latents are stochastic, and the model is trained to predict the next time steps of a target from its latest latents, an objective inspired by predictive coding. The same model can be used in two ways:

• Unsupervised: with the recording itself as the target, SPARKS is a variational autoencoder, to explore the structure of the data without labels.
• Supervised: with a behavioural variable or a stimulus as the target, it is a variational information bottleneck, to decode that variable with high accuracy.

One model can embed several recording sessions, with different numbers of neurons, in a single latent space, and the Hebbian coefficients it computes can be analysed to infer directed functional interactions between neurons. SPARKS works with electrophysiology (spike trains) and calcium imaging data.

About this library

sparks-ai is written in PyTorch and supports Python 3.9 to 3.13. It provides a SPARKS model, high-level train and evaluate functions to get started quickly, and data loaders for the datasets of the paper. Every part of the model is a regular torch.nn.Module, so if you know PyTorch you can write your own training loop, decoder or front end. The library is released under the GPL-3.0-or-later licence.

What's new in 0.3

Version 0.3 is a major rewrite and a breaking release: code written for 0.2 has to be updated (see Migrating from 0.2), and models trained with 0.2 cannot be loaded. The main changes are:

• Whole sequences in one pass. The model takes [batch, time, neurons] spikes and returns latents and predictions at every time step. The Hebbian traces are computed with a parallel scan instead of a Python loop over time.
• Long recordings in chunks. chunk_length processes a recording in consecutive chunks, carrying the recurrent state across them (truncated backpropagation through time), so memory use stays flat.
• A Perceiver bottleneck. The neurons of any session are compressed into bottleneck_dim latent slots, so everything after the session's Hebbian layer is shared by all sessions. New sessions can be added to a trained model with add_session.
• A learned autoregressive prior over the latents (kl_type='sequential'), a new Hebbian rule for calcium imaging, optional shared attention or Mamba blocks, training utilities (parameter groups, learning-rate schedule, gradient clipping), co-smoothing, and an interface to replace Hebbian attention with your own front end.
• Data loaders that all return a DataSplit of train and test datasets and dataloaders.

The full list is in the changelog.

Getting started

Installation

We recommend installing SPARKS in a fresh environment, for example with conda:

conda create -n sparks python=3.12
conda activate sparks
pip install "sparks-ai>=0.3"

or with a virtual environment:

python -m venv .venv && source .venv/bin/activate
pip install "sparks-ai>=0.3"

This installs SPARKS with PyTorch. To use a GPU, first install the PyTorch build that matches your CUDA version, following the instructions on pytorch.org, then install sparks-ai. On Apple silicon, you can try the 'mps' device.

To work on the code itself, install it from source in editable mode:

git clone https://github.com/FrancisCrickInstitute/sparks.git
cd sparks
pip install -e ".[dev,all]"

See CONTRIBUTING.md to run the tests and contribute.

Optional extras

The base install only contains what the model needs. Optional extras add the dependencies of the data loaders and of the demo notebooks:

For example, pip install "sparks-ai[nlb,demos]" lets you run the monkey reaching demo. If you use a feature whose dependency is missing, the error message names the extra to install.

Mamba blocks (sparks.models.MambaBlock) need the mamba-ssm package, which is compiled against your PyTorch and CUDA versions. Install ninja, then run pip install causal-conv1d mamba-ssm --no-build-isolation.

Quickstart

In this example, SPARKS learns to decode a (random) two-dimensional behaviour from the spikes of 50 neurons:

import torch
import sparks

# 32 trials of 200 time bins (1 ms) from 50 neurons: [trials, time, neurons]
spikes = (torch.rand(32, 200, 50) < 0.05).float()
behaviour = torch.rand(32, 200, 2)  # e.g. a 2D position in [0, 1] at each time bin
loader = torch.utils.data.DataLoader(torch.utils.data.TensorDataset(spikes, behaviour), batch_size=8, shuffle=True)

model = sparks.SPARKS(n_neurons_per_session=50,
                      embed_dim=32,          # each neuron's embedding
                      latent_dim=2,          # the latent space
                      bottleneck_dim=4,      # latent slots the neurons are compressed into
                      hebbian_config=sparks.HebbianAttentionConfig(tau_s=0.5, dt=0.001),  # STDP timescale and bin size, in s
                      output_dim_per_session=2)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
loss_fn = torch.nn.BCEWithLogitsLoss()

for epoch in range(10):
    sparks.train(model, [loader], loss_fn, optimizer, beta=1e-3)  # beta weighs the KL term

_, latents, predictions = sparks.evaluate(model, [loader], act=torch.sigmoid)
print(latents[0].shape, predictions[0].shape)  # torch.Size([32, 200, 2]) twice: one per trial and time bin

A few things to note:

• train runs one epoch. It takes a list of dataloaders, one per recording session, even when there is a single session.
• evaluate returns the loss (here None, since no loss function was given), and lists of latents and predictions with one tensor per session, on the CPU.
• Since the targets lie in [0, 1], the decoder is trained with a binary cross-entropy on its logits, and act=torch.sigmoid maps its outputs back to [0, 1].
• To train on a GPU, pass device=torch.device('cuda') to SPARKS (or call model.to('cuda')). Batches are moved to the model's device for you.

The demos show complete workflows on real data.

Core concepts

Data format

• Spikes are float tensors of shape [batch, time, neurons], time first. For electrophysiology, these are binned spike trains, typically binary (0 or 1) at a fine resolution such as 1 ms. For calcium imaging, they can be deconvolved, binarised spikes (as in the MEC demo) or the fluorescence signal itself (see calcium data).
• Targets have time first too: [batch, time, features] for continuous variables, or [batch, time] integer labels for classification.
• Dataloaders must yield (spikes, targets) pairs. In unsupervised mode the targets are ignored, but a pair is still expected: TensorDataset(spikes, spikes) works.
• All trials in a batch must have the same number of time bins. Cut or pad trials of different lengths, or use batch_size=1.
• The bin size dt and the STDP time constant tau_s are given to HebbianAttentionConfig, both in seconds.

Architecture

1. Hebbian attention layer. At every time step, each pair of neurons (i, j) has a coefficient that is potentiated when neuron i spikes shortly before neuron j, and depressed when it spikes shortly after, through pre- and post-synaptic eligibility traces that decay with time constant tau_s. Each neuron's row of coefficients is projected by a learned value matrix into an embed_dim-dimensional embedding. Several heads with different time constants can run in parallel. Each session has its own Hebbian layer, sized to its neurons.
2. Perceiver bottleneck. After a feed-forward block, the embeddings of the N neurons, plus a learned identity embedding per neuron, are read out by bottleneck_dim learned queries, giving a bottleneck_dim × embed_dim representation whose size does not depend on N. Everything from here on is shared by all sessions.
3. Shared blocks (optional, none by default). Causal attention blocks with rotary embeddings, or Mamba blocks, operating over time.
4. Latents. Two linear heads give the mean and log-variance of a Gaussian posterior over the latent_dim-dimensional latents. During training, latents are sampled from it; in evaluation mode, the model returns its mean.
5. Decoder. At each time step, the decoder reads a window of the last tau_p latents and predicts the next tau_f targets. This is the predictive causally conditioned objective of the paper: longer windows encourage temporally coherent latents.

Supervised and unsupervised learning

The same model is trained differently depending on the target:

• Supervised: the dataloaders' targets are the variable to decode, and output_dim_per_session is its number of features (or classes).
• Unsupervised: pass unsupervised=True to train and evaluate, and set output_dim_per_session to the number of neurons. The model then predicts its own inputs.

n_neurons = 50
model = sparks.SPARKS(n_neurons_per_session=n_neurons, embed_dim=32, latent_dim=2, bottleneck_dim=4,
                      hebbian_config=sparks.HebbianAttentionConfig(tau_s=0.5, dt=0.001),
                      output_dim_per_session=n_neurons)  # reconstruct every neuron
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
loss_fn = torch.nn.BCEWithLogitsLoss()  # binary spikes: Bernoulli likelihood on the logits

sparks.train(model, [loader], loss_fn, optimizer, beta=1e-3, unsupervised=True)
_, latents, spike_probs = sparks.evaluate(model, [loader], unsupervised=True, act=torch.sigmoid)

The two are complementary: unsupervised learning explores the structure of the data without imposing any task variable, while supervised learning reveals how the variables recorded during the experiment are encoded.

Key hyperparameters

Configuration

Hebbian attention: HebbianAttentionConfig

The front end of the encoder is configured with a HebbianAttentionConfig, passed to SPARKS(hebbian_config=...):

from sparks import HebbianAttentionConfig

fixed = HebbianAttentionConfig(tau_s=0.5, dt=0.001)                   # SPARKS-fixed (the default 'light' config)
learned = HebbianAttentionConfig(tau_s=0.5, dt=0.001, config='full')  # SPARKS-learned
multi = HebbianAttentionConfig(tau_s=[0.01, 0.1, 1.0], dt=0.001)      # three heads, one per timescale
calcium = HebbianAttentionConfig(tau_s=1.0, dt=1 / 30, data_type='calcium')

Fixed or learned? The light configuration only learns the value projection, keeps one trace per neuron rather than per pair, and is the faster and lighter of the two. In the paper, it performed best on small datasets, while the full configuration, which learns 4N² extra parameters, was better on large populations and long recordings, when there is enough data to fit them.

Calcium data. With data_type='calcium', the calcium traces are used directly, without deconvolving them into spikes, since a calcium trace is already a decaying memory of past spikes. The coefficient of each pair of neurons, i → j, integrates the difference of their traces x_i and x_j, with a leak of time constant tau_s, and is clipped to [min_attn_value, max_attn_value]:

a_ij,t = clip(e^−dt/τ_s a_ij,t−1 + x_i,t − x_j,t, min_attn_value, max_attn_value)

The calcium layer only learns its value projection, and config has no effect on it. If you do have deconvolved spikes, the default electrophysiology layer also works well, as in the MEC demo.

Shared blocks: SharedConfig

Between the Hebbian attention block and the latents, you can stack sequence blocks shared by every session, operating over time. There are none by default.

from sparks import SharedConfig

model = sparks.SPARKS(n_neurons_per_session=50, embed_dim=32, latent_dim=2, bottleneck_dim=4,
                      shared_config=SharedConfig(n_layers=2, n_heads=4, dropout=0.1),
                      output_dim_per_session=2)

The blocks operate on vectors of size embed_dim × bottleneck_dim, which must be divisible by n_heads, with an even width per head.

Decoders

By default, SPARKS builds an MLP decoder with one hidden layer of width 4 × tau_p × latent_dim, and one output head per session of size output_dim_per_session (times tau_f). You can instead pass any decoder with the decoder argument. sparks.models provides three:

from sparks.models import LinearDecoder, MLPDecoder, TransformerDecoder

latent_dim, tau_p = 4, 5
in_dim = latent_dim * tau_p  # the decoder reads the last tau_p latents

linear = LinearDecoder(in_dim=in_dim, output_dim_per_session=2)
mlp = MLPDecoder(in_dim=in_dim, hidden_dims=[64, 64], output_dim_per_session=2)
transformer = TransformerDecoder(embed_dim=in_dim, output_dim_per_session=2, n_layers=1, n_heads=1)

model = sparks.SPARKS(n_neurons_per_session=50, embed_dim=32, latent_dim=latent_dim, bottleneck_dim=4,
                      tau_p=tau_p, decoder=linear)

A custom decoder is any module whose forward(x, session_id) maps the windows of latents, of shape [batch, time, tau_p × latent_dim], to predictions of shape [batch, time, output_dim × tau_f]:

import torch.nn as nn


class GRUDecoder(nn.Module):
    """Reads the windows of latents with a GRU."""

    def __init__(self, in_dim, hidden_dim, output_dim):
        super().__init__()
        self.gru = nn.GRU(in_dim, hidden_dim, batch_first=True)
        self.readout = nn.Linear(hidden_dim, output_dim)

    def forward(self, x, session_id):  # x: [batch, time, tau_p * latent_dim]
        return self.readout(self.gru(x)[0])  # [batch, time, output_dim * tau_f]


model = sparks.SPARKS(n_neurons_per_session=50, embed_dim=32, latent_dim=latent_dim, bottleneck_dim=4,
                      tau_p=tau_p, decoder=GRUDecoder(in_dim=latent_dim * tau_p, hidden_dim=64, output_dim=2))

Data

Your own data

Any dataloader that yields (spikes, targets) batches works, such as the TensorDataset of the quickstart. For a recording cut into trials, you can also subclass sparks.data.BaseDataset, which documents the expected shapes and records the bin size, and wrap it in dataloaders with make_split:

import numpy as np
from sparks.data import BaseDataset, make_split


class MyRecording(BaseDataset):
    """A recording cut into trials, with a behavioural variable to decode."""

    def __init__(self, spikes, behaviour, dt):
        self.spikes = torch.as_tensor(spikes, dtype=torch.float32)        # [trials, time, neurons]
        self.behaviour = torch.as_tensor(behaviour, dtype=torch.float32)  # [trials, time, features]
        self.dt = dt                                                      # bin size, in seconds

    def __len__(self):
        return len(self.spikes)

    @property
    def n_neurons(self):
        return self.spikes.shape[-1]

    def get_spikes(self, index):
        return self.spikes[index]     # [time, neurons]

    def get_target(self, index):
        return self.behaviour[index]  # [time, features]


my_spikes = np.random.rand(40, 300, 30) < 0.05  # replace with your recording
my_behaviour = np.random.rand(40, 300, 2)
data = make_split(MyRecording(my_spikes[:32], my_behaviour[:32], dt=0.001),
                  MyRecording(my_spikes[32:], my_behaviour[32:], dt=0.001), batch_size=8)
train_dataset, test_dataset, train_loader, test_loader = data

make_split returns a DataSplit, a named tuple of the train and test datasets and their dataloaders (the training one shuffled).

Built-in loaders

sparks.data has loaders for the datasets of the paper. Each returns a DataSplit, or a list of them with one per session.

from sparks.data.nlb.area2_bump import make_area2_bump_data

data = make_area2_bump_data('/path/to/000127/sub-Han/sub-Han_desc-train_behavior+ecephys.nwb',
                            y_keys='hand_pos', mode='prediction', batch_size=32)
print(data.train_dataset.n_neurons, data.train_dataset.target_dim, data.train_dataset.dt)

The Allen loaders read files prepared from an AllenSDK cache by a script in the repository. On a first run, it downloads the sessions' data, which takes hundreds of GB and several hours:

python scripts/preprocess_allen.py /path/to/allen_cache natural_movie_one --areas VISp

The stimulus can be gratings_flashes, spontaneous or natural_movie_one. See the documentation of the sparks.data.allen module for the layout of the prepared files.

Training

Train and evaluate

sparks.train(model, train_dls, loss_fn, optimizer, ...) trains the model for one epoch over a list of dataloaders, one per session. Its options are keyword-only:

Batches alternate between sessions, in a random order; shorter dataloaders are cycled until the longest one is exhausted. To control every step yourself, sparks.training.train_on_batch trains on a single batch.

sparks.evaluate(model, test_dls, loss_fn=None, ...) runs the model in evaluation mode and returns (loss, latents, predictions). The loss is summed over batches (None without loss_fn); latents and predictions are lists with one [trials, time, ...] tensor per session, on the CPU. It takes session_ids, unsupervised, burnin and chunk_length like train, and act, a function applied to the decoder's outputs before the loss and before they are returned (e.g. torch.sigmoid).

Choosing a loss

The loss is called as loss_fn(predictions, targets), so any PyTorch loss works. Common choices are:

from sparks.training import TemporalCrossEntropyLoss

n_classes = 10
labels = torch.randint(n_classes, (32, 200))  # [trials, time]: one class per time bin
class_loader = torch.utils.data.DataLoader(torch.utils.data.TensorDataset(spikes, labels), batch_size=8)

model = sparks.SPARKS(n_neurons_per_session=50, embed_dim=32, latent_dim=4, bottleneck_dim=4,
                      output_dim_per_session=n_classes)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
loss_fn = TemporalCrossEntropyLoss(label_smoothing=0.1)  # smoothing keeps the latents from growing without bound

sparks.train(model, [class_loader], loss_fn, optimizer, beta=1e-3)
_, _, logits = sparks.evaluate(model, [class_loader])
accuracy = (logits[0].argmax(-1) == labels).float().mean()

Recommended recipe

The demos train SPARKS with AdamW, a learning rate that warms up then follows a cosine decay, and the sequential KL. build_param_groups sorts the parameters into weight-decay groups (biases, norms and embeddings are not decayed) and, for the sequential KL, creates the learned prior, which is trained by the same optimizer:

from sparks.training import build_param_groups, get_lr_multiplier

model = sparks.SPARKS(n_neurons_per_session=50, embed_dim=32, latent_dim=2, bottleneck_dim=4,
                      hebbian_config=sparks.HebbianAttentionConfig(tau_s=0.5, dt=0.001),
                      output_dim_per_session=2)
n_epochs = 10
loss_fn = torch.nn.BCEWithLogitsLoss()

# Weight-decay groups, and the learned prior p(z_t | z_{t-1})
param_groups, transition = build_param_groups(model, weight_decay=0.01, kl_type='sequential')
optimizer = torch.optim.AdamW(param_groups, lr=1e-3)

# Linear warmup over the first 5% of steps, then cosine decay
total_steps = n_epochs * len(loader)
scheduler = torch.optim.lr_scheduler.LambdaLR(
    optimizer, lambda step: get_lr_multiplier(step, warmup_steps=int(0.05 * total_steps),
                                              total_steps=total_steps, min_lr=1e-4))

for epoch in range(n_epochs):
    sparks.train(model, [loader], loss_fn, optimizer, scheduler=scheduler,
                 beta=1e-3, kl_type='sequential', transition=transition)

Long recordings

The Hebbian layer keeps an N × N matrix of coefficients for every time step, and backpropagation goes through the whole sequence, so memory grows with the length of the recording. With chunk_length, sequences are processed in consecutive chunks: the recurrent state (Hebbian traces, Perceiver memory, and the decoder's window of past latents) is carried from one chunk to the next, and each chunk gets its own backward pass and optimizer step (truncated backpropagation through time). Memory then depends on the chunk length only, and the forward dynamics are unchanged.

sparks.train(model, [loader], loss_fn, optimizer, beta=1e-3, chunk_length=50)
_, latents, predictions = sparks.evaluate(model, [loader], chunk_length=50)  # same outputs as without chunks

When you call the model yourself, the same mechanism is available with carry_state=True. Call reset_state() before the first chunk of a new sequence:

long_recording = (torch.rand(1, 1000, 50) < 0.05).float()  # [1, time, neurons]

model.eval()
model.reset_state()
with torch.no_grad():
    chunks = [model(chunk, session_id=0, carry_state=True)[0] for chunk in long_recording.split(250, dim=1)]
latents = torch.cat(chunks, dim=1)  # [1, 1000, latent_dim]

Calling model(spikes, session_id) returns (z, predictions, mu, logvar), each with shape [batch, time, ...].

Multiple sessions

Training on several sessions

One model can embed recordings from several sessions or animals in a single latent space, even when they have different neurons and different numbers of them. Each session has its own Hebbian attention layer and neuron identity embeddings; everything else is shared. Give the number of neurons and an id for each session, and one dataloader per session:

spikes_a = (torch.rand(32, 200, 50) < 0.05).float()  # session 'mouse_a': 50 neurons
spikes_b = (torch.rand(24, 200, 80) < 0.05).float()  # session 'mouse_b': 80 neurons
loader_a = torch.utils.data.DataLoader(torch.utils.data.TensorDataset(spikes_a, torch.rand(32, 200, 2)), batch_size=8)
loader_b = torch.utils.data.DataLoader(torch.utils.data.TensorDataset(spikes_b, torch.rand(24, 200, 2)), batch_size=8)

model = sparks.SPARKS(n_neurons_per_session=[50, 80],
                      id_per_session=['mouse_a', 'mouse_b'],
                      embed_dim=32, latent_dim=2, bottleneck_dim=4,
                      hebbian_config=sparks.HebbianAttentionConfig(tau_s=0.5, dt=0.001),
                      output_dim_per_session=[2, 2])  # one decoder output head per session
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
loss_fn = torch.nn.BCEWithLogitsLoss()

sessions = ['mouse_a', 'mouse_b']
sparks.train(model, [loader_a, loader_b], loss_fn, optimizer, session_ids=sessions, beta=1e-3)
_, latents, predictions = sparks.evaluate(model, [loader_a, loader_b], session_ids=sessions, act=torch.sigmoid)

When every session predicts the same variable, a single output head can be shared with output_dim_per_session=2, joint_decoder=True, as in the multi-session demo.

Adding a session to a trained model

add_session gives a new session its own Hebbian attention layer and, unless the decoder is joint, its own output head; it reuses everything else. Rebuild the optimizer so that it includes the new parameters, then fine-tune on the new session. In the paper, a new session reached its final accuracy within about 50 fine-tuning epochs, without loss of accuracy on the sessions the model was trained on.

spikes_c = (torch.rand(16, 200, 65) < 0.05).float()  # a new session with 65 neurons
loader_c = torch.utils.data.DataLoader(torch.utils.data.TensorDataset(spikes_c, torch.rand(16, 200, 2)), batch_size=8)

model.add_session(n_neurons=65, session_id='mouse_c', new_output_dim=2)  # leave out new_output_dim with a joint decoder
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)

for epoch in range(5):
    sparks.train(model, [loader_c], loss_fn, optimizer, session_ids=['mouse_c'], beta=1e-3)

The new session can have its own Hebbian settings, e.g. model.add_session(..., tau_s=0.1, dt=0.01); any setting you leave out is inherited from the model.

Analysis

Latents and predictions

evaluate returns the latents and predictions at every time step of every trial. In evaluation mode, the latents are the mean of the posterior, so they are deterministic. With tau_f > 1, the predictions at each time step hold the next tau_f targets one after the other, so the prediction for the current time step comes first:

model = sparks.SPARKS(n_neurons_per_session=50, embed_dim=32, latent_dim=2, bottleneck_dim=4,
                      tau_p=10, tau_f=5, output_dim_per_session=2)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
sparks.train(model, [loader], torch.nn.BCEWithLogitsLoss(), optimizer, beta=1e-3)

_, latents, predictions = sparks.evaluate(model, [loader], act=torch.sigmoid)
print(predictions[0].shape)              # [32, 200, 10]: 5 future steps of 2 features
current_step = predictions[0][..., :2]   # the prediction for the current time step

You can also decode arbitrary latents, bypassing the encoder, to explore the latent space: model.generate(z, session_id) with z of shape [batch, time, latent_dim].

Hebbian coefficients

The coefficients computed by the Hebbian attention layer capture directed, spike-timing relationships between neurons. In the paper, averaging them over pairs of brain areas recovered the known functional hierarchy of the mouse visual cortex. get_hebbian_coeffs_on_batch runs a batch through the Hebbian layer only and returns the coefficients:

from sparks.training import get_hebbian_coeffs_on_batch

model = sparks.SPARKS(n_neurons_per_session=50, embed_dim=32, latent_dim=2, bottleneck_dim=4,
                      hebbian_config=sparks.HebbianAttentionConfig(tau_s=0.01, dt=0.001),
                      output_dim_per_session=2)
trials = spikes[:4]

coeffs = get_hebbian_coeffs_on_batch(model, trials, session_id=0)  # [trials, time, N, N]

# For long recordings: walk through the trials in chunks and keep only the last time step
final = get_hebbian_coeffs_on_batch(model, trials, session_id=0, chunk_length=50, last_only=True)  # [trials, 1, N, N]
coupling = final[:, 0].mean(0)  # [N, N]: entry (i, j) grows when neuron i fires before neuron j

With several Hebbian heads (a list of tau_s), a leading head dimension is added. Coefficients are non-zero for any pair of active neurons, partly because of chance co-activations. As in the paper, you can remove this rate-driven part by subtracting the coefficients obtained after shuffling each neuron's spikes in time:

def shuffle_time(spikes):
    """Permutes each neuron's spikes in time, independently: keeps its rate, removes timing relations."""
    order = torch.rand(spikes.shape).argsort(dim=1)  # [trials, time, neurons]
    return spikes.gather(1, order)


baseline = torch.stack([get_hebbian_coeffs_on_batch(model, shuffle_time(trials), session_id=0, last_only=True)
                        for _ in range(10)]).mean(0)[:, 0].mean(0)
corrected = coupling - baseline  # the part of the coefficients due to spike timing

The coefficients depend on the Hebbian parameters, not on the rest of the model: with the light configuration, they only depend on tau_s, w_plus and alpha, while with the full configuration, they also reflect the decays and weights learned for each pair during training. With calcium data, they only depend on tau_s and the bounds, and since they leak with time constant tau_s, the last time step reflects the end of a recording rather than all of it.

Co-smoothing

Co-smoothing measures how well the activity of held-out neurons, which the encoder never sees, can be predicted from the latents. evaluate_cosmoothing fits a readout per held-out neuron on training latents and scores it on test latents, in bits per spike (0 for a model that only captures the neurons' mean rates):

from sparks.training import bits_per_spike, evaluate_cosmoothing

recording = (torch.rand(32, 200, 60) < 0.05).float()
held_out = torch.arange(50, 60)  # neurons the encoder never sees
inputs, held_out_spikes = recording[..., :50], recording[..., held_out]
input_loader = torch.utils.data.DataLoader(torch.utils.data.TensorDataset(inputs, inputs), batch_size=8)

model = sparks.SPARKS(n_neurons_per_session=50, embed_dim=32, latent_dim=4, bottleneck_dim=4,
                      output_dim_per_session=50)
# ... train the model with unsupervised=True ...
_, latents, spike_probs = sparks.evaluate(model, [input_loader], unsupervised=True, act=torch.sigmoid)

X = latents[0].reshape(-1, latents[0].shape[-1]).numpy()  # [trials * time, latent_dim]
Y = held_out_spikes.reshape(-1, len(held_out)).numpy()     # [trials * time, held-out neurons]
n_train = 4000                                             # first time bins to fit, the rest to test
score, _ = evaluate_cosmoothing(X[:n_train], Y[:n_train], X[n_train:], Y[n_train:], metric='bernoulli')

# Bits per spike of the model's own reconstruction of its inputs
reconstruction = bits_per_spike(spike_probs[0], inputs, distribution='bernoulli')

metric can be 'poisson', 'bernoulli', 'gaussian' or 'r2'. The MEC loader can hold out a random fraction of neurons for you, with make_mec_data(..., p_left_out=0.1) and dataset.left_out_neurons.

Saving and loading

Save the model's state_dict. To load it, build a model with the same arguments, add any session that was added after training, then load the weights:

torch.save(model.state_dict(), 'sparks.pt')

model = sparks.SPARKS(...)  # the same arguments as the saved model
model.add_session(...)      # the same sessions, if any were added after building it
model.load_state_dict(torch.load('sparks.pt'))

Advanced

Custom front ends

The Hebbian attention block can be replaced by any other front end, for instance to test what Hebbian attention contributes on your data, while keeping the rest of the model unchanged. Any module implementing the sparks.models.FrontEndBlock interface can be used, selected with HebbianAttentionConfig(block_class=...). Its constructor receives n_neurons_per_session (a dict from session ids to numbers of neurons), embed_dim, bottleneck_dim and the HebbianAttentionConfig fields, and its forward maps a session's spikes to vectors of size bottleneck_dim × embed_dim:

from sparks.models import FrontEndBlock


class LinearFrontEnd(FrontEndBlock):
    """Projects each time bin's spikes straight to the shared stream: no pairwise structure."""

    def __init__(self, n_neurons_per_session, embed_dim, bottleneck_dim, **config):  # config: unused fields
        super().__init__()
        self.out_dim = embed_dim * bottleneck_dim
        self.proj = nn.ModuleDict({sid: nn.Linear(n, self.out_dim) for sid, n in n_neurons_per_session.items()})

    def forward(self, x, session_id, carry_state=False):  # x: [batch, time, neurons]
        return self.proj[str(session_id)](x)                # [batch, time, bottleneck_dim * embed_dim]

    def add_session(self, n_neurons, session_id, **overrides):
        self.proj[str(session_id)] = nn.Linear(n_neurons, self.out_dim).to(next(self.parameters()))


model = sparks.SPARKS(n_neurons_per_session=50, embed_dim=32, latent_dim=2, bottleneck_dim=4,
                      hebbian_config=sparks.HebbianAttentionConfig(block_class=LinearFrontEnd),
                      output_dim_per_session=2)

A front end with recurrent state should store it with the helpers of CarriedState, its base class, so that carry_state and chunk_length work. The test suite of the repository has a minimal example that does (tests/front_ends.py), and the ablation controls of the paper are complete ones.

Ablation controls

The control models of the paper, which replace Hebbian attention with other modules to test what it contributes, are not part of sparks-ai. They are front ends defined in the controls package of the sparks-controls repository, along with the scripts of the paper's control runs:

The first two keep the rest of the Hebbian attention block (one embedding of each neuron per tau_s, the feed-forward network and the Perceiver readout), and only replace the spike-timing coefficients of each pair of neurons with other ones. The last three replace both the Hebbian attention block and the Perceiver: the smoothed firing rates of every neuron, one per tau_s, are read out by a linear layer specific to each session, then the sequence model, shared by every session, writes directly to the latent stream. From the root of the sparks-controls repository, a control plugs in like any other front end:

from controls import GRUControlBlock  # run from the root of the sparks-controls repository

control = sparks.SPARKS(n_neurons_per_session=50, embed_dim=32, latent_dim=2, bottleneck_dim=4,
                        hebbian_config=sparks.HebbianAttentionConfig(block_class=GRUControlBlock, tau_s=0.5, dt=0.001),
                        output_dim_per_session=2)

Tips

• Memory. The Hebbian layer stores an N × N matrix per time step and trial, so memory scales with batch size × time × N². Reduce the batch size, or use chunk_length, for large populations or long trials. The full configuration also learns 4N² parameters per session.
• Reconstructing sparse spikes. The built-in decoders' output heads start at a probability of 0.5. For unsupervised learning on sparse spike trains, start them at the mean firing probability instead, with e.g. LinearDecoder(..., base_rate=train_spikes.mean().item()).
• Warm-up. The Hebbian traces start at zero at the beginning of each sequence. If the first time steps are hard to predict, leave them out of the loss with burnin.
• Gradient clipping. Gradients are clipped to max_grad_norm=5 by default. A RuntimeWarning tells you if clipping rescales every step; use sparks.training.measure_grad_norm, or train_on_batch(..., return_grad_norms=True), on a few batches to choose a better value.
• Sampling latents at evaluation. In evaluation mode the latents are the posterior mean. To draw samples instead, use z, _, mu, logvar = model(x, session_id) then model.encoder.reparametrize(mu, logvar, sample=True).
• Multi-GPU. carry_state and chunk_length are not supported with nn.DataParallel; use DistributedDataParallel or a single GPU.

Migrating from 0.2

The main changes to update 0.2 code are below. The changelog has the complete tables, including the data loaders and modules that moved.

Reproducing the paper

The scripts that reproduce the results of the paper are in the sparks-examples repository, and those of the paper's ablation controls in the sparks-controls repository. In the paper, SPARKS-fixed and SPARKS-learned correspond to config='light' and config='full'.

Citation and help

If you use SPARKS, please cite:

@article{skatchkovsky2024sparks,
  title   = {A Biologically Inspired Attention Model for Neural Signal Analysis},
  author  = {Skatchkovsky, Nicolas and Glazman, Natalia and Egea-Weiss, Alexander and Sadeh, Sadra and Iacaruso, Florencia},
  journal = {bioRxiv},
  year    = {2024},
  doi     = {10.1101/2024.08.13.607787}
}

For bugs and questions, please open an issue on GitHub, or contact us by email.