POCO: Scalable Neural Forecasting through Population Conditioning

TL;DR

POCO addresses the challenge of forecasting neural activity across sessions and individuals by uniting a simple univariate forecaster with a population encoder that conditions predictions on brain-wide state. Using FiLM conditioning and a Perceiver-IO–based tokenization of population activity, POCO scales to large neural populations and learns unit embeddings that align with brain-region structure without anatomical labels. Across five diverse calcium-imaging datasets, it achieves state-of-the-art cellular-resolution forecasting and demonstrates rapid adaptation to new sessions through embedding fine-tuning, while revealing meaningful functional organization in the learned embeddings. These results advance scalable neural foundation-model ideas for cross-subject forecasting and offer practical insights for designing adaptable neurotechnologies.

Abstract

Predicting future neural activity is a core challenge in modeling brain dynamics, with applications ranging from scientific investigation to closed-loop neurotechnology. While recent models of population activity emphasize interpretability and behavioral decoding, neural forecasting-particularly across multi-session, spontaneous recordings-remains underexplored. We introduce POCO, a unified forecasting model that combines a lightweight univariate forecaster with a population-level encoder to capture both neuron-specific and brain-wide dynamics. Trained across five calcium imaging datasets spanning zebrafish, mice, and C. elegans, POCO achieves state-of-the-art accuracy at cellular resolution in spontaneous behaviors. After pre-training, POCO rapidly adapts to new recordings with minimal fine-tuning. Notably, POCO's learned unit embeddings recover biologically meaningful structure-such as brain region clustering-without any anatomical labels. Our comprehensive analysis reveals several key factors influencing performance, including context length, session diversity, and preprocessing. Together, these results position POCO as a scalable and adaptable approach for cross-session neural forecasting and offer actionable insights for future model design. By enabling accurate, generalizable forecasting models of neural dynamics across individuals and species, POCO lays the groundwork for adaptive neurotechnologies and large-scale efforts for neural foundation models. Code is available at https://github.com/yuvenduan/POCO.

Figures (20)

• Figure 1: Model Architecture. (A) POCO combines a univariate MLP forecaster (orange part) and a population encoder that conditions the MLP (blue part). This is a schematic for illustration only; traces and feature maps shown are not actual model input, output, or embedding. (B) The population encoder is adapted from POYO azabou2023unified. We split the trace of each neuron into several tokens, encode the tokens with POYO, and then use unit embedding to query the conditioning parameters. See the Method section for more details.
• Figure 2: POCO maintains accuracy advantage over time and benefits from longer context. (A) MSE increases when forecasting longer into the future. Results are shown for two different datasets, see Figure \ref{['step_vs_performance']} for additional datasets. Error bars show SEM of 3 random seeds. (B) Model performance improves with longer context. (C) Sample prediction traces produced by POCO, where the first $C = 48$ steps are given to the model as context.
• Figure 3: POCO performance improves with longer recordings and more sessions. (A) Prediction score vs. training recording length (x-axis in log scale) for two different datasets. Models were trained using increasing portions of each session’s data. Error bars show SEM across 3 random seeds. (B) We split all sessions in one dataset into $n$ approximately equal partitions and train one model on each partition, then take the average of model prediction scores across all partitions. Average prediction scores vs the number of splits $n$ is shown for two datasets.
• Figure 4: Multi-session POCO improves when individuals are similar; POCO can quickly adapt to new sessions (A) Performance gain of multi-session POCO compared to single-session POCO on synthetic data for different values of $\eta$. Larger $\eta$ means individuals are less similar. We randomly generated 16 cohorts, each with 16 individuals. Each blue cross represents a cohort. Error bars are SEM. (B) Validation loss curve of fine-tuning pre-trained POCO (Pre-POCO) and training POCO, NLinear, and MLP from scratch. We also compared full-finetuning with only tuning the embedding. Dashed gray lines represent the copy baseline. Error shades represent SEM for 3 random seeds. Two sample sessions from two datasets are shown here; see Figure \ref{['finetune_curves']} for more sessions.
• Figure 5: POCO learns meaningful unit embeddings without supervision. (A) UMAP mcinnes2018umap visualization of unit embeddings after training POCO on the first 512 PCs in a zebrafish dataset. (B) Visualization of unit embedding after training POCO on the mice dataset, where neurons are colored by the brain region. One sample session is shown for (A) and (B), see Figure \ref{['embedding']} for more sessions. (C) Normalized average cosine similarity of unit embeddings between each pair of regions. Each row is normalized to $[0, 1]$ and then averaged across 4 runs. Patterns are consistent for different seeds (Figure \ref{['embedding_dist']}). See the Appendix for more details.