LFADSEvalTrainer#

class ts2.models.single_session.LFADSEvalTrainer(cfg, rank, world_size)[source]#

Bases: ts2.ts2_eval_trainer.TS2EvalTrainer

TS2 trainer for LFADS.

Replaces the base reconstruction loss with the LFADS ELBO:

L = recon + l2_ramp * (l2_gen_scale * l2_gen + l2_con_scale * l2_con)

  • kl_ramp * (kl_ic_scale * kl_ic + kl_co_scale * kl_co)

All three groups are per-sample sums averaged over the batch, which makes the KL scales true ELBO weights (1.0 is the unmodified ELBO). Weighting a per-element mean reconstruction against per-sample summed KL terms is what produces posterior collapse.

training_step(X, y)[source]#

Compute loss for one training batch.

Override this to customize masking, auxiliary losses, or any model-specific training logic. The base implementation is a plain forward pass against the benchmark-fixed target.

Parameters:

  • X (dict) – Input batch dict containing model_inputs.

  • y (dict) – Target batch dict containing values (benchmark-fixed spike counts).

Returns:

Scalar loss.

Return type:

torch.Tensor