Time-Dependent VAE for Building Latent Representations from Visual Neural Activity with Complex Dynamics

7. How exactly the train/test splits are done?

As described in lines 302-308, each visual stimulus is presented for multiple trials, and we divide all trials into training, validation, and test sets, rather than dividing them with respect to the time. In training, positive samples are selected only within the same trial and do not involve trials in the validation and test sets. Therefore, we ensure that none of the validation/test responses were shown during training.

[1] Tyler D Marks and Michael J Goard. Stimulus-dependent representational drift in primary visual cortex. Nature communications, 2021.

[2] Mohammad Bashiri, et al. A flow-based latent state generative model of neural population responses to natural images. NeurIPS, 2021.

[3] Jaivardhan Kapoor, et al. Latent Diffusion for Neural Spiking Data. arXiv, 2024.

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

1. Reproducubility might be an issue.

We have provided our code in the supplementary material, including the model, training and testing. We use a one-layer GRU implemented by pytorch for each latent space, and the dimensions of the hidden states of the GRU are the same as the dimensions of the latent variables. We will provide an explicit figure of the model architecture and more detailed descriptions in the revision.

2. Generalisability and robustness of the model is an issue.

This problem may be due to the fact that neural responses elicited by complex natural visual stimuli show large variability across subjects [1] and we have discussed this limitation in lines 533-537. However, as suggested by the reviewer umrd, we choose five other scenes that elicit the weakest responses, while the experiments in our work are conducted on five scenes that elicit the strongest responses. We conduct two experiments, one with these new scenes only and one with all ten scenes. As shown in Tables R1 and R2, our model performs best in most cases, demonstrating its robustness to some extent. Besides, we will evaluate our model on different visual stimuli (e.g. gratings) in the future.

Table R1: The decoding results (%) on five new scenes.

Table R2: The decoding results (%) on all ten scenes.

3. Some references are missing.

Thank you for your suggestion. These references are a good complement to our related work. The first [2] combined a deep network for stimulus-driven activity and a flow-based generative model to capture the stimulus-conditioned variability, providing new insights into latent state construction. The second [3] introduced structured state-space layers into the autoencoder to project discrete high-dimensional spiking data into continuous time-aligned latent states, simultaneously achieving inference of low-dimensional latent states and realistic conditional generation of neural spiking data. While the former uses visual stimuli to predict the distribution of neural population responses and the latter aims to reconstruct realistic neural activity, our work focuses on building latent representations from neural activity and decoding visual stimuli. We will add these discussions in the revision.

4. RNNs are often tricky to train.

We have found that our model is prone to overfitting, especially when dealing with long sequential data. We have largely mitigated this problem by adjusting the learning rate and applying the techniques of learning rate decay and early stopping, which enables our model to achieve the best performance.

5. Did you train a new separate model per mouse?

Since each mouse is implanted with Neuropixel probes separately in biological experiments, there is no guarantee that the neural sites recorded from each mouse are aligned, and the number of recorded neurons is different. Therefore, we train separate models for each mouse.

6. About Figure 3.

We think this is similarly related to the large variability in neural activity across subjects and trials (as discussed in Response 2). Specifically, on some trials, Mouse 1's neural activity is easily distinguishable when it is focusing its attention on the visual stimulus when , but on other trials, its neural activity may be dominated by its own state rather than driven by stimuli. As shown in Figure 7 in the Appendix, our model makes a good separation between stimulus-based clusters, especially for Mouse 2, suggesting that analysis of subject-to-subject and trial-to-trial variability is a meaningful direction for future investigation.

Thanks a lot for clarifications!

One more question about train-val-test split. Did I got correctly that responses for the repeats of the same stimuli can appear in both train and test? (If I show image A at time 10 and 20, the responses for the first presentation can go to test and for the second - to train?)

And as you train a model per animal, have you thought about some adapter layer to unify the dimensions and then keep the shared encoder-decoder structure, reported in the paper? What are the limitations of this approach? Otherwise, the latent states between mice are not necessary easily comparable due to the potential misalignment between the models. (See [1] for more misalignment details)

[1] Moschella, L., Maiorca, V., Fumero, M., Norelli, A., Locatello, F., & Rodolà, E. (2022). Relative representations enable zero-shot latent space communication. arXiv preprint arXiv:2209.15430.

8. The ablation studies on the RNNs and latent spaces.

