- The paper introduces CalM, a self-supervised framework that decouples task-specific modeling with an innovative tokenization scheme and an autoregressive transformer for calcium imaging data.
- It utilizes a two-stage pipeline with a neural quantizer based on VQ-VAE and a dual-axis transformer, achieving superior performance in population forecasting and behavior decoding compared to specialized baselines.
- The study demonstrates emergent interpretability in neuronal embeddings, with robust cross-session transferability and up to 20.9% improvement in yaw velocity decoding R² on held-out data.
Self-Supervised Foundation Model for Calcium-Imaging Population Dynamics
Overview
The paper proposes CalM, a self-supervised foundation model targeting scalable representation learning from large-scale, multi-animal calcium imaging datasets. CalM decouples task-specific modeling by pretraining on purely neuronal calcium traces and introduces a dedicated tokenization scheme and an autoregressive transformer backbone facilitating transfer across diverse downstream objectives such as population dynamics forecasting and behavior decoding. The framework achieves strong empirical performance relative to specialized baselines and exhibits emergent interpretability properties in its learned representations, supporting broad applicability in functional neural analysis.
Model Architecture and Training Paradigm
CalM is structured as a two-stage pipeline:
- Neural Quantizer (NQ):
- Utilizes VQ-VAE to discretize single-neuron calcium traces into tokens.
- Employs convolutional segmentation followed by transformer encoding with RoPE for sequence order, ensuring context-aware feature extraction.
- Discrete tokens are derived by nearest-neighbor lookup in a learnable codebook, with reconstruction handled via a transformer-based decoder.
- Codebook regularization (entropy maximization, orthogonality, and dead code revival) and an auxiliary autoregressive head optimize for temporal predictability and prevent mode collapse.
- Dual-Axis Transformer (DAT):
- Ingests tokenized population activity across trials, factorizing attention along neural (bidirectional, per timestep) and temporal (causal, per neuron) axes.
- Incorporates neuron and session embeddings for identity and context conditioning, crucial for multi-session/multi-animal robustness.
- Scheduled sampling and neighborhood replacement during training mitigate exposure bias and quantization error propagation.
- Primary objective is cross-entropy over token prediction; auxiliary losses foster regularization.
Downstream tasks are supported via frozen pre-trained backbones and task-specific heads (linear or nonlinear GLU-style) trained for forecasting and behavior decoding.
Experimental Evaluation
Dataset
Comprehensive evaluation was performed using both simulated RNN-driven calcium traces and the public dataset from [tseng2022shared], comprising 8 mice, 286 sessions, over 273,000 neurons across six regions during a decision-making task. Trials were partitioned for rigorous held-in/held-out testing.
Forecasting
- Population dynamics forecasting: CalM consistently outperforms specialized baselines and is competitive with POCO in multi-session settings. Importantly, CalM accommodates variable-length horizons flexibly (unlike POCO's fixed window), and operates without task-specific forecasting supervision after self-supervised pretraining.
- Strong numerical results include comparative evaluations on Pearson correlation and full autoregressive rollout, demonstrating the backbone's transferability and robustness to session/animal variability.
Behavior Decoding
- The fine-tuned behavior head yields superior decoding results versus GLM, TCN, and neural forecasting models. On held-in data, improvements over POYO+ reach up to 20.9% in yaw velocity decoding R2, with an average increase of 10.1%.
- On held-out sessions, CalM retains a 7.2% advantage without updating backbone weights, highlighting its capacity for generalization and transfer.
Interpretability
- Neuronal Embedding Structure: PCA and LDA analyses show functional segregation of cue- and choice-tuned neurons, with linear axes corresponding to orthogonal gradients in task variable encoding.
- Low-dimensional Dynamics: CalM's forecasted traces project onto ground-truth principal component axes with higher correlation than POCO and significantly outperform shuffled controls, demonstrating preservation of intrinsic manifold structure.
Ablation studies verify the necessity of all regularization components and training strategies, particularly scheduled sampling, for robust temporal forecasting.
Implications and Future Directions
CalM establishes a scalable, self-supervised paradigm for foundation modeling in systems neuroscience, facilitating principled transfer across animals, sessions, and tasks. The tokenization scheme and autoregressive backbone enable both generative and discriminative objectives, reducing reliance on task-specific engineering and enabling efficient downstream adaptation. Emergent interpretability in representations indicates potential for functional discovery beyond predictive accuracy.
Future research directions include:
- Extending CalM to multimodal neural data (combining, e.g., electrophysiology and calcium imaging) to further contextualize behavioral outputs.
- Enabling direct backbone fine-tuning in raw-trace space for enhanced precision in forecasting tasks.
- Unifying architecture for end-to-end learning, potentially incorporating online adaptation mechanisms for real-time neural interface applications.
- Systematic investigation of representation spaces using nonlinear and information-theoretic probes to elucidate emergent properties and manifold geometry.
Implications span both practical neuroscience workflows (accelerating cross-subject transfer and reducing annotation burden) and theoretical efforts in understanding population dynamics and neural coding.
Conclusion
CalM demonstrates effective self-supervised pretraining for large-scale calcium imaging, outperforming specialized models on both forecasting and decoding tasks across single and multi-session, multi-animal datasets. The framework leverages a novel discrete tokenization method and dual-axis transformer architecture to capture structured dynamics and enable transferable representations. Linear analyses confirm the interpretability and functional organization within the learned embeddings, supporting scalable neural discovery. CalM thus represents a strong step toward general-purpose foundation models for population-level neuroscience.