Disentangling 3D Animal Pose Dynamics with Scrubbed Conditional Latent Variables

We thank the reviewer for their thoughtful insights. We are glad the reviewer recognizes the importance of our research directions in neuroscience. Some of the comments in this review will take some time to address. However, we would like to submit this initial rebuttal to initiate discussion. By the end of the author-reviewer discussion period, all comments will be fully addressed. We thank the reviewer for helping us improve our paper.

Key claims lack adequate empirical or theoretical support. For example, L52-53…

We understand that the language was confusing here. The evidence of the claim in L52-53 is shown in the results in Figure 2A-D. In these panels, we show that C-VAE latent variables are not invariant to the conditioned factor and contribute significantly to the latent structure. We will clarify this language.

Similarly, Line 60 lacks evidence...

To clarify, we assert the opposite: “certain nuisance variables, like speed, are not always completely independent of meaningful behavioral representations”. Our results in Fig 3F-I demonstrate this. When complete independence is assumed when not the case, as in SC-VAE-MI, performance decreases compared to models with more relaxed assumptions (SC-VAE-MALS). Furthermore, we did not intend to argue that speed is always a nuisance variable. However, in many cases, variables like speed can lead to over-segmentation in an action space and less interpretable representations (see [1]).

Several statements are overly vague …

We thank the reviewer for this suggestion. We will expand our Appendix B to include more justifications behind our parameter choices.

The model is evaluated against only a single baseline, despite references to multiple adversarial models…

In case of misunderstanding, we would like to highlight that the SC-VAE-GR model in our paper is a baseline using an existing adversarial method [2]. Furthermore, we intended one of our primary contributions to be defining explicit functional forms for the scrubbers. This bypasses the need to train and tune separate adversarial neural networks (ANN) as existing methods do. In Fig 3, we show the benefits of using explicit functions instead of nondescript ANNs for scrubbing.

Line 65 refers to "weak supervision"...

We understand that the term weak supervision can be vague. We based our definitions on those in papers like [3] in which our approach falls into the category of “restricted labeling”. A fully supervised approach would use action class labels during training. However, we are open to changing our language to something like “partial supervision” to make it clear that we are supervising with a subset of features.

The comparison between SC-VAE and C-VAE is not entirely fair…

Indeed, many current methods of disentanglement in behavioral neuroscience do not apply direct optimization similar to C-VAE [4-6]. We introduced this comparison exactly to illustrate the advantage of doing so. As explained above, we also compared to a different baseline where there is direct optimization of latent variables (SC-VAE-GR).

[1] Costacurta, J. et al. Neurips (2022). [2] Ding, Z. et al. CVPR (2020). [3] Shu, R. et al. ICLR (2020) [4] Whiteway, M.R. et al. PLoS Comp Bio (2021). [5] Shi, C. et al. Neurips (2021) [6] Yi, D. et al. eLife (2023)

We thank the reviewer for their thorough and thoughtful comments. We are glad the reviewer found our work appropriate for the ICLR audience and recognized our efforts to increase the practicality of learning disentangled representations in neurobehavioral data. Some of the comments in this review will take some time to address. However, we would like to submit this initial rebuttal to initiate discussion. By the end of this discussion period, all comments will be fully addressed. We thank the reviewer for helping us improve our paper.

On Costacurta et al. 2022 …

Despite similar motivations between our work and Costacurta et al. 2022, we found some challenges that made comparison nontrivial. First, their method was developed for depth camera recordings, whereas we used 3D keypoint trajectories. The adaptation to another modality presents a nontrivial challenge (indeed there was an entire Nature Methods paper published on this [1]). Second, their method is a time series clustering method, whereas our goal was to learn a low-dimensional feature space useful for many downstream tasks beyond clustering. Furthermore, our methods are complementary, not mutually exclusive. In principle, the methods developed in Costacurta et al. 2022 can be used to cluster our low-dimensional representation similar to [2]. Finally, as you have already stated, our goal was to allow flexibility for disentangling kinematics beyond speed.

That being said, we agree that baselines are important. We would like to emphasize that we include three baselines that correspond to many methods in the literature: VAE, C-VAE, and SC-VAE-GR. In the context of behavioral disentanglement, VAE corresponds to [2]; C-VAE to [3-5]; SC-VAE-GR to [6].

Is there a separate model fit for each of the nuisance variables …

For the consistency of comparison in this paper, we trained models from scratch and separate models were fit for each of the nuisance variable scrubbers. We can include some preliminary results in the Appendix from models trained on scrubbing both heading direction and average speed which show more pronounced benefits due to scrubbing. It will be interesting in future work to explore scrubbing with inter-dependencies between nuisance variables in more detail. For example, heading direction and average speed are not completely independent in our dataset. This can be seen in (Fig. 3D, right) where there are four heading direction peaks within walking behaviors. This is because mice tend to walk along the four walls of the arena in our dataset.

The MMD analysis…

We will rewrite this section to make our motivation more clear. Previous works have used MMD to quantify distances in VAE latent spaces, a notable example in neuroscience is [7]. In our analyses, the MMD ratio is similar to a measured effect size between healthy and diseased conditions. It captures the size of the difference between behavioral distributions relative to the within-condition differences across animals. When the variability across animals decreases (when animal identity is scrubbed), we gain a more precise estimate of the disease effect. We chose to use MMD over alternative measures of distribution difference, e.g., Jensen-Shannon Divergence, because MMD captures differences without discretization.