As suggested by the reviewer, we conduct ablation studies on the RNNs and latent spaces. As shown in Tables R9 and R10, the non-recurrent version and the two versions without content/style latent space almost all perform worse than our original model, suggesting that the RNNs and the two latent spaces play a critical role.

Table R9: The ablation results (%) for the RNNs and latent spaces in natural scene decoding experiments.

Table R10: The ablation results (%) for the RNNs and latent spaces in natural movie frame decoding experiments.

Besides, Swap-VAE, as one of the best models we compare in our work, includes the swap operation and uses an alignment loss computed on positive sample pairs which plays a similar role to the contrastive loss. Its performance is slightly worse than or comparable to the non-recurrent version of our model, demonstrating the importance of RNNs in introducing time dependence.

5. Training on all the mice? Or training on some of them and testing on the left-out individuals?

In biological experiments, each mouse is implanted with Neuropixel probes separately, so there is no guarantee that the neural sites recorded from each mouse are aligned, and the number of recorded neurons is different (Table 6 in the Appendix). Therefore, we treat each mouse as a single dataset. As suggested by the reviewer, we sample neurons from all mice evenly to form a dataset (All Mice dataset) with a similar number of neurons to the single mouse dataset. The results in Tables R3 and R4 show that our model performs better than other models on this dataset, but the performance is overall lower than that on the single mouse dataset, suggesting variability in neural activity to the same visual stimuli across mice.

Table R3: The results (%) on the All Mice dataset in natural scene decoding experiments.

Table R4: The results (%) on the All Mice dataset in natural movie frame decoding experiments.

In addition, we sample the same number of neurons in each mouse. We then train our model on one mouse and test it on the other. As shown in Tables R5 and R6, for the same mouse in the test, there is a significant drop in performance for models trained on other mice. This may be due to the fact that it is hard to align the recorded neurons and there is a large variation in response patterns across mice.

Table R5: The results (%) of our model trained and tested on different mouse datasets in natural scene decoding experiments.

Table R6: The results (%) of our model trained and tested on different mouse datasets in natural movie frame decoding experiments.

6. The number of samples for each dataset.

We report the number of samples for each dataset in Table R7, where the size of the mouse dataset under natural scene stimuli is somewhat small. In training, we use a small learning rate and early stopping to ensure that the models are not overfitting.

Table R7: The number of samples for each dataset.

7. The number of parameters of each model.

The number of model parameters is roughly proportional to the number of input neurons, so here we only report the number of parameters for the Mouse 1 dataset (Table R8), and we will include other results in the Appendix. Due to incorporating the recurrent modules, our model has slightly more parameters than some models and is comparable to CEBRA.

Table R8: The number of model parameters for the Mouse 1 dataset.

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

1. About the novelty.

We agree that it is important to take advantage of the visual modality when dealing with visual neural data. Indeed, our work focuses on revealing the intrinsic correlation between neural activity and visual stimuli. As described in lines 243-246, we decode visual stimuli using low-dimensional latent representations transformed from neural activity and assess the clarity of latent temporal structures to explain neural dynamics under different types of visual stimuli. In Tables 2 and 3 of the main text, we report the decoding performance for natural scene/movie frame classification, showing that our model outperforms other baselines. In Figures 4 and 5, we exhibit the latent structures and analyzed their characters corresponding to visual stimuli. Ablation studies are also based on the decoding performance of visual stimuli.

Our lack of emphasis on the evaluation metric on neural datasets may have led to your misunderstanding. We will describe the evaluation in more detail in the revision.

2. About the writing.

We apologize for the unclear expression or grammatical flaws in our manuscript. We will modify the sentence in line 269 to "others even struggle to split four clusters" and the word in line 482 to "contrastive learning". What's more, we will re-check and improve our writing.

3. Why are the reconstruction and KL losses computed on the positive samples as well as the other samples?

These loss terms are motivated by the emphasis on the positive samples. As suggested by the reviewer, we perform an ablation study to remove these losses computed on the positive samples and report the results in Tables R1 and R2. The results show that these loss terms contribute less to performance, suggesting that the contrastive loss may be sufficient to emphasize the positive samples. We will consider removing these loss terms to speed up model training in the future.

Table R1: The ablation results (%) for losses on the positive samples in natural scene decoding experiments.

Table R2: The ablation results (%) for losses on the positive samples in natural movie frame decoding experiments.

