LFADS

class ts2.models.single_session.LFADS(ic_enc_dim=64, ci_enc_dim=64, ic_dim=64, gen_dim=200, fac_dim=40, co_dim=4, con_dim=64, ci_lag=1, dropout=0.05, cd_rate=0.3, unit_cd_rate=0.1, cd_pass_rate=0.0, cell_clip=5.0, ic_post_var_min=0.0001, log_rate_clamp=8.0, ic_prior_variance=0.1, co_prior_tau=10.0, co_prior_nvar=0.1, forecast_grad_scope='window')

Bases: core.model.BaseModel

LFADS for TS2 co-smoothing and forecasting [Keshtkaran et al., 2022, Pandarinath et al., 2018].

Consumes binned spike counts (B, T, N) and returns log firing rates of the same shape for the benchmark’s Poisson NLL.

During training the input is corrupted in the shape of the task’s hold-out and the reconstruction gradient flows only through the corrupted entries (see _mask_input()). Val and test arrive already stripped by the dataset.

Parameters:

• ic_enc_dim (int) – Hidden size per direction of the initial-condition encoder.

• ci_enc_dim (int) – Hidden size per direction of the controller-input encoder.

• ic_dim (int) – Dimension of the initial-condition latent g0.

• gen_dim (int) – Generator hidden size.

• fac_dim (int) – Number of latent factors.

• co_dim (int) – Dimension of the inferred inputs u_t (0 disables the controller).

• con_dim (int) – Controller hidden size.

• ci_lag (int) – Lag on the controller input, so the controller does not see the current timestep of the data it must explain.

• dropout (float) – Rate applied to input, encoder output, controller input and generator state.

• cd_rate (float) – Coordinated-dropout rate: single input entries hidden from the encoder and used for the reconstruction gradient. Co-smoothing only.

• unit_cd_rate (float) – Rate of hiding a unit’s whole row for the window, the shape the co-smoothing hold-out takes at eval. Co-smoothing only.

• cd_pass_rate (float) – Fraction of visible entries that also pass gradient.

• cell_clip (float) – Clip value for generator and controller states.

• ic_post_var_min (float) – Floor on the initial-condition posterior variance.

• log_rate_clamp (float) – Clamp on output log-rates, keeping the Poisson NLL from overflowing early in training.

• ic_prior_variance (float) – Variance of the fixed-variance prior over g0.

• co_prior_tau (float) – Initial AR(1) time constant of the prior over u_t, in bins.

• co_prior_nvar (float) – Initial AR(1) process variance of the prior over u_t.

• forecast_grad_scope (Literal['tail', 'window']) – Forecasting only, ignored for co-smoothing. Whether the reconstruction gradient comes from the held-out tail (“tail”) or the whole window (“window”).

link_datasets(train_dataset, val_dataset=None, 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 (IBLBrainWideBenchTS2) – Training split, the one to size from.

• val_dataset (Optional[IBLBrainWideBenchTS2]) – Validation split.

• test_dataset (Optional[IBLBrainWideBenchTS2]) – Test split, or None when the run scores none.

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, return_posterior=False)

Run the sequential VAE.

Parameters:

• spikes (Tensor) – (B, T, N) binned spike counts.

• return_posterior (bool) – Also return the KL terms, L2 penalties and gradient mask. A forward argument rather than a second method so training still goes through DDP.forward.

Return type:

Tensor | dict

Returns:

(B, T, N) log firing rates, or a dict with those under log_rates plus the training quantities.

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.