Vector Quantized Representations for Efficient Hierarchical Delineation of Behavioral Repertoires

We thank the reviewer for appreciating the novelty of our proposed method and applications. For your comments on:

• in equation (3), can you elaborate more on the embedding and commitment terms? intuitively what do they mean and how are they derived?

Great question and we are happy to elaborate on this. Given the non-differentiability of nearest-neighbor lookup operation from the embedding space, the gradient is directly copied from decoder inputs z_q after quantization to the encoder outputs z(x) (straight-through estimate). A pseudo PyTorch code for this process is given by z_q = z + (z_q - z).detach()

In this way, the reconstruction loss (first term in equation) can only update the decoder and encoder, not embeddings in the codebook with no gradients passing through. We then included the second term, embedding loss, to optimize the embedding vectors toward the encoder outputs z(x) ia L2 loss. In practice, this loss was replaced by exponential moving average (EMA) to update the codebook embeddings as moving averages of encoder outputs, which provided better performance empirically.

The last term, commitment loss, can be intuitively understood as strategies constraining the volume of the embedding space and encouraging the encoder outputs to ‘commit’ to a specific embedding vector rather than growing arbitrarily, as not controlled by the reconstruction term alone. This also ensures that our straight-through gradient estimates remain concrete.

To summarize, the first reconstruction term optimizes the encoder and decoder; the second embedding term (implemented by EMA) optimizes the embedding vectors in the codebook; the last commitment term optimizes the encoder only. The commitment term is scaled by a hyperparameter alpha in the final objective.

• How robust is the method to variations in the number and granularity of discrete codes? In practice, how do you determine the number of codes in each level?

Thank you for the thoughtful comments. We have included new experiments in Table. 1 that examines the changes in downstream performance on motion synthesis with different code granularities. Setting the downsampling rate as 16 frames (320 ms) has yielded overall the best synthesis fidelity and diversity, which is consistent with previous works about the average duration of rodent behaviors (Wiltschko et al. 2015, Marshall et al. 2021). In addition, we have included experiments varying the number of top- and bottom- codes in Supp. Table. 4 in the updated manuscript (also in response to all reviewers).

The number of codes in each level are hyperparameters to tune, but we do deliberately set the number of codes in the top level to be smaller than that in the bottom level, similar to Singh et al. 2019, encouraging disentanglement in the hierarchy. Although it is challenging to form a systematic procedure for tuning the number of codes, as the plausible number of behavioral motifs is more data-driven, we monitored the changes in codebook perplexity (the fraction of inactive codes) during training and examined whether motion sequences derived from the learned motifs are plausible post training.

• What are the quantitative comparison results against the KPMS benchmark?

We agree with this comment and have included additional quantitative analyses that compare our method with KPMS (please refer to the overall response to all reviewers). We did not compare with KPMS on the downstream task of motion synthesis, as their model is not designated for such tasks.

We would like to point out that comprehensive, quantitative comparison between behavioral analysis methods has been challenging due to lack of consensus over behavior definitions and benchmarks. To the best of our knowledge, there exists no public continuously annotated rodent behavior datasets (notice the CalMS21 Dataset we have used in the paper is a 2D mouse dataset with annotations on dyadic social behaviors alone). Moreover, the inter-labeler behavior annotation agreement remains low even among well-trained human experts (Segalin et al. 2021) and these labels are inherently biased by human limited understanding of rodent behaviors.

References:

• Singh et al (2019). Finegan: Unsupervised hierarchical disentanglement for fine-grained object generation and discovery.
• Wiltschko et al (2015). Mapping sub-second structure in mouse behavior.
• Marshall et al. (2021). Continuous whole-body 3D kinematic recordings across the rodent behavioral repertoire.
• Segalinet al. (2021). The Mouse Action Recognition System (MARS) software pipeline for automated analysis of social behaviors in mice.

We appreciate the reviewer’s constructive feedback for better positioning the motivations and contributions of the proposed method, as well as strengthening the paper with more comprehensive quantitative comparisons. Regarding your comments on:

• Unfortunately, it’s not clear whether there is sufficient technical novelty for the ICLR community ...

Thank you for the constructive criticism, and we agree that it is important to ensure that our work is a good fit for ICLR and we believe it is. To your point about how the novelty for neuroscience suggests that a different venue could be a better fit:

• ICLR has always been a venue for papers on neuroscience applications (also explicitly listed as a subfield in the ICLR Call for Papers), for example, Eyjolfsdottir et al. (2017) for learning recurrent representations for animal behaviors, Li et al. (2022) for animal pose estimation and Molano-Mazon M et al. (2018) for neural activity synthesis, among many more. Papers specifically targeting improving and benchmarking animal behavioral analyses, as we do, have also appeared in other AI/ML conferences including ICML, for instance, Sun et al. (2023) that introduced a multi-species benchmark for learned animal behavioral representations, and NeurIPS, such as Costacurta et al. (2022), Shi et al. (2021) and Batty et al. (2019), to only mention a few. In particular, the methodology presented in Costacurta et al. is a modification to (Keypoint-)MoSeq, the method we have compared to in the paper.
• Our approach is broadly relevant to neuroscience, extending beyond the scope of behavioral neuroscience. The ultimate goal of neuroscience is to understand how the brain generates behavior, which is, after all, the brain’s principal output. Until recently, analyzing behavior was just not as tractable as analyzing neural activity, as we lacked methods for rigorously and efficiently quantifying the behavior itself. This has changed, and now behavioral analysis methods, and ultimately computational models that relate neural activity to behavior, are key areas of interest and development in computational neuroscience.

To your point about whether the architectural/learning improvements are relevant to the broader ICLR reader:

• First, please see our above points regarding neuroscience applications in ICLR; at a minimum, the neuro community at ICLR will appreciate this work.