4. What are these properties in the synthetic datasets?

As described in lines 259-261, the first property is the different categories of visual stimuli that visual neural activity corresponds to. Consequently, the non-temporal dataset is generated from several sets of labels (four clusters). The second property is the time dependence of neural activity. Taking this into account, the temporal dataset is constructed using the Lorenz system.

Dear Reviewer 35tP,

Thank you for your valuable time and constructive comments. As the discussion period will end soon, we hope to discuss whether your concerns have been addressed and would appreciate it if you could reconsider the rating. If you have any further questions, we would be happy to discuss them with you and do our best to address them.

Best,

The authors

We are very sorry that we didn't update our manuscript timely. We have included all the clarifications and additional results in our revision according to the reviewers' suggestions. Furthermore, we would like to address your remaining concerns.

1. About reconstructing neural activity and classifying stimuli.

First, we would like to emphasize that our work focuses on constructing low-dimensional high-quality latent representations from neural activity to reveal the intrinsic correlation between neural activity and visual stimuli. Second, VAE-based latent variable models (LVMs) aim to build latent representations, while reconstructing neural activity is not the research objective, but rather an objective of model training. Third, the use of latent representations for classification tasks is an important quantification of the quality of latent representations, which is a common approach in studies oriented to the motor cortex [1-3] and the visual cortex [4-5]. Last but not least, directly reconstructing visual stimuli is another research objective and usually requires neural activity with much larger datasets, such as fMRI [6]. In conclusion, our work follows the common analytical approach in the field, and our model achieves the highest performance on widely used evaluation methods.

2. A small version of TiDeSPL-VAE.

As suggested by the reviewer, we build a small version of our model (TiDeSPL-VAE-small) with fewer trainable parameters than Swap-VAE. Since the introduction of the recurrent module inevitably increases the number of parameters, we reduce the parameters of the encoder and decoder to achieve this (see Appendix E of the revision for the number of parameters). As shown in Tables R1 and R2, the small version of our model still outperforms Swap-VAE in most cases, suggesting that the performance improvement mainly stems from the design of our model rather than the larger number of parameters. We have included these results in Table 2 and 3 of our revised manuscript.

Table R1: The decoding results (%) in natural scene decoding experiments.

Table R2: The decoding results (%) in natural movie frame decoding experiments.

[1] Chethan Pandarinath, et al. Inferring single-trial neural population dynamics using sequential auto-encoders. Nature methods, 2018.

[2] Ding Zhou and Xue-Xin Wei. Learning identifiable and interpretable latent models of highdimensional neural activity using pi-vae. NeurIPS, 2020.

[3] Ran Liu, et al. Drop, swap, and generate: A self-supervised approach for generating neural activity. NeurIPS, 2021.

[4] Steffen Schneider, Jin Hwa Lee, and MackenzieWeygandt Mathis. Learnable latent embeddings for joint behavioral and neural analysis. Nature, 2023.

[5] Kendrick N Kay, et al. Identifying natural images from human brain activity. Nature, 2008.

[6] Jiaxuan Chen, et al. Rethinking Visual Reconstruction: Experience-Based Content Completion Guided by Visual Cues. ICML, 2023.

Thank you for recognizing our work. More importantly, we sincerely appreciate your valuable comments and suggestions that have helped us improve our work.

7. The scientific rationale for the design choice.

The overall design of our model is not an arbitrary overlay of these advanced technologies, but is scientifically motivated. First of all, we have given a conceptual interpretation to the two latent variables in the case of visual neural activity analysis. As described in lines 64-67, content latent variables correspond to the component of neural activity driven by the current visual stimulus, while style latent variables correspond to the neural dynamics influenced by the internal state of the organism. Based on the interpretation of the two latent variables, the overall objective is to construct high-quality low-dimensional latent variables. Using VAE and the contrastive loss both serve this objective. We clarify each term of our loss function point by point below.

• The reconstruction loss: The basic requirement for high-quality latent variables is to effectively reconstruct the input spike data.

• The constrastive loss: As we assume that content latent variables correspond to the component of neural activity driven by the current visual stimuli, we use self-supervised contrastive learning to make them more relevant to visual stimuli. Specifically, the positive sample is offset by several time steps from a given sample and both correspond to the same or similar visual stimuli. Consequently, the content latent variables of the positive pairs are required to be closer, which is achieved by the contrastive loss. The swap operation helps to further bring the content latent variables of the positive sample pairs closer together.

