RRRDecoder

class pretrain.models.RRRDecoder(temporal_rank=10, bin_size=0.02)

Bases: core.model.BaseModel

Reduced-rank decoder mapping a window of binned spikes to behavior [Zhang et al., 2026].

Parameters:

• temporal_rank (int) – rank R of the factorization, the model’s main regularizer. Constrained to R <= min(min_N, T).

• bin_size (float) – spike bin width in seconds; T = context_window / bin_size.

Us

(N_sess, R) neuron -> latent projections, keyed by session id so a checkpoint survives a change in which sessions are included.

V

(R, T, out_dim) read-out. Shared, and the only tensor that transfers.

bs

(out_dim,) intercepts, keyed by session id.

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)

Transfer the shared basis V.

Us/bs index the pretraining sessions’ neurons and baselines, so they keep their fresh initialization and are fit on this session.

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(spikes, session_index)

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:

Tensor

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.