POSSM

class pretrain.models.POSSM(*, backbone='gru', bin_width=0.05, bin_step=0.05, num_latents=1, dim=256, depth=0, dim_head=64, cross_heads=1, self_heads=8, ffn_dropout=0.2, lin_dropout=0.4, atn_dropout=0.0, rnn_dim=512, num_rnn_layers=4, rnn_dropout=0.2, bidirectional=True, output_ca_ctx_lim=3, emb_init_scale=0.02, t_min=0.0001, t_max=10.0, dec_t_min=0.01, dec_t_max=10.0, task_vocab=None, finetune_enable=False)

Bases: core.model.BaseModel

POSSM (POYO + State Space Model) for the IBL benchmark [Ryoo et al., 2025].

Combines a perceiver-based per-bin encoder with a sequential backbone for processing neural spike data in temporal bins. Uses core.nn.MultitaskReadout for output projection so that the same model can be used for single-task eval or multi-task pretraining.

This differs from the paper in one respect: only the GRU backbone is implemented here, not the S4D or Mamba variants.

Architecture:

1. Input spikes are binned into temporal intervals.

2. Per-bin perceiver encoder compresses spikes into latent tokens via cross-attention, optionally refined by self-attention processing layers.

3. A sequential backbone processes the per-bin latent representations.

4. Decoder cross-attention with bounded causal context produces output query embeddings.

5. MultitaskReadout projects to task-specific output dimensions.

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

Parameters:

• backbone (str) – Sequential backbone type; "gru" is the only one implemented.

• bin_width (float) – Width of each temporal bin (seconds).

• bin_step (float) – Step between consecutive bins (seconds).

• num_latents (int) – Number of latent tokens per bin.

• dim (int) – Hidden dimension of all embeddings.

• depth (int) – Number of self-attention processing layers in the per-bin encoder.

• dim_head (int) – Dimension of each attention head.

• cross_heads (int) – Number of cross-attention heads.

• self_heads (int) – Number of self-attention heads.

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

• lin_dropout (float) – Dropout rate for linear layers.

• atn_dropout (float) – Dropout rate for attention.

• rnn_dim (int) – Backbone hidden dimension.

• num_rnn_layers (int) – Number of backbone layers.

• rnn_dropout (float) – Backbone dropout rate.

• bidirectional (bool) – If True (default), use a bidirectional GRU. Note this makes the per-bin representation non-causal, adjust output_ca_ctx_lim if real-time decoding is required.

• output_ca_ctx_lim (int) – Number of past bins visible to decoder cross-attention.

• emb_init_scale (float) – Embedding initialization scale.

• t_min (float) – Min time period for encoder rotary embeddings.

• t_max (float) – Max time period for encoder rotary embeddings.

• dec_t_min (float) – Min time period for decoder rotary embeddings.

• dec_t_max (float) – Max time period for decoder rotary embeddings.

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_multitask_readout(readout_specs)

Configure MultitaskReadout over a set of tasks.

Each task in readout_specs must be present in self.task_vocab so that its task_emb row and MultitaskReadout head id are stable across regimes (multi-task pretrain -> single-task finetune).

Parameters:

readout_specs (dict[str, TS1ReadoutSpec]) – mapping task_name -> TS1ReadoutSpec.

configure_readout(readout_spec)

Convenience wrapper around configure_multitask_readout().

Accepts a single ReadoutSpec (single-task), a list of them, or a {task_name: ReadoutSpec} dict. The container types are tested first, so the spec itself is matched structurally and never by its suite’s class.

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 POSSM 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.

Prepares per-bin spike inputs and multitask output queries for all configured readout specs. Each output query carries a decoder index identifying which readout head it belongs to.

Return type:

dict

forward(*, spike_unit_index, spike_timestamps, spike_type, input_mask, n_intervals, latent_index, latent_timestamps, output_timestamps, output_decoder_index, output_bin_index, output_session_index, unflatten_output=True, return_dict=False)

Forward pass of the POSSM model.

Parameters:

• spike_unit_index – per-bin spike unit indices, (B, n_intervals, n_in)

• spike_timestamps – per-bin spike timestamps (bin-relative)

• spike_type – per-bin spike token types

• input_mask – per-bin attention mask (True = valid token)

• n_intervals – number of bins per sample (or int when constant)

• latent_index – latent token indices, (B, n_latent)

• latent_timestamps – latent token timestamps (bin-relative)

• output_timestamps – output query timestamps, (B, n_out)

• output_decoder_index – per-query readout id, (B, n_out)

• output_bin_index – per-query 1-based bin index, (B, n_out)

• output_session_index – per-query session index, (B, n_out)

• unflatten_output (bool) – if True and the (single) eval task is timestep-level, returns (B, n_out, dim_out); sequence-level tasks return (B, dim_out). Used for compatibility with TS1EvalTrainer.

• return_dict (bool) – if True, returns the raw multitask dict {task_id: tensor} from MultitaskReadout. Used by multi-task training.

compute_bin_index(timestamps)

Map query timestamps to 1-based bin indices using the model’s bin structure (matches the boundary convention in input_fn()).

Parameters:

timestamps (Tensor) – float tensor of any shape; values in [0, context_duration].

Return type:

Tensor

Returns:

Long tensor of the same shape with 1-based bin indices.

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.