Paper: Vector Quantized Representations for Efficient Hierarchical Delineation of Behavioral Repertoires (ICLR-2024)

Reviewer comment (Theory):
In equation (3), can you elaborate more on the embedding and commitment terms? Intuitively what do they mean and how are they derived?

Author response:
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.

Reviewer comment (Experiments):
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?

Author response:
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.

Reviewer comment (Evaluation):
What are the quantitative comparison results against the KPMS benchmark?

Author response:
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.

Reviewer comment (Novelty):
Unfortunately, it’s not clear whether there is sufficient technical novelty for the ICLR community. The application of VQ-VAE to animal behavioural data specifically may be novel, but other than the multi-level encoding, it is not clear whether there are architectural/learning improvements that may be relevant to the broader ICLR reader. The listed contributions indicate potential novelty in behavioural neuroscience (ie. that this model simplifies the behavioural pipeline), which could perhaps suggest a different venue might be a better fit.

Author response:
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).

Reviewer comment (Novelty):
There is quite a lot of work on existing work on learning movement primitives or behavioural decomposition via latent variable models in embodied and robotics domains. The work could better be positioned in this context, and architectural design choices could be better justified.

Author response:
Thank you for suggesting these relevant references. We agree that these methods are motivating, particularly Merel et al. (2019b) that describes behavioral representations of an artificial rodent model by applying classic analysis methods in neuroscience and Rao et al. 2022 that does offline learning of hierarchical, discrete motor skills. While we acknowledge that our knowledge in RL-oriented skill learning may be limited and biased, the choice of discrete representations versus encoding motion sequences into continuous latent dynamics is better aligned with the longstanding goal of constructing formal taxonomical structures for better describing and understanding animal behaviors (similar ideas also explored in a more RL-oriented setting such as Paulius et al. 2020). The coding scheme we have deployed would better support real-world investigations into behavioral patterns and perturbations, as well as causal relationships across behaviors of different individuals in an information theoretic sense (Ferrer‐i‐Cancho R et al. 2013). Moreover, our studies on naturalistic behaviors in animals are largely disjoint from the goal-directed manipulation of body movements in most RL optimization settings and many findings in related RL literature may not be readily applicable to behavioral analyses. While it would be valuable to offer a comprehensive review of behavioral understanding and analysis in both real organisms and embodied agents, elaborating on this is likely beyond the scope of this paper and deviates from our major focus.

Reviewer comment (Experiments):
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).

Author response:
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.

Reviewer comment (Novelty):
There is quite a lot of work on learning movement primitives (eg. Paraschos et al, 2013), or behavioural decomposition via latent variable models in embodied and robotics domains. Merel et al (2019a); Bohez et al (2022) apply hierarchical latent variable models to motion-capture data from humans and other mammals, and Merel et al (2019b) specifically studies learned representations of simulated rodent behaviour. Other works leverage latent variable representations of offline behavioural data in robotics (eg. Singh et al, 2021), including hierarchical discrete representations that can decompose data into discrete motifs that can execute / synthesize meaningful behaviours (Rao et al, 2022).

Author response:
Thank you for suggesting these relevant references. We agree that these methods are motivating, particularly Merel et al. (2019b) that describes behavioral representations of an artificial rodent model by applying classic analysis methods in neuroscience and Rao et al. 2022 that does offline learning of hierarchical, discrete motor skills. While we acknowledge that our knowledge in RL-oriented skill learning may be limited and biased, the choice of discrete representations versus encoding motion sequences into continuous latent dynamics is better aligned with the longstanding goal of constructing formal taxonomical structures for better describing and understanding animal behaviors (similar ideas also explored in a more RL-oriented setting such as Paulius et al. 2020). The coding scheme we have deployed would better support real-world investigations into behavioral patterns and perturbations, as well as causal relationships across behaviors of different individuals in an information theoretic sense (Ferrer‐i‐Cancho R et al. 2013). Moreover, our studies on naturalistic behaviors in animals are largely disjoint from the goal-directed manipulation of body movements in most RL optimization settings and many findings in related RL literature may not be readily applicable to behavioral analyses. While it would be valuable to offer a comprehensive review of behavioral understanding and analysis in both real organisms and embodied agents, elaborating on this is likely beyond the scope of this paper and deviates from our major focus.

Reviewer comment (Theory):
Related to these points, it’s not clear why specific architectural choices were made. For example, encoding entire trajectories into a single latent code scales poorly with dimensionality of the inputs and length of the sequence, and many of the approaches from the previous point use sequential latent variable models to better model embodied temporal data.

Author response:
We appreciate the feedback and apologize if there is any miscommunication with respect to how the proposed model encodes trajectories. The latent variables are still sequential, while a long trajectory is encoded as a long sequence of latent codes. We updated Figure 1 to make this more clear.  Due to the downsampling that occurs from input sequences to latent code, each latent code captures approximately 320ms of the trajectory (out of ~30 minute recordings).

Reviewer comment (Writing):
Broken citation reference in the first paragraph.

Author response:
Thank you for catching this and we have fixed it.

Reviewer comment (Presentation):
Having the related work as a final section reads a bit awkwardly, as it feels like an afterthought. Consider moving it to at least before the final discussion / conclusions, and ideally before the method itself to provide some scaffolding and context for the contributions and claims.

Author response:
We agree that it would be helpful to rearrange the Related Work section and thus provide more of an orientation with respect to the motivations and contributions. Please refer to the updated manuscript for changes that have been made.

Reviewer comment (Experiments):
It would be better to compare the proposed method with more advanced baseline approaches to demonstrate its effectiveness. There should also be more ablative analysis to illustrate the effectiveness of each component of the model.

Author response:
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. 

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.

Reviewer comment (Writing):
There is a missing citation on the first page.

Author response:
Thank you for pointing this out; have fixed it.

Reviewer comment (Novelty):
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.

Author response:
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.