Predictive Representation Learning for Partially Observed Neural Dynamics
Abstract
In long-term multi-session neural recordings, the same latent population dynamics are often observed through changing subsets of neural channels due to electrode drift, signal instability, or session-dependent recording variability. This creates a partially observed system-identification problem: each session provides only an incomplete view of a shared dynamical system, and missing channels are not merely absent inputs but unobserved variables coupled to the same latent dynamics. Existing approaches often rely on observation-pattern-specific mapping, which treats missing channels passively and leaves the shared latent dynamics weakly constrained. We propose CHAMA (CHannel-Aware Masked Attention), a mask-conditioned predictive representation learning framework for identifying shared neural dynamics under heterogeneous channel availability. CHAMA introduces a single global observation interface conditioned on channel-availability masks, allowing observed and unobserved channels to jointly constrain latent dynamics. The key idea is bidirectional: missing channels impose structural constraints on the latent representation, while the learned latent dynamics enable principled recovery of missing activity. To improve identifiability under sparse observations, CHAMA aggregates causal temporal windows, effectively trading temporal context for missing spatial information. Across synthetic dynamical systems and real multi-session neural recordings, CHAMA improves missing-channel completion, forecasting, latent dynamical fidelity, and downstream behavioral decoding over conditional adapter, zero-padding, and masked-autoencoding baselines. These results show that dynamics-aware recovery preserves globally consistent and behaviorally relevant structure beyond point-wise reconstruction accuracy. The source code is available at: \url{https://anonymous.4open.science/r/CHAMA-D691}.