• Second, despite our focus on applications to behavior and computational neuroscience, our approach is novel at learning the inherent hierarchy among motion primitives using multi-level encoding with a relatively simple, computationally economic framework. This approach could be applicable to either type of sequential data. The state-of-the-art performance on motion synthesis relative to recent methods in human motion synthesis also suggest the benefits for deploying such discrete representations in generative tasks. These tasks are not limited to forecasting or generation of animal movements, but also establishing mappings with other behavioral modalities, such as vocal activities in rodents (languages/speech for animals).

• The analysis could be more thorough; as it stands there is only one table of quantitative results, for motion synthesis, and the method is compared mostly to quite weak baselines (fully-connected MLP, GRU).

We appreciate and agree with the constructive feedback and have therefore added quantitative comparisons with Keypoint-MoSeq (please refer to the response to all reviewers). For the task of motion synthesis, we have included experiments with more recent and stronger baselines, for instance MotionDiffuse (Zhang et al. 2022), as well as internal comparisons on our proposed method with different codebook sizes and temporal granularities. We hope that your concerns have been addressed in this revision. We would also like to reiterate that motion synthesis is not the primary focus of this paper, but applications tangential to the analysis of animal behaviors.

Thank you very much for the valuable feedback that has helped us to improve the manuscript. Regarding the comments on:

• Missing/failed citation in first paragraph

Thank you for pointing this out and we apologize for the error. We have fixed it in the latest version.

• How is the quantization "codebook" initialized and updated? This is not clear to me.

Prior to training, the codebook is initialized with zeros. During training, as we had mentioned on Method page 3 in the first version of manuscript, the codebook is updated with exponential moving average (EMA), where each embedding vector is updated as the moving average of encoder outputs. This protocol was first described in the original Van Den Oord et al. VQ-VAE paper and further investigated by Zhang et al. 2023.

• An ablation showing the benefit of using the proposed quantization would be helpful.

We agree, thank you for the suggestion. While quantization has the clear benefit of directly inferring the discrete behavioral representations currently used by the computational neuroscience field, we agree that it is important to objectively quantify these benefits. Our previous motion synthesis results actually spoke a bit to this – all external methods we have compared to use a continuous latent representation. But in the revision we include a new comparison to a hierarchical continuous VAE (in the behavioral segmentation tables, see our overall response). We have also attempted a hierarchical, continuous VAE (same architecture but replacing quantization with a standard Gaussian VAE bottleneck, results not shown) but found it challenging to clearly separate different behaviors from post hoc clustering of sample latents, resulting in high perplexity but extremely low average code usage. In this context, quantization plays a role in better structuring the learned latent space to be more interpretable. We elaborate the discussion here in the related question below.

• Were there no other SOTA models to compare against? The evaluations seem a bit lacking. Additional applications, comparison models and a detailed ablation would help here as well as more detail on limitations and failure cases.

We appreciate the thoughtful feedback and agree that it would be helpful to strengthen the results by more quantitative comparisons. For motion synthesis, we have now added comparisons with

• A diffusion-based framework MotionDiffuse (Zhang et al. 2022).
• Action2Motion (Guo et al. 2020), in addition to the existing GRU VAE baseline.

and internal comparisons with

• Reconstructed motion sequences by VQ-VAE.
• Synthesized motion sequences with different downsampling rates.

Please refer to the updated manuscript Table. 1 for detailed results.

We sincerely thank the reviewer for the thoughtfulness concerning the assessment of the effectiveness of the work. For the questions listed in the weakness section,

• It would be better to compare the proposed method with more advanced baseline approaches to demonstrate its effectiveness.

We appreciate the feedback and have now significantly expanded our comparisons. In addition to new quantitative comparisons for behavioral representations/segmentation, we now compare to a larger set of recent methods for motion synthesis, including ACTOR, Action2Motion and MotionDiffusion (Table 1 in revised text). Notably, our method outperforms even the strong MotionDiffuse baseline, among others. While beyond the scope of this paper, in the future it will be interesting to explore whether these VQ advantages extend to human motion synthesis benchmarks.

• There should also be more ablative analysis that illustrate the effectiveness of each component of the model.

Agreed. We have now performed ablative analysis for both behavioral segmentation and motion syntheses, testing the impact of the hierarchical latents, quantization, and codebook size. Details are included in Supp. Table 2 and 3 (same as reported in ‘response to all reviewers’) and Table 1 in the updated manuscript. Notice that we have also attempted a hierarchical, continuous VAE (same architecture but replacing quantization with a standard Gaussian VAE bottleneck, results not shown) but found it challenging to clearly separate different behaviors from post hoc clustering of sample latents, resulting in high perplexity but extremely low average code usage. In this context, quantization plays a role in better structuring the learned latent space to be more interpretable.

• There is a missing citation on the first page.

Thank you for pointing this out; have fixed it.

• The major innovations seem not very clear. It would be better to clearly state the major novelty of the proposed method and indicate its advantages over existing methods in the literature. The related work section is suggested to be refined and moved to an earlier place for readers to understand the context of the field.

We agree with the reviewer that the major novelty and advantages of the proposed method would be better conveyed to the target audience if the Related Work section is positioned earlier and modified accordingly. Corresponding changes have been updated in the current manuscript, with additional modifications in the Introduction section for better readability.

Our approach is capable of delineating animal behavioral repertoire from motion capture data, in a more computationally efficient and human-interpretable way than established methods and is novel at finding the inherent type/subtype hierarchy among motion primitives without post hoc manual annotation and grouping. . The hierarchical representations recapitulate the standard behavioral analysis tasks as existing techniques, while quantitatively capturing and grouping finer-scale kinematic details, as well as enabling novel tasks.