Sa-SVAE: A Shared and Aligned Structured Variational Autoencoder for Extracting Behaviorally Relevant and Preserved Neural Dynamics Across Animals

Abstract

Understanding the preserved behaviorally-relevant neural dynamics across individuals when performing similar tasks presents a critical challenge. Current methods typically focus on analyzing subject-specific neural dynamics or employing post-training alignment to adapt latent dynamics across sessions and individuals. Yet, establishing a shared latent space that effectively captures the continuous nature of behavioral data remains elusive. In this study, we introduce sa-SVAE, a Shared and Aligned Structural Variational AutoEncoder that integrates neural recordings from multiple subjects and uncovers the shared, behaviorally-relevant latent dynamics, facilitating the prediction of corresponding behaviors through a universal decoder. Utilizing a Structured Variational AutoEncoder (SVAE), our approach infers nonlinear latent factors and learns tractable dynamics driven by behavior on a circuit-level manifold. We employ contrastive learning to align low-dimensional, behaviorally-relevant geometries across subjects, thereby preserving the integrity of neural representations linked to specific behaviors across different sessions and subjects. This alignment enables the development of a unified behavior decoder that outperforms previous methods. Our model demonstrates robust decoding of task-relevant behaviors by capturing these preserved latent dynamics, underscoring the factors essential for cross-subject generalization. This study highlights the potential for building a universal behavior decoder and provides neuroscience insights into preserved and behaviorally constrained neural representations.