SPARKS ← All demos Documentation Download this notebook (GitHub)

Decoding hand position from somatosensory cortex

This notebook trains SPARKS to predict a monkey's hand position from the spikes of neurons in its somatosensory cortex (area 2), using the Area2_Bump dataset of the Neural Latents Benchmark (NLB). On each trial, the monkey reaches toward one of eight targets. For an overview of the data, see the NLB tutorial.

You will need:

• the nlb and demos extras: pip install "sparks-ai[nlb,demos]";
• the dataset, DANDI 000127, which you can download with the DANDI client: dandi download DANDI:000127.

Training takes about 4 minutes on a recent GPU.

import matplotlib.pyplot as plt
import numpy as np
import torch
from sklearn.metrics import r2_score
from tqdm import tqdm

import sparks
from sparks.data.nlb.area2_bump import UNIQUE_ANGLES, make_area2_bump_data
from sparks.models import LinearDecoder
from sparks.training import build_param_groups, get_lr_multiplier

Data

make_area2_bump_data cuts each active reach into a 600 ms trial, from 100 ms before movement onset to 500 ms after, in 1 ms bins. With mode='prediction', the target at each time bin is the hand position 40 ms later, scaled to [0, 1]. The NLB 'train' trials form the train set, and its 'val' trials the test set.

data_path = '/path/to/000127/sub-Han/sub-Han_desc-train_behavior+ecephys.nwb'  # change to your local path

data = make_area2_bump_data(data_path, y_keys='hand_pos', mode='prediction', batch_size=141)
train_dataset, test_dataset = data.train_dataset, data.test_dataset

print(f'{len(train_dataset)} train and {len(test_dataset)} test trials, {train_dataset.n_neurons} neurons')
spikes, hand_position = train_dataset[0]
print('spikes:', tuple(spikes.shape), '(time, neurons); targets:', tuple(hand_position.shape), '(time, x/y)')

141 train and 52 test trials, 65 neurons
spikes: (600, 65) (time, neurons); targets: (600, 2) (time, x/y)

Model

• The Hebbian attention layer tracks the correlations between neurons through STDP with a 0.5 s time constant (tau_s), on the data's 1 ms time bins (dt).
• Six causal attention blocks shared over time (SharedConfig) refine its output, which is compressed into a 2-dimensional latent at each time bin.
• A linear decoder reads the hand position from the last tau_p = 10 latents.

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')  # on Apple silicon, you can try 'mps'
torch.manual_seed(0)

latent_dim = 2
tau_p = 10  # number of past latents the decoder reads

model = sparks.SPARKS(n_neurons_per_session=train_dataset.n_neurons,
                      embed_dim=32,
                      latent_dim=latent_dim,
                      bottleneck_dim=2,
                      tau_p=tau_p,
                      hebbian_config=sparks.HebbianAttentionConfig(tau_s=0.5, dt=train_dataset.dt),
                      shared_config=sparks.SharedConfig(n_layers=6, n_heads=4),
                      decoder=LinearDecoder(in_dim=latent_dim * tau_p, output_dim_per_session=train_dataset.target_dim),
                      device=device)

Training

• The latents are regularised toward a learned autoregressive prior (kl_type='sequential'), with a weight beta. build_param_groups creates this prior and sorts the parameters into AdamW's weight-decay groups.
• The learning rate warms up over the first 5% of steps, then follows a cosine decay.
• Since the hand positions lie in [0, 1], the decoder is trained with a binary cross-entropy on its logits; act=torch.sigmoid maps its outputs back to positions at evaluation.

We score the predictions on the test trials with the $R^2$, from movement onset onward.

n_epochs = 200
beta = 1e-4

param_groups, transition = build_param_groups(model, weight_decay=0.01, kl_type='sequential')
optimizer = torch.optim.AdamW(param_groups, lr=0.01)
total_steps = n_epochs * len(data.train_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))
loss_fn = torch.nn.BCEWithLogitsLoss()

movement_onset = 100  # the time bin of movement onset
true_positions = test_dataset.target_data.numpy()  # [trials, time, 2]


def r2_after_onset(predictions):
    return r2_score(true_positions[:, movement_onset:].reshape(-1, 2),
                    predictions[:, movement_onset:].reshape(-1, 2), multioutput='variance_weighted')


pbar = tqdm(range(n_epochs))
for epoch in pbar:
    sparks.train(model, [data.train_loader], loss_fn, optimizer, scheduler=scheduler,
                 beta=beta, kl_type='sequential', transition=transition)

    if (epoch + 1) % 10 == 0:
        _, _, predictions = sparks.evaluate(model, [data.test_loader], act=torch.sigmoid)
        pbar.set_description(f'test R2: {r2_after_onset(predictions[0].numpy()):.3f}')

test R2: 0.947: 100%|██████████| 200/200 [01:45<00:00,  1.89it/s]

Results

evaluate returns, for each session, the latents and the predictions on every test trial, at every time bin. Averaged over the trials of each reach direction, the 2-dimensional latents trace out the hand's trajectories.

_, latents, predictions = sparks.evaluate(model, [data.test_loader], act=torch.sigmoid)
latents, predictions = latents[0].numpy(), predictions[0].numpy()  # the only session
print(f'Test R2: {r2_after_onset(predictions):.3f}')

colors = plt.cm.hsv(np.linspace(0, 1, len(UNIQUE_ANGLES), endpoint=False))
fig, axs = plt.subplots(1, 3, figsize=(18, 5.5))
for direction, color in enumerate(colors):
    trials = test_dataset.directions == direction
    for ax, trajectories in zip(axs, [true_positions, predictions, latents]):
        mean_trajectory = trajectories[trials, movement_onset:].mean(axis=0)
        ax.plot(mean_trajectory[:, 0], mean_trajectory[:, 1], color=color, linewidth=3,
                label=f'{UNIQUE_ANGLES[direction]:.0f}°')

for ax, title in zip(axs, ['Hand position', 'Predicted hand position', 'Latents']):
    ax.set_title(title)
    ax.set_xticks([])
    ax.set_yticks([])
axs[2].set_xlabel('Neural axis 1')
axs[2].set_ylabel('Neural axis 2')
axs[0].legend(title='Reach direction', fontsize=8)
plt.show()

Test R2: 0.947