Evaluating Latent Variable Models of Neural Population Activity: GPFA Fitting & Multi-Metric Evaluation

This repo fits Gaussian Process Factor Analysis (GPFA) to single-unit spiking data from extracellular recordings, in order to extract low-dimensional latent trajectories that summarize population-level neural dynamics on a trial-by-trial basis, and evaluates how good those latent trajectories actually are using multiple complementary metrics.

Evaluating Latent Variable Models

Fitting GPFA and picking a dimensionality via leave-neuron-out cross-validation (as in GPFA_data.ipynb) answers one specific question: which dimensionality captures shared structure in the population's activity without overfitting to noise? That's the approach from the original GPFA paper (Yu et al., 2009), and it's a good starting point — but it's only one lens on model quality.

Neural Latents Benchmark '21 (Pei, Ye, Zoltowski et al., 2021) introduces a more standardized evaluation framework for latent variable models (LVMs) of neural population activity, built around several complementary metrics rather than a single score. The table below summarizes those metrics and how each is used in this repo:

Metric	What it tells us	Why we need it
Leave-neuron-out prediction error (MSE)	How well the trajectory captures correlated firing rates across the population.	A fundamental, distance-based measure of predictive ability that works for both probabilistic and non-probabilistic methods.
Co-smoothing (bits per spike)	How well the model fits the raw neural data.	Ensures the model isn't just "hallucinating" structure that isn't there.
Behavioral decoding (R²)	How well the latent signals relate to actual movement/behavior.	Shows the latents are biologically grounded, not just math that happens to fit the spikes.
PSTH matching	How well the model captures stereotyped, repeated response patterns.	Checks whether the model reconstructs the trial-averaged ("PSTH") response well.
Forward prediction	How well the model predicts future activity.	Tests whether the model has learned the dynamics governing how the brain's state evolves over time, not just its instantaneous structure.

Metric-specific notebooks in this evaluation suite:

• Cross-validated held-out-trial neural reconstruction error (MSE) — GPFA_data.ipynb
• Co-smoothing evaluation using held-out bits per spike — cosmoothing_bps.ipynb
• Behavioral decoding from GPFA latent trajectories — behavioral_decoding.ipynb
• VLMD_decoding.ipynb isn't a new metric — it reuses the behavioral-decoding metric above, applied to a different latent variable model (VLMD instead of GPFA), so its results are directly comparable to behavioral_decoding.ipynb's.

PSTH matching and forward prediction aren't implemented in this repo yet. No single metric here is meant to stand alone as "the" measure of model quality.

Notebooks in this repo

GPFA_analysis.ipynb — fit GPFA and select dimensionality

Walks through the full pipeline for a single session:

1. Load session data — spikes and trial event times for a recording session
2. Unit selection — keep well-isolated ("good") units with enough spikes to support GPFA fitting
3. Build GPFA-ready trials — convert raw spike times into per-trial, per-neuron spike trains aligned to trial onset
4. Cross-validate latent dimensionality — fit GPFA across a range of candidate dimensionalities and score each by leave-neuron-out prediction error, with a held-out test set for a final unbiased evaluation
5. Fit the final model at the selected dimensionality and save results
6. Visualize latent trajectories — single-trial and trial-averaged dynamics, 3D trajectories, latent traces alongside spike rasters, and the GPFA loading matrix
7. Relate latents to behavior — correlate each latent dimension with measured behavioral variables (e.g. licking, whisking, pupil size, respiration)

behavior_decoding.ipynb — behavioral decoding evaluation

Asks a more direct question than the correlation check in step 7 above: can a simple linear decoder actually predict behavior from the GPFA latents, cross-validated across trials and swept across latent dimensionality? See Evaluating Latent Variable Models below for how this fits into the broader evaluation picture.

1. Load session data, select units, build GPFA trials — same data-preparation steps as GPFA_data.ipynb
2. Bin spike trains for GPFA at the chosen BIN_SIZE_MS
3. Single-behavior decoding walkthrough (lick rate) — bin lick rate into GPFA time bins, then for each candidate x_dim: fit GPFA on training trials, transform held-out test trials, and score a cross-validated Ridge regression decoder (R²) predicting lick rate from the latents
4. Multi-behavior decoding sweep — repeat the same procedure for every measured behavior (wheel velocity, whisking, pupil area/radius, respiration, face motion) at once, reusing each fold's fitted GPFA model across behaviors, and plot cross-validated R² vs. dimensionality for all behaviors together

