SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Model	Sequence-level subtasks	Frame-level subtasks
Day (↓)	Time (↓)	Strain	Lights	Approach	Chase	Close	Contact	Huddle	O/E	O/G	O/O	Watching
PCA	0.09416	0.09445	51.60	54.65	0.86	0.14	49.27	37.87	12.71	0.21	0.60	0.53	6.65
TVAE	0.09403	0.09442	52.98	56.80	1.07	0.45	59.33	44.77	21.96	0.27	0.83	0.62	10.20
T-BYOL	0.09362	0.09373	60.95	65.05	1.68	0.72	62.48	48.19	18.52	0.35	0.96	0.82	17.77
T-BERT	0.09262	0.09276	78.63	68.84	1.80	0.87	70.22	55.84	30.24	0.51	1.40	1.12	17.27
TS2Vec	0.09380	0.09422	57.12	65.60	1.29	0.66	59.53	46.13	24.74	0.35	1.09	0.74	12.37
T-Perceiver	0.09322	0.09323	69.81	69.68	1.57	1.27	60.84	47.81	28.32	0.41	1.16	0.86	16.42
T-GPT	0.09269	0.09384	64.45	65.39	1.73	0.64	69.05	55.78	23.80	0.46	1.12	1.05	17.86
T-PointNet	0.09275	0.09320	66.01	67.15	2.56	4.57	70.68	55.96	21.23	0.84	2.79	2.32	15.61
BAMS	0.09094	0.08989	88.23	72.00	2.74	1.89	67.22	53.43	31.43	0.59	1.61	1.57	18.15
⊥ short-term	0.09288	0.09294	61.33	66.34	1.80	1.15	66.58	52.60	25.34	0.39	1.09	0.98	16.74
⊥ long-term	0.09174	0.09037	86.49	70.91	2.10	1.06	61.99	49.12	29.09	0.45	1.32	1.08	13.98
Table 2: Linear readouts of mouse behavior. We report the BAMS against the top performing models in the MABe 2022 challenge. All numbers are reported in (59) except for TS2Vec and T-BYOL, which we produce. The scores show the performance of the linear readouts across 13 different tasks. Mean-squared error (MSE) is used in the case of tasks 1 and 2, since they are continuously labeled. In the rest of the subtasks, which are binary (yes/no), F1-scores are used. The best-performing models are those with low MSE scores and high F1-scores, are highlighted in bold.
and  $\mathbf{z}_{t,3}$ , for each mouse respectively and concatenate the embeddings for all three individuals to build a joint embedding for evaluation. The model is trained for 500 epochs using the Adam optimizer with a learning rate of  $10^{-3}$ . More details can be found in Appendix B.
Interaction loss. To learn additional features that are useful in predicting animal interactions for multi-agent settings, we introduce a simple auxiliary loss to predict the distances between the trio at time  $t$ . Our input features do not include any information about the global position of the mice in the arena, so the model can only rely on the inherent behavior and movement of each individual mouse to draw conclusions about their proximity. Thus, we build a network  $h$  that takes in the embeddings of two mice  $i$  and  $j$  and predicts the distance  $d_{i,j}$  between them as follows,
$$
\mathcal {L} _ {\mathrm {a u x}} = \left\| h \left(\mathbf {z} _ {t, i}, \mathbf {z} _ {t, j}\right) - \mathrm {d} _ {i, j} \right\| _ {2} ^ {2}. \tag {6}
$$
This penalty is added to the loss in Equation 5 to encourage learning of shared features across the different individual embeddings.
Evaluation protocol. To evaluate the performance of the model in detecting frame- and sequence-level behaviors, we compute representations  $\mathbf{z}_{t,1}$ ,  $\mathbf{z}_{t,2}$  and  $\mathbf{z}_{t,3}$  for each mouse respectively, which we then aggregate into a single mouse triplet embedding using two different pooling strategies. First, we apply average pooling to get  $\mathbf{z}_{t,\mathrm{avg}}$ . Second, we apply max pooling and min pooling, then compute the difference to get  $\mathbf{z}_{t,\mathrm{minmax}} = \mathbf{z}_{t,\mathrm{max}} - \mathbf{z}_{t,\mathrm{min}}$ . Both aggregated embeddings are concatenated into a 128-dim embedding for each frame in the sequence. Evaluation of each one of the 13 tasks, is performed by training a linear layer on top of the frozen representations, producing a final F1 score or a mean squared error depending on the task.
3.2.2 Results
We compare our model against PCA, trajectory VAE (TVAE) (13), TS2Vec (69), BYOL (17), and the top performing models in the MABe2022 challenge (61) which are adapted respectively from Perceiver (23), GPT (7), PointNet (49) and BERT (14). All methods are trained using some form of reconstruction-based objective, with both T-Perceiver and T-BERT using masked modeling. Both T-BERT and T-PointNet also supplement their training with a contrastive learning objective using positives from the same sequence. It is also important to note that the training set labels from two tasks (Lights and Chase) are made publicly available, and are used as additional supervision in the T-Perceiver and T-BERT models. We do not use any supervision for BAMS.
Our model achieves a new state-of-the-art result on the MABe Multi-Agent Behavior 2022 - Mouse Triplets Challenge, as can be seen in Table 2. We rank first overall based upon our performance on all tasks, and show impressive boosts in performance on global sequence-level tasks where we find a  $22\%$  improvement in the Strain task and  $5\%$  improvement in the Light task. In the frame-level tasks, we remain competitive with other approaches, and rank 1st on 3 out of 9 of the frame-level subtasks; this is in contrast to the other top performing model that explicitly models the interaction between mice by introducing hand-crafted pairwise features. This is outside the scope of this work, as we do not focus on social interactions beyond predicting the distance between mice.
We observe that our proposed method results in significant gaps on sequence-level tasks. In particular, we observe a marked improvement on the prediction of strain over the second-place model at  $78.63\%$ , with BAMS yielding a  $10\%$  improvement in accuracy at  $88.23\%$ . These big improvements suggest that we might have identified and addressed a critical problem in behavior modeling. In the next section, we conduct a series of ablations to further dissect our model's performance.
Multi-timescale embedding separation also enables us to probe our model for timescale-specific features. In Table 2, we report the decoding performance with short-term embeddings and long-term embeddings respectively. We find that sequence-level behavioral factors are better revealed in the long-term space, while the frame-level factors are more distinct in the short-term space. That being said, decoding performance is still best when using both timescales.
By default, BAMS is pre-trained with all available trajectory data. We also test BAMS in the inductive setting, where we only pre-train using the training split (only 1800 out of 5336) of the dataset. Results are reported in Appendix B.5. We find that the performance drop is modest, and that BAMS trained in this setting still beats all other methods.
3.2.3 Ablations
To understand the role of the different proposed components in enabling us to achieve state-of-the-art performance, we conduct a series of ablations on the MABe benchmark that we report in Table 3. First, we compare the performance of BAMS when trained with the traditional reconstruction-based objective (multi-step sequential prediction). BAMS without the HoA prediction objective, performs comparably with many of the top-entries, though performance is ultimately improved when using this novel learning objective. We note that we had to reduce the number of prediction frames to 10, as the training fails if we go beyond. With our HoA objective, we can use 30
Hist of actions	Multi-timescale
Seq MSE	Seq F1	Frame F1
✓	✓	0.090415	80.12	19.85
✓	0.093100	68.30	19.46
✓	0.090717	78.11	19.45
✓	✓	0.092483	73.37	19.52
Table 3: Ablations on the Mouse Triplet Dataset. We report the average sequence-level MSE and F1-score, and the average frame-level F1-score.
frames, which strongly suggests that this loss can be stably applied over longer time horizons when compared to other losses.
Next, we analyze the role of the multi-timescale bootstrapping in improving the quality of the representation. When removing the bootstrapping objective, we find a  $2\%$  drop in the sequence-level averaged F1-score, as well as modest drops in frame-level performance. This suggests that this learning objective may help to resolve global and intermediate-scale features. We also ablate the multi-timescale component, i.e. we perform bootstrapping but in the same space, by using a single TCN that spans the receptive field of the two used originally. This results in sub-optimal performance, which emphasizes the idea that having different objectives in different spaces is critical to prevent interference and provide ideal performance. We report additional ablations in Appendix B.4, including an ablation of the interaction module where we find that it is not critical for good performance.
3.3 Experiments on Fruit Fly Groups Dataset
3.3.1 Experimental Setup and Tasks
The fruit fly groups dataset is the second dataset in the MABe benchmark. It consists of tracking data of a group of 9 to 11 flies interacting in a small dish (Figure 5). Tracking data consists of 19 keypoints on each fly body and wings, and is recorded at a much higher frame rate compared to the mice (150 vs. 30 fps). Precise neural activity manipulations are performed on certain neurons which, when activated, induce certain types of behavior including courtship, avoidance and female aggression (61). Additionally, the groups of flies are differentiated by various genetic mutations and tagged by sex. This along with other behavioral factors provide us with 50 different subtasks, both frame-level and sequence-level, that can be used to evaluate the learned representations.
Figure 5: Sample frame from the Fruit Fly Groups Dataset.
In this set of experiments, we are interested in testing whether our proposed method generalizes to a dataset from another organism, and whether it can work out-of-the-box with minimal expert
intervention. In particular we chose to not extract any features manually and only use the provided data, we use the default hyperparameters found in the previous experiment and do not perform hyperparameter tuning, and finally we do not use an interaction module. Note that we follow the same evaluation procedure established previously.
3.3.2 Results
We compare our model against PCA, TVAE, TS2Vec, and the top performing models in the fly challenge: T-GPT, and T-Perceiver. Our model achieves state-of-the-art performance on the fly dataset, as seen in Table 4. BAMS outperforms other models on both frame-level and sequence-level sub-tasks, and we note a significant boost in the average frame-level F1 score. This result further demonstrates the generalizability of our approach to new datasets and scaling to an even larger numbers of animals.
Model	All F1	Seq F1	Frame F1
PCA	42.5	23.0	45.2
TVAE	37.0	22.2	39.0
T-Perceiver	44.8	19.7	48.2
T-GPT	45.8	24.5	48.7
BAMS	48.2	25.4	51.1
Table 4: Linear readouts of fly behavior. We report the average performance of various models on both frame-level, and sequence-level subtasks.
4.1 Animal behavior analysis
Pose estimation and animal tracking. Recently, there has been a recent democratization of automated methods for pose estimation and animal tracking that has made it possible to conduct large scale behavioral studies in many scientific domains. These tools abstract behavior trajectories from video recordings, and facilitate the modeling of behavioral dynamics. Most pipelines for analyzing animal behavior consist of three key steps (36): (1) pose estimation (60; 67; 15; 48; 40), (2) spatial-temporal feature extraction (35; 10), and (3) quantification and phenotyping of behavior (43; 39; 66). In our work, we consider the analysis of behavior after pose estimation is performed. However, one could imagine using our multi-timescale bootstrapping approach for representation learning in video analysis, where self-supervised losses have been proposed recently for keypoint discovery (60; 64; 24).
Disentanglement of animal behavior in videos. In recent work (54; 32; 63), two distinct disentangled behavioral embeddings are learned from video, separating non-behavioral features (context, recording condition, etc.) from the dynamic behavioral factors (pose). This has even been applied to situations with multiple individuals, performing disentanglement on each individual (21). This is performed by training two encoders, typically variational autoencoders (VAEs) (29), on either multi-view dynamic information or a single image. These encoders create explicitly separated embeddings of behavioral and context features. In comparison to this work, rather than seperating behavior from context, our model considers the explicit separation of behavioral embeddings across multiple timescales, and considers the construction of a global embedding that is consistent over long timescales.
Modeling social behavior. For social and multi-animal datasets, there are a number of other challenges that arise. Simba (45) and MARS (52) have similar overall workflows for detecting keypoints and pose of many animals and classifying social behaviors. More recently, a semi-supervised approach TREBA has been introduced (59) for building behavior embeddings using task programming. TREBA is built on top of the trajectory VAE (13), a variational generative model for learning representations of physical trajectories in space. In our work, we do not consider a reconstruction objective but a future prediction objective, in addition to bootstrapping the behavior representations at different timescales.
4.2 Representation learning for sequential data
The self-supervised learning (SSL) framework has gained a lot of popularity recently due to its impressive performance in many domains (14; 9; 12). Many SSL methods are built based on the concept of instance-specific alignment loss: Different views of each datapoint are created based on pre-selected augmentations, and the views that are produced from the same datapoint are treated as positive examples; while the views that are produced from different datapoints are treated as negative examples. While contrastive methods like SimCLR (9) utilize both positive examples and negative examples to guide the learning, BYOL (17) proposes a framework in which augmentations of a
sample are brought closer together in the representation space through a predictive regression loss. Recent work (18; 50; 44; 2) applies BYOL to learn representations of sequential data. In such cases, neighboring samples in time are considered to be positive examples of each other, assuming temporal smoothness of the semantics underlying the sequence. The model is trained such that neighboring samples in time are mapped close to each other in the representation space. However, these methods use a single scale to define similarities unlike our method.
Recently, self-supervised methods such as TS2vec (69) learn self-supervised representations for sequential data by generating positive sub-sequence views through more complex temporal augmentations that can be integrated at instance or local level. Positive sub-sequences are temporally contrasted with other representations within the same sequence as well as contrasted with representations of other available sequences. However, unlike our method, the same space of representations are used for learning these multi-scale representations. We learn a separate space of representations for different time scales to encourage disentanglement.
The idea of using multi-scale feature extractors can be found in representation learning. In (50), a video representation learning framework, two different encoders process a narrow view and a broad view respectively. The narrow view corresponds to a video clip of a few seconds, while the broad view spans a larger timescale. The objective, however, is different to ours, as the narrow and broad representations are brought closer to each other, in the goal of encoding their mutual information. This strategy is also used in graph representation learning where a local-neighborhood of node is compared to its global neighborhood (19; 22). Our method differs in that we bootstrap the embeddings at different timescales separately, this is important to maintain the fine granularity specific to each timescale, thus revealing richer information about behavioral dynamics.
Behavior is likely to be driven by a number of factors that can unfold over different timescales. Thus, having ways to model behavior and discover differences in behavioral repertoires or actions at many scales could provide insights into individual differences, and help, for example, detect signatures of cognitive impairment (66). We make steps towards addressing these needs by proposing BAMS, a novel self-supervised approach that learns representations for behavioral data at different timescales.
Our analysis on realworld datasets centers on two different datasets in a multi-agent benchmark, both with very different dimensionalities and variable tasks. Despite this, our approach is designed to model a wide range of behaviors, whether it's individual, social, structured or naturalistic. More multi-scale and multi-task benchmarks like MABe are needed in order properly evaluate and guide the development of tools for general behavior analysis.
Currently, our model has a relatively simple way of (optionally) modeling interactions between animals in the multi-agent setting. Despite this, we show really competitive performance in multi-agent analysis without using handcrafted interaction features like T-PointNet or T-BERT (61). Moving forward, it would be interesting to develop new ways to learn features of multi-animal interactions, especially in open environments where animals might come in and out of frame or get occluded.
Our experiments highlight the truly multi-scale nature of BAMS and show that our method can learn to distinguish global as well as temporally local behaviors. Currently, we separate short-term from long-term dynamics through a contrastive loss and separation of the information into two latent spaces for each scale. Although this approach appears to be effective, we don't modify the loss explicitly to disentangle the two latent spaces. By providing additional incentives for the model to separate short from long-term dynamics we hope to improve the interpretability of the model.
When extending our model from mice to flies, we found that it was possible to use the same overall model architecture and hyperparameters, despite the major differences in datasets and underlying tasks. Thus, given the robustness of the method, we imagine that it can be utilized in the analysis of other species and even more diverse types of behavior (42; 26). By combining our modeling approach with methods for self-supervised video keypoint discovery (57), we could further extend BAMS to raw video data without needing the intermediate step of pose estimation. We also look forward to the opportunities provided by simultaneous recordings of the brain and behavior (30; 47) and other multi-modal sources of input that could be leveraged to further study how behavior unfolds across different timescales.
We would like to thank Mohammad Gheshlaghi Azar and Remi Munos for their feedback on the work. This project was supported by NIH award 1R01EB029852, NSF awards IIS-2039741 and IIS-2146072, as well as generous gifts from the Alfred Sloan Foundation, the McKnight Foundation, and the CIFAR Azrieli Global Scholars Program.
[67] Anqi Wu, Estefany Kelly Buchanan, Matthew Whiteway, Michael Schartner, Guido Meijer, Jean-Paul Noel, Erica Rodriguez, Claire Everett, Amy Norovich, Evan Schaffer, et al. Deep graph pose: a semi-supervised deep graphical model for improved animal pose tracking. Advances in Neural Information Processing Systems, 33:6040-6052, 2020. 9
[68] Yuzhe Yang, Xin Liu, Jiang Wu, Silviu Borac, Dina Katabi, Ming-Zher Poh, and Daniel McDuff. Simper: Simple self-supervised learning of periodic targets. arXiv preprint arXiv:2210.03115, 2022. 1
[69] Zhihan Yue, Yujing Wang, Juanyong Duan, Tianmeng Yang, Congrui Huang, Yunhai Tong, and Bixiong Xu. Ts2vec: Towards universal representation of time series. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 36, pages 8980-8987, 2022. 4, 7, 10, 18
A Experimental details: Simulated Quadrupeds
A.1 Data generation
Simulation details. We record a total of 5182 trajectories. 2756 were generated for robots of type ANYmal B and 2426 sequences were generated for robots of type ANYmal C. These are quadruped robots, which means that they have four legs. Each leg has 3 degrees of freedoms - hip, shank and thigh. The position and velocities of these degrees of freedom for all 4 legs were recorded. This results in 24 features for each robot. Robots are generated while traversing an procedurally generated environment with different terrain types and traversal difficulty, as show in Figure 6. We only keep trajectories that correspond to a successful traversal.
B
Figure 6: Visualization of different simulated environments. (A) Screenshots from the simulator showing a robot walking down some stairs, and a view of the terrain landscape. (B) Visualization of the different terrain sections, characterized by a terrain type and different levels of difficulty. Terrain are made more difficult to traverse by either making them more rough or have steeper slopes.
Ascendir
Descending stairs
Tasks. To evaluate the representation quality of our model, we use multi-task probes that correspond to different long-term and short-term behavioral factors.
- Robot type: the robot can either be of type "ANYmal B" or "ANYmal C". These robots have the same degrees of freedom and tracked joints but differ by their morphology. This is a sequence-level task.
- Linear velocity: the command of the robot is a constant velocity vector. The amplitude of the velocity dictates how fast the robot is commanded to traverse the environment. A higher velocity would translate into more clumps and more risk-taking behavior. This is a sequence-level task.
- Terrain type: the environment is generated with multiple segments of five terrain types that are categorized as: flat surfaces, pits, hills, ascending and descending stairs. This is a frame-level task.
- Terrain slope: the slope of the surface the robot is walking on. This is a frame-level task.
- Terrain difficulty: the different terrain segments have different difficulty levels based on terrain roughness or steepness of the surface. This is a frame-level task.
Why this dataset. Simulation-based data collection enables access to information that is generally inaccessible or hard to acquire in a real-world setting. Unlike noisy measurements coming from the camera-based feature extractor in the case of the mouse dataset, physics engines do not suffer from the problem of noise. Instead, they provide accurate ground-truth information about the creature and the world state free of charge. Access to such information is at times critical for scrutinizing the capabilities of the learning algorithms.
A.2 Visualizing differences between short-term and long-term embeddings
In Figure 7, we visualize how the short-term and long-term embeddings evolve over time, for a single sample sequence. We note a clear difference in the smoothness in the two timescales. In the short-term embeddings, we note a clear block structure corresponding to different blocks of behavior that span a few seconds, while in the long-term embeddings the representation is more stable over
time. This suggests that, as expected, the bootstrapping objectives are forming representations with different levels of granularity.
Figure 7: Visualization of the short-term and long-term embeddings. We visualize for a single sequence how the short-term and long-term embeddings evolve over time.
B Experimental details: Mouse Triplet
B.1 Feature extraction
Each mouse in the arena is tracked using 12 anatomically defined keypoints. We process these keypoints to extract 36 different features characterizing each mouse individually, similar to (52). We separate the keypoints into two different areas, the head and the body, for each we extract different measures of displacement, that we express in the frame of the mouse, i.e. these features are invariant to the pose of the mouse relative to the arena. These features include: