SPARKS ← All demos Documentation Download this notebook (GitHub)

One model for several mice: decoding a natural movie

SPARKS can embed recordings from several sessions, with different neurons and different numbers of them, in a single latent space. This notebook trains one model on the responses of primary visual cortex (VISp) of three mice to a 30 s natural movie, from the Allen Brain Observatory Visual Coding Neuropixels dataset. The model decodes, at each time bin, which of the movie's 900 frames is showing. A fourth mouse is then added to the trained model.

You will need:

• the allen and demos extras: pip install "sparks-ai[allen,demos]";
• the preprocessed natural-movie responses. To make them, clone the sparks repository and run python scripts/preprocess_allen.py /path/to/allen_visual natural_movie_one --areas VISp, where /path/to/allen_visual is the directory of your AllenSDK cache. On a first run, this downloads the sessions' data, which takes hundreds of GB and several hours.

Training takes about 8 minutes on a recent GPU.

import matplotlib.pyplot as plt
import numpy as np
import torch
from tqdm import tqdm

import sparks
from sparks.data.allen import get_session_ids_with_data, make_session_movies_data
from sparks.models import LinearDecoder
from sparks.training import TemporalCrossEntropyLoss, build_param_groups, get_lr_multiplier

Data

The movie was shown 20 times in each session. make_session_movies_data returns one DataSplit per session. With the default block='first', each split trains on 9 of the first 10 presentations and tests on the 10th. With mode='prediction', the target at each time bin is the index of the frame showing. Sessions are identified by their Allen session ids.

data_path = '/path/to/allen_visual'  # change to your AllenSDK cache directory
dt = 0.01  # time bins of 10 ms

session_ids = get_session_ids_with_data(data_path, 'natural_movie_one', areas=('VISp',))
train_session_ids, new_session_id = list(session_ids[:3]), session_ids[3]

splits = make_session_movies_data(data_path, train_session_ids, areas=('VISp',), mode='prediction', dt=dt, batch_size=9)
for session_id, split in zip(train_session_ids, splits):
    print(f'Session {session_id}: {split.train_dataset.n_neurons} neurons, '
          f'{len(split.train_dataset)} train and {len(split.test_dataset)} test presentations')

Session 715093703: 58 neurons, 9 train and 1 test presentations
Session 719161530: 52 neurons, 9 train and 1 test presentations
Session 721123822: 31 neurons, 9 train and 1 test presentations

Model

Each session gets its own Hebbian attention layer, sized to its neurons. Everything after it is shared: the model compresses any number of neurons into bottleneck_dim slots, then into 16-dimensional latents. A single linear decoder (joint_decoder=True), also shared by every session, classifies the frame from the last tau_p = 6 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 = 16
tau_p = 6
n_frames = 900

model = sparks.SPARKS(n_neurons_per_session=[split.train_dataset.n_neurons for split in splits],
                      id_per_session=train_session_ids,
                      embed_dim=32,
                      latent_dim=latent_dim,
                      bottleneck_dim=8,
                      tau_p=tau_p,
                      hebbian_config=sparks.HebbianAttentionConfig(tau_s=5., dt=dt),
                      decoder=LinearDecoder(in_dim=latent_dim * tau_p, output_dim_per_session=n_frames, joint_decoder=True),
                      device=device)

Training

sparks.train runs one epoch over the sessions' dataloaders, alternating between sessions. The frame is classified with a cross-entropy at every time bin (TemporalCrossEntropyLoss); label smoothing keeps the latents from growing without bound. The latents are regularised toward a learned autoregressive prior (kl_type='sequential').

A prediction counts as correct when it is within 1 s (30 frames) of the frame showing.

n_epochs = 200
beta = 0.1

param_groups, transition = build_param_groups(model, weight_decay=0.01, kl_type='sequential')
optimizer = torch.optim.AdamW(param_groups, lr=0.002)
total_steps = n_epochs * len(splits)  # one batch per session and epoch
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 = TemporalCrossEntropyLoss(label_smoothing=0.1)

frame_shown = splits[0].test_dataset.frame_indices  # [T]: the frame showing at each time bin