• The KL divergence: Style latent variables correspond to the neural dynamics influenced by the internal state of the organism, which contains a lot of intrinsic noise and variability. Therefore, we build them as random values from a parameterized distribution and apply a time-dependent prior distribution to guide them.

8. Your method is more complicated than those you benchmark, and your ablation studies hint at this as well.

As discussed in the Responses 4 and 7, the design of our model is complex but has a clear conceptual interpretation and scientific motivation. The number of trainable parameters in our model is slightly higher than some models and lower than CEBRA. In Table 4, the decrease in model performance when conducting ablation studies on the loss function (which is larger only for decoding natural movie in Mouse 2) exactly demonstrates the validity of our scientifically reasonable design, rather than showing that complexity drives performance. What's more, in the ablation studies on the model structure, we have provided results of replacing GRU with vanilla RNN (simple MLP), showing that our model still outperforms baselines in most cases.

5. About the benchmark datasets.

First, we have not only benchmarked on natural scenes but also on the natural movie. The decoding score metric follows that proposed by CEBRA (line 431). Second, as described in lines 259-264, the synthetic data follows some previous work. The non-temporal dataset was used in Swap-VAE and the temporal dataset was used in LFADS.

However, we agree that we should test our model on more existing benchmarks, as suggested by the reviewer. First, we evaluate our model on the Allen Brain Natural Movie 1 presented in CEBRA. In contrast to our use of the neural activity of each mouse independently as a dataset, CEBRA sampled neurons from the same brain regions in all mice to form the dataset. The experiment setup and model training fully follow the demo provided by CEBRA and the description in their paper. As shown in Table R2, our model outperforms CEBRA on datasets with different numbers of neurons sampled. However, we observe that our model's advantage decreases with more input neurons, which may be due to the fact that CEBRA's encoder has more layers than ours (we have only two), resulting in a lack of feature encoding ability in the neuron dimension of our model. Besides, these results are obtained in the case of a sample with 4 time steps (as presented in the CEBRA demo). We find that CEBRA reaches a performance of more than 90% in the case of a sample with 40 time steps (Figure 5c in their paper). After training with a sample of 40 time steps, our model achieves a performance of 94.62±0.48.

Table R2: The decoding scores (%) on the Allen Brain Natural Movie 1 (V1 data) presented in CEBRA.

We further evaluate our model on the synthetic dataset presented in CEBRA. The results show that our model (92.98±0.02) also outperforms CEBRA (90.60±0.04). What's more, although the design of our model is motivated by the properties of visual neural activity, we evaluate it on motor data presented in CEBRA. We compared our model with three CEBRA models (CEBRA-Position, CEBRA-Target and CEBRA-Time) which use positional, target and time information as auxiliary variables to select positive and negative samples, respectively. As shown in Table R3, our model outperforms CEBRA-Time, which also uses only time information like ours, and CEBRA-Target, which uses target information, in decoding position. However, our model performs worse in decoding direction.

Table R3: The decoding scores (%) on the motor data presented in CEBRA.

Decoding motor behavior information with latent variable models has been long and widely explored. However, there are fewer studies using latent variable models to decode visual stimuli, and CEBRA made a great step. Given that the task of the Sensorium Challenge is to predict neural activity using visual stimuli, we believe that our work can make a contribution to the benchmark of the task of decoding visual stimuli with LVMs.

6. The performance of the separate embeddings.

In Tables 5 and 8 of the main text, we have extracted the content and style latent representations from the full representations and evaluated their decoding performance respectively. The results show that the content latent representations perform better than the style latent representations, which is consistent with our conceptual interpretation of the two, i.e. that the content representations are more relevant to visual stimuli.

4. About the complexity.

In terms of model structure and loss function design, we admit that our model is more complex. However, each component has a clear conceptual interpretation and motivation (as described in Section 3). For example, we choose GRU to address the lack of long-time dependence inherent in RNNs, which is also used in LFADS. In addition, we have performed ablation studies replacing GRU with vanilla RNN (Tables 4 and 7 in the main text), and the results show that even with RNN, our model still achieves the highest performance in most cases. For the hyperparameters, we have searched and adjusted them for all models to ensure that their losses converge after training and they achieve the best performance.