cosmoothing.ipynb — co-smoothing evaluation (bits per spike)

Asks whether the fitted GPFA model explains the raw spiking data itself — specifically, whether its latents generalize to neurons it never saw during fitting, not just new trials. See Evaluating Latent Variable Models below for how this fits into the broader evaluation picture.

1. Load session data, select units, build GPFA trials — same data-preparation steps as the other notebooks
2. Held-in/held-out neuron split and train/test trial split — held-in neurons are used to fit GPFA and a latent→rate bridge; held-out neurons are only ever scored
3. Diagnostics — firing-rate and spike-count distributions for held-in vs. held-out neurons, to check the split is reasonable and see where low-firing neurons sit
4. Bridge method — fit GPFA on held-in neurons, bridge latents → held-in rates via linear regression, then bridge held-in rates → each held-out neuron's spikes via a Poisson GLM; score held-out predictions in bits/spike relative to a constant-rate null model
5. Cross-validate latent dimensionality by mean held-out bits/spike, with a final unbiased score on a held-out test set

7. Robustness check — how sensitive the co-smoothing score is to removing low-firing held-in neurons, used to choose a firing-rate cutoff for the final evaluation

VLMD_decoding.ipynb — behavioral decoding from an alternative model (VLMD)

Repeats the behavioral-decoding evaluation from behavioral_decoding.ipynb, but on latents from a different latent variable model — Variational Latent Mode Decomposition (VLMD) — loaded from a precomputed grid-search run on disk, rather than fit live with GPFA. This is a model-comparison notebook: it asks whether an alternative to GPFA produces latents that are at least as behaviorally meaningful, using the same decoding metric as behavioral_decoding.ipynb so the two are directly comparable.