def accuracy(predictions, tolerance=30):
    """The fraction of time bins whose predicted frame is within `tolerance` frames of the right one."""
    return ((predictions.argmax(dim=-1) - frame_shown).abs() <= tolerance).float().mean().item()


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

    if (epoch + 1) % 10 == 0:
        _, _, predictions = sparks.evaluate(model, [split.test_loader for split in splits], session_ids=train_session_ids)
        pbar.set_description('test accuracy: ' + ', '.join(f'{accuracy(p):.2f}' for p in predictions))

test accuracy: 0.86, 0.69, 0.92: 100%|██████████| 200/200 [03:10<00:00,  1.05it/s]

A shared latent space

We plot the first three of the 16 neural axes of the latent space over time, for every session's test presentation. The three mice's curves follow each other: the model represents the movie the same way, whichever mouse's neurons it reads.

_, latents, predictions = sparks.evaluate(model, [split.test_loader for split in splits], session_ids=train_session_ids)
for session_id, session_predictions in zip(train_session_ids, predictions):
    print(f'Session {session_id}: test accuracy {accuracy(session_predictions):.2f}')

latents = [session_latents[0, tau_p:].numpy() for session_latents in latents]  # the test presentation, after the first tau_p bins
time = np.arange(tau_p, len(frame_shown)) * dt

fig, axs = plt.subplots(3, 1, figsize=(12, 7), sharex=True)
for session_id, session_latents in zip(train_session_ids, latents):
    for axis, ax in enumerate(axs):
        ax.plot(time, session_latents[:, axis], linewidth=1, label=f'Session {session_id}')
for axis, ax in enumerate(axs):
    ax.set_ylabel(f'Neural axis {axis + 1}')
axs[0].legend(loc='upper right', fontsize=8)
axs[-1].set_xlabel('Time in the movie (s)')
plt.show()

Session 715093703: test accuracy 0.86
Session 719161530: test accuracy 0.69
Session 721123822: test accuracy 0.92

Adding a new mouse

add_session gives a new session its own Hebbian attention layer, sized to its neurons, and reuses everything else. Since the decoder is shared, the new session needs no output head of its own. We then fine-tune the model on the new session's data. The optimizer is rebuilt so that it includes the new session's parameters.

(new_split,) = make_session_movies_data(data_path, [new_session_id], areas=('VISp',), mode='prediction', dt=dt, batch_size=9)
model.add_session(n_neurons=new_split.train_dataset.n_neurons, session_id=new_session_id)

_, _, predictions = sparks.evaluate(model, [new_split.test_loader], session_ids=[new_session_id])
print(f'Session {new_session_id} ({new_split.train_dataset.n_neurons} neurons), before fine-tuning: '
      f'test accuracy {accuracy(predictions[0]):.2f}')

Session 732592105 (108 neurons), before fine-tuning: test accuracy 0.02

n_epochs_finetuning = 100

param_groups, _ = build_param_groups(model, weight_decay=0.01)
param_groups[0]['params'] += list(transition.parameters())  # keep training the prior
optimizer = torch.optim.AdamW(param_groups, lr=0.002)
total_steps = n_epochs_finetuning * len(new_split.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))

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

    if (epoch + 1) % 10 == 0:
        _, _, predictions = sparks.evaluate(model, [new_split.test_loader], session_ids=[new_session_id])
        pbar.set_description(f'test accuracy: {accuracy(predictions[0]):.2f}')

test accuracy: 0.78: 100%|██████████| 100/100 [01:02<00:00,  1.59it/s]

_, _, predictions = sparks.evaluate(model, [new_split.test_loader], session_ids=[new_session_id])
predicted_frame = predictions[0][0].argmax(dim=-1)
print(f'Session {new_session_id}, after fine-tuning: test accuracy {accuracy(predictions[0]):.2f}')

fig, ax = plt.subplots(figsize=(8, 4))
ax.plot(np.arange(len(frame_shown)) * dt, frame_shown, color='grey', linewidth=3, label='Frame shown')
ax.scatter(np.arange(len(frame_shown)) * dt, predicted_frame, color='tab:orange', s=2, label='Predicted frame')
ax.set_xlabel('Time (s)')
ax.set_ylabel('Frame')
ax.legend()
plt.show()

Session 732592105, after fine-tuning: test accuracy 0.78