Exploring Behavior-Relevant and Disentangled Neural Dynamics with Generative Diffusion Models

Dear Reviewer HCYS,

Thank you for your recognition of the work and insightful questions. We provide clarifications and new results that we have generated to address your questions below.

This work primarily relies on a behavior LVM for neural activity disentanglement, which can be influenced by the newly introduced total correlation (TC) penalty term with 𝛽 weighted. The impact of this penalty should be considered through ablation studies.

We appreciate your attention to this aspect. The following are our ablation studies with various $\beta$ values. Each experiment condition is repeated with 5 runs, and their mean and variance are listed.

Table 1: Hyper-parameter $\beta$ Investigation on the SSp-Left region of Session-1.

Table 2: Hyper-parameter $\beta$ Investigation on the MOs-Left region of Session-1.

From the results, we observed that setting $\beta = 1$ or $\beta = 2$ resulted in lower disentanglement of the latent subspace. However, increasing $\beta = 8$ led to substantial degradation in neural reconstruction quality. Hence, we determined that $\beta = 4$ offers a balanced trade-off between disentanglement and reconstruction. These Tables will be added in the Appendix of the revised manuscript.

There is significant research combining behavioral labels with latent variables of neural activity [1,2], considering the effects of different disentanglement algorithms on behavior encoding would further enhance this study.

This is an insightful point. We would like to note that the listed two research works [1, 2] combining behavioral labels for neural latent variable learning do not focus on the disentanglement of the neural subspace. The total-correlation (TC) term proposed here is promising to be combined with these two approaches [1, 2] to achieve a more disentangled latent subspace. We added the TC term to their loss functions, and the results are reported in the following Tables 3 and 4. Each experiment condition is repeated with 5 runs, and their mean and variance are listed.

Table 3: Total-correlation term effect analyses on the SSp-Left region of Session-1.

Table 4: Total-correlation term effect analyses on the MOs-Left region of Session-1.

After incorporating the total-correlation term, we observe that the disentanglement performance of these baseline neural LVMs improves. These results will also be added to the Appendix of the revised manuscript. We would like to emphasize that the disentangled neural LVM module is not our primary contribution. Instead, the main modeling and scientific contribution of our work both lies in the video diffusion modeling (VDM) module of BeNeDiff, which interprets the neural dynamics of each disentangled latent factor in a generative manner. The VDM module is crucial for clearly interpreting the neural dynamics of each latent factor, demonstrating specificity to behaviors of interest (e.g., paw-x-axis movement, jaw movement) and aligning with ground-truth behavioral trajectories. Moreover, the VDM module can be generalized to all the aforementioned neural LVM baselines with the total-correlation term for the discovery of neural dynamics in the learned neural subspace. We can also include more generated video results in an anonymized link upon request.

Refs:

[1] Learning identifiable and interpretable latent models of high-dimensional neural activity using pi-VAE. (Zhou et al., 2020)

[2] Learnable latent embeddings for joint behavioural and neural analysis. (Schneider et al., 2023)

Dear Reviewer CLKJ,

Thank you for your valuable comments. We would like to make the following clarifications. Hopefully these will resolve most of your concerns, and they can be taken into account when deciding the final review score.

It is unclear to me if the behavior video generation part is trained together with the disentangled subspace learning. It seems true based on the figure, but didn't see how they are optimized together from text.

We would like to note that given the learned disentangled neural subspace and trajectories (from the neural latent variable model), the behavioral video (generation) diffusion model (VDM) module serves as an expressive ML tool for interpreting and visualizing complex factor-wise neural dynamics. Hence, the behavioral VDM is trained after the neural LVM in two steps, as we need high-quality neural latents for training the classifier (neural encoder) in VDM. This procedure is also in line with the conventions of previous studies [1] for scientific interpretation.

Regarding Figure 2, it describes the neural dynamics interpretation phase of BeNeDiff rather than the model training phase. In this phase, behavior videos are generated with guidance from the neural encoder. In the revised manuscript, we will also enhance Figure 2 for clearer visualization.

It is interesting to learn the disentangled neural space. Some visualization of the learned latent space is beneficial. Does the learned latent trajectory show some clustering property, or somehow organized, could you visualize it? It could be better the paper include more analyze on the learned latent trajectory and see if it could provide some scientific insight.

We thank the reviewer for this practical suggestion. In Figure 1(B) of the attached PDF, in each subfigure, we select two neural latent factors at a time and plot their 2D trajectories. Latent-factor-3 is informed by the "Paw-x" behavior-of-interest. We observe that the learned dimension axes are disentangled, but the behavior dynamics are not readable from the trajectories. Figure 1(C) depicts the temporal evolution of all six latent factors, which is also difficult to interpret in terms of the behavior dynamics encoded by the neural population. In Figure 1(D) of the attached PDF, the full neural latent trajectories of trials (dimension reduction by PCA to 3D) are also hard to directly interpret or identify any clustering structure.