We think our current MMD analysis is a simple proof of principle that illustrates a potential benefit of our approach while being rooted in a metric used by prior work in the field. Rather than affecting binary healthy vs. diseased classification, increased separation between these two conditions (as quantified by MMD) is better considered as a way to reveal subtler phenotypic differences that would otherwise be obscured by animal-animal variability, and to potentially increase sensitivity for detecting the effects of additional experimental manipulations (e.g., optogenetics, ablations, pharmaceuticals, or therapies). Our partial lesion analyses (Fig 4B) are an example of this.

…transgenic lines where a single subject will fall into one or the other category…

Here our identity scrubbing paradigm makes sense for our dataset in which we have recordings of every subject both before and after inducing PD symptoms. However, as you say in your described scenario with unpaired data, full identity scrubbing may be less appropriate. In that case, SC-VAE can be used in principle to scrub identity only within each condition to enhance the disease vs healthy effect size.

[1] Weinreb, C. et al. Nature Methods (2024) [2] Luxem, K. et al. Nature Comm (2022) [3] Whiteway, M.R. et al. PLoS Comp Bio (2021) [4] Shi, C. et al. Neurips (2021) [5] Yi, D. et al. eLife (2023) [6] Ding, Z. et al. CVPR (2020) [7] Goffinet, J. et al. eLife (2021)

We thank the reviewer for their consideration and suggestions. We are glad that the reviewer recognized our motivations for more interpretable and researcher-centric methods of deep behavioral analysis, as well as the applicability of some of our ideas to other domains.

Generalizability: ….

This is a very thoughtful point, and we agree that this is essential to consider. First, as the reviewer observed, works like Locatello et. al. 2019b have shown that unsupervised disentanglement is difficult, if not impossible, to guarantee. Much of the disentanglement literature has since moved towards introducing assumptions or weak supervision as we do (Shu et al. 2019; “Weakly supervised disentanglement with guarantees”). In our paper, our fully unsupervised baseline (VAE) consistently performs less favorably when compared to the SC-VAE models.

Second, we would like to clarify that our intention was not to completely replace unsupervised behavioral representation learning. Our goal was to introduce approaches that complement an exploratory analysis of complex behavioral datasets. Here our nuisance factors (speed, heading, and identity) are easy to access in practice for most datasets. In Fig. 4, we show that by reducing the structure of nuisance identity effects, scrubbing enhances the subtle structure about PD itself (lesion severity).

Figures …

We thank the reviewer for these suggestions for improving our figures. We agree that they will enhance our presentation and will implement these changes in the final paper.

In terms of pose estimation, I often hear about DeepLabCut…

Including DeepLabCut, there are now numerous 2D and 3D animal pose estimation tools available which have enabled a wide variety of neuroscience studies (see L31-33 of our paper). We use these pose tracks as the features which are input into our (SC-)VAE models. So the contribution of our work is downstream of DeepLabCut-style analysis (i.e., it interprets the output of tools like DeepLabCut).

In our contributed dataset here, we chose to use a pose estimation tool called DANNCE. For a more in-depth discussion of the differences, please see the DANNCE paper (Dunn et al. 2021).

We will incorporate these suggested discussions. We thank the reviewer for helping us improve our paper.

We thank the reviewer for their time and thoughtfulness. We are encouraged to see that the reviewer appreciated the impact and timeliness of our work for neuroscience and found our results to be thorough and compelling in this application.

Some more discussion of the Parkinsonian dataset would be appreciated….

We thank the reviewer for this suggestion. In case this was missed, some details are already included in Appendix B and C. However, we will expand these sections to include more details and clarifications about our datasets, model architectures, and hyperparameters. We will also include references to these sections in the main text.

For the PD results in Fig. 4, we used a larger dataset (36 animals) in which animals were recorded before and after PD induction (via 6-OHDA injection). Otherwise, the data were collected in the same manner as the other figures. For the results in Fig. 4, we implemented animal identity scrubbing as described in Sec 4.2 “Animal Identity.” We reused the model architectures and most of the hyperparameters from the previous sections. The exception was the weight $\lambda$ on the scrubbing loss, which was re-tuned on the larger dataset (i.e., to minimize the decodability of $\mathbf{v}$ from $\mathbf{z}$ as in Fig. 2A-C).

Why do the authors think that the linear scrubbing improved the conditioned sequence best …

This is a great question. Assuming that the reviewer is referencing the results in Table 1, we were not trying to suggest that linear scrubbers are definitively the best at this task. Note that the models used to produce the results in Table 1 and Fig. 3 had their hyperparameters tuned based on performance in Fig. 2A-C. It is possible that if tuned specifically for the motion synthesis task, the order of performance would change slightly.

Our intent was to assess the difference between models without adversarial disentanglement (C-VAE), models using adversarial neural networks (SC-VAE-GR), and models which do not require adversarial neural networks (SC-VAE-MALS, -MAQS, and -MI). For our purpose, the larger-scale differences between these three categories of approaches were more important than the relative ordering of -MALS, -MAQS, and -MI. We show that even the simplest linear scrubber (-MALS) performs favorably to the neural network scrubber (-GR) with the additional advantage of not requiring the specification of a second neural network and its hyperparameters. We will update our writing in Sec 4.3 to clarify these points.

If any additional questions or concerns arise, we will be happy to address them during this discussion period. In our revised paper, we will implement these changes and those suggested by other reviewers. Again, we thank the reviewer for helping us improve our paper.

We thank the reviewers once again for their helpful comments thus far. With the deadline approaching, please let us know if there are any unresolved concerns or questions that you would like us to address.