Time-Evolving Dynamical System for Learning Latent Representations of Mouse Visual Neural Activity

We thank the reviewer for the pointed and constructive comments. We will do our best to address the comments and answer the questions.

R1. References and baselines.

We sincerely appreciate the reviewer's valuable suggestions regarding additional Latent Variable Model (LVM) references, and we will gladly incorporate a discussion of these works into our revised manuscript. However, we respectfully contend that these suggested models may not be more appropriate for direct comparison with our method than the baselines we have already employed.

First, our work proposes a novel visual dynamical system that leverages self-supervised contrastive learning to disentangle neural activity components related to visual stimuli from those influenced by internal brain states. A core characteristic of our approach is the absence of any supervised signals (e.g., external behavioral variables) in the model's training. Given these specific model characteristics, we selected our current baselines. LFADS is a classic and widely recognized dynamical system, providing a robust benchmark for temporal neural data. pi-VAE incorporates external supervised information (and supporting discrete variables) to construct interpretable latent variables, contrasting with our self-supervised objective. Swap-VAE utilizes contrastive learning to construct separated latent variables, aligning with our disentanglement goal, although it does not explicitly model temporal dynamics. CEBRA also employs contrastive learning for latent variable construction. Critically, the work on CEBRA has conducted extensive analyses of neural activity in the mouse visual cortex and our experimental setup for visual neural activity largely follows its established methodology. This direct alignment in an experimental context makes CEBRA a particularly pertinent baseline.

Second, most of the models mentioned by the reviewer are either oriented towards other research topics or possess characteristics that render them less applicable to our specific context. For instance, models such as LDS and PLINE primarily focus on modeling the dynamics of synthetic data or motor behaviors in the motor cortex. Similarly, PSID, DFINE and DPAD are designed to separate behavior-relevant and behavior-irrelevant neural dynamics using supervised information, which differs from our self-supervised approach. MM-GPVAE is tailored for multimodal inputs, while CREIMBO and POYO focus on modeling population dynamics across brain regions or individuals. Finally, models like RP-GSSM and DCCA are primarily concerned with learning low-dimensional representations of images or videos, which fall outside the scope of our neural activity analysis.

Nevertheless, as suggested by the reviewer, we incorporate PSID, DFINE, and DCCA into our experiments on visual neural activity for comparison. Notably, the supervised versions of PSID and DFINE leverage continuous, high-temporal-resolution behavioral variables as supervised information to separate behavior-relevant and behavior-irrelevant neural dynamics. This approach presents significant challenges for application to the disentanglement of visual-related and visual-unrelated neural activity in our context. Specifically, biological experiments often utilize static visual stimuli, or dynamic stimuli with temporal resolutions considerably lower than that of neural activity. Furthermore, natural visual stimuli are high-dimensional and complex signals. The various methods of constructing visual information for supervision introduce strong human-induced priors that could influence the nature of the constructed latent variables. For example, using pixel values, low-order visual features, or high-order semantic information for supervision corresponds to visual information encoded by different regions of the visual pathway. Therefore, we adopt the unsupervised versions of PSID and DFINE. As presented in Table r1, our model consistently demonstrates superior performance compared to these models.

Table r1: The results (%) for decoding natural movie frames.

R2. The low performance of CEBRA in Table 1.

We think that CEBRA demonstrates stronger performance in distinguishing between coarse-grained categories but less capability for finer-grained distinctions. First, this is suggested by Figure 8 in the Appendix, where, despite CEBRA's capacity to differentiate some categories through contrastive learning, considerable overlap is observed in the neural latent representations corresponding to other categories. This limitation is further evinced in the movie decoding experiment (Figure 4B), where CEBRA's performance declines more significantly as the constraint window narrows (i.e., an increased number of movie frame categories).

To corroborate this view, we further perform an experiment on static scene decoding. We train the models on neural data of the five natural scenes that elicit the strongest neural responses and evaluate their performance on the five-classification task. As shown in Table r2, although CEBRA's performance remains inferior to our model, it achieves a relatively good level, supporting our hypothesis regarding its better performance on coarser distinctions.

Furthermore, we note that the lower performance of pi-VAE is confined to Mouse 1, with its performance on other mice being at an intermediate level, thus ruling out the possibility of a bug.