1. Setup and session loading
2. Load precomputed VLMD latents — reconstruct the full latent trace from a saved VLMD run's modes_list.npy and related arrays
3. Inspect the VLMD run directory (diagnostic) — sanity-check array shapes and list sibling hyperparameter runs from the same grid search
4. Align VLMD latents to trial structure — VLMD is fit on the full continuous recording, so its latent trace has to be time-aligned and cropped into per-trial segments (unlike GPFA's latents, which come pre-segmented by trial)
5. Decode all behaviors from VLMD latents — same cross-validated Ridge regression / R² procedure as behavioral_decoding.ipynb, but at VLMD's fixed latent dimensionality rather than a swept range

Requirements

• Python 3.11
• elephant (spike train analysis + GPFA)
• neo and quantities
• numpy, scipy, scikit-learn, matplotlib

pip install elephant neo quantities numpy scipy scikit-learn matplotlib

Both notebooks also depend on two local modules from this repo, imported via a path relative to the notebook's location:

• data_paths.py — a DATA_PATH dict mapping session keys to .mat file paths
• Analysis_library/ — a local package providing file_loading.SessionData (for loading a session's spikes/trial data) and analysis.get_good_units (for filtering to well-isolated units)

Update data_paths.py to point at your own data before running.

Usage

GPFA_data.ipynb

1. Set SESSION_KEY in the "Load Session Data" section to the session you want to analyze.
2. Adjust analysis parameters as needed:
   • BIN_SIZE_MS — GPFA time-bin width (default 30 ms; consider 10 ms for finer temporal resolution at the cost of more compute)
   • PRE_TRIAL_S / POST_TRIAL_S — trial window around each trial event
   • MIN_SPIKES — minimum total spikes for a unit to be included
   • X_DIMS — candidate latent dimensionalities to cross-validate over
   • N_FOLDS — number of cross-validation folds
3. Run the notebook top to bottom. The cross-validation section is the slowest step (it fits len(X_DIMS) * N_FOLDS + 2 GPFA models).
4. Fitted models and results are saved to ISS/gpfa_results.pkl.

For the behavior-correlation section at the end, you'll need a binned_behaviors dict (behavior name → array of shape [trials, time bins], aligned to the same trial windows and bin size used above) — this notebook doesn't build it, so load or construct it from your session's behavioral traces first.

behavioral_decoding.ipynb

1. Set SESSION_KEY and the binning/trial-window parameters the same way as in GPFA_data.ipynb (these are independent notebooks, so parameters aren't shared automatically — keep them consistent if you want comparable results).
2. Adjust X_DIMS to the dimensionalities you want to sweep for decoding, and N_DECODE_FOLDS / the KFold settings for cross-validation.
3. Run top to bottom. The multi-behavior sweep (Section 5) is the slowest step — it fits len(X_DIMS) * n_folds GPFA models, each reused to decode every behavioral variable.
4. Behavioral signals are pulled directly from s.data["behavior"] and s.data["physiology"] — if your session's data structure uses different keys, update the behavior_signals dict in Section 5 accordingly.

cosmoothing_bps.ipynb

1. Set SESSION_KEY and binning/trial-window parameters the same way as in the other notebooks (again, independent notebooks — keep parameters in sync for comparable splits).
2. Adjust X_DIMS and N_FOLDS for the dimensionality cross-validation, and MIN_TEST_SPIKES (Section 6.3) if you want a different low-spike cutoff for evaluated held-out neurons.
3. Run top to bottom. Section 7's CV loop is the slowest step — it fits one GPFA model, one linear regression, and one Poisson GLM per held-out neuron, for every fold and every x_dim.
4. After the CV curve and final per-neuron plot (Section 7.1), Section 8 runs a robustness sweep at a fixed dimensionality (FIXED_X_DIM, default 10 — update this if your CV picks a different best x_dim) to check how sensitive the co-smoothing score is to low-firing held-in neurons, then applies a chosen cutoff (CHOSEN_CUTOFF_PERCENT) for a final filtered score.

VLMD_decoding.ipynb

1. Update STORAGE_VLMD_PATH and RUN_DIR in Section 3 to point at your own VLMD grid-search output (see Section 4's diagnostic cell for listing other available runs/hyperparameter configurations to swap in).
2. Set SESSION_KEY and keep PRE_TRIAL_S / POST_TRIAL_S / BIN_SIZE_MS in Section 5 consistent with the run's parameters (VLMD's own bin_size hyperparameter should match BIN_SIZE_MS here for the trial alignment to be meaningful).
3. Run top to bottom. This notebook is much faster than the other three — there's no model fitting, just loading, alignment, and one decoding pass per behavior.
4. Compare the printed vlmd_results R² scores against behavioral_decoding.ipynb's per-behavior results at a comparable dimensionality to see how VLMD stacks up against GPFA.

Known limitations / TODOs

• The final GPFA model in GPFA_data.ipynb isn't currently saved in a form that supports inference on new data without refitting.
• Raster plots aren't yet shown alongside every latent trajectory plot.
• behavioral_decoding.ipynb and GPFA_data.ipynb each rebuild gpfa_data and binned_trials independently — if you're evaluating the same session across metrics, this means the two notebooks may not be using an identical train/test trial split unless RANDOM_SEED and fold settings are kept in sync.
• cosmoothing_bps.ipynb's FIXED_X_DIM (used for the robustness sweep in Section 8) is hardcoded rather than automatically pulled from the CV result in Section 7 — update it manually if cross-validation picks a different best dimensionality.

Output

• ISS/gpfa_results.pkl (from GPFA_data.ipynb) — pickled dict containing the fitted GPFA model, estimated parameters, selected x_dim, cross-validation errors, unit selections, and the analysis settings used to produce them.
• behavioral_decoding.ipynb doesn't currently save results to disk — CV R² scores per behavior/dimensionality live in the behavior_scores and all_results dicts in-notebook, and are surfaced via the summary plots.
• cosmoothing_bps.ipynb doesn't currently save results to disk either — CV bits/spike per dimensionality lives in cv_results, final per-neuron scores in final_bps_per_neuron, and the robustness sweep in percentile_results, all surfaced via the in-notebook plots.
• VLMD_decoding.ipynb doesn't save results to disk — per-behavior CV R² scores live in the vlmd_results dict and are printed at the end of the notebook.