To sum up, the learned neural subspace exhibits certain disentangled properties, but the visualization results are not intuitive or readable enough for experimentalists. Therefore, the VDM module in BeNeDiff is crucial for clearly interpreting the neural dynamics of each latent factor, which demonstrates specificity to behaviors of interest and aligns with ground-truth behavioral trajectories (e.g., Latent-factor-3 is aligned with "Paw-x" in Figure 1(E)). Please refer to Section 1.2 of the global response for detail.

Additional ablation study on different contribution of the loss function could be performed.

Please refer to Section 1.4 of the global response for a detailed discussion of this point.

(1) How training and testing data are split? Are they session wise or random? (2) How the method can be generalized to different sessions or animals?

(1) We train separate models for each session independently. Within each session, we use an 80/20 split for training and testing, employing cross-validation to ensure model robustness.

(2) This is an interesting point. The generalizability of each method is a vital consideration. For the neural data part, the difficulty arises from the varying number of observed neurons across different sessions. Additionally, the same brain region can exhibit varying encoding for the same behavior of interest across different animals. One possible solution to mitigate the issues is to add a linear probing layer to perform neural align [2, 3] from various sessions and animals into a unified subspace, then can be adpated to the BeNeDiff framework. Meanwhile, the VDM module can generalize well to multi-sessions and multi-animals settings by fitting a unified neural encoder for the aligned neural data.

(1) For session 3.2, 𝑐 is better to be described as it is first mentioned in eq.5. (2) Does the number of class is equals to feature dimension D and the same as the number of the behavior classes"?

(1) Thank you for this practical suggestion. We will move the description of class label $\mathbf{c}$ from line 147 to after Eq. (5).

(2) Yes. The number of class labels $\mathbf{c}$ is the same as the latent factor dimension $D$, as well as the number of behaviors of interest.

To get figure 4, which latent variable do you use, 𝑍^ or 𝑍?

We use each latent factor of $\mathbf{Z}$ for behavior decoding, and get the results in Figure 4.

For fig.6, and fig.7, could you show the difference of the ground truth frame with the same behavior type as comparison.

We apologize for the confusion and would like to clarify that the results in Figures 6 and 7 are all derived from analyzing the behavior-related neural dynamics of one target neural latent factor using two baseline latent manipulation methods and BeNeDiff. The ground truth frames from behavior trials in the dataset do not specifically activate one neural latent factor, so the frame differences between consecutive images result in entangled behaviors. Including them would not provide a fair comparison. We have, however, provided several ground-truth behavior trials in the supplementary materials for reference.

We look forward to further discussion, and are happy to answer any questions that may arise.

Refs:

[1] Cortical discovery using large scale generative models. (Luo et al., 2023)

[2] Stabilizing brain-computer interfaces through alignment of latent dynamics. (Karpowicz et al., 2022)

[3] Using adversarial networks to extend brain computer interface decoding accuracy over time. (Ma et al., 2023)

Dear Reviewer CLKJ,

We thank you for the valuable response. We agree that incorporating behavior labels can be helpful for learning a disentangled neural subspace. However, we note that our neural LVM module in BeNeDiff maintains a high neural data reconstruction rate (as shown in Table 1 of the manuscript), thereby preserving the neural dynamics in the latent trajectories.

It is true that we can train a probabilistic behavior classifier (a behavior decoder from neural data $\mathbf{X}$ to behavior labels $\mathbf{U}$) to approximate the conditional distribution $p(\mathbf{U} \mid \mathbf{X})$ and derive latent variables $\mathbf{Z}'$ from it. The $\mathbf{Z}'$ would also be disentangled. However, such $\mathbf{Z}'$ lacks physical or semantic meaning, as there is no reconstruction of the neural data, making it unable to interpret neural dynamics or provide any scientific insights. We further trained a diffusion model guided by $\mathbf{Z}'$ and presented the visualization results. We have sent the AC a comment linking to an anonymous repository containing these synthesized video results. In accordance with this year’s NeurIPS guidelines, please have the AC forward it to you.

We observe that while the synthesized videos by activating each latent factor does show some specificity to a particular behavior of interest, the overall video results are quite chaotic and do not align with the ground-truth behavioral trajectories. The behavior movements are overly drastic and contain many unnatural distortions compared to the videos generated by the neural latent trajectories in BeNeDiff. These videos will be included in the supplementary materials of the revised manuscript. Additionally, we have included the videos generated by BeNeDiff and the ground-truth behavior videos in the anonymous repository provided to the AC for comparison.

Looking ahead, our future research aims to generalize BeNeDiff to settings where the latent (feature) dimension is larger than the number of behavior classes. We hope that these responses and the above rebuttals address your concerns and will be considered when determining the final review score.

Dear Reviewer o3cJ,

Thank you for your detailed and constructive comments. We would like to make the following clarifications. Hopefully these will resolve most of your concerns, and they can be taken into account when deciding the final review score.