Table r2: The results (%) for decoding five natural scenes.

R3. The evaluation for internal latent variables.

While our current findings indicate a low correlation between internal latent variables and visual stimuli, we acknowledge the absence of quantitative metrics in this work to definitively verify if these variables reflect an animal's internal state. To address this limitation, we will utilize more comprehensive datasets (e.g., incorporating behavioral and physiological recordings) for robust evaluation in the future.

R4. The data requirements.

This represents a trade-off: our model, being more complex than linear alternatives, incurs higher computational costs when applied to smaller datasets. However, for larger datasets, where linear models are prone to underfitting, our model demonstrates superior scalability and the capacity to capture more intricate relationships.

We thank the reviewer for the notable and perceptive comments. We will do our best to address the reviewer's concerns and answer the questions in the following.

R1. Problem motivation, the core methodological contribution and the gap in prior work.

Problem motivation and scientific objective: LVMs have proven highly influential in studies for the motor cortex, primarily by extracting low-dimensional neural embeddings that encode motor features and capture motor-related neural dynamics. However, LVMs have limited application in the analysis of visual neural activity. In our work, the scientific objective is to capture low-dimensional, stimulus-relevant embeddings from visual neural activity, particularly in the context of naturalistic visual stimuli. More specifically, we aim to leverage LVMs to disentangle neural components related to visual stimuli from those influenced by internal states. Through this disentanglement, we seek to further explain the observed variability in visual neural activity across subjects and across cortical regions.

The core methodological contribution: We have provided a concise approach description in lines 45-50. Here, we make a more detailed clarification. First, we introduce temporal structures to explicitly establish temporal relationships for latent variables in a causal order. This design allows the latent variables to evolve over time, capturing the temporal dependencies inherent in neural activity. Second, we implement a split structure and design distinct loss functions to build two specialized sets of latent variables: external latent representations are trained to capture stimulus-relevant components of visual neural activity and internal latent representations are developed to reflect the animal's dynamic internal states.

The gap in prior work: Existing studies via LVMs on visual neural activity primarily focus on simple visual stimuli and do not investigate extracting meaningful neural embeddings under naturalistic visual stimuli. More importantly, there was a lack of LVMs capable of splitting latent representations to explicitly differentiate between external stimulus-relevant and internal state-related components in visual neural activity in a self-supervised manner. For detailed differences in structure and design from previous models, please refer to R9 to Reviewer RtPx.

R2. Model architecture and methodology of model training.

Model architecture: Our model's core comprises external and internal latent variables, representing distinct visual neural activity components, alongside state variables for temporal processing. These variables interdepend through modules, fully described in the manuscript; here, we highlight design motivations.

• Many component dependencies are adapted from the classic VRNN.
• In the encoder, both external and internal latent variables are conditioned on input neural activity and their respective RNN state variables, crucial for causal temporal capture.
• The internal latent variable's prior is conditioned solely on state variables, promoting internal state spontaneity.
• The decoder reconstructs neural activity conditioned on both latent and internal state variables, with internal state variables enriching neural dynamics for more precise reconstruction.
• Finally, external RNN state variables depend only on current input and previous state, while internal state variables also depend on latent variables, reflecting how external visual stimuli influence an animal's internal states.

Loss functions: The objective of multiple training terms within our designed loss functions is to achieve disentanglement between external and internal latent variables, compelling them to reflect distinct components of visual neural activity.

• As in lines 132–141, contrastive learning is applied to the external latent variables. It encourages clustering external latent variables corresponding to sample pairs with small temporal offsets (i.e., under similar visual stimuli), associating these variables with the visual input. Concurrently, a time-dependent prior distribution exhibiting spontaneity guides the internal latent variables via KL divergence. The differential application of these distinct loss functions to the two sets of latent variables is instrumental in promoting their disentanglement.
• The swapping operation (detailed in R1 to Reviewer NU2a) further reinforces this disentanglement. It compels the model to strengthen the distinct roles of the external and internal latent variables, thereby enhancing their disentangled representation.

For the reason for the baseline choices, please refer to R1 to Reviewer RfuT.

R3. The experiments of synthetic datasets.

