MtM

class pretrain.models.MtM(hidden_dim, bin_size, encoder_num_layers, encoder_dim_feedforward, encoder_num_heads, encoder_dropout, encoder_activation, pre_encoder_dropout, post_encoder_dropout, shared_proj_input_dim, shared_proj_ffn_mult=2, use_custom_init=False, tfixup_scale_base=None, tfixup_v_scale_factor=None, finetune_enable=False)

Bases: core.model.BaseModel

Multi-task masked transformer over binned spikes [Zhang et al., 2024].

Reference implementation: IBL_MtM_model.

Transformer encoder-based self-supervised multi-task model for neural population dynamics.

Extend NDT with multi-task masking objectives:

• co-smooth: predict a masked neuron from the activity of other neurons.

• causal: predict future time steps from past context.

• inter-region: predict one region’s activity from other regions.

• intra-region: predict one neuron from other neurons in the same region.

Returns predicted firing rates and neural latents.

Note: MtM intentionally has a lot of overlap with the NDTStitch implementation (i.e., the DRY code principle is not respected). This design choice is motivated by the desire to make an individual model fully understandable for a user without having to navigate between models.

We document/highlight the key differences between NDTStitch and MtM.

Removed:

• Nothing. The NDTStitch structure is kept whole: per-session in/out stitchers, session embedding, encoder and T-Fixup init are unchanged.

Added:

• A session-shared projection after the per-session in_stitcher: Linear -> Softsign -> Linear, taking shared_proj_input_dim to hidden_dim. The per-session layer stays a single linear, so the width added here is paid once and not once per session.

• Mask token added in the sequence: one learned embedding per mask mode, prepended before the encoder and dropped after it.

• Region tokenization: input_fn also emits a regions array, which the inter-region and intra-region objectives need to pick their target region.

Updated:

• in_stitcher: num_units -> shared_proj_input_dim, rather than num_units -> hidden_dim, since the shared projection now produces hidden_dim.

• encode() and forward() take a mask_mode, and encode() drops one extra leading token when a mode is active.

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.

configure_readout(readout_spec)

Size the readout head for the task being evaluated.

Called after link_datasets(). The spec carries dim and num_timesteps, so a head is sized without branching on it.

Parameters:

readout_spec (ReadoutSpec) – See ReadoutSpec.

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.

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.

forward(spikes, positions, session_tokens, mask_mode=None)

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.

custom_init()

T-Fixup init for the encoder; see core.nn.tfixup_init_().