Explicit disentanglement of neural descriptive factors using the query-based neural activity prediction (qNAP) model

Deep learning models have shown significant efficacy in predicting neural activity, offering promising avenues for understanding complex neural dynamics. However, disentangling key explanatory factors---crucial for achieving generalizable and interpretable neural activity predictions---remains underexplored. Such disentanglement is essential for downstream simulations and for elucidating inter-individual variability in neural responses. To address this gap, we introduce the Query-based Neural Activity Prediction (qNAP) model, a novel framework that leverages an attention mechanism to explicitly disentangle intrinsic factors such as batch effects, neuronal types, receptive field locations, and other relevant descriptors. Our model can learn both categorical and continuous variables, including intra-type factors, and effectively trains on limited datasets---requiring as few as a dozen samples to learn a new factor. This approach allows for the exploration and manipulation of neural activity predictions by editing the query. Moreover, the qNAP model can predict neural activities in response to complex stimuli, enhancing its applicability to real-world neural data. A key advantage of our approach is its versatility across various neural network architectures, extending beyond convolutional neural networks, and its ability to incorporate additional explanatory factors to account for individual differences. This enhances the model's modularity and flexibility. By effectively disentangling neural descriptive factors, our work contributes to more generalizable and interpretable neural activity prediction models.