We have provided the detailed generation procedure for the two synthetic datasets in Appendix C. Specifically, the non-temporal dataset is constructed using several label clusters to represent distinct visual stimulus categories, while the temporal dataset is designed to characterize the temporal relationships inherent in neural activity.

These synthetic experiments are crucial for evaluating the fundamental ability of LVMs to construct accurate latent variables—a standard preliminary step in this field. The verification of our model's core contribution and its efficacy in extracting latent variables related to actual visual stimuli is subsequently accomplished through experiments with real visual neural activity datasets.

R4. About the decoding task.

For decoding tasks and coarse-grained (within +/-1 second) movie frame classification, we largely follow CEBRA's established approach, as it has conducted extensive analysis of mouse visual cortex activity. This alignment facilitates fair baseline comparisons. In addition, our model consistently outperforms others in finer-grained (the minimum is +/-0.2 second) frame classification, as shown in Figure 4B.

Nevertheless, this approach has limitations with longer movies or larger scene datasets. Future work will explore more generalized metrics for evaluating latent variable-visual stimuli relationships, such as comparing the similarity between latent variables and low- or high-level visual features.

For k-value selection, empirical values were derived considering dataset size, latent variable dimension, and category count. An odd $k$ value was chosen to reduce decision uncertainty.

R5. Evaluation of spike reconstruction and internal latent variables.

We evaluate spike reconstruction on the test set for models other than CEBRA. Our model consistently demonstrated superior performance, as indicated by lower Poisson negative log-likelihood values (Tables r1, r2).

We acknowledge the current lack of quantitative metrics for internal latent variables. While preliminary results suggest low correlation with visual stimuli, future work will require more comprehensive datasets to rigorously validate whether these variables accurately reflect an animal's internal state.

Table r1: The reconstruction results for the scene dataset.

Table r2: The reconstruction results for the movie dataset.

R6. The improvement in written.

We will revise those sentences mentioned by the reviewers in the revised manuscript.

• L38-39: Given the significant challenge of elucidating the correlation between neural activity and visual stimuli, the development of robust models for extracting vision-related latent representations is critically imperative.

• L48: External latent representations aim to capture stimulus-relevant components within visual neural activity, while internal latent representations reflect dynamic internal states influencing an animal's sensory capabilities.

• L117: High-quality latent representations, as compressed forms of original neural activity, must effectively reconstruct their input.

R7. Why use the contrastive loss?

Please refer to R2.

R8. Whether the resulting objective still maintains a valid ELBO?

The ELBO derivation hinges on introducing an approximate posterior for latent variables and utilizing variational inference. Contrastive learning, applied to external latent variables under the approximate posterior's input-conditioned constraint, acts as a regularizer without conflict. Its non-negative loss term also preserves the ELBO inequality.

R9. The variation in decoding performance across mice.

We have further investigated this phenomenon in Sections 4.5 and 4.6. Our findings suggest that it may stem from variations in the internal states of different mice at distinct stages, which, in turn, influence their capacity for encoding visual information.

R10. Limitations.

Without simultaneous internal state recordings, interpreting the functional roles of inferred internal latent variables is limited. Future work will analyze datasets incorporating both visual neural activity and state monitoring. Subsequently, we will extend this approach to other brain regions and experimental paradigms to investigate broader principles of biological coding mechanisms.

We thank the reviewer for the thoughtful and detailed comments. We will do our best to address the comments and provide detailed responses point by point, below.

R1. About new insights.

Our work offers two key insights into visual information processing mechanisms in the mouse visual cortex. First, the encoding capabilities for visual information are influenced by the animals' internal state, potentially accounting for variability in neural activity both across individuals and within the same individual over time. Second, under long-duration continuous visual stimuli, the neural activity of mice becomes more focused on encoding visual stimuli and is less influenced by their internal state. Further details on these points are elaborated in R4 and R5.

We acknowledge, however, that certain phrasing, such as claims of providing "new insights" into the intrinsic correlation between neural activity and visual stimuli, was imprecise. This should be revised to state that our model captures a greater degree of correlation between neural activity and visual stimuli. We will implement these necessary modifications in the revised manuscript.

R2. About the deterministic external latent variables.

We agree that factors beyond neuronal noise and internal brain states contribute to trial-to-trial variability. However, spontaneous neural activity, which is highly correlated with internal states, is recognized as a predominant source of cortical variability [1].

