POYO

class pretrain.models.POYO(*, 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, finetune_enable=False)

Bases: core.model.BaseModel

Transformer-based model for neural decoding from spike trains [Azabou et al., 2023].

Adapted from torch_brain.

1. Input tokens are constructed by combining unit embeddings, token type embeddings,

   and time embeddings for each spike in the sequence.

2. The input sequence is compressed using cross-attention, where learnable latent

   tokens (each with an associated timestamp) attend to the input tokens.

3. The compressed latent token representations undergo further refinement through

   multiple self-attention processing layers.

4. Query tokens are constructed for the desired outputs by combining session

   embeddings, and output timestamps.

5. These query tokens attend to the processed latent representations through

   cross-attention, producing outputs in the model’s dimensional space (dim).

6. Finally, a task-specific linear layer maps the outputs from the model dimension

   to the appropriate output dimension.

Parameters:

• context_duration – Maximum duration of the input spike sequence (in seconds)

• modality_spec – A torch_brain.registry.ModalitySpec specifying readout properties

• latent_step (float) – Timestep of the latent grid (in seconds)

• num_latents_per_step (int) – Number of unique latent tokens (repeated at every latent step)

• dim (int) – Hidden dimension of the model

• depth (int) – Number of processing layers (self-attentions in the latent space)

• dim_head (int) – Dimension of each attention head

• cross_heads (int) – Number of attention heads used in a cross-attention layer

• self_heads (int) – Number of attention heads used in a self-attention layer

• ffn_dropout (float) – Dropout rate for feed-forward networks

• lin_dropout (float) – Dropout rate for linear layers

• atn_dropout (float) – Dropout rate for attention

• emb_init_scale (float) – Scale for embedding initialization

• t_min (float) – Minimum timestamp resolution for rotary embeddings

• t_max (float) – Maximum timestamp resolution for rotary embeddings

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.

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)

Input function used to convert Data into model inputs for the POYO model.

This input function can be called as a transform. If you are applying multiple transforms, make sure to apply this one last.

This code runs on CPU. Do not access GPU tensors inside this function.

Return type:

dict

forward(*, input_unit_index, input_timestamps, input_token_type, input_seqlen, output_session_index, output_timestamps=None)

Forward pass of the POYO model.

The model processes input spike sequences through its encoder-processor-decoder architecture to generate task-specific predictions.

Parameters:

• input_unit_index (Tensor) – \((S_{in},)\) int, unit index per input token.

• input_timestamps (Tensor) – \((S_{in},)\) float, spike timestamp per input token.

• input_token_type (Tensor) – \((S_{in},)\) int, token type id per input token.

• input_seqlen (Tensor) – \((B,)\) int, number of input tokens per sample.

• output_session_index (Tensor) – \((S_{out},)\) int, session index per output query.

• output_timestamps (Optional[Tensor]) – \((S_{out},)\) float, timestamp per output query, or None.

Return type:

Tensor

Returns:

\((S_{out}, D_{out})\) float. Callers that need a batch axis reshape it against the target they already hold.

classmethod create_search_space(trial, cfg)

Map out the model’s Optuna search space.

Call trial.suggest_*; the names suggested become the keys process_tunable_params() receives.

Parameters:

• trial (Trial) – Optuna trial to register suggestions on.

• cfg (DictConfig) – The run config, for values the space depends on.

classmethod process_tunable_params(tune_params)

Turn suggested hyperparameters into config overrides.

Runs before the config is filled, so this is where a suggestion is mapped onto the config path it sets (batch_size_log2 -> batch_size), a value is derived from another, or a default is supplied for something not being tuned.

Parameters:

tune_params (dict) – The names create_search_space() suggested, with their values.

Return type:

dict

Returns:

The overrides to apply to the config. The default returns them unchanged.