NEDS

class pretrain.models.NEDS(bin_size, hidden_dim, tokenizer_dropout, encoder_num_layers, encoder_num_heads, encoder_dropout, encoder_ffn_factor, modalities=['spikes', 'choice', 'reward', 'stimulus_contrast', 'whisker_motion_energy', 'wheel_speed', 'right_paw_speed', 'left_paw_speed', 'licking_rate'], finetune_enable=False)

Bases: core.model.BaseModel

Multimodal masked transformer over spikes and behavior [Zhang et al., 2025].

Reference implementation: NEDS.

This differs from the reference in one respect: the per-session stitchers, in and out, are single linear layers. The two-layer MLP the reference uses was dropped as not scalable: its hidden layer is paid once per session and per modality, which on the output side alone comes to 250M parameters over the 423 pretraining sessions, next to a 9.5M encoder.

A run is made unimodal by its modalities and the masker’s mask_types, not by a mode flag: that is how ts1.models.pretrained.NEDSEvalTrainer decodes.

Not tested at any context_length other than the default (1.0); see Variable context.

link_datasets(train_dataset, val_dataset, test_dataset=None)

Build whatever the model sizes from the datasets.

Called before configure_readout(), so what is read here sizes the readout.

Parameters:

• train_dataset (IBLBrainWideBench2026) – Training split, the one to size from.

• val_dataset (IBLBrainWideBench2026) – Validation split.

• test_dataset (Optional[IBLBrainWideBench2026]) – Test split, or None when the run scores none.

load_ckpt(ckpt)

Copy pretrained weights out of a checkpoint into this model.

Read only the weights: the trainer restores optimizer and epoch state itself.

Parameters:

ckpt (dict) – The loaded checkpoint, as written by the pretraining run.

input_fn(data)

Convert one trial’s data into the model’s inputs.

Runs per item on the dataloader workers, as a dataset transform, so it returns tensors rather than batches.

Parameters:

data – One trial, exposing the interval key configure_readout() set.

Returns:

Dict whose model_inputs entry is splatted into forward(). Extra top-level keys are the model’s own to read in its trainer.

forward(model_inputs, keep_masks, modality_masks)

Run the model on one collated batch.

Takes whatever input_fn() put under model_inputs, splatted in as keyword arguments, and returns what the trainer’s loss reads.