To comprehensively model variability involving external inputs, it would be necessary to incorporate auxiliary information, such as neural activity from other cortical regions (e.g., encoding other sensory inputs, behaviors, or memories). Given that this work analyzes neural activity within visual regions, an attempt to rashly model such externally driven variability (e.g., by pre-specifying the variance of external latent variables) risks both an ineffective representation of this variability and a detrimental impact on the modeling of internal state-related latent variables.

Therefore, to maintain our focus on variability arising from internal brain states, we have opted to simplify the modeling of external latent variables as deterministic values.

R3. About the inputs of decoder.

We admit that the choice of state variables as input to the decoder is primarily motivated by engineering considerations. First, following the architectural dependencies established in the VRNN, we introduce internal state variables in the decoder to further enrich the dynamic information, extending beyond what is captured by internal latent variables, thereby facilitating a more accurate reconstruction of neural activity. Second, since the external latent variables are deterministic functions of external state variables, introducing external state variables into the decoder yields less information gain. Therefore, we omitted this inclusion to mitigate computational complexity.

R4. How computing RSM is "facilitating the unraveling of information processing mechanisms in the visual cortex"?

External latent variables reflect neural activity components related to visual stimuli; consequently, their performance in decoding these stimuli cannot surpass the inherent validity of the visual information encoded by the neural activity itself. Therefore, observed differences in decoding scores across individual mice are more plausibly attributed to differences in the visual information encoded by their respective visual neural activity, rather than being a direct consequence of the external latent variables. More abstractly, these variations may be linked to the visual sensory capabilities of different mice at various times.

Building upon this premise, our analysis of the RSM reveals distinct patterns of internal state changes (as reflected by the internal latent variables) across mice, as detailed in lines 255-260 and further elucidated in R7 to Reviewer NU2a. We thus hypothesize that such inter-individual differences in sensory performance may be influenced by the internal state of each animal, potentially serving as a significant factor contributing to inter-individual variability. This investigation into the influence of internal states on sensory performance supports our broader aim of exploring and explaining the mechanisms of visual information processing across different individuals.

R5. About the Figure 4F.

As suggested by the reviewer, we compute the RSM for external latent variables and internal latent variables separately. As expected, the RSM for external latent variables consistently exhibits a clear diagonal structure, even for different trials. In contrast, the RSM for internal latent variables lacks such a structure, displaying no consistent organization across trials. These findings further support the hypothesized distinct roles of the external and internal latent variables in encoding different aspects of neural activity. We will incorporate these additional results into the revised manuscript.

R6. The interpretation of the decoding scores across areas.

First, our conclusions regarding hierarchical organization are primarily drawn from observed differences in decoding performance between primary and high-level cortical regions (lines 310-313). We wish to clarify that the discussion concerning the lowest decoding performance in VISrl (a multisensory brain region) serves merely as a potential explanation for this specific observation, rather than as an argument for hierarchical organization. Second, we offer the following explanation for the superior decoding performance observed in VISp. The neural datasets used in our experiments include categories with similar semantic content (e.g., various natural scenes featuring the same animals, or movie frames that are temporally distant but depict similar characters). Consequently, high-level cortical regions (VISpm and VISam), which primarily encode high-level visual features such as semantic information, may struggle to distinguish these subtle categorical differences. Conversely, the primary visual cortex (VISp), which encodes low-level visual features like texture and orientation, is better equipped to discriminate these fine-grained distinctions.

Nevertheless, we acknowledge that our initial conclusions regarding functional hierarchy were overly strong. The current results primarily suggest that the visually encoded information is heterogeneous across different mouse cortical regions. We will temper our conclusions accordingly in the revised manuscript. Furthermore, the interactive effects of hierarchical structure and multimodal integration represent a compelling avenue for future exploration, likely necessitating a synergistic analysis of neural activity across multiple sensory cortices.

R7. About the potential breakthrough.

It appears there may be a misunderstanding regarding the "potential breakthrough" referenced in lines 330-331. We wish to clarify that this phrase pertains to a promising direction for future research—specifically, the challenge of constructing stable, high-quality latent representations across diverse experimental conditions. It is not presented as a contribution of our work. In fact, the observed variability in our model's decoding performance across different mice is a recognized limitation of this work, underscoring the need for further investigation into robust modeling.