In terms of model trainable parameters, we report the number of parameters for the Mouse 1 dataset for example (Table R1), as the number of model parameters is roughly proportional to the number of input neurons. Due to incorporating the recurrent modules, our model has slightly more parameters than some models and is comparable to CEBRA. In terms of running time, since our model only uses current and past information, it can only process neural data serially, and the running time of our model is indeed longer compared to models without time dependence and CEBRA, which uses time convolution that can be operated in parallel.

Table R1: The number of model parameters for the Mouse 1 dataset.

We thank the reviewer for the clear and pointed comments. We will try our best to address the reviewer's concerns and answer the questions in the following.

1. About the Motivation.

We would like to elaborate on some of the statements and clarify our motivation. First, what is natural reality and antecedent time dependence? Neural activity at the current time evolves only from past neural activity but does not depend on future neural activity. Therefore, modeling time dependence and constructing latent representations in a natural way should use only current and past neural activity.

Based on this motivation, we state that LFADS and CEBRA treat temporal relationships unnaturally since they both use future information to construct latent representations of the current time. As discussed in lines 251-254, although LFADS uses dynamical systems at the generator/decoder, its encoder uses bidirectional RNNs. Consequently, its decoder constructs latent variables at all times using future information, and even includes temporal relationships that evolve from the future to the past. As for CEBRA, it takes the surrounding time points centered on the current point and uses time convolution to encode temporal features, resulting in the inclusion of future information and the absence of chronological relationships.

In contrast, our work uses unidirectional RNNs to construct latent representations that depend only on current and past neural activity, which is an equally important contribution as our design for splitting latent variables. Therefore, we state that our model builds latent representations progressively along the chronological order. However, we agree that realistic motor behavior encoding-decoding is an unsolved problem and we will check the grammar in our manuscript.

2. About the related works.

Thank you for your suggestion. We will include discussions about task-relevant and task-irrelevant models to our related work. Besides, we want to address that we have benchmarked several task/behavior-relevant models, such as pi-VAE and CEBRA, as well as task-irrelevant models, such as LFADS and Swap-VAE. Unfortunately, our work focuses on spiking activity recorded from the mouse visual cortex under natural scene/movie stimuli, thus we did not notice the study in [2], which provides an accurate and efficient method for learning multiscale dynamical models for multimodal spike-LFP recordings of a monkey performing a reach-and-grasp movement task. (BTW, the link to ref [1] can not be accessed). We will evaluate this model in our experiments in the future.

3. About the training settings of the other models.

As described in lines 318-319 and Appendix E, we have provided detailed experiment setups for baseline models. For the latent variable dimension, we choose 128 for fair comparisons, because it was used in Swap-VAE and CEBRA and produced the highest task performance. We have also conducted ablation studies on the dimension (Figure 5), showing that the model performance saturates gradually as the dimension of latent variables increases and suggesting that it is sufficient to choose a dimension at reasonable intervals. For the training iteration, we adjust it based on the size of the dataset and ensure that the losses for our model and baseline models all converge.

The input time sequence processing of LFADS and CEBRA has been presented in Appendix E. In each experiment, the length of their input samples is consistent with our model, 5 time points for natural scenes (lines 313-314) and 4 time points for natural movie (lines 425-426). Besides, we have also provided the offset time, ±3 for natural scenes (lines 314-315) and ±2 for natural movie (lines 425-426).

Dear Reviewer XzNA,

Thank you for your valuable time and constructive comments. As the discussion period will end soon, we hope to discuss whether your concerns have been addressed and would appreciate it if you could reconsider the rating. If you have any further questions, we would be happy to discuss them with you and do our best to address them.

Best,

The authors

My final comment is that I re-read the PDF at its current state.

I still believe the introduction and related works need improvement.

I also still believe benchmarking on the NLB would make a lot of sense and insure fair comparisons.

The natural movie 1 decoding table doesn’t show cebra-behavior, which likely outperforms your method (and is what those authors used).

The overall novelty of the work is minimal as it’s effectively a time extension to swapVAE, as other reviewers have pointed out.

I also disagree with your “natural” statements - to expand, you use recurrency (RNN) - you then also use “future and past” information. And people use RNNs as models of the brain, especially for HC and M1 because it’s doing forward planning, meaning neural activity will not be time aligned with many internal brain states, and forward trajectories are often simulated. The HC (hippocampus) literature if full of examples of forward “preplay” and “replay” of simulated future behaviors. Thus, arguing that generalists algorithms like LFADS and CEBRA are unnatural just doesn’t make sense to me and hurts your presentation.

