MLP#

class ts1.models.single_session.MLP(bin_size=0.02, depth=2, hidden_dim=64, dropout=0.2, activation='relu', batch_norm=True)[source]#

Bases: core.model.BaseModel

Multi-layer perceptron mapping binned spike counts to task outputs.

Notation: \(B\) = batch size, \(T_{in}\) = input time bins, \(N\) = units, \(D_{out}\) = task output dim, \(T_{out}\) = output time steps, \(D\) = final hidden dim.

configure_readout() must be called before inference; it fixes \(D_{out}\) and the output shape.

  1. input_fn(): bin raw spikes into \((T_{in}, N)\) then flatten to \((T_{in} \cdot N,)\).

  2. forward(): pass \((B, T_{in} \cdot N)\) through the MLP to \((B, D)\), apply the readout head, and reshape to \((B, 1, D_{out})\) or \((B, T_{out}, D_{out})\) depending on the target resolution.

Parameters:

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

  • depth (int) – Number of hidden layers.

  • hidden_dim (int) – Width of the first hidden layer \(D_0\); each successive layer halves the width: \(D_0, D_0/2, \ldots\). The final layer has width \(D = D_0 / 2^{depth-1}\).

  • dropout (float) – Dropout probability applied after each activation.

  • activation (Literal['relu', 'gelu', 'tanh']) – Pointwise non-linearity.

  • batch_norm (bool) – If True, insert BatchNorm1d after each linear layer.

configure_readout(readout_spec)[source]#

Fix \(D_{out}\) and build the linear readout head.

The output shape depends on the target resolution:

  • Sequence-level (\(T_{out}=1\)): linear \(D \to D_{out}\), reshaped to \((B, 1, D_{out})\).

  • Timestep-level: \(D_{out}\) is expanded by \(T_{out}\), so linear \(D \to D_{out} \cdot T_{out}\), reshaped to \((B, T_{out}, D_{out})\).

Parameters:

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

input_fn(data)[source]#

Bin and flatten spikes.

Parameters:

data (Data) – Trial data containing raw spike times and unit metadata.

Return type:

dict[str, Tensor]

Returns:

Dict with model_inputs.spikes of shape \((T_{in} \cdot N,)\).

forward(spikes)[source]#

Map flattened binned spikes to task predictions.

Parameters:

spikes (Tensor) – \((B, T_{in} \cdot N)\) flattened binned spike counts.

Return type:

Tensor

Returns:

\((B, 1, D_{out})\) for sequence-level tasks or \((B, T_{out}, D_{out})\) for timestep-level tasks.

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.