Foundation model for neural spiking data using multi-task masking (MtM) to translate across population, region, and single-neuron levels. Enables zero-shot and few-shot brain decoding across multiple brain areas. Activation triggers: neural translator, foundation model spiking, MtM, multi-task masking, IBL dataset, brain decoding.
Foundation model for neural spiking data using multi-task masking (MtM) to translate across population, region, and single-neuron levels. Enables zero-shot and few-shot brain decoding across multiple brain areas. Activation triggers: neural translator, foundation model spiking, MtM, multi-task masking, IBL dataset, brain decoding.
Neural Dynamics Universal Translator Foundation
A foundation model for neural spiking data that seamlessly "translates" across all spatial scales of the brain through multi-task masking self-supervised learning.
deftrain_mtm_model(model, dataloader, epochs=100, lr=1e-4):
"""Train multi-task masking model"""
optimizer = torch.optim.Adam(model.parameters(), lr=lr)
criterion = nn.MSELoss()
modes = ['temporal', 'neuron', 'region']
mask_ratio = 0.15for epoch inrange(epochs):
total_loss = 0for batch in dataloader:
x = batch['activity'] # (batch, time, neurons)
batch_size, T, N = x.shape
# Randomly select masking mode
mode = np.random.choice(modes)
# Create mask based on modeif mode == 'temporal':
mask_indices = create_temporal_mask(batch_size, T, N, mask_ratio)
elif mode == 'neuron':
mask_indices = create_neuron_mask(batch_size, T, N, mask_ratio)
else: # region
mask_indices = create_region_mask(batch_size, T, N, mask_ratio)
# Forward pass
reconstructed = model(x, mode, mask_indices)
# Compute loss only on masked positions
loss = criterion(reconstructed[mask_indices], x[mask_indices])
# Backward pass
optimizer.zero_grad()
loss.backward()
optimizer.step()
total_loss += loss.item()
if epoch % 10 == 0:
print(f"Epoch {epoch}, Loss: {total_loss / len(dataloader):.4f}")
return model
defcreate_temporal_mask(batch_size, T, N, ratio):
"""Create temporal masking indices"""
n_mask = int(T * ratio)
mask_t = torch.randperm(T)[:n_mask]
mask_indices = torch.cartesian_prod(
torch.arange(batch_size),
mask_t,
torch.arange(N)
)
return mask_indices[:, 0], mask_indices[:, 1], mask_indices[:, 2]
Step 4: Few-Shot Adaptation
classFewShotAdapter(nn.Module):
"""Linear probe for few-shot downstream tasks"""def__init__(self, encoder, output_dim):
super().__init__()
self.encoder = encoder
# Freeze encoderfor param inself.encoder.parameters():
param.requires_grad = Falseself.classifier = nn.Linear(encoder.embed_dim, output_dim)
defforward(self, x):
# Extract features using frozen encoderwith torch.no_grad():
features = self.encoder(x, mode='neuron', mask_indices=None)
# Classifyreturnself.classifier(features.mean(dim=1)) # Pool over timedeffew_shot_adapt(model, support_set, n_shots=5):
"""
Adapt model with few labeled examples
Args:
model: Pretrained MtM model
support_set: Dict with 'activity' and 'labels'
n_shots: Number of examples per class
"""# Create linear probe
adapter = FewShotAdapter(model, output_dim=n_classes)
# Train only classifier on support set
optimizer = torch.optim.Adam(adapter.classifier.parameters(), lr=1e-3)
for epoch inrange(50): # Few epochs for few-shot
logits = adapter(support_set['activity'])
loss = F.cross_entropy(logits, support_set['labels'])
optimizer.zero_grad()
loss.backward()
optimizer.step()
return adapter
Applications
1. Brain-Computer Interfaces (BCI)
Zero-shot decoding on new subjects
Reduced calibration time for neural prosthetics
Cross-subject motor imagery classification
2. Cross-Region Neural Analysis
Study information flow between brain regions
Identify region-specific neural codes
Map distributed neural computation
3. Behavior Prediction
Decode behavioral states from neural activity
Predict decision-making processes
Analyze cognitive task performance
4. Neurological Disorder Research
Compare neural dynamics across patient populations
Identify biomarkers for brain disorders
Track disease progression
Benchmarks
Task
Metric
Performance
Single-neuron prediction
R²
0.72
Region-level prediction
R²
0.68
Forward prediction (200ms)
MAE
0.15
Behavior decoding (choice)
Accuracy
82%
Cross-animal generalization
R²
0.61
Pitfalls
Data Quality: Model performance heavily depends on spike sorting quality
Temporal Resolution: Requires high temporal resolution recordings (≥1kHz sampling)
Recording Stability: Assumes consistent electrode placement across sessions
Animal Variability: May require fine-tuning for animals with significant anatomical differences
Computational Cost: Large transformer models require significant GPU memory
Related Skills
brain-dit-fmri-foundation-model
neurostorm-fmri-foundation
reve-eeg-foundation
spike-mllm-multimodal-spiking
meta-learning-in-context-brain-decoding
References
@article{zhang2024universal,
title={Towards a "universal translator" for neural dynamics at single-cell, single-spike resolution},
author={Zhang, Yizi and Wang, Yanchen and Jim{\'e}nez Benet{\'o}, Donato and Wang, Zixuan and Azabou, Mehdi and Richards, Blake and Winter, Olivier and others},
journal={arXiv preprint arXiv:2407.14668},
year={2024}
}
