POYOPlus

class pretrain.models.POYOPlus(*, latent_step, num_latents_per_step=64, dim=512, depth=2, dim_head=64, cross_heads=1, self_heads=8, ffn_dropout=0.2, lin_dropout=0.4, atn_dropout=0.0, emb_init_scale=0.02, t_min=0.0001, t_max=2.0627, task_vocab=None, attn_impl='nested')

Bases: core.model.BaseModel

POYO+ over chained spike tokens [Azabou et al., 2025].

attn_impl picks the kernel behind every attention layer. The default nested runs on stock torch; xformers needs the optional xformers extra (uv pip install -e ".[train,xformers]") and is worth installing for anything long-running. Its main win is memory – a substantially lower peak, which is often what decides whether a batch size fits on one GPU – with a smaller speedup on top. The two compute the same attention, so checkpoints are interchangeable: switching backends needs no retraining and no state-dict surgery.

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

link_datasets(train_dataset, val_dataset, test_dataset=None, finetune_enable=False)

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.

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 (Data) – One trial, exposing the interval key configure_readout() set.

Return type:

dict

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(*, input_unit_index, input_timestamps, input_token_type, input_seqlen, latent_index, latent_timestamps, latent_seqlen, output_session_index=None, output_timestamps=None, output_decoder_index=None, output_batch_index=None, return_dict=True, unpack_output=False, unflatten_output=True)

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.

Return type:

dict[str, Tensor] | Tensor