I will not be raising my score. I encourage you to reformulate your presentation, fairly show results, and resubmit. I think there is value in the work, but it’s not well presented in the current form.

Thanks for your feedback.

In fact, CEBRA only used "offset1-model" in the experiments of synthetic data which has no temporal relationships. For all neural datasets, CEBRA used "offset10-model" or "offset40-model" (as described on page 13 of their paper). As shown in https://github.com/AdaptiveMotorControlLab/CEBRA/blob/main/cebra/models/model.py, the former takes the surrounding 10 time points centered on the current point and the latter takes the surrounding 40 time points. Therefore, we think our description is correct. Besides, CEBRA uses time convolution to encode temporal features, which captures the causal relationships among all time points without considering the chronological relationships. Although time convolution is a good method for extracting temporal relationships, it is indeed less natural compared to our model.

Thanks for your pointed feedback.

1. We will improve our introduction and related work.

2. Although our work focuses on visual brain regions and the NLB is based on motor brain regions, we will test our model on the NLB in the future. More importantly, we believe that our existing experimental comparisons are fair enough, since for the other benchmark models we have ensured that the training losses converge and their optimal performance is obtained through a grid search of the hyperparameters.

3. Indeed, the results we report is CEBRA-Behavior, since we directly follow their demo.

4. We not only novelly introduce time dependence, but also make a new conceptual interpretation for the split latent representations in the case of visual neural activity analysis. For fair comparisons, we have built a small version of our model with fewer trainable parameters than Swap-VAE. The results demonstrate the superiority of our model over Swap-VAE.

5. First, although we use RNNs, they do not contain future information, but only past and current information, as shown by all the formulas in our manuscript. Second, we agree that neural activity may not be time aligned with brain states or behaviors. However, in terms of the time dependence within neural activity itself, "forward planning" and "forward preplay" just show that the brain uses past information to predict future information, rather than directly using future information, which is not the same thing. "replay" also uses the information from the past. Therefore, the "natural" statements point at both the time dependence and non-use of future information.

In conclusion, we will improve our manuscript, but we believe that we have made fair and sufficient comparisons and our statements about "naturally incorporating time dependence" are correct.

9. About the performance difference.

The large performance difference may be due to the fact that neural responses elicited by complex natural visual stimuli show large variability across subjects [5]. It is difficult for latent variable models to consistently construct high-quality stimulus-relevant latent representations across conditions and this requires further exploration. We have discussed this limitation in lines 533-537.

Besides, we think you may have a misunderstanding about Table 4. We didn't remove any mice in the ablation studies. The four columns of results are for Mouse 1 and Mouse 2 under natural scenes and natural movie respectively. The results for the remaining three mice are shown in Table 7 in the Appendix.

10. About the mouse visual hierarchy.

Our work focuses on revealing the intrinsic correlation between neural activity and visual stimuli and puts neurons from all brain regions together as inputs. Therefore, we could not analyze the hierarchy of the mouse visual cortex from the existing results. However, this is a meaningful and constructive suggestion. In the future, we will try to extract latent variables from each brain region separately, and the transformation relationship of latent variables between different brain regions may be able to reveal the mouse visual hierarchy.

[1] Chethan Pandarinath, et al. Inferring single-trial neural population dynamics using sequential auto-encoders. Nature methods, 2018.

[2] Ding Zhou and Xue-Xin Wei. Learning identifiable and interpretable latent models of highdimensional neural activity using pi-vae. NeurIPS, 2020.

[3] Ran Liu, et al. Drop, swap, and generate: A self-supervised approach for generating neural activity. NeurIPS, 2021.

[4] Steffen Schneider, Jin Hwa Lee, and MackenzieWeygandt Mathis. Learnable latent embeddings for joint behavioral and neural analysis. Nature, 2023.

[5] Tyler D Marks and Michael J Goard. Stimulus-dependent representational drift in primary visual cortex. Nature communications, 2021.

4. How well are the spikes reconstructed by the autoencoder?

