CEBRA#

class ts1.models.single_session.CEBRA(output_dimension=8, bin_size=0.02, mode='time', max_iterations=1000, batch_size=512, learning_rate=0.0003, temperature=1.0, model_architecture='offset10-model', verbose=True, num_hidden_units=1024, time_offsets=10, temperature_mode='auto')[source]#

Bases: core.model.BaseModel

CEBRA-based neural decoding model on single-session [Schneider et al., 2023].

Note

Only single-session inference is supported.

Notation: \(B\) = batch size, \(T_{in}\) = input time bins, \(N\) = units, \(D_{emb}\) = CEBRA embedding dimension (output_dimension), \(D_{out}\) = task output dim, \(D\) = MLP hidden dim (num_hidden_units, shared with the CEBRA encoder).

  1. fit(): train the CEBRA encoder on the eval training split. The trainer then calls collect_embeddings to run cebra_model.transform on all splits and cache the results in model.embeddings.

  2. configure_readout(): build the two-layer MLP head and fix \(D_{out}\).

  3. input_fn(): look up cached CEBRA embeddings to get \((T_{in}, D_{emb})\) or \((1, D_{emb})\) for sequence-level tasks.

  4. forward(): pass embeddings through the MLP readout to produce \((B, T_{out}, D_{out})\).

Parameters:

  • output_dimension (int) – CEBRA embedding dimension \(D_{emb}\).

  • bin_size (float) – Width of each time bin in seconds.

  • mode (str) – Contrastive objective, "time" uses temporal labels only; "behavior" conditions on behavioral labels (and enables hybrid mode for continuous labels).

  • max_iterations (int) – Number of CEBRA training iterations.

  • batch_size (int) – Mini-batch size for CEBRA training.

  • learning_rate (float) – CEBRA optimiser learning rate.

  • temperature (float) – Contrastive loss temperature (used when temperature_mode="constant").

  • model_architecture (str) – CEBRA model architecture string (e.g. "offset10-model").

  • verbose (bool) – Whether to print CEBRA training progress.

  • num_hidden_units (int) – MLP hidden dimension \(D\).

  • time_offsets (int) – Number of time offsets for the temporal objective.

  • temperature_mode (str) – How to set the contrastive temperature ("auto" or "constant").

fit(session, behavior_labels=None)[source]#

Train the CEBRA encoder on a single session.

In "time" mode the model uses a temporal contrastive objective with synthetic time labels. In "behavior" mode it conditions on behavior_labels and additionally enables the hybrid (time + behavior) objective for continuous labels. After this call the trainer runs collect_embeddings to cache cebra_model.transform outputs in model.embeddings before configure_readout() is invoked.

Parameters:

  • session (ndarray) – \((T_{in}, N)\) float spike-rate array for the session.

  • behavior_labels (Optional[ndarray]) – \((T_{in}, ...)\) behavioral labels required when mode="behavior". Pass continuous floats to enable hybrid mode; integer labels disable it.

configure_readout(readout_spec)[source]#

Fix \(D_{out}\) and build the two-layer MLP readout head.

The output shape depends on the target resolution:

  • Sequence-level (\(T_{out}=1\)): input_fn() mean-pools embeddings to \((1, D_{emb})\) before the MLP.

  • Timestep-level: input_fn() returns \((T_{in}, D_{emb})\); the MLP is applied at each step to yield \((B, T_{in}, D_{out})\).

The readout is a two-layer MLP: \(D_{emb} \to D \xrightarrow{\text{GELU}} \text{Dropout}(0.2) \to D_{out}\).

Parameters:

readout_spec (ReadoutSpec) – Task specification carrying \(D_{out}\) and the target resolution.

input_fn(data)[source]#

Look up pre-computed CEBRA embeddings for a single trial.

Parameters:

data (Data) – Trial data; must expose the interval key specified in configure_readout() and a single trial index into self.embeddings.

Returns:

  • model_inputs.embeddings: \((T_{in}, D_{emb})\) float tensor for timestep-level tasks, or \((1, D_{emb})\) after mean-pooling for sequence-level tasks.

Return type:

Dict with

forward(embeddings)[source]#

Map CEBRA embeddings to task predictions via the MLP readout.

Parameters:

embeddings (Tensor) – \((B, T_{out}, D_{emb})\) float tensor of CEBRA embeddings, where \(T_{out}=1\) for sequence-level tasks or \(T_{out}=T_{in}\) for timestep-level tasks.

Return type:

Tensor

Returns:

\((B, T_{out}, D_{out})\) task predictions.

classmethod create_search_space(trial, cfg)[source]#

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)[source]#

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.