NDTSuperv#

class ts1.models.single_session.NDTSuperv(bin_size, max_spikes, spike_readin, hidden_dim, encoder_num_heads, encoder_ffn_factor, encoder_dropout, encoder_activation, encoder_num_layers, pre_encoder_dropout, post_encoder_dropout, use_custom_init=False, initrange=0.1, tfixup_scale_base=None, tfixup_v_scale_factor=None)[source]#

Bases: core.model.BaseModel

Single-session Transformer encoder for supervised neural decoding [Ye and Pandarinath, 2021].

Reference implementation: neural-data-transformers.

Note

Only single-session inference is supported.

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

link_datasets() must be called first to build the spike embedding, positional embedding, and Transformer encoder (their sizes depend on \(N\)).

configure_readout() must then be called to fix \(D_{out}\) and the output shape.

  1. input_fn(): bin raw spikes into \((T_{in}, N)\) and generate position indices.

  2. forward(): embed \((B, T_{in}, N)\) to \((B, T_{in}, D)\), add positional embeddings, encode with a Transformer, and project to \((B, 1, D_{out})\) or \((B, T_{in}, D_{out})\) depending on the target resolution.

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

Parameters:

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

  • max_spikes (int) – Maximum spike count per bin (used by the count embedding).

  • spike_readin (Literal['none', 'count', 'linear']) – How binned spikes enter the model (see link_datasets()).

  • hidden_dim (int) – Transformer hidden dimension \(D\) (overridden by spike_readin "none" or "count").

  • encoder_num_heads (int) – Number of attention heads.

  • encoder_ffn_factor (int) – Feed-forward sublayer width multiplier relative to \(D\).

  • encoder_dropout (float) – Dropout inside each encoder layer.

  • encoder_activation (str) – Activation of the feed-forward sublayer.

  • encoder_num_layers (int) – Number of Transformer encoder layers.

  • pre_encoder_dropout (float) – Dropout applied to embeddings before the encoder.

  • post_encoder_dropout (float) – Dropout applied to encoder output.

  • use_custom_init (bool) – If True, apply T-Fixup weight initialisation (see custom_init()).

  • initrange (float) – Uniform init range for non-Transformer weights.

  • tfixup_scale_base (Optional[float]) – T-Fixup base scale factor (required when use_custom_init=True).

  • tfixup_v_scale_factor (Optional[float]) – Additional scale applied to value weights (required when use_custom_init=True).

Build session-specific components from the training dataset.

Must be called before configure_readout(). Reads \(N\) from the dataset and builds:

  • Spike readin: maps \((B, T_{in}, N)\) to \((B, T_{in}, D)\), where \(D\) depends on spike_readin:

    • "none": identity cast; forces \(D = N\).

    • "count": per-unit count lookup of size 2, concatenated; forces \(D = 2N\).

    • "linear": linear projection; keeps the configured \(D\).

  • Positional embedding: lookup table of shape \((T_{in}, D)\).

  • Transformer encoder: \(d_{model} = D\), \(d_{ff} = D \cdot encoder\_ffn\_factor\).

Parameters:
Return type:

None

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\)): project \(D \to D_{out}\) at each time step, then mean-pool to \((B, 1, D_{out})\).

  • Timestep-level: project \(D \to D_{out}\) at each time step, output \((B, T_{in}, D_{out})\).

Parameters:

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

input_fn(data)[source]#

Bin spikes and generate position indices.

Parameters:

data – Trial data containing raw spike times and unit metadata.

Returns:

  • model_inputs.spikes: \((T_{in}, N)\) int spike counts.

  • model_inputs.positions: \((T_{in},)\) int position indices.

Return type:

Dict with

forward(spikes, positions)[source]#

Map binned spikes to task predictions.

Parameters:

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

  • positions (Tensor) – \((T_{in},)\) int position indices.

Return type:

Tensor

Returns:

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

custom_init()[source]#

Apply the same weight initialisation scheme used in NDT [Ye and Pandarinath, 2021].

Combines the T-Fixup strategy [Huang et al., 2020] with the stabilisation technique from [Nguyen and Salazar, 2019] to scale Transformer weights for better optimisation without a warm-up schedule.

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.