R8. What is the contribution of each of these losses?

Our approach achieves disentanglement between external and internal latent variables through multiple training objectives.

First, as described in lines 132–141, contrastive learning is applied to the external latent variables. This mechanism encourages external latent variables corresponding to sample pairs with small temporal offsets (i.e., under similar visual stimuli) to cluster together, thereby establishing an association between these variables and the visual input. Concurrently, a time-dependent prior distribution with spontaneity is used to guide the internal latent variables via the KL divergence. The application of these distinct loss functions to the two types of latent variables is instrumental in promoting their disentanglement.

Second, the swapping operation (details in R1 to Reviewer NU2a) further reinforces this disentanglement. This operation involves exchanging the external latent variables between positive sample pairs while preserving their internal latent variables. This manipulation specifically aims to further reduce the distance between the external latent variables of positive pairs. By doing so, it compels the model to isolate and strengthen the distinct roles of the external and internal latent variables, thereby enhancing their disentangled representation.

R9. The main differences between our model and alternative models.

We have provided a brief description of the differences in Section 4.3. Here, we elaborate on the main differences between our proposed model and alternatives:

• pi-VAE: Unlike our self-supervised approach, it incorporates additional supervised information (specifically, visual stimulus categories) for latent variable construction. Additionally, pi-VAE does not model the temporal dynamics inherent in neural activity.

• Swap-VAE: Similar to pi-VAE, Swap-VAE does not explicitly model the temporal information of neural activity, operating under an assumption of temporal independence.

• LFADS and CEBRA: A key distinguishing feature is that neither LFADS nor CEBRA performs a decomposition into external and internal latent variables. Moreover, the temporal relationships modeled by these approaches are not strictly causal. The former employs bidirectional RNNs to capture temporal relationships, while the latter uses temporal convolutional kernels to extract time information.

[1] Spontaneous behaviors drive multidimensional, brainwide activity. Carsen Stringer, et al. Science 2019.

We thank the reviewer for the constructive and detailed comments. We will do our best to address the comments and answer the questions. Below are our detailed responses.

R1. Architectural choice and swapping operation.

Our model leverages the architecture of sequential VAEs, with many component dependencies derived from the classic model VRNN. Specifically, in the encoder, the external latent variables $z^{(e)}$ and internal latent variables $z^{(i)}$ are conditioned on the input neural activity, as well as the respective state variables ($h^{(e)}$ and $h^{(i)}$) of two distinct RNNs. This conditioning facilitates the capture of temporal information in a causal order. The prior distribution of $z^{(i)}$ is conditioned solely on the state variables. In the decoder, the reconstructed neural activity is conditioned on the $z$ and $h^{(i)}$. The deliberate inclusion of $h^{(i)}$ here is intended to enrich the information of neural dynamics, providing a more precise reconstruction. Finally, regarding the update for the RNN state variables, $h^{(e)}$ depend solely on the input neural activity and their values from the previous time step. Differently, $h^{(i)}$ depend not only on the input neural activity and their previous time-step values, but also on the latent variables. This design choice reflects a consideration that animals' internal states are susceptible to the influence of current external visual stimuli.

The swapping operation for details is as follows: For a given sample and its corresponding positive sample, we exchange their external latent variables while keeping their internal latent variables unchanged. This process yields two new latent variables. We then decode neural firing rates using these new latent variables and compute the additional reconstruction loss. The intuition is this: our model is trained to reduce the distance between the external latent variables of positive sample pairs. Consequently, if the model has learned an effective external latent space, it should still accurately reconstruct neural activity even after this swap. This robust reconstruction, even with interchanged external latent variables from positive pairs, significantly enhances the training process by imposing strong regularization and encouraging more meaningful latent representations.

We will include the necessary clarifications in our revised version.

R2. How the model was applied to the non-temporal dataset.

For the model input of non-temporal datasets, each sample can be conceptualized as a sequence comprising a single time step (as described in Appendix F). In the context of contrastive learning for such data, the approach of defining positive pairs through temporal shifting is inapplicable. Instead, we adopt a strategy provided in CEBRA, selecting positive samples based on the distance of their respective labels (Appendix C), thereby implementing a form of supervised contrastive learning.

R3. The sentence on line 232.