Each of the primary contributions to neural LVM type models (the use of total correlation and the diffusion modeling) are not sufficiently benchmarked against competing approaches.

We would like to clarify that the contribution of BeNeDiff is two-fold: (1) The use of the total correlation term in the behavior-informed neural latent variable model (LVM) to induce a disentangled neural latent subspace. While this module is not our main modeling focus, it is effective compared to baselines (please refer to Table 1 of the attached PDF for the results). (2) The main novelty and contribution of our paper is the video diffusion modeling (VDM) module, which serves as an interpretation tool for visualizing the neural dynamics of each learned disentangled latent factor. The generated behavioral videos of VDM conditioned on activating one learned latent factor over time demonstrate that each factor shows specificity to a ground-truth behavior of interest (e.g., paw-x-axis movement, jaw movement) and aligns with ground-truth behavioral trajectories. The qualitative performance comparison of this VDM module against baselines is presented in Figures 6 and 7 of the manuscript, with detailed analyses in Sections 3.2.1 and 5.3.

If the authors wanted to strengthen the case that their model is an important novel contribution for behaviorally-relevant latent disentangling, they could compare how total correlation VAE compares to other methods where behavioral labels are used to help disentangle the latents. One important comparison that I would be curious to see would be the pi-VAE [1].

Including the suggested pi-VAE [1], we have added the following three baselines for extensive comparisons. Please refer to Section 1.3 of the global response for the baseline descriptions and tables with results. We also note that the VDM module of BeNeDiff can generalize to interpret the neural dynamics of the latent factors learned by all the neural LVM baselines mentioned above.

The authors could also focus on how using video diffusion is a novel application in this setting. Again, some further discussion and emphasis as to how the diffusion methodology solves an important problem in neuroscience that perhaps existing methods cannot solve would strengthen this paper.

This is a highly valid point. We would like to emphasize that the scientific question in neuroscience we aim to answer in this paper is "how can we enable in-depth exploration of neural latent factors with video behavior recorded, revealing interpretable neural dynamics associated with behaviors". Therefore, the modeling goal is to develop a machine learning tool for visualizing neural dynamics encoded by each learned latent factor. This involves mapping the neural latent factors $\mathbf{Z}$ to behavior videos $\mathbf{Y}$, activating a single latent factor $\mathbf{z}^{(d)}$, and observing how the manipulation affects $\mathbf{Y}$. The induced changes in the videos reveal the dynamics encoded by $\mathbf{z}^{(d)}$.

Previous manipulation methods [2, 3] often change the target $\mathbf{z}^{(d)}$ while fixing non-target subspaces to arbitrary values. However, setting arbitrary values without knowing the true distributions of non-target subspaces can lead to unnatural distortions in generated videos, complicating the interpretation and visualization of genuine animal behavioral dynamics.

BeNeDiff proposes using a video diffusion model (VDM) module to explore disentangled neural dynamics in a generative manner. Specifically, the VDM module employs a neural encoder (Eq. (7) of the manuscript) to guide the generation of videos that maximize variance in the neural trajectory of target latent factor $\mathbf{z}^{(d)}$ while minimizing variance in other factors' trajectories. This approach ensures that the generated behavioral videos predominantly reflect the neural dynamics of the target latent factor. It maintains naturalness by avoiding assumptions about specific values in non-target subspaces, thereby preventing the generation of videos with unnatural distortion, which manifests the modeling-wise novelty of BeNeDiff. As shown in Figures 6 and 7 of the manuscript, for each latent factor, the VDM module provides interpretable quantifications of its neural dynamics with specificity to the behaviors of interest. These generated video results are also available in the 'BeNeDiff-Generated-Video' folder of the supplementary materials.

Though the authors highlight the application of the model to this specific dataset as a contribution, this alone would be more appropriate for a more experimentally-focused venue. For neurips, the focus should be on the model advancements and their applications.

We would like to politely disagree with the statement that the application of our proposed method is specific to this dataset. Please refer to Section 1.1 of the global response for details on how BeNeDiff can be generalized to multiple datasets. For the model advancements, please refer to the answer of the above point. We also note that the diffusion model has been employed as an interpretation tool in previous neuroscience works, resulting in publication [4] at NeurIPS.

The authors only minimally discussed the limits of the ability of their model.

Due to space limitations, please refer to Section 1.5 of the global response for detail.

We look forward to further discussion, and are happy to answer any questions that may arise.

Refs:

[1] pi-VAE. (Zhou et al., 2020)

[2] Partitioning variability in animal behavioral videos using ss-vaes. (Whiteway et al., 2021)

[3] Classifier-free diffusion guidance. (Ho et al., 2022)

[4] Cortical discovery using large scale generative models. (Luo et al., 2023)

Dear Reviewer o3cJ,

We sincerely appreciate your positive evaluation of our method's contribution and additional benchmarking. We apologize for the confusion in the Introduction section and will clarify the model's generalizability across datasets when discussing the major contributions in the revised manuscript. We thank you once again for the valuable feedback and suggestions.