Since the output of models is the reconstruction of the firing rate, we compute the root mean square error (RMSE) between the PSTH of the real spike data and the reconstructed firing rate of each model over time as a metric to evaluate the reconstruction performance, which has been used in related work [2-3]. As shown in Tables R3 and R4, our model performs slightly worse than Swap-VAE but better than LFADS and pi-VAE, suggesting that our model is effective in reconstructing neural activity.

Table R3: The reconstruction results (%) on mouse neural dataset under natural scene stimuli.

Table R4: The reconstruction results (%) on mouse neural dataset under natural movie stimuli.

5. Have you considered using only the contrastive objective, i.e. without the decoder and VAE loss?

As described in Response 2, all the loss terms are based on our conceptual interpretation of the latent variables in relation to visual neural activity, and all play a crucial role in the construction of high-quality latent variables. Indeed, one of the models used for comparison in our work, CEBRA, is the model that uses only the encoder and the contrastive loss and our model outperforms it in almost all experiments.

6. About the natural scenes.

The five scenes that elicit the strongest average responses are eagle, lion, cat, bird, and salamander. These are all animals, some of which are natural enemies of mice, so it makes sense that they would elicit the strongest responses. As suggested by the reviewer, we choose five other scenes that elicit the weakest responses. They are forests, withered leaves, mountains, geese, and reeds, which also makes sense. We conduct two experiments, one with these new scenes only and one with all ten scenes. As shown in Tables R5 and R6, our model performs best in most cases, demonstrating its robustness.

Table R5: The decoding results (%) on five new scenes.

Table R6: The decoding results (%) on all ten scenes.

Compared to the experiments on the scenes that elicit the strongest responses, the performance on the scenes that elicit the weakest responses is lower overall, suggesting that it is more difficult for models to build stimulus-relevant latent variables from neural activity with low signal-to-noise.

We will provide some examples of natural scenes in the revision.

7. About the natural movie.

The natural video is a 30-second clip extracted from a movie. Most of the film consists of perspective changes of background objects (streets, buildings, parked cars, etc.) with a small amount of character movement. There are few scenes of fast motion. The criterion for correct decoding follows the previous publication, CEBRA [4], as described in line 431 of our manuscript.

We will provide the natural movie in the supplementary material.

8. About the tSNE visualizations.

We try different hyperparameters (perplexity in the range [5, 50] and early_exaggeration in the range [12, 24]) and the properties explained in our work of embedding are stable. Besides, this visualization method has been used in related work [3-4] and can effectively cluster similar samples. Therefore, we believe that it is safe to interpret these nonlinear embeddings. We will include some visualization with different hyperparameters in the Appendix of the revision.

2. The clarifications of objective function.

The overall objective is to construct high-quality low-dimensional latent variables based on our interpretation of the two latent states (see Response 1). The VAE style ELBO derived from the likelihood of spike data and the contrastive loss both serve this objective rather than being considered separately. We clarify each term of our loss function point by point below.

• The reconstruction loss: The basic requirement for high-quality latent variables is to effectively reconstruct the input spike data. Since directly reconstructing spike data poses difficulties for model training, it is commonly assumed that the input data follow a Poisson distribution, and the output of models is the reconstruction of the firing rate of the Poisson distribution, as described in line 134. This approach and formulation has been widely used in related work [1-3].

• The constrastive loss: As we assume that content latent states correspond to the component of neural activity driven by the current visual stimuli, we use self-supervised contrastive learning to make them more relevant to visual stimuli. Specifically, the positive sample is offset by several time steps from a given sample and both correspond to the same or similar visual stimuli. Consequently, the content latent states of the positive pairs are required to be closer, which is achieved by the contrastive loss. The offset is less than the length of the sample to ensure that the positive sample pairs overlap, providing a strong time constraint. As for the sentence in line 231, we use formulas for clear explanation. For the latent variables of a given sample $z=[z^{(c)},z^{(s)}]$ and the positive sample $z\_{pos}=[z^{(c)}\_{pos},z^{(s)}\_{pos}]$, we swap their content latent states and keep the style latent states, resulting in $\hat{z}=[z^{(c)}\_{pos},z^{(s)}]$ and $\hat{z}\_{pos}=[z^{(c)},z^{(s)}\_{pos}]$. The swapped latent variables are used to compute new reconstructed firing rates $\hat{r}$ and $\hat{r}\_{pos}$, and the additional reconstruction loss $L\_P(x,\hat{r})+L\_P(x_{pos},\hat{r}\_{pos})$, which help to further bring the content latent states of the positive sample pairs closer together.