We apologize for the imprecision in the original statement on line 232. Our intended meaning was that, while both our model and LFADS exhibited substantial performance degradation, the performance reduction observed in Swap-VAE and CEBRA was comparatively less severe (although CEBRA's decline was still notable). We will revise the sentence to clarify this distinction.

R4. Statistical tests for the sentence on line 272.

Following the reviewer's suggestions, we conducted a one-sided t-test comparing the performance of our model against baseline models, with the alternative hypothesis stating that our model's mean performance is superior. The results (t-statistic > 10; p-value < 1e-5) demonstrate that our model consistently exhibits significantly higher performance than all other models across nearly all mice, with the sole exception being the comparison with CEBRA in Mouse 3. The intent of the description on line 272 is to emphasize the differential magnitude of performance gains, specifically, that our model's advantage over LFADS is greater than that over the other models.

R5. The experiment for mouse cortical regions.

In the current results (Figure 5), performance variations across brain regions, under differing visual stimuli, exhibit similar patterns. This at least indicates heterogeneous and specialized visual encoding information across each brain region. However, we acknowledge that the initial conclusion regarding a sequential functional hierarchy was overly strong, and we will temper this statement in the revised version.

Following the reviewer's suggestion, we fine-tuned a model pre-trained on one region to another. This involves training only the first encoder layer and the last decoder layer, while freezing all other parameters. Tables r1 and r2 present these results, with rows indicating the pre-trained region and columns the fine-tuned region. We observe that even when pre-trained in different brain regions, the performance of the fine-tuned model shows no significant difference. The performance is lower than that of a model trained directly on the target region. This phenomenon suggests that neural activity across different brain regions may share similar latent spaces but encode distinct visual information.

Table r1: The fine-tuned results (%) for decoding natural scenes.

Table r2: The fine-tuned results (%) for decoding natural movie frames.

R6. Ablation study about the negative samples.

In response to the reviewer's suggestions, we investigate an alternative negative sampling strategy: drawing negative samples from the same sequence but with longer time offsets. Specifically, for the scene dataset, we use a 100 ms offset for negative samples (compared to 30 ms for positive samples). For the movie dataset, negative samples are generated with offsets of approximately 330 ms and 165 ms (versus an approximate 16 ms offset for positive samples). As detailed in Tables r3 and r4, this negative sampling approach results in a slight performance decline, with a larger decrease observed when the negative sample offset is smaller. These results indicate the feasibility of this alternative strategy, yet they also highlight a trade-off: selecting appropriate offsets for different datasets increases the computational cost associated with hyperparameter search. Our current methodology, which involves randomly selecting negative samples from the entire training dataset, emerges as a more generalized and efficient approach. Besides, we carefully review the code of CEBRA, confirming that it adopts the same random sampling strategy.

Table r3: The ablation results (%) for decoding natural scenes.

Table r4: The ablation results (%) for decoding natural movie frames.

R7. Do the authors expect that this internal state is represented by their internal latent variable?

In fact, the RSM of mouse 1 in Figure 3B has already shown regular changes in the internal state, i.e., the two redder blocks in the top-left and bottom-right corners, as discussed in lines 255–258 of the text. Upon checking the RSM of the internal latent variables, we observe that this pattern is even more pronounced. Subsequently, a quantitative assessment of the similarity between two blocks of trials, performed using the Pearson correlation coefficient, reveals that the similarity for Mouse 1 (0.248) is much higher than that for Mouse 2 (0.022). This finding provides further corroboration for our statement that visual sensory performance is influenced by internal states.

R8. Why does the decoding score concatenate vectors in the scenes, but average them in the movies?

The presentation duration of static scenes (250 ms) necessitates the preservation of dynamic changes within their latent representations for effective stimulus discrimination. Consequently, a concatenation operation is employed to retain maximal temporal information. In contrast, the frame rate of the movie is 30 Hz, meaning that each movie frame lasts only about 33 ms. Using an averaging operation for the latent representation of a single movie frame does not result in much information loss but can reduce the computational complexity of KNN algorithms.

R9. Limitations.

The main limitation is the inability to quantitatively analyze internal latent variables reflecting internal states. Moreover, as noted in lines 332-334, extending our method to other brain regions and species represents a promising avenue for future exploration.