• The KL divergence: Style latent states correspond to the neural dynamics influenced by the internal state of the organism, which contains a lot of intrinsic noise and variability. Therefore, we build them as random values from a parameterized distribution and apply a time-dependent prior distribution to guide them. To avoid excessive fluctuations over time, we use the L2 norm of the mean and log-variance of the prior distribution (assumed to be Gaussian) as a regularization loss term.

The ablation results in Section 4.4 demonstrate that these loss terms contribute to the construction of high-quality latent variables. We will include these clarifications of our objective function in the revision and make it more readable.

3. The results are difficult to interpret because you compare methods that are all quite sophisticated.

In the experiments on the mouse visual neural dataset, we have provided the decoding scores of a simple linear method, PCA (Tables 2 and 3 in the main text). We apply PCA to neural activity and obtain embeddings with the same dimensions as latent variables of network models for the decoding experiments. We have also made comparisons with vanilla VAE ($\beta$-VAE) as a simple network model. In addition, we present the results of PCA on the two synthetic datasets here (88.71±0.01 for the synthetic non-temporal dataset and 20.76±0.08 for the synthetic temporal dataset).

We are sorry for the unclear description of data selection in line 298. In the analysis, we actually selected 5 mice (5 sessions) that have the highest number of recorded neurons, and these recorded neurons are evenly spread across six visual areas. We will revise the description in the revision.

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

1. The justification for the design of factorized latent states.

The design of two latent states in the previous publication is motivated by specific motor behaviors. In our work, we introduce time dependence into the two latent states by recurrent modules to improve their construction and have given them a new conceptual interpretation in the case of visual neural activity analysis. As described in lines 64-67, content latent states correspond to the component of neural activity driven by the current visual stimulus, while style latent states correspond to the neural dynamics influenced by the internal state of the organism.

In Tables 5 and 8 of the main text, we have extracted the content and style latent states from the full states and evaluated their decoding performance respectively. In addition, we train two models without each of the latent states as suggested by the reviewer, and report the results in Tables R1 and R2. These results suggest that the factorized latent states are important and that the content latent states perform better than the style latent states, which is consistent with our conceptual interpretation of the two, i.e. that the content states are more relevant to visual stimuli.

Table R1: The ablation results (%) for the latent states in natural scene decoding experiments.

Table R2: The ablation results (%) for the latent states in natural movie frame decoding experiments.

Dear Reviewer umrd,

Thank you for your valuable time and constructive comments. As the discussion period will end soon, we hope to discuss whether your concerns have been addressed and would appreciate it if you could reconsider the rating. If you have any further questions, we would be happy to discuss them with you and do our best to address them.

Best,

The authors

Dear Reviewer umrd,

Thank you for your valuable comments and insightful suggestions. We have updated our manuscript to include the clarifications and additional results. As the extended discussion period draws to a close, we hope to discuss whether your concerns have been addressed and would be very grateful if you could reconsider the rating of our work. If you have any further suggestions or comments, we would be delighted to address them.

Best,

The authors

Thanks for your pointed feedback.

1. The results of Tables 5 and 10 in our manuscript show that the performance of content latent representations is significantly higher than that of style latent representations, suggesting that it is the content latent representations that mainly contribute to decoding. The passable decoding performance of style latent representations may be due to the fact that visual stimuli necessarily affect the organism’s internal state. These results do not run counter to our interpretations of the two latent representations. In future work, we will evaluate our model on other visual neural datasets that include behavior data and validate that style latent representations are more relevant to internal states.

2. We agree that the contrastive loss that competes with the VAE loss may result in worse reconstruction performance. However, the main goal of our work is to construct high-quality latent representations and our model has achieved it. We will explore how to address the limitation of reconstruction to aid our model in being good at both.

3. As we have shown that the model without the contrastive loss (applied to content latent variables) performs worse (Tables 4 and 9), we perform another ablation study of a model without the KL loss (applied to style latent variables). As shown in Tables R1 and R2 (we report results of two mice due to time constraints), the results demonstrate that the KL loss also contributes to decoding performance. We will trade off the two terms and investigate their effect in future work.

Table R1: The ablation results (%) for the latent states in natural scene decoding experiments.

Table R2: The ablation results (%) for the latent states in natural movie frame decoding experiments.