# A study of animal action segmentation algorithms across supervised,
unsupervised, and semi-supervised learning paradigms

Ari Blau
Department of Statistics
Columbia University &Evan S Schaffer
Icahn School of Medicine
Mount Sinai &Neeli Mishra
Zuckerman Institute
Columbia University &Nathaniel J Miska
Sainsbury Wellcome Centre
University College London &The International Brain Laboratory
International Brain Laboratory &Liam Paninski
Department of Statistics
Zuckerman Institute
Columbia University &Matthew R Whiteway
Zuckerman Institute
Columbia University a list of consortium members appears in Appendix B

###### Abstract

Action segmentation of behavioral videos is the process of labeling each frame
as belonging to one or more discrete classes, and is a crucial component of
many studies that investigate animal behavior. A wide range of algorithms
exist to automatically parse discrete animal behavior, encompassing
supervised, unsupervised, and semi-supervised learning paradigms. These
algorithms –- which include tree-based models, deep neural networks, and
graphical models –- differ widely in their structure and assumptions on the
data. Using four datasets spanning multiple species –- fly, mouse, and human
–- we systematically study how the outputs of these various algorithms align
with manually annotated behaviors of interest. Along the way, we introduce a
semi-supervised action segmentation model that bridges the gap between
supervised deep neural networks and unsupervised graphical models. We find
that fully supervised temporal convolutional networks with the addition of
temporal information in the observations perform the best on our supervised
metrics across all datasets.

##  1 Introduction

Action segmentation of video data is a process that classifies discrete animal
behaviors given a set of spatiotemporal video features. It is an indispensable
tool for quantifying natural animal behavior across a range of experimental
paradigms, from spontaneous behaviors in open arenas to complex social
interactions [1, 2, 3]. This procedure begins with the collection of raw
behavioral data during an experiment, typically with video cameras or motion-
capture sensors. The next step is to reduce the dimensionality of the video
data using pose estimation [4, 5, 6], autoencoders [7, 8], or other techniques
[9, 10]. Finally, an action segmentation model finds discrete behaviors from
the low-dimensional representation (Fig. 1).

Supervised action segmentation (Fig. 2A) requires the experimenter to annotate
a subset of frames that contain behaviors of interest, such as grooming or
sniffing. Then a classifier is trained to match each frame (or its low-
dimensional representation) with the corresponding label [11, 12, 13, 14, 15,
16, 17, 18, 10, 19, 20, 21]. As the scale of behavioral data continues to grow
[22, 3], it becomes infeasible to densely label behaviors in every video.
Therefore, it is crucial to develop action segmentation models that perform
well with sparsely labeled data. Furthermore, we would like to develop
techniques that also take advantage of the vast amounts of unlabeled data.

Unsupervised action segmentation models (Fig. 2B) are a complementary approach
that do not require any hand labels [23, 24, 25, 26, 27, 28] and perform
clustering on the low-dimensional behavioral representation [29]. These
unsupervised models are more scalable than their supervised counterparts since
they do not require annotations for training. Another benefit of this approach
is its ability to discover new behaviors that are not pre-defined by the
experimenter [29, 2]. However, there may be certain behaviors of particular
interest to downstream users, and unsupervised models cannot guarantee that
they cluster these behaviors accurately.

Figure 1: Overview of the action segmentation pipeline. Raw sensor data (e.g.
video) is collected, then features are extracted (e.g. pose estimates), then
an action segmentation model is trained to map those features to a discrete
behavioral class for each frame.

The recent proliferation of approaches to animal action segmentation raises
obvious practical questions: how do these approaches compare, and what are
their trade-offs? To aid in this endeavor, we introduce an action segmentation
model that bridges the gap between supervised deep neural networks and
unsupervised graphical models through the use of semi-supervised state space
models (Fig. 2C). Graphical models, particularly switching linear dynamical
systems, are an increasingly ubiquitous approach to unsupervised animal action
segmentation [24, 25, 30, 31, 28]. These models posit that each discrete
action is defined by a set of linear dynamics, and infer which set of dynamics
best describes the observations at each time point [29]. We utilize such a
model and bias it towards interpretable solutions by providing a small set of
hand labels during training. Our inference procedure leads to a deep neural
network classifier that allows us to make direct comparisons with fully
supervised classifiers. We refer to the resulting model and inference
technique as a semi-supervised switching linear dynamical system (S3LDS).

Another important consideration for the animal action segmentation problem is
the behavioral representation used as input to a given model. Much previous
work with supervised models uses a host of quantities derived from pose
estimates, like distances and angles between different keypoints [18, 19, 21].
Much of the unsupervised literature simply uses the pose estimates themselves
(perhaps converting to ego-centric coordinates first) and allow the model to
perform the necessary transformations [27, 28] (though see [23, 26] for
approaches that involve feature engineering). We would also like to know how
simple choices about the behavioral representation – such as whether or not
temporal information contained in velocity or acceleration measurements is
included – impacts the performance of these models.

To address these questions we use an array of models, including S3LDS, to
evaluate four behavioral datasets: a head-fixed fly spontaneously behaving on
an air-supported ball [32], a mouse freely moving around an open arena [18], a
head-fixed mouse performing a perceptual decision-making task from the
International Brain Laboratory [33], and a human gait dataset (HuGaDB) [34].
For the first three datasets, we experiment with both position features
(static pose information) and position-velocity features (including the
velocity of each static keypoint). The HuGaDB data is comprised of three
dimensional accelerometer and gyroscope data collected from inertial
measurement units (IMUs).

It is challenging to comprehensively compare the performance of such a wide
range of models, and we therefore limit ourselves to evaluating how much the
outputs of the models align with human-annotated behaviors. This supervised
metric will of course miss many of the nuances provided by these models,
especially those which are fully unsupervised. We find that the S3LDS offers
improvements over supervised models when using the position features as input,
but these advantages disappear once we include the velocity features. In the
latter case, a fully supervised temporal convolutional network (TCN) performs
best across all datasets [35, 32, 36]. We also compare the S3LDS and fully
supervised models to Keypoint-MoSeq [28], a closely related yet fully
unsupervised action segmentation method, and find that the models trained with
discrete behavior labels learn an improved latent representation, as measured
by several supervised cluster quality metrics.

Figure 2: Overview of action segmentation models. A: Top: Graphical model for
supervised classification. Both discrete states $y_{t}$ and poses
$\mathbf{x}_{t}$ are observed. Bottom: Inference network for the supervised
model. We use a window of observed behavioral features for state prediction.
B: Top: Graphical model for an unsupervised recurrent switching dynamical
system. The set of discrete states $\\{y_{t}\\}$ and continuous latents
$\\{\mathbf{z}_{t}\\}$ are unobserved. Bottom: The inference network uses a
window of observed behavioral features to create a deterministic hidden
representation $\mathbf{h}_{t}$ (purple arrows); this is then used to predict
the continuous latents $\mathbf{z}_{t}$ (blue arrows) and discrete latents
$y_{t}$ (red arrows). Note that the purple and red arrows together define a
classifier for the discrete state at each time step. C: Graphical model and
inference network for a semi-supervised recurrent switching dynamical system.
A subset of the discrete states are observed. During inference, the observed
discrete state is used for the inference of $\mathbf{z}_{t}$ when possible.

##  2 Results

We first provide a brief overview of the different types of supervised models
that we compare, and establish which one we use as our baseline. We explore
how performance is affected by the choice of behavioral features used by the
various models. Next, we introduce our proposed semi-supervised framework,
S3LDS, which we compare to our supervised model using supervised
classification metrics. We also perform several ablation experiments
demonstrating which aspects of the semi-supervised model are most important
for performance. Finally, we compare the supervised and semi-supervised models
to Keypoint-MoSeq [28], an unsupervised model closely related to the S3LDS.

###  2.1 Comparing supervised baselines

We experiment with three fully supervised models: a temporal convolution
network (TCN) [37, 38], random forests [39], and XGBoost [40]. TCNs have been
competitive on human action segmentation and recognition benchmarks [41, 42,
43]. The two tree-based methods are commonly used as supervised classifiers
for animal action segmentation [19, 21].

Another common question in the literature is how to transform the data before
model fitting. Previous work has proposed various hand-engineered features
[23, 26, 19, 21] and learned features [9, 10]. While a thorough
characterization of this choice is beyond the scope of this study, we test two
scenarios: modeling with position features (no temporal information) versus
position-velocity features (contains temporal information).

We find the TCN achieves the best performance with both feature types (Fig.
S1). We also find that temporal information in the features is useful for all
methods. We select the TCN to use as our baseline supervised model, not only
because of its superior performance, but also because of its compatibility to
function as an inference network in the deep graphical models we introduce
next.

###  2.2 Introducing a semi-supervised action segmentation model

Next we move to semi-supervised models, which utilize both labeled and
unlabeled data. We focus on the switching linear dynamical system (SLDS) and
develop a new amortized variational inference scheme. The SLDS is commonly
used to model complex nonlinear behavior in dynamical systems [44, 45, 46].
Several studies have used the SLDS (and the related auto-regressive hidden
Markov model, or ARHMM) to model animal behavior [24, 25, 30, 47, 7, 48, 8,
31, 49, 28]. The motivation to construct our semi-supervised model from the
SLDS is two-fold: (1) the SLDS is typically utilized as an unsupervised model,
and hence allows us to make easy connections between these two learning
paradigms; (2) the amortized variational inference scheme will result in a
classifier that can be directly compared to the supervised TCN model above.

The standard setup for the generative model of an SLDS is as follows (Fig.
2B): let $y_{t}\in\\{0,\ldots,K-1\\}$ be the discrete behavioral state at time
$t$ with Markovian dynamics, i.e. the state at time $t$ is only dependent on
the state at time $t-1$. The observed data $\mathbf{x}_{t}$ are modeled using
piecewise-linear dynamics in a latent space denoted by $\mathbf{z}_{t}$:

$\displaystyle p_{\theta}(y_{t}|y_{t-1})$ $\displaystyle=$
$\displaystyle\textrm{Cat}(y_{t}|\pi_{\theta}({y_{t-1}}))$ (1) $\displaystyle
p_{\theta}(\mathbf{z}_{t}|\mathbf{z}_{t-1},{y_{t}})$ $\displaystyle=$
$\displaystyle\mathcal{N}(\mathbf{z}_{t}|A_{y_{t}}\mathbf{z}_{t-1}+b_{{y}_{t}},Q_{y_{t}})$
(2) $\displaystyle p_{\theta}(\mathbf{x}_{t}|\mathbf{z}_{t},y_{t})$
$\displaystyle=$
$\displaystyle\mathcal{N}(\mathbf{x}_{t}|C_{y_{t}}{\mathbf{z}_{t}}+d_{y_{t}},S_{y_{t}}),$
(3)

where $t\in\\{1,2,\dots,T\\}$. $A_{y_{t}}$ and $b_{y_{t}}$ define the latent
linear dynamics associated with state $y_{t}$; $C_{y_{t}}$ and $d_{y_{t}}$
represent the linear mapping from the continuous latents to the observations;
and $Q_{y_{t}}$ and $S_{y_{t}}$ represent noise covariance matrices of the
latent and observed data. In other words, at each time point, the model
selects one set of the $K$ linear dynamics (which each define a discrete
action class; Eq. 1) that operate in a continuous latent space, evolves the
latents forward in time using those linear dynamics (Eq. 2), then projects
those latents to the observations (the behavioral features; Eq. 3).

The recurrent SLDS (rSLDS) builds on the SLDS by conditioning $y_{t}$ on
$\mathbf{z}_{t-1}$ in addition to $y_{t-1}$ [50] (Fig. 2B). This recurrence
allows the model to more flexibly switch between discrete states based on the
trajectory through the latent space, rather than just static switching
probabilities. The new generative model therefore replaces Eq. 1 with

$p_{\theta}(y_{t}|y_{t-1},\mathbf{z}_{t-1})=\textrm{Cat}(y_{t}|\pi_{\theta}(y_{t-1},\mathbf{z}_{t-1})).$
(4)

In both the SLDS and recurrent SLDS, it is assumed the discrete action classes
are all unobserved. We propose to allow the model access to labels for the
discrete action classes for a small number of time points (Fig. 2C, top),
resulting in a semi-supervised SLDS (S3LDS). We develop an efficient and
flexible amortized variational inference strategy that incorporates both
labeled and unlabeled data. Importantly, the proposed inference strategy leads
to the definition of an approximate posterior
$q_{\phi}(y_{t}|\mathbf{x}_{1:T})$ which also serves as a classifier [51],
predicting the discrete action class $y_{t}$ from the observed data
$\mathbf{x}_{1:T}$ (Fig. 2C, bottom). We implement this approximate posterior
using the TCN model of the previous section, allowing us to directly assess
the benefits of training this classifier network with both labeled and
unlabeled data. For more details, see the Methods section.

###  2.3 Semi-supervised and supervised model comparisons

We start by comparing supervised and semi-supervised models using a supervised
classification metric, and report several ablation studies along the way that
highlight the role of various components of the S3LDS. Due to our choice of
the TCN for the S3LDS inference network, we also use a TCN with the exact same
architecture as our supervised baseline.

We begin with the head-fixed fly dataset (Fig. 3A), which contains five
annotated behaviors: still, walk, front groom, back groom, and abdomen move
(Fig. 3B,C). We extract pose estimates from the videos using Lightning Pose
and the Ensemble Kalman Smoother post-processor [6], and use these position
features as input to the TCN and S3LDS. We train models with an increasing
number of labeled videos (five networks trained per condition, each using a
random subset of labeled videos). The S3LDS uses all available unlabeled
frames. We find that performance, as measured by F1 (the harmonic mean of
accuracy and precision; higher is better, with a maximum value of 1.0)
increases with increasing labels. The S3LDS outperforms the TCN across all
labeled data amounts (Fig. 3D, solid lines) using these position features,
demonstrating how the addition of unlabeled frames can improve supervised
classification performance.

While modeling the raw pose estimates is intuitive, the SLDS framework can
accommodate any type of behavioral features, and has been used, for example,
with depth cameras [24, 25, 47, 48, 31, 49], latent representations from
autoencoders [7, 8], and other behavioral features [30]. While our goal here
is not to determine the optimal behavioral representation for our models, we
find that simply concatenating the pose vector $\mathbf{x}_{t}$ with its
velocity $\mathbf{x}_{t}-\mathbf{x}_{t-1}$ leads to a new behavioral
representation
$\tilde{\mathbf{x}}_{t}=\left[\mathbf{x}_{t};\mathbf{x}_{t}-\mathbf{x}_{t-1}\right]$
of position-velocity features which outperforms the position features
$\mathbf{x}_{t}$ in both model types (Fig. 3D, dashed lines). With this new
behavioral representation, the performance improvements due to unlabeled
frames disappear.

Figure 3: Supervised vs semi-supervised results for the head-fixed fly. A:
Example frame of the fly, overlaid with pose markers. B: Proportion of each
labeled behavior in the training dataset. C: Sample of ground truth labels,
along with predictions from both the TCN and the S3LDS models. Below is a
subset of the corresponding features used as inputs to the models. D: F1
scores for the TCN and S3LDS models. We show results for the position features
(solid lines) as well as the position-velocity features (dashed lines). Adding
velocity improves performance for both models. Error bars represent the
standard deviation of the F1 scores over five subsamples of the training data.
E: Confusion matrices for the TCN and S3LDS models. F: Average entropy of the
false positives (left) and true positives (right) for both models. In both
cases entropies are larger for the S3LDS model. These results are consistent
with the other datasets as shown in Fig. S2. Panels E and F only show results
from the models trained with position-velocity features.

To further investigate this, we next look at a snippet of model predictions
using $\tilde{\mathbf{x}}_{t}$ (Fig. 3C). One observation is that S3LDS makes
mistakes across a wider range of behavior classes. We compute the confusion
matrices for each model across all test videos, but do not see any obvious
structural differences between the models (Fig. 3E). However, another
observation is that the S3LDS predictions contain some rapid state switches.
One cause of these rapid switches would be two or more classes whose
probabilities are similar, and small amounts of noise can lead to switches
between these states. To test this hypothesis, we computed the entropy of the
discrete state probabilities for each action class. The models are evenly
matched for the false positives (Fig. 3F). However, for true positives the
entropy for the S3LDS is higher across all classes, indicating the semi-
supervised model is not as confident as the supervised one.

We find similar patterns in the remaining datasets: the S3LDS can outperform
the TCN when using position features (freely moving mouse), but the TCN
performs best across all datasets when incorporating temporal information with
the position-velocity features (Fig. 4, Fig. S2).

The finding that inclusion of temporal information in the observations leads
to improved classification performance raises a related question: are the
temporal components of the S3LDS necessary for accurate action segmentation?
To answer this question we perform an ablation study. We start with the semi-
supervised Gaussian Mixture Deep Generative Model (GMDGM) [51], a static
analogue of the S3LDS. The GMDGM was developed to model static data such as
images, and as such neither the approximate posteriors nor the generative
model take temporal information into account. As in the S3LDS, we allow for
some of the discrete states to be observed (Fig. S3A). To bridge the gap
between the GMDGM and the S3LDS we introduce the GMDGM-TCN, which uses the
static generative model of the GMDGM and the temporal context-aware
approximate posterior of the S3LDS (Fig. S3B). The GMDGM and GMDGM-TCN thus
highlight the effects of ablating different aspects of the dynamic structure
in the S3LDS.

Figure 4: Supervised vs semi-supervised results across datasets. Conventions
as in Fig. 3. As in the head-fixed fly, we find that using position-velocity
features improves performance over the position features across both model
types, and in all datasets the TCN performs best. A: Results on the freely
moving mouse dataset. Rather than using the raw markers, we compute the
features introduced in [18]. These features compute transformations on the raw
markers, including distances and angles between different groups of keypoints.
B: Results on the head-fixed mouse dataset. C: Results on the HuGaDB dataset.
The data is collected from sensors that already contain velocity data, so we
only use one set of features.

Across all datasets, we find the GMDGM performs much worse than the temporal
models when the observations do not contain temporal information (Fig. S4).
Except for the fly dataset, the GMDGM performance is improved upon by adding
temporal information into the inference network (the GMDGM-TCN).
Interestingly, adding temporal information to the generative model (the S3LDS)
does not meaningfully improve performance across any of the datasets,
regardless of the behavioral features used. We conclude that incorporating
temporal information in the observations and the approximate posterior provide
complementary improvements in model performance, and incorporating temporal
information into the generative model does not provide additional gains.

Finally, we look at the role of nonlinearities in the S3LDS. In our
implementation we use linear transformations for the recurrent state
transitions (Eq. 4) and the continuous latent space dynamics (Eq. 2). We also
implement a semi-supervised nonlinear dynamical system (S3NLDS) where we
replace each of these linear transformations with a one-hidden-layer dense
neural network (Methods). The flexibility of our amortized variational
inference scheme allows for such changes to the generative model with no
corresponding changes in the inference. We find that – at least for these
datasets – there is no benefit to using nonlinearities, regardless of the
behavioral representation (Fig. S5).

###  2.4 Adding sparse labels to an unsupervised method modifies the latent
space

We now turn to comparisons between the S3LDS and a closely related
unsupervised model, keypoint-MoSeq [28]. Keypoint-Moseq is also in the class
of SLDS models, but differs in several key aspects. Importantly, keypoint-
MoSeq uses a different prior on the discrete states that allows the model to
choose the number of states from the data; this also requires a very different
inference approach than the one developed for the S3LDS. Despite these
differences (among others), we are interested in comparing the discrete states
uncovered by these models.

We fit keypoint-Moseq on the position-velocity fly features, and find 37
discrete states. We then match each of these states to the labeled action
class that contains the most overlap in the training data, and find good
overlap with the ground truth labels on held-out test data (Fig. 5A). We
quantify this overlap with the same F1 score as before, and find that
keypoint-MoSeq achieves competitive results, but unsurprisingly the TCN and
S3LDS – which both have access to the hand labels during training – perform
better on this supervised metric (Fig. 5B).

It is also possible to investigate the effects of hand labels on the
continuous latent space of each model type. The latent embeddings are computed
by passing the observations through the trained encoders for the TCN and
S3LDS, and performing inference of the continuous latents in keypoint-MoSeq.
We first project the latents of each model in 2D space using UMAP [52] (Fig.
5C). The points with corresponding hand labels are colored, revealing that
similar behaviors are clustered together across all model types. To quantify
this observation, we next perform k-means clustering in the original latent
space for each model. We then compute a cluster homogeneity score [53] that
measures the extent to which the k-means clusters contain data points from a
single behavior class. The S3LDS achieves a higher score than keypoint-MoSeq,
demonstrating that the presence of hand labels during training does in fact
shift the continuous latent space to be more aligned with the labeled behavior
classes (Fig. 5D). The TCN, which is purely supervised, scores much higher
than the other two models since its only objective is to properly separate the
hand labels. These observations are replicated across the other datasets (Fig.
6) and with various other cluster quality metrics (Fig. S6). We show results
on the non-temporal features in the supplemental material (Figs. S7, S8, S9).

We close by noting the F1 score and cluster quality metrics are all based on
hand labels; a purely unsupervised method like keypoint-MoSeq is not trained
to optimize these types of metrics. Rather, keypoint-MoSeq (and related
methods) finds action classes defined by statistical regularities in the
behavioral dynamics. Our finding here is that these statistical regularities
may not match up to human-defined behavioral classes as precisely as models
trained specifically for this task.

Figure 5: Supervised and semi-supervised latent spaces more closely align with
labels than unsupervised latents. All models use position-velocity features
and all available training videos for the head-fixed fly dataset. A: The top
row shows a segment of ground truth labels. The following two rows show
predictions from the TCN and S3LDS models. The third row shows the state
outputs of keypoint-MoSeq (KPM), aligned to the ground truth class with
highest overlap on the training data. The final row shows the raw state
outputs of keypoint-MoSeq. B: F1 scores for the TCN, S3LDS and KPM models.
Error bars represent the standard deviation of the F1 scores over five trained
models (different initialization seeds). C: 2D UMAP embedding of continuous
latents colored by discrete labels for three different models. D: The addition
of hand labels produces more homogeneous clusters in the models’ latent
spaces. Error bars represent the standard deviation of the cluster scores over
five models. We use a range of cluster numbers to show that cluster scores are
not biased by cluster size.  Figure 6: Keypoint-MoSeq performance on non-fly
datasets: position-velocity features. Models are trained with position-
velocity features for all of the datasets. The mice datasets (panels A and B)
use position and velocity information constructed from the pose estimates, the
HuGaDB dataset uses inertial sensor data (panel C). Other conventions as in
Fig. 5. As in the fly dataset, we find the TCN, which is purely supervised,
achieves the highest alignment of the latent space with the ground truth
labels as measured by the cluster homogeneity score.

##  3 Discussion

We surveyed a range of animal action segmentation models spanning supervised,
semi-supervised, and unsupervised learning paradigms. Across four behavioral
datasets, we established that supervised TCN models trained with both static
and dynamic pose information performed the best on supervised classification
metrics. We also proposed the S3LDS, a semi-supervised method for animal
action segmentation that incorporates a small amount of discrete labels. With
the S3LDS we were able to selectively ablate different aspects of the model
and assess how they impacted performance; we found that including temporal
information in the inference network and the observations is critical, but
temporal information in the generative model did not lead to performance
improvements. Furthermore, we compared these models to the unsupervised
keypoint-MoSeq, and showed how the inclusion of labeled data during training
shapes the latent spaces of these models to be more aligned with labeled
classes.

A straightforward but important conclusion from our work is that the
behavioral representation used for action segmentation – whether supervised,
unsupervised, or something in between – plays a crucial role in the
performance of the algorithm. [18] did not use raw pose estimation outputs for
action segmentation, but rather computed a set of derived quantities like
distances and angles between different keypoints. We confirmed that this
representation is much more performant than the raw keypoints themselves (data
not shown), but also found that additionally using the velocities of these
features led to further improvements. But what are the “best” set of features
for any given dataset? The problem of featurization becomes even more complex
in the case of social interactions; for example, both the MARS [19] and SimBA
[21] packages compute hundreds of behavioral features within and between
animals as input to their classifiers. There has been some work on automating
this process. [9] use large amounts of unlabeled data to build an effective
feature extractor from pose data, and then use simple, fully supervised
classifiers for the final stage. Similarly, [10] use a set of unsupervised
feature extractors to bypass pose estimation and directly compute abstract
behavioral representations from video frames, which are then fed into fully
supervised classifiers.

An important direction for future work is to leverage the S3LDS framework for
semi-unsupervised learning [51, 54]; that is, where we provide a small number
of labels for some behavioral categories, but also provide the model with
capacity to discover completely novel classes simultaneously. This scenario
allows experimenters to precisely quantify known behaviors of interest while
also capturing previously unknown behaviors in the data. Another avenue for
future work involves modeling the social interactions between multiple
animals. Recent work has modeled the interactions between multiple brain
regions using rSLDS-like models [55, 56], and the flexibility of our approach
– as well as the relative ease with which we can label discrete animal
behaviors compared to discrete brain states – make the S3LDS an interesting
starting point for such modeling efforts.

##### Acknowledgements

We thank Peter Latham and Matteo Carandini for helpful discussions. This work
was supported by the following grants: Gatsby Charitable Foundation GAT3708,
NIH U19NS123716, NSF 1707398, The NSF AI Institute for Artificial and Natural
Intelligence, Simons Foundation and the Wellcome Trust 216324. The funders had
no role in study design, data collection and analysis, decision to publish, or
preparation of the manuscript.

##### Code availability

The code for the TCN, S3LDS, and GMDGM is available at
https://github.com/ablau100/daart under the GPL-3.0 license.

##### Data availability

The pose estimates and discrete state annotations for the freely moving mouse
dataset are available at https://github.com/ETHZ-INS/DLCAnalyzer; the raw
videos are available at https://zenodo.org/records/3608658.

To access the raw videos, pose estimates, and wheel traces for the IBL head-
fixed mouse dataset, see the documentation at https://int-brain-
lab.github.io/ONE/FAQ.html#how-do-i-download-the-datasets-cache-for-a-
specific-ibl-paper-release and use the tag 2024_Q2_Blau_et_al.

The IMU sensor data for the HuGaDB dataset can be found at
https://www.dropbox.com/scl/fi/jsq0mr26rthrzy64vkkjc/HuGaDB-v2.zip?rlkey=101j8lvdktdejm105cf9fpisi&e=1&dl=0.

The remaining data will be publicly released upon acceptance of this
manuscript for publication, or upon reasonable request.

##  4 Methods

###  4.1 Datasets

##### Head-fixed fly.

A head-fixed fly engaged in spontaneous behaviors on an air-supported ball
[32]. A camera captured a side view of the fly at 70 Hz. We tracked eight
points across the legs and abdomen using five supervised Lightning Pose
networks, followed by Ensemble Kalman Smoothing [6].

We consider five behavior categories: still, walk, front groom, back groom,
and abdomen move. We labeled behaviors in chunks of 50 to several hundred
contiguous time points using the DeepEthogram GUI [10]. Flies often engage in
behaviors for longer than 50 frames, so the selected chunks did not contain
any transitions from one behavior to another. We use five labeled videos for
training, and five for testing. The dataset contains 1.01M frames, only 16.6k
of which are labeled ($\sim$1.6%; Table S1). We train the models on 2, 3, 4,
and 5 of the training videos to see how our models perform with different
amounts of labeled data. For the models trained with 5 videos, we run each
model five times with different random seeds. For the other models, we
subsample different permutations of the training videos, i.e. for 2 videos, we
take 5 random permutations of 2 videos from the list of all 5 training videos.

##### Freely moving mouse.

In this publicly available dataset, a mouse freely moved around an open arena
[18]. A camera captured a top-down view of the mouse at 25 Hz. Thirteen points
were tracked across the tail, body, and head of the mouse. Instead of using
the raw pose data for observations, we use a set of 22 features computed from
the pose data based on distances, angles, and areas, following [18]. This
provides a position- and orientation-invariant representation that we did not
need to account for in the head-fixed fly dataset.

We consider three behavior categories: unsupported rearing, supported rearing
(when the mouse uses the wall of the arena), and grooming. In the rest of the
frames, there is no specific action or behavior displayed, so we do not train
any of the models to identify this “other” category. We only consider the
three specific behaviors described above in our analyses. We use 10 labeled
videos for training, and 10 for testing. The dataset contains 280k frames, 83k
of which are labeled ($\sim$29.5%; Table S2). We train our models on 4, 6, 8,
and 10 videos. We choose the permutations of these videos similar to the head-
fixed fly videos.

##### Head-fixed mouse.

This is a publicly available dataset from the International Brain Lab [33]. A
head-fixed mouse performed a perceptual decision-making task, using a wheel to
indicate its decision [57, 33]. Two cameras – “left” (60 Hz) and “right” (150
Hz) – capture roughly orthogonal side views of the mouse’s face and upper
trunk during each session [58], though we only use the left view. Multiple
points are tracked across the paws and face, though for this work we only
consider a single tracked point on the paw closest to the camera. In addition
to the paw coordinates, we also use the 1D wheel velocity as input to the
models. See example videos from one of the sessions here:
https://viz.internationalbrainlab.org/app?dset=bwm&pid=94fcff55-2da2-4366-a2c7-2f58c05b54dc&tid=57&cid=598&qc=0#trialviewer.

We consider four behavioral categories for this single paw: still, moving
without turning the wheel, moving while turning the wheel, and grooming. We
labeled behaviors using the DeepEthogram GUI [10]. We use five labeled videos
for training, and five for testing. The dataset contains 1M frames, 14k of
which are labeled ($\sim$1.4%; Table S3). We train the models on 2, 3, 4, and
5 of the training videos, as described for the other datasets.

##### Human Gait Database.

The Human Gait Database (HuGaDB) is a publicly available action segmentation
dataset that contains lower limb activities such as walking, running, and
sitting for a total of 18 subjects [34]. HuGaDB measured movements from six
inertial measurement units (IMUs) with a 60 Hz sampling frequency. The IMUs
were placed on each individual’s right and left feet, shins, and thighs. Each
inertial sensor tracks the acceleration data and gyroscope data on each of the
three x,y,z axes.

We consider eight behavioral categories: walking, running, going up, going
down, sitting, sitting down, standing up, and standing. The original dataset
contained four additional behaviors that we excluded from our analysis. Two of
those behaviors, “bicycling” and “sitting in car,” were not present in the
publicly available dataset. The IMU data for the classes “up by elevator” and
“down by elevator” did not seem to contain measurable variability. The dataset
contains 517k frames, all of which are labeled (Table S4). We train our model
on 100, 250, 500, and 1000 labeled frames from each class. For each number of
training frames, we randomly select the specified number of labeled frames,
and remove the rest of the labels for training.

###  4.2 Semi-supervised linear dynamical system (S3LDS)

####  4.2.1 S3LDS model formulation

Our starting point is the switching linear dynamical system (SLDS) defined in
Eqs. 1-4. The specific model that we implement further generalizes the above
model by allowing the transitions between discrete states to be recurrent, and
the observation mapping to be nonlinear, leading to a modified recurrent
switching linear dynamical system (rSLDS) model [59, 56]:

$\displaystyle p_{\theta}(\mathbf{x}_{1:T},y_{1:T},\mathbf{z}_{1:T})$
$\displaystyle=$ $\displaystyle
p_{\theta}(\mathbf{x}_{1}|\mathbf{z}_{1})p_{\theta}(\mathbf{z}_{1})p(y_{1})\prod_{t=2}^{T}p_{\theta}(\mathbf{x}_{t}|\mathbf{z}_{t})p_{\theta}(\mathbf{z}_{t}|\mathbf{z}_{t-1},y_{t})p_{\theta}(y_{t}|y_{t-1},\mathbf{z}_{t-1})$
(5) $\displaystyle p_{\theta}(\mathbf{x}_{t}|\mathbf{z}_{t})$ $\displaystyle=$
$\displaystyle\mathcal{N}(\mathbf{x}_{t}|g_{\theta}(\mathbf{z}_{t}),S)$ (6)
$\displaystyle p_{\theta}(\mathbf{z}_{t}|\mathbf{z}_{t-1},{y_{t}})$
$\displaystyle=$
$\displaystyle\mathcal{N}(\mathbf{z}_{t}|A_{y_{t}}\mathbf{z}_{t-1}+b_{{y}_{t}},Q_{y_{t}})$
(7) $\displaystyle p_{\theta}(y_{t}|y_{t-1},\mathbf{z}_{t-1})$
$\displaystyle=$
$\displaystyle\text{Cat}(y_{t}|\text{Softmax}(R_{y_{t-1}}\mathbf{z}_{t-1}+r_{{y}_{t-1}}))$
(8) $\displaystyle p(\mathbf{z}_{1})$ $\displaystyle=$
$\displaystyle\mathcal{N}(\mathbf{z}_{1}|0,\mathbb{I})$ (9) $\displaystyle
p(y_{1})$ $\displaystyle=$ $\displaystyle\text{Cat}(y_{1}|\pi),$ (10)

where $R_{y_{t}}$ and $r_{y_{t}}$ define the recurrent transition associated
with state $y_{t}$, and $g_{\theta}$ is a one-hidden-layer dense neural
network that is shared across all discrete states. For our investigation of
nonlinearities (Fig. S5), we replaced each of the linear transformations in
Eqs. 7 and 8 with a one-hidden-layer dense neural network.

####  4.2.2 S3LDS inference and learning

Switching models contain discrete latent variables, and therefore we cannot
simply use the reparameterization trick commonly employed in variational
inference schemes. One option is to use a continuous relaxation of the
discrete variables such as the Gumbel Softmax [60], Concrete [61], or REBAR
[62]. However, in practice this approach has not been successful with
switching nonlinear models (SNLDS) [59]. Both [59, 56] use the exact posterior
of the discrete states under the generative model, and perform inference using
the forward-backward algorithm. We take a different approach and marginalize
out the discrete state such that we obtain a soft mixture of dynamics at each
time step, similar to the Kalman VAE [63].

We define the approximate posteriors for this model in a way that leads to a
classifier $q_{\phi}(y_{t}|\mathbf{x}_{1:T})$, which is trained on both
labeled and unlabeled data [51]. We also incorporate temporal information in
our approximate posteriors to more closely match the generative model [64,
65]. We propose to do so with a temporal convolution network [37, 38]. To
begin, let us assume we have no observed labels $y_{t}$; we define a joint
posterior over $\\{y_{1:T},\mathbf{z}_{1:T}\\}$ that factorizes over time:

$q_{\phi}(y_{1:T},\mathbf{z}_{1:T}|\mathbf{x}_{1:T})=\prod_{t=1}^{T}q_{\phi}(y_{t},\mathbf{z}_{t}|\mathbf{x}_{\mathcal{T}_{t}})$
(11)

where $\mathcal{T}_{t}$ defines a set of time points, for instance
$\\{1,\ldots,T\\}$ if we condition on all observations, or
$\\{t-\tau,\ldots,t+\tau\\}$ if we condition on a window of observations
centered at time $t$. We further factorize
$q_{\phi}(y_{t},\mathbf{z}_{t}|\mathbf{x}_{\mathcal{T}_{t}})$ as

$\displaystyle q_{\phi}(y_{t},\mathbf{z}_{t}|\mathbf{x}_{\mathcal{T}_{t}})$
$\displaystyle=$ $\displaystyle
q_{\phi}(y_{t}|\mathbf{x}_{\mathcal{T}_{t}})q_{\phi}(\mathbf{z}_{t}|\mathbf{x}_{\mathcal{T}_{t}},y_{t})$
(12) $\displaystyle q_{\phi}(y_{t}|\mathbf{x}_{\mathcal{T}_{t}})$
$\displaystyle=$
$\displaystyle\text{Cat}(y_{t}|\pi_{\phi}(\mathbf{x}_{\mathcal{T}_{t}}))$ (13)
$\displaystyle q_{\phi}(\mathbf{z}_{t}|\mathbf{x}_{\mathcal{T}_{t}},y_{t})$
$\displaystyle=$
$\displaystyle\mathcal{N}(\mathbf{z}_{t}|\mu_{\phi}(\mathbf{x}_{\mathcal{T}_{t}},y_{t}),\text{diag}(\sigma_{\phi}^{2}(\mathbf{x}_{\mathcal{T}_{t}},y_{t}))).$
(14)

This combination of amortization and variational approximation results in the
distribution $q_{\phi}(y_{t}|\mathbf{x}_{\mathcal{T}_{t}})$ which can be
interpreted as a classifier that can efficiently predict the discrete behavior
state after training the model is complete.

There has been much work on developing inference strategies for rSLDS models
[50, 66, 67]. Our contribution here is the development of an efficient and
flexible amortized variational inference strategy that incorporates both
labeled and unlabeled data. We note that the proposed inference strategy leads
to the definition of an approximate posterior
$q_{\phi}(y_{t}|\mathbf{x}_{\mathcal{T}_{t}})$. This approximate posterior –
which is trained on both labeled and unlabeled data – also serves as a
classifier, predicting the discrete action class $y_{t}$ from the observed
data $\mathbf{x}_{1:T}$ using a neural network [51]. Existing inference
methods for rSLDS models do not use labeled data or an amortized inference
network [50, 66, 67]. We use a TCN [37, 38] as our model backbone, as
motivated in a previous section. This inference approach allows us to directly
compare classifiers that are trained with labeled data only (i.e., supervised)
or with both labeled and unlabeled data (i.e., semi-supervised).

##### ELBO with all labels observed.

Let us first consider the variational lower bound assuming all $y_{t}$ are
observed (details in Appendix A). We will use $\ \tilde{}\ $ to denote that a
variable is sampled from its approximate posterior using the
reparameterization trick [68]:

$\displaystyle\mathcal{L}_{l}(\mathbf{x}_{1:T},y_{1:T})=$
$\displaystyle\sum_{t=1}^{T}{\log p(\mathbf{x}_{t}|\tilde{\mathbf{z}}_{t}})$
$\displaystyle-\text{KL}\left[q(\mathbf{z}_{1}|\mathbf{x}_{\mathcal{T}_{1}},y_{1})||p(\mathbf{z}_{1})\right]+\log
p(y_{1})$
$\displaystyle-\sum_{t=2}^{T}{\text{KL}\left[q(\mathbf{z}_{t}|\mathbf{x}_{\mathcal{T}_{t}},y_{t})||p(\mathbf{z}_{t}|\tilde{\mathbf{z}}_{t-1},y_{t})\right]}$
$\displaystyle+\sum_{t=2}^{T}\log p(y_{t}|y_{t-1},\tilde{\mathbf{z}}_{t-1}).$
(15)

Note that Eq. 4.2.2 does not explicitly contain a term that looks like a
classifier; we will address this in a later section.

##### ELBO with all labels unobserved.

Next, let us derive the variational lower bound assuming all $y_{t}$ are
unobserved (details in Appendix A).

$\displaystyle\mathcal{L}_{u}(\mathbf{x}_{1:T})=$
$\displaystyle\sum_{t=1}^{T}\sum_{k=1}^{K}\Big{(}q(y_{t}=k|\mathbf{x}_{\mathcal{T}_{t}}){\log
p(\mathbf{x}_{t}|{\tilde{\mathbf{z}}}_{t}^{k}})\Big{)}$
$\displaystyle-\sum_{k=1}^{K}\Big{(}q(y_{1}=k|\mathbf{x}_{\mathcal{T}_{t}}){\text{KL}\left[q(\mathbf{z}_{1}|\mathbf{x}_{\mathcal{T}_{1}},y_{1}=k)||p(\mathbf{z}_{1})\right]}\Big{)}-\text{KL}\left[q(y_{1}|\mathbf{x}_{\mathcal{T}_{1}})||p(y_{1})\right]$
$\displaystyle-\sum_{t=2}^{T}\sum_{k=1}^{K}\sum_{k^{\prime}=1}^{K}\Big{(}q(y_{t}=k|\mathbf{x}_{\mathcal{T}_{t}})q(y_{t-1}=k^{\prime}|\mathbf{x}_{\mathcal{T}_{t-1}})$
$\displaystyle\quad\quad{\text{KL}\left[{q(\mathbf{z}_{t}|\mathbf{x}_{\mathcal{T}_{t}},y_{t}=k)}||{p(\mathbf{z}_{t}|\tilde{\mathbf{z}}_{t-1}^{k^{\prime}},y_{t}=k)}\right]}\Big{)}$
$\displaystyle-\sum_{t=2}^{T}\sum_{k=1}^{K}\Big{(}q(y_{t-1}=k|\mathbf{x}_{\mathcal{T}_{t-1}}){\text{KL}\left[q(y_{t}|\mathbf{x}_{\mathcal{T}_{t}})||{p(y_{t}|y_{t-1}=k,\tilde{\mathbf{z}}_{t-1}^{k})}\right]}\Big{)}$
(16)

where $\tilde{\mathbf{z}}_{t}^{k}\sim
q(\mathbf{z}_{t}|\mathbf{x}_{\mathcal{T}_{t}},y_{t}=k)$ and
$\tilde{\mathbf{z}}_{t-1}^{k}\sim
q(\mathbf{z}_{t-1}|\mathbf{x}_{\mathcal{T}_{t-1}},y_{t-1}=k)$. A nice property
of this lower bound is that each of the individual terms in the unlabeled ELBO
(Eq. 4.2.2) reduces to the corresponding term in the labeled ELBO (Eq. 4.2.2)
when the output of $q(y_{t}|\mathbf{x}_{\mathcal{T}_{t}})$ is a one-hot vector
(i.e. observed), which we exploit when computing the full semi-supervised ELBO
(more below).

##### Classification loss.

Note that Eq. 4.2.2 contains terms of the form
$q(y_{t}|\mathbf{x}_{\mathcal{T}_{t}})$, which maps observations to the
discrete state. This part of the approximate posterior is a classifier that
only appears in the unlabeled ELBO, and hence does not utilize labeled data
for learning. Therefore, we can improve classification performance by
including a cross entropy loss term for our model to learn the classes. We
weight this loss with a hyperparameter $\alpha$, chosen through cross-
validation.

##### Total loss function.

The total loss function for the general semi-supervised case will be a
combination of Eqs. 4.2.2 and 4.2.2 depending on which subset of time points
include observed discrete labels. In practice we compute
$q(y_{t}|\mathbf{x}_{\mathcal{T}_{t}})$ for all time points, then replace the
probability vector with a one-hot vector when the discrete label is observed,
and then compute Eq. 4.2.2. In addition, we compute the classification loss
for all observed discrete labels. See Appendix A for full details.

###  4.3 Gaussian mixture deep generative model (GMDGM)

For one of our baselines we implement the Gaussian mixture deep generative
model (GMDGM) proposed in [51], which is similar to the S3LDS but without
temporal dependencies (Fig. S3A). This is a modified version of the semi-
supervised deep generative model proposed in [69].

####  4.3.1 GMDGM model formulation

The generative model is defined as

$\displaystyle p_{\theta}(\mathbf{x},y,\mathbf{z})$ $\displaystyle=$
$\displaystyle p_{\theta}(\mathbf{x}|\mathbf{z})p_{\theta}(\mathbf{z}|y)p(y)$
(17) $\displaystyle p_{\theta}(\mathbf{x}|\mathbf{z})$ $\displaystyle=$
$\displaystyle\mathcal{N}(\mathbf{x}|g_{\theta}(\mathbf{z}),R)$ (18)
$\displaystyle p_{\theta}(\mathbf{z}|y)$ $\displaystyle=$
$\displaystyle\mathcal{N}(\mathbf{z}|f_{\theta}(y),\text{diag}(\sigma_{\theta}^{2}(y)))$
(19) $\displaystyle p(y)$ $\displaystyle=$ $\displaystyle\text{Cat}(y|\pi).$
(20)

This model defines an explicit clustering mechanism by conditioning the
latents $\mathbf{z}$ on the discrete label $y$, and defining a categorical
prior over $y$. Therefore, even in the absence of labeled data, the GMDGM will
attempt to cluster the data.

####  4.3.2 GMDGM inference and learning

Like the S3LDS we perform approximate inference in the GMDGM model.
Furthermore, we structure the approximate posteriors in a way that leads to a
classifier $q_{\phi}(y|\mathbf{x})$, which is trained on both labeled and
unlabeled data.

$\displaystyle q_{\phi}(y,\mathbf{z}|\mathbf{x})$ $\displaystyle=$
$\displaystyle q_{\phi}(y|\mathbf{x})q_{\phi}(\mathbf{z}|\mathbf{x},y)$ (21)
$\displaystyle q_{\phi}(y|\mathbf{x})$ $\displaystyle=$
$\displaystyle\text{Cat}(y|\pi_{\phi}(\mathbf{x}))$ (22) $\displaystyle
q_{\phi}(\mathbf{z}|\mathbf{x},y)$ $\displaystyle=$
$\displaystyle\mathcal{N}(\mathbf{z}|\mu_{\phi}(\mathbf{x},y),\text{diag}(\sigma_{\phi}^{2}(\mathbf{x},y))).$
(23)

For the GMDGM-TCN model we use the same approximate posteriors as the S3LDS
model, defined in Eqs. 12-14 (Fig. S3B).

The variational lower bound for labeled data is

$\displaystyle\mathcal{L}_{l}(\mathbf{x},y)$ $\displaystyle=$
$\displaystyle\mathbb{E}_{q(\mathbf{z}|\mathbf{x},y)}\left[\log\frac{p(\mathbf{x}|\mathbf{z})p(\mathbf{z}|y)p(y)}{q(\mathbf{z}|\mathbf{x},y)}\right]$
(24) $\displaystyle=$
$\displaystyle\mathbb{E}_{q(\mathbf{z}|\mathbf{x},y)}\left[\log
p(\mathbf{x}|\mathbf{z})\right]+\log
p(y)-\text{KL}[q(\mathbf{z}|\mathbf{x},y)||p(\mathbf{z}|y)]$
$\displaystyle\approx$ $\displaystyle\log
p(\mathbf{x}|\tilde{\mathbf{z}})+\log
p(y)-\text{KL}[q(\mathbf{z}|\mathbf{x},y)||p(\mathbf{z}|y)]$

The variational lower bound for unlabeled data is

$\displaystyle\mathcal{L}_{u}(\mathbf{x})$ $\displaystyle=$
$\displaystyle\mathbb{E}_{q(y,\mathbf{z}|\mathbf{x})}\left[\log\frac{p(\mathbf{x}|\mathbf{z})p(\mathbf{z}|y)p(y)}{q(y,\mathbf{z}|\mathbf{x})}\right]$
(25) $\displaystyle=$
$\displaystyle\mathbb{E}_{q(y,\mathbf{z}|\mathbf{x})}\left[\log\frac{p(\mathbf{x}|\mathbf{z})p(\mathbf{z}|y)p(y)}{q(y|\mathbf{x})q(\mathbf{z}|\mathbf{x},y)}\right]$
$\displaystyle=$
$\displaystyle\sum_{y}q(y|\mathbf{x})\left[\mathbb{E}_{q(\mathbf{z}|\mathbf{x},y)}\left[\log\frac{p(\mathbf{x}|\mathbf{z})p(\mathbf{z}|y)p(y)}{q(y|\mathbf{x})q(\mathbf{z}|\mathbf{x},y)}\right]\right]$
$\displaystyle=$
$\displaystyle\sum_{y}q(y|\mathbf{x})\left[\mathbb{E}_{q(\mathbf{z}|\mathbf{x},y)}\left[\log\frac{p(\mathbf{x}|\mathbf{z})p(\mathbf{z}|y)p(y)}{q(\mathbf{z}|\mathbf{x},y)}\right]-\log
q(y|\mathbf{x})\right]$ $\displaystyle=$
$\displaystyle\sum_{y}q(y|\mathbf{x})\left[\mathcal{L}_{l}(\mathbf{x},y)-\log
q(y|\mathbf{x})\right]$ $\displaystyle=$
$\displaystyle\sum_{y}q(y|\mathbf{x})\mathcal{L}_{l}(\mathbf{x},y)-H(q(y|\mathbf{x}))$
$\displaystyle=$
$\displaystyle\sum_{k=1}^{K}q(y=k|\mathbf{x})\mathcal{L}_{l}(\mathbf{x},k)-H(q(y|\mathbf{x})).$

As with the S3LDS, note that the term $q(y|\mathbf{x})$, which we will use as
a final classifier, only appears in $\mathcal{L}_{u}(\mathbf{x})$ and is
therefore not trained with labeled data. We again address this by adding the
cross entropy loss between the ground truth label and the distribution
provided by $q(y|\mathbf{x})$ to the lower bound of the labeled data, so that

$\displaystyle\mathcal{L}=\mathbb{E}_{\mathbf{x}^{(l)},y^{(l)}\sim\mathcal{D}_{l}}\left[\mathcal{L}_{l}(\mathbf{x}^{(l)},y^{(l)})+\alpha\log
q(y^{(l)}|\mathbf{x}^{(l)})\right]+\mathbb{E}_{\mathbf{x}^{(u)}\sim\mathcal{D}_{u}}\left[\mathcal{L}_{u}(\mathbf{x}^{(u)})\right],$
(26)

where $\mathcal{D}_{l}$ is the labeled dataset and $\mathcal{D}_{u}$ is the
unlabeled dataset.

###  4.4 Model implementation details

####  4.4.1 Temporal Convolutional Network backbone

We used the same dilated TCN backbone for the supervised model, GMDGM models,
and S3LDS. For all models and datasets we used 2 dilation blocks, 4 temporal
lags per convolution layer, and 32 filters per convolution layer. Each
dilation block consisted of a sequence of 2 sub-blocks (1D convolution layer
$\rightarrow$ leaky ReLU nonlinearity $\rightarrow$ dropout with
probability=0.10), as well as a residual connection between the input and
output of the dilation block. The dilation of the convolutional filters starts
with 1 for the first dilation block, then increases by a factor of 2 for each
additional dilation block. This results in a larger temporal receptive field
as the model gets deeper, allowing for learning of longer range dependencies
[70]. For our specific parameters this results in outputs of the network at
time $t$ being dependent on time points $\\{t-24,\ldots,t,\ldots,t+24\\}$, or
in other words a 0.7 s receptive field for the fly data (70 Hz), a 1.96 s
receptive field for the freely moving mouse data (25 Hz), and a 0.82 s
receptive field for the head-fixed mouse and human data (both 60 Hz). We
experimented with different values of these hyperparameters but found the
results to be robust to specific choices.

####  4.4.2 S3LDS model

The S3LDS model used two TCN networks as described above: one for the
approximate posterior of the discrete states
$q(y_{t}|\mathbf{x}_{\mathcal{T}_{t}})$, and one for the approximate posterior
of the continuous latents
$q(\mathbf{z}_{t}|\mathbf{x}_{\mathcal{T}_{t}},\tilde{y}_{t})$. We used a one-
hidden-layer dense neural network for all the generative model nonlinearities
in the S3LDS and S3NLDS. The number of hidden units was always equal to the
dimensionality of the continuous latent space. To determine this, we performed
PCA on the observed data and chose the number of dimensions that explained 95%
of the variance.

####  4.4.3 Model training

To train the supervised TCN, GMDGM, and S3LDS models we used the Adam
optimizer [71] with a learning rate of $0.0001$. Each batch contained 8
sequences of 1000 frames each. We trained the models for 500 epochs. For the
semi-supervised models, we tested a range of possible $\alpha$s (the
hyperparameter on the supervised classification loss). We chose $\alpha=100$
for all of the datasets since that value generalized well across datasets. We
anneal the two KL losses and $\log{y}$ (for the labeled data) by linearly
increasing the weight from 0 to 1 over 100 epochs.

###  4.5 Tree-based methods

For the tree-based methods we used the same input features as the other
models. We used a temporal window of features to match the input to the first
layer of the TCN, i.e. the prediction at time $t$ is made from the
concatenated behavioral features from $\\{t-4,\ldots,t,\ldots,t+4\\}$.

To train the random forest models we used the function RandomForestClassifier
from the sklearn package with the following arguments: n_estimators=6000,
max_features=‘sqrt’, criterion=‘entropy’, min_samples_leaf=1, bootstrap=True
following [21].

To train the XGBoost models we used the function XGBClassifier from the
xgboost package with the following arguments: n_estimators=2000, max_depth=3,
learning_rate=0.1, objective=‘multi:softprob’, eval_metric=‘mlogloss’,
tree_method=‘hist’, gamma=1, min_child_weight=1, subsample=0.8,
colsample_bytree=0.8 following [19].

###  4.6 Keypoint-MoSeq

For keypoint-MoSeq, we use both the position features (Fig. S7) as well as the
position-velocity features (Figs. 5, 6) as inputs for the head-fixed fly,
freely moving mouse, and head-fixed mouse datasets. For the HuGaDB data, we
use the IMU features. We used the keypoint_moseq library to train the models.
We initialized the model with ARHMM run for fifty iterations followed by
training the full keypoint-MoSeq model for 500 additional iterations. We set
kappa=1e4 for all models.

## References

  * [1] David J Anderson and Pietro Perona  “Toward a science of computational ethology”  In _Neuron_ 84.1 Elsevier, 2014, pp. 18–31
  * [2] Talmo D Pereira, Joshua W Shaevitz and Mala Murthy  “Quantifying behavior to understand the brain”  In _Nature neuroscience_ Nature Publishing Group, 2020, pp. 1–13
  * [3] Lukas Ziegler, Oliver Sturman and Johannes Bohacek  “Big behavior: challenges and opportunities in a new era of deep behavior profiling”  In _Neuropsychopharmacology_ 46.1 Nature Publishing Group, 2021, pp. 33–44
  * [4] Alexander Mathis et al.  “DeepLabCut: markerless pose estimation of user-defined body parts with deep learning”  In _Nature neuroscience_ 21.9 Nature Publishing Group, 2018, pp. 1281–1289
  * [5] Talmo D Pereira et al.  “Fast animal pose estimation using deep neural networks”  In _Nature methods_ 16.1 Nature Publishing Group US New York, 2019, pp. 117–125
  * [6] Dan Biderman et al.  “Lightning Pose: improved animal pose estimation via semi-supervised learning, Bayesian ensembling and cloud-native open-source tools”  In _Nature Methods_ Nature Publishing Group US New York, 2024, pp. 1–13
  * [7] Eleanor Batty et al.  “BehaveNet: nonlinear embedding and Bayesian neural decoding of behavioral videos”  In _Advances in Neural Information Processing Systems_ 32, 2019
  * [8] Matthew R Whiteway et al.  “Partitioning variability in animal behavioral videos using semi-supervised variational autoencoders”  In _PLoS computational biology_ 17.9 Public Library of Science San Francisco, CA USA, 2021, pp. e1009439
  * [9] Jennifer J. Sun et al.  “Task Programming: Learning Data Efficient Behavior Representations”  arXiv, 2020
  * [10] James P. Bohnslav et al.  “DeepEthogram: a machine learning pipeline for supervised behavior classification from raw pixels”  In _eLife_ eLife Sciences Publications Ltd, 2021
  * [11] Hueihan Jhuang et al.  “Automated home-cage behavioural phenotyping of mice”  In _Nature communications_ 1.1 Nature Publishing Group, 2010, pp. 1–10
  * [12] Mayank Kabra et al.  “JAABA: interactive machine learning for automatic annotation of animal behavior”  In _Nature methods_ 10.1 Nature Publishing Group, 2013, pp. 64
  * [13] Gregory Kramida et al.  “Automated mouse behavior recognition using VGG features and LSTM networks”  In _Proc. Vis. Observ. Anal. Vertebrate Insect Behav. Workshop (VAIB)_ , 2016, pp. 1–3
  * [14] Van Anh Le and Kartikeya Murari  “Recurrent 3D Convolutional Network for Rodent Behavior Recognition”  In _ICASSP 2019 - 2019 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)_ , 2019, pp. 1174–1178
  * [15] Kartikeya Murari  “Recurrent 3D Convolutional Network for Rodent Behavior Recognition”  In _ICASSP 2019-2019 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)_ , 2019, pp. 1174–1178  IEEE
  * [16] Ngoc G Nguyen et al.  “Applying deep learning models to mouse behavior recognition”  In _Journal of Biomedical Science and Engineering_ 12.2 Scientific Research, 2019, pp. 183–196
  * [17] Elsbeth A Dam, Lucas PJJ Noldus and Marcel AJ Gerven  “Deep learning improves automated rodent behavior recognition within a specific experimental setup”  In _Journal of neuroscience methods_ 332 Elsevier, 2020, pp. 108536
  * [18] Oliver Sturman et al.  “Deep learning-based behavioral analysis reaches human accuracy and is capable of outperforming commercial solutions”  In _Neuropsychopharmacology_ 45.11, 2020, pp. 1942–1952
  * [19] Cristina Segalin et al.  “The Mouse Action Recognition System (MARS) software pipeline for automated analysis of social behaviors in mice”  In _Elife_ 10 eLife Sciences Publications Limited, 2021, pp. e63720
  * [20] Christopher J Gabriel et al.  “BehaviorDEPOT is a simple, flexible tool for automated behavioral detection based on markerless pose tracking”  In _Elife_ 11 eLife Sciences Publications Limited, 2022, pp. e74314
  * [21] Nastacia L Goodwin et al.  “Simple Behavioral Analysis (SimBA) as a platform for explainable machine learning in behavioral neuroscience”  In _Nature Neuroscience_ Nature Publishing Group US New York, 2024, pp. 1–14
  * [22] Alex Gomez-Marin et al.  “Big behavioral data: psychology, ethology and the foundations of neuroscience”  In _Nature neuroscience_ 17.11 Nature Publishing Group, 2014, pp. 1455–1462
  * [23] Gordon J Berman, Daniel M Choi, William Bialek and Joshua W Shaevitz  “Mapping the stereotyped behaviour of freely moving fruit flies”  In _Journal of The Royal Society Interface_ 11.99 The Royal Society, 2014, pp. 20140672
  * [24] Alexander B Wiltschko et al.  “Mapping sub-second structure in mouse behavior”  In _Neuron_ 88.6 Elsevier, 2015, pp. 1121–1135
  * [25] Matthew J Johnson et al.  “Composing graphical models with neural networks for structured representations and fast inference”  In _Advances in neural information processing systems_ 29, 2016
  * [26] Alexander I Hsu and Eric A Yttri  “B-SOiD, an open-source unsupervised algorithm for identification and fast prediction of behaviors”  In _Nature communications_ 12.1 Nature Publishing Group UK London, 2021, pp. 5188
  * [27] Kevin Luxem et al.  “Identifying behavioral structure from deep variational embeddings of animal motion”  In _Communications Biology_ 5.1 Nature Publishing Group UK London, 2022, pp. 1267
  * [28] Caleb Weinreb et al.  “Keypoint-MoSeq: parsing behavior by linking point tracking to pose dynamics”  In _Nature Methods_ 21 Epub, 2024, pp. 1329–1339
  * [29] Sandeep Robert Datta et al.  “Computational neuroethology: a call to action”  In _Neuron_ 104.1 Elsevier, 2019, pp. 11–24
  * [30] E Kelly Buchanan, Akiva Lipschitz, Scott W Linderman and Liam Paninski  “Quantifying the behavioral dynamics of C. elegans with autoregressive hidden Markov models”  In _Workshop on Worm’s neural information processing at the 31st conference on neural information processing systems_ , 2017
  * [31] Julia Costacurta et al.  “Distinguishing discrete and continuous behavioral variability using warped autoregressive HMMs”  In _Advances in Neural Information Processing Systems_ 35, 2022, pp. 23838–23850
  * [32] Evan S Schaffer et al.  “The spatial and temporal structure of neural activity across the fly brain”  In _Nature Communications_ 14.1 Nature Publishing Group UK London, 2023, pp. 5572
  * [33] IBL et al.  “A brain-wide map of neural activity during complex behaviour”  In _bioRxiv_ Cold Spring Harbor Laboratory, 2023, pp. 2023–07
  * [34] Roman Chereshnev and Attila Kert’esz-Farkas  “HuGaDB: Human Gait Database for Activity Recognition from Wearable Inertial Sensor Networks”  In _International Conference on Analysis of Images, Social Networks and Texts_ , 2017, pp. 131–141  Springer
  * [35] Jennifer J Sun et al.  “The multi-agent behavior dataset: Mouse dyadic social interactions”  In _arXiv preprint arXiv:2104.02710_ , 2021
  * [36] Claire P Everett et al.  “Coordination and persistence of aggressive visual communication in Siamese fighting fish”  In _bioRxiv_ Cold Spring Harbor Laboratory, 2024, pp. 2024–04
  * [37] Colin Lea, Rene Vidal, Austin Reiter and Gregory D Hager  “Temporal convolutional networks: A unified approach to action segmentation”  In _European Conference on Computer Vision_ , 2016, pp. 47–54  Springer
  * [38] Colin Lea et al.  “Temporal convolutional networks for action segmentation and detection”  In _proceedings of the IEEE Conference on Computer Vision and Pattern Recognition_ , 2017, pp. 156–165
  * [39] Leo Breiman  “Random Forests”  In _Machine Learning_ 45.1 USA: Kluwer Academic Publishers, 2001, pp. 5–32
  * [40] Tianqi Chen and Carlos Guestrin  “Xgboost: A scalable tree boosting system”  In _Proceedings of the 22nd acm sigkdd international conference on knowledge discovery and data mining_ , 2016, pp. 785–794
  * [41] Tae Soo Kim and Austin Reiter  “Interpretable 3d human action analysis with temporal convolutional networks”  In _2017 IEEE conference on computer vision and pattern recognition workshops (CVPRW)_ , 2017, pp. 1623–1631  IEEE
  * [42] Yazan Abu Farha and Jurgen Gall  “Ms-tcn: Multi-stage temporal convolutional network for action segmentation”  In _Proceedings of the IEEE/CVF conference on computer vision and pattern recognition_ , 2019, pp. 3575–3584
  * [43] Benjamin Filtjens, Bart Vanrumste and Peter Slaets  “Skeleton-based action segmentation with multi-stage spatial-temporal graph convolutional neural networks”  In _IEEE Transactions on Emerging Topics in Computing_ IEEE, 2022
  * [44] Guy A Ackerson and King-Sun Fu  “On state estimation in switching environments”  In _IEEE Transactions on Automatic Control_ , 1970
  * [45] Chaw-Bing Chang and Michael Athans  “State estimation for discrete systems with switching parameters”  In _IEEE Transactions on Aerospace and Electronic Systems_ , 1978
  * [46] James D Hamilton  “Analysis of time series subject to changes in regime”  In _Journal of econometrics_ , 1990
  * [47] Jeffrey E Markowitz et al.  “The striatum organizes 3D behavior via moment-to-moment action selection”  In _Cell_ 174.1 Elsevier, 2018, pp. 44–58
  * [48] Alexander B. Wiltschko et al.  “Revealing the structure of pharmacobehavioral space through motion sequencing”  In _Nature Neuroscience_ 23.11, 2020
  * [49] Jeffrey E Markowitz et al.  “Spontaneous behaviour is structured by reinforcement without explicit reward”  In _Nature_ 614.7946 Nature Publishing Group UK London, 2023, pp. 108–117
  * [50] Scott W. Linderman et al.  “Recurrent switching linear dynamical systems”  arXiv, 2016
  * [51] Matthew Willetts, Stephen Roberts and Chris Holmes  “Semi-Unsupervised Learning: Clustering and Classifying using Ultra-Sparse Labels”  In _2020 IEEE International Conference on Big Data (Big Data)_ , 2020, pp. 5286–5295  IEEE
  * [52] Leland McInnes, John Healy and James Melville  “Umap: Uniform manifold approximation and projection for dimension reduction”  In _arXiv preprint arXiv:1802.03426_ , 2018
  * [53] Andrew Rosenberg and Julia Hirschberg  “V-measure: A conditional entropy-based external cluster evaluation measure”  In _Proceedings of the 2007 joint conference on empirical methods in natural language processing and computational natural language learning (EMNLP-CoNLL)_ , 2007, pp. 410–420
  * [54] Padraig Davidson et al.  “Semi-unsupervised Learning: An In-depth Parameter Analysis”  In _German Conference on Artificial Intelligence (Künstliche Intelligenz)_ , 2021, pp. 51–66  Springer
  * [55] Joshua Glaser et al.  “Recurrent switching dynamical systems models for multiple interacting neural populations”  In _Advances in neural information processing systems_ 33, 2020, pp. 14867–14878
  * [56] Orren Karniol-Tambour et al.  “Modeling state-dependent communication between brain regions with switching nonlinear dynamical systems”  In _The Twelfth International Conference on Learning Representations_ , 2024
  * [57] IBL et al.  “Standardized and reproducible measurement of decision-making in mice”  In _Elife_ 10 eLife Sciences Publications, Ltd, 2021, pp. e63711
  * [58] IBL et al.  “Video hardware and software for the International Brain Laboratory”  In _figshare_ , 2022  DOI: 10.6084/m9.figshare.19694452
  * [59] Zhe Dong, Bryan Seybold, Kevin Murphy and Hung Bui  “Collapsed amortized variational inference for switching nonlinear dynamical systems”  In _International Conference on Machine Learning_ , 2020, pp. 2638–2647  PMLR
  * [60] Eric Jang, Shixiang Gu and Ben Poole  “Categorical reparameterization with gumbel-softmax”  In _arXiv preprint arXiv:1611.01144_ , 2016
  * [61] Chris J Maddison, Andriy Mnih and Yee Whye Teh  “The concrete distribution: A continuous relaxation of discrete random variables”  In _arXiv preprint arXiv:1611.00712_ , 2016
  * [62] George Tucker et al.  “Rebar: Low-variance, unbiased gradient estimates for discrete latent variable models”  In _Advances in Neural Information Processing Systems_ 30, 2017
  * [63] Marco Fraccaro, Simon Kamronn, Ulrich Paquet and Ole Winther  “A disentangled recognition and nonlinear dynamics model for unsupervised learning”  In _Advances in neural information processing systems_ 30, 2017
  * [64] Rahul G Krishnan, Uri Shalit and David Sontag  “Deep kalman filters”  In _arXiv preprint arXiv:1511.05121_ , 2015
  * [65] Rahul Krishnan, Uri Shalit and David Sontag  “Structured inference networks for nonlinear state space models”  In _Proceedings of the AAAI Conference on Artificial Intelligence_ 31.1, 2017
  * [66] Scott Linderman et al.  “Bayesian learning and inference in recurrent switching linear dynamical systems”  In _Artificial Intelligence and Statistics_ , 2017, pp. 914–922  PMLR
  * [67] Josue Nassar, Scott W Linderman, Monica Bugallo and Il Memming Park  “Tree-structured recurrent switching linear dynamical systems for multi-scale modeling”  In _arXiv preprint arXiv:1811.12386_ , 2018
  * [68] Diederik P Kingma and Max Welling  “Auto-encoding variational bayes”  In _arXiv preprint arXiv:1312.6114_ , 2013
  * [69] Durk P Kingma, Shakir Mohamed, Danilo Jimenez Rezende and Max Welling  “Semi-supervised learning with deep generative models”  In _Advances in Neural Information Processing Systems_ , 2014, pp. 3581–3589
  * [70] Fisher Yu and Vladlen Koltun  “Multi-scale context aggregation by dilated convolutions”  In _arXiv preprint arXiv:1511.07122_ , 2015
  * [71] Diederik P Kingma and Jimmy Ba  “Adam: A method for stochastic optimization”  In _arXiv preprint arXiv:1412.6980_ , 2014

##  5 Supplementary figures

Figure S1: Temporal convolutional networks outperform tree-based methods in
supervised classification. We compare the TCN network to XGBoost and random
forest models. We compare models trained on position features (solid lines),
as well as on models trained on position-velocity features (dashed lines). For
the position feature models, the TCN is the best model on all three datasets
that use markers (A, B, C). For the position-velocity features, the TCN is the
best model for the freely moving mouse and head-fixed mouse datasets (B, C).
The TCN and XGBoost models perform similarly for the head-fixed fly dataset
(A), and for the HuGaDB dataset, the TCN and random forest perform similarly
(D). Given that the TCN performs better or similarly across all four datasets,
we select this as the backbone model throughout the rest of the model
comparisons in the paper.

Figure S2: Semi-supervised models have less confident predictions than their
supervised counterparts. The first two columns show the confusion matrices for
the TCN and S3LDS models, respectively. The latter two columns display the
average entropy of the false positives (left) and true positives (right) for
both models. In both cases entropies are larger for the S3LDS model,
indicating higher uncertainty in the state distributions. All results are from
the models trained on position-velocity features.

Figure S3: S3LDS model ablations. A: Graphical model for the Gaussian Mixture
Deep Generative Model (GMDGM) [51]. The GMDGM does not contain temporal
dependencies in the generative model (top) or the inference network (bottom).
B: The GMDGM-TCN, which does not contain temporal dependencies in the
generative model (top) but does in the inference model (bottom), where we use
a window of observed behavioral features for state prediction. C: The S3LDS
contains temporal dependencies in both the generative model (top) and
inference network (bottom). Same schematic as Fig. 2C.

Figure S4: The role of temporal information in semi-supervised model
performance. Temporal ablations comparing S3LDS (temporal inference and
generative model), GMDGM-TCN (temporal inference only), and GMDGM (non-
temporal) models. The non-temporal model performs much worse than the other
two across all datasets, especially when using non-temporal features.  Figure
S5: The role of nonlinearities in semi-supervised model performance. We
compare our S3LDS model to a model that uses one-hidden-layer dense neural
networks for the dynamics and recurrent transitions (S3NLDS). The
nonlinearities do not improve model performance on any dataset.

Figure S6: The latent spaces of models are shaped by the extent to which
labels are used in training: position-velocity features. Cluster metrics for
models trained with position-velocity features. These cluster metrics –
including normalized mutual info score, adjusted rand score, and purity score
– complement the cluster homogeneity score in Fig. 5. The addition of hand
labels produces more homogeneous clusters in the models’ latent spaces.

Figure S7: Keypoint-MoSeq performance on non-fly datasets: position features.
Comparisons between the models using only position features. Conventions as in
Fig. 6. We find the TCN, which is purely supervised, achieves the highest
alignment of the latent space with the ground truth labels as measured by the
cluster homogeneity score.  Figure S8: Keypoint-MoSeq performance: position-
velocity features vs position features. Comparisons between the keypoint-Moseq
models trained on position-velocity features vs position features. We find
that keypoint-Moseq performs best with the position features on the head-fixed
fly dataset, and the position-velocity features on the two mouse datasets.

Figure S9: Cluster metrics for model latent spaces: position features. Cluster
metrics for models trained with position features. Conventions as in Fig. S6.

##  6 Supplementary tables

Experiment ID | Total frames | Walk | Still | Front groom | Back groom | Abdomen move
---|---|---|---|---|---|---
Train 2019_08_07_fly2 | 50000 | 300 | 300 | 100 | 300 | 591
2019_08_08_fly1 | 94960 | 300 | 300 | 300 | 150 | 594
2019_08_20_fly2 | 65108 | 300 | 300 | 300 | 300 | 1067
2019_10_10_fly3 | 141042 | 300 | 300 | 300 | 300 | 478
2019_10_14_fly3 | 140929 | 300 | 300 | 300 | 300 | 318
Test 2019_06_26_fly2 | 45000 | 300 | 300 | 300 | 300 | 706
2019_08_14_fly1 | 124144 | 300 | 300 | 300 | 300 | 693
2019_08_20_fly3 | 73055 | 300 | 300 | 300 | 300 | 405
2019_10_14_fly2 | 139925 | 300 | 300 | 300 | 300 | 101
2019_10_21_fly1 | 142554 | 300 | 300 | 300 | 300 | 0
Totals | 1016717 | 3000 | 3000 | 2800 | 2850 | 4953

Table S1: The number of labeled frames per behavior for the head-fixed fly dataset. | Experiment ID | Total frames | Supported rear | Unsupported rear | Groom
---|---|---|---|---|---
Train | OFT_5 | 14000 | 3267 | 1208 | 387
| OFT_6 | 14000 | 2424 | 1090 | 428
| OFT_11 | 14000 | 1900 | 1036 | 117
| OFT_12 | 14000 | 2354 | 2213 | 202
| OFT_14 | 14000 | 3165 | 2071 | 140
| OFT_15 | 14000 | 2038 | 1433 | 154
| OFT_16 | 14000 | 1773 | 826 | 493
| OFT_23 | 14000 | 1708 | 263 | 396
| OFT_24 | 14000 | 2130 | 2312 | 298
| OFT_38 | 14000 | 2564 | 2598 | 513
Test | OFT_39 | 14000 | 2159 | 2062 | 570
| OFT_41 | 14000 | 2381 | 1479 | 178
| OFT_43 | 14000 | 3642 | 338 | 125
| OFT_44 | 14000 | 2314 | 1069 | 139
| OFT_49 | 14000 | 2638 | 1120 | 142
| OFT_50 | 14000 | 1550 | 873 | 386
| OFT_51 | 14000 | 1799 | 782 | 442
| OFT_52 | 14000 | 2949 | 1605 | 547
| OFT_54 | 14000 | 2587 | 2012 | 367
| OFT_58 | 14000 | 2468 | 1858 | 577
| Totals | 280000 | 47810 | 28248 | 6601
Table S2: The number of labeled frames per behavior for the freely moving
mouse dataset.

| Experiment ID | Total frames | Still | Move | Wheel turn | Groom
---|---|---|---|---|---|---
Train | danlab_DY_009_2020-02-27-001 | 100000 | 497 | 316 | 374 | 300
| danlab_DY_018_2020-10-15-001 | 100000 | 400 | 268 | 241 | 580
| hoferlab_SWC_061_2020-11-23-001 | 100000 | 320 | 261 | 494 | 265
| mrsicflogellab_SWC_058_2020-12-11-001 | 100000 | 433 | 357 | 358 | 298
| wittenlab_ibl_witten_26_2021-01-27-002 | 100000 | 624 | 235 | 353 | 360
Test | churchlandlab_CSHL045_2020-02-27-001 | 100000 | 306 | 368 | 522 | 469
| cortexlab_KS020_2020-02-06-001 | 100000 | 378 | 340 | 322 | 208
| hoferlab_SWC_043_2020-09-15-001 | 100000 | 469 | 363 | 306 | 513
| mrsicflogellab_SWC_052_2020-10-22-001 | 100000 | 338 | 179 | 494 | 242
| wittenlab_ibl_witten_27_2021-01-21-001 | 100000 | 335 | 35 | 536 | 50
| Totals | 1000000 | 4100 | 2722 | 4000 | 3285

Table S3: The number of labeled frames per behavior for the IBL head-fixed
mouse dataset.

| Experiment ID | Total frames | Walking | Running | Going up | Going down | Sitting | Sitting down | Standing up | Standing
---|---|---|---|---|---|---|---|---|---|---
Train | subject 05 | 56199 | 23034 | 2154 | 9820 | 10246 | 4215 | 834 | 1052 | 4844
| subject 07 | 53571 | 23533 | 4111 | 5767 | 6460 | 5000 | 1164 | 1273 | 6263
| subject 09 | 40066 | 12254 | 3712 | 7374 | 5954 | 3278 | 1038 | 1227 | 5229
| subject 12 | 52005 | 14343 | 5796 | 7683 | 7805 | 6843 | 908 | 1085 | 7542
| subject 14 | 55366 | 23365 | 5529 | 8829 | 6199 | 4771 | 920 | 900 | 4853
Test | subject 06 | 67747 | 33386 | 11334 | 5367 | 6214 | 3153 | 922 | 1008 | 6363
| subject 08 | 49093 | 17327 | 4081 | 7138 | 6262 | 4710 | 1101 | 1509 | 6965
| subject 11 | 45137 | 16202 | 3727 | 7768 | 5797 | 4471 | 1257 | 1329 | 4586
| subject 13 | 52053 | 17971 | 5374 | 10969 | 6942 | 3896 | 1095 | 1314 | 4492
| subject 17 | 45547 | 12579 | 1862 | 8723 | 6130 | 6867 | 1276 | 1235 | 6875
| Totals | 516784 | 193994 | 47680 | 79438 | 68009 | 47204 | 10515 | 11932 | 58012

Table S4: The number of labeled frames per behavior for the human gait
database.

##  Appendix A Variational lower bound of S3LDS model

###  A.1 All labels observed

We first derive the variational lower bound, assuming all $y_{t}$ are
observed. For clarity we drop the dependence on $\theta$ and $\phi$.

$\displaystyle\log p(\mathbf{x}_{1:T},y_{1:T})$ $\displaystyle\geq$
$\displaystyle\int_{\mathbf{z}_{1:T}}q(\mathbf{z}_{1:T}|\mathbf{x}_{1:T},y_{1:T})\log\frac{p(\mathbf{x}_{1:T}|\mathbf{z}_{1:T})p(\mathbf{z}_{1:T},y_{1:T})}{q(\mathbf{z}_{1:T}|\mathbf{x}_{1:T},y_{1:T})}$
$\displaystyle=$
$\displaystyle\mathbb{E}_{q(\mathbf{z}_{1:T}|\mathbf{x}_{1:T},y_{1:T})}\left[\log
p(\mathbf{x}_{1:T}|\mathbf{z}_{1:T})\right]+\mathbb{E}_{q(\mathbf{z}_{1:T}|\mathbf{x}_{1:T},y_{1:T})}\left[\log\frac{p(\mathbf{z}_{1:T},y_{1:T})}{q(\mathbf{z}_{1:T}|\mathbf{x}_{1:T},y_{1:T})}\right].$

For the first term, recall that both
$q(\mathbf{z}_{1:T}|\mathbf{x}_{1:T},y_{1:T})$ and
$p(\mathbf{x}_{1:T}|\mathbf{z}_{1:T})$ factorize across time:

$\displaystyle\mathbb{E}_{q(\mathbf{z}_{1:T}|\mathbf{x}_{1:T},y_{1:T})}\left[\log
p(\mathbf{x}_{1:T}|\mathbf{z}_{1:T})\right]$ $\displaystyle=$
$\displaystyle\int_{\mathbf{z}_{1:T}}\prod_{t=1}^{T}q(\mathbf{z}_{t}|\mathbf{x}_{\mathcal{T}_{t}},y_{t})\log\prod_{t=1}^{T}p(\mathbf{x}_{t}|\mathbf{z}_{t})$
(27) $\displaystyle=$
$\displaystyle\sum_{t=1}^{T}\int_{\mathbf{z}_{1:T}}\prod_{t=1}^{T}q(\mathbf{z}_{t}|\mathbf{x}_{\mathcal{T}_{t}},y_{t})\log
p(\mathbf{x}_{t}|\mathbf{z}_{t})$ $\displaystyle=$
$\displaystyle\sum_{t=1}^{T}\mathbb{E}_{q(\mathbf{z}_{t}|\mathbf{x}_{\mathcal{T}_{t}},y_{t})}\left[\log
p(\mathbf{x}_{t}|\mathbf{z}_{t})\right].$

For the second term,

$\displaystyle\mathbb{E}_{q(\mathbf{z}_{1:T}|\mathbf{x}_{1:T},y_{1:T})}\left[\log\frac{p(\mathbf{z}_{1:T},y_{1:T})}{q(\mathbf{z}_{1:T}|\mathbf{x}_{1:T},y_{1:T})}\right]=$
$\displaystyle\int_{\mathbf{z}_{1:T}}\prod_{t=1}^{T}q(\mathbf{z}_{t}|\mathbf{x}_{\mathcal{T}_{t}},y_{t})\log\frac{p(\mathbf{z}_{1})p(y_{1})\prod_{t=2}^{T}p(\mathbf{z}_{t}|\mathbf{z}_{t-1},y_{t})p(y_{t}|y_{t-1},\mathbf{z}_{t-1})}{\prod_{t=1}^{T}q(\mathbf{z}_{t}|\mathbf{x}_{\mathcal{T}_{t}},y_{t})}$
$\displaystyle=$
$\displaystyle\int_{\mathbf{z}_{1:T}}\prod_{t=1}^{T}q(\mathbf{z}_{t}|\mathbf{x}_{\mathcal{T}_{t}},y_{t})\log\frac{p(\mathbf{z}_{1})p(y_{1})}{q(\mathbf{z}_{1}|\mathbf{x}_{\mathcal{T}_{1}},y_{1})}$
$\displaystyle+\sum_{t=2}^{T}\int_{\mathbf{z}_{1:T}}\prod_{t=1}^{T}q(\mathbf{z}_{t}|\mathbf{x}_{\mathcal{T}_{t}},y_{t})\log\frac{p(\mathbf{z}_{t}|\mathbf{z}_{t-1},y_{t})p(y_{t}|y_{t-1},\mathbf{z}_{t-1})}{q(\mathbf{z}_{t}|\mathbf{x}_{\mathcal{T}_{t}},y_{t})}$
$\displaystyle=$
$\displaystyle\int_{\mathbf{z}_{1}}q(\mathbf{z}_{1}|\mathbf{x}_{\mathcal{T}_{1}},y_{1})\log\frac{p(\mathbf{z}_{1})p(y_{1})}{q(\mathbf{z}_{1}|\mathbf{x}_{\mathcal{T}_{1}},y_{1})}$
$\displaystyle+\sum_{t=2}^{T}\int_{\mathbf{z}_{t-1}}\int_{\mathbf{z}_{t}}q(\mathbf{z}_{t-1}|\mathbf{x}_{\mathcal{T}_{t-1}},y_{t-1})q(\mathbf{z}_{t}|\mathbf{x}_{\mathcal{T}_{t}},y_{t})\log\frac{p(\mathbf{z}_{t}|\mathbf{z}_{t-1},y_{t})p(y_{t}|y_{t-1},\mathbf{z}_{t-1})}{q(\mathbf{z}_{t}|\mathbf{x}_{\mathcal{T}_{t}},y_{t})}$
$\displaystyle=$ $\displaystyle\log
p(y_{1})-\text{KL}\left[q(\mathbf{z}_{1}|\mathbf{x}_{\mathcal{T}_{1}},y_{1})||p(\mathbf{z}_{1})\right]-\sum_{t=2}^{T}\mathbb{E}_{q(\mathbf{z}_{t-1}|\mathbf{x}_{\mathcal{T}_{t-1}},y_{t-1})}\left[\text{KL}\left[q(\mathbf{z}_{t}|\mathbf{x}_{\mathcal{T}_{t}},y_{t})||p(\mathbf{z}_{t}|\mathbf{z}_{t-1},y_{t})\right]\right]$
$\displaystyle+\sum_{t=2}^{T}\mathbb{E}_{q(\mathbf{z}_{t-1}|\mathbf{x}_{\mathcal{T}_{t-1}},y_{t-1})}\left[\log
p(y_{t}|y_{t-1},\mathbf{z}_{t-1})\right].$ (28)

Now we combine Eqs. A.1 and A.1 to get the full variational lower bound:

$\displaystyle\log p(\mathbf{x}_{1:T},y_{1:T})\geq$
$\displaystyle\sum_{t=1}^{T}\mathbb{E}_{q(\mathbf{z}_{t}|\mathbf{x}_{\mathcal{T}_{t}},y_{t})}\left[\log
p(\mathbf{x}_{t}|\mathbf{z}_{t})\right]$
$\displaystyle-\text{KL}\left[q(\mathbf{z}_{1}|\mathbf{x}_{\mathcal{T}_{1}},y_{1})||p(\mathbf{z}_{1})\right]+\log
p(y_{1})$
$\displaystyle-\sum_{t=2}^{T}\mathbb{E}_{q(\mathbf{z}_{t-1}|\mathbf{x}_{\mathcal{T}_{t-1}},y_{t-1})}\left[\text{KL}\left[q(\mathbf{z}_{t}|\mathbf{x}_{\mathcal{T}_{t}},y_{t})||p(\mathbf{z}_{t}|\mathbf{z}_{t-1},y_{t})\right]\right]$
$\displaystyle+\sum_{t=2}^{T}\mathbb{E}_{q(\mathbf{z}_{t-1}|\mathbf{x}_{\mathcal{T}_{t-1}},y_{t-1})}\left[\log
p(y_{t}|y_{t-1},\mathbf{z}_{t-1})\right]$ (29) $\displaystyle\approx$
$\displaystyle\sum_{t=1}^{T}{\log p(\mathbf{x}_{t}|\tilde{\mathbf{z}}_{t}})$
$\displaystyle-\text{KL}\left[q(\mathbf{z}_{1}|\mathbf{x}_{\mathcal{T}_{1}},y_{1})||p(\mathbf{z}_{1})\right]+\log
p(y_{1})$
$\displaystyle-\sum_{t=2}^{T}{\text{KL}\left[q(\mathbf{z}_{t}|\mathbf{x}_{\mathcal{T}_{t}},y_{t})||p(\mathbf{z}_{t}|\tilde{\mathbf{z}}_{t-1},y_{t})\right]}$
$\displaystyle+\sum_{t=2}^{T}\log p(y_{t}|y_{t-1},\tilde{\mathbf{z}}_{t-1}).$
(30)

###  A.2 All labels unobserved

Next, we derive the variational lower bound for the more complicated case
where we assume all $y_{t}$ are unobserved.

$\displaystyle\log p(\mathbf{x}_{1:T})\geq$
$\displaystyle\int_{\mathbf{z}_{1:T}}\int_{y_{1:T}}q(\mathbf{z}_{1:T},y_{1:T}|\mathbf{x}_{1:T})\log\frac{p(\mathbf{x}_{1:T}|\mathbf{z}_{1:T})p(\mathbf{z}_{1:T},y_{1:T})}{q(\mathbf{z}_{1:T},y_{1:T}|\mathbf{x}_{1:T})}$
$\displaystyle=$
$\displaystyle\int_{\mathbf{z}_{1:T}}\int_{y_{1:T}}q(\mathbf{z}_{1:T},y_{1:T}|\mathbf{x}_{1:T})\log
p(\mathbf{x}_{1:T}|\mathbf{z}_{1:T})$
$\displaystyle+\int_{\mathbf{z}_{1:T}}\int_{y_{1:T}}q(\mathbf{z}_{1:T},y_{1:T}|\mathbf{x}_{1:T})\log\frac{p(\mathbf{z}_{1:T},y_{1:T})}{q(\mathbf{z}_{1:T},y_{1:T}|\mathbf{x}_{1:T})}$

Again, recall that both $q(\mathbf{z}_{1:T}|\mathbf{x}_{1:T},y_{1:T})$ and
$p(\mathbf{x}_{1:T}|\mathbf{z}_{1:T})$ factorize across time. We now expand
the first term:

$\displaystyle\int_{\mathbf{z}_{1:T}}\int_{y_{1:T}}q(\mathbf{z}_{1:T},y_{1:T}|\mathbf{x}_{1:T})\log
p(\mathbf{x}_{1:T}|\mathbf{z}_{1:T})=$
$\displaystyle\int_{\mathbf{z}_{1:T}}\int_{y_{1:T}}\prod_{t=1}^{T}q(y_{t}|\mathbf{x}_{\mathcal{T}_{t}})q(\mathbf{z}_{t}|\mathbf{x}_{\mathcal{T}_{t}},y_{t})\log\prod_{t=1}^{T}p(\mathbf{x}_{t}|\mathbf{z}_{t})$
$\displaystyle=$
$\displaystyle\sum_{t=1}^{T}\int_{\mathbf{z}_{1:T}}\int_{y_{1:T}}\prod_{t=1}^{T}q(y_{t}|\mathbf{x}_{\mathcal{T}_{t}})q(\mathbf{z}_{t}|\mathbf{x}_{\mathcal{T}_{t}},y_{t})\log
p(\mathbf{x}_{t}|\mathbf{z}_{t})$ $\displaystyle=$
$\displaystyle\sum_{t=1}^{T}\int_{\mathbf{z}_{t}}\int_{y_{t}}q(y_{t}|\mathbf{x}_{\mathcal{T}_{t}})q(\mathbf{z}_{t}|\mathbf{x}_{\mathcal{T}_{t}},y_{t})\log
p(\mathbf{x}_{t}|\mathbf{z}_{t})$ $\displaystyle=$
$\displaystyle\sum_{t=1}^{T}\mathbb{E}_{q(y_{t}|\mathbf{x}_{\mathcal{T}_{t}})q(\mathbf{z}_{t}|\mathbf{x}_{\mathcal{T}_{t}},y_{t})}\left[\log
p(\mathbf{x}_{t}|\mathbf{z}_{t})\right].$ (31)

For the second term,

$\displaystyle\int_{\mathbf{z}_{1:T}}\int_{y_{1:T}}q(\mathbf{z}_{1:T},y_{1:T}|\mathbf{x}_{1:T})\log\frac{p(\mathbf{z}_{1:T},y_{1:T})}{q(\mathbf{z}_{1:T},y_{1:T}|\mathbf{x}_{1:T})}$
$\displaystyle=$
$\displaystyle\int_{\mathbf{z}_{1:T}}\int_{y_{1:T}}\prod_{t=1}^{T}q(y_{t}|\mathbf{x}_{\mathcal{T}_{t}})q(\mathbf{z}_{t}|\mathbf{x}_{\mathcal{T}_{t}},y_{t})\log\frac{p(\mathbf{z}_{1})p(y_{1})\prod_{t=2}^{T}p(\mathbf{z}_{t}|\mathbf{z}_{t-1},y_{t})p(y_{t}|y_{t-1},\mathbf{z}_{t-1})}{\prod_{t=1}^{T}q(y_{t}|\mathbf{x}_{\mathcal{T}_{t}})q(\mathbf{z}_{t}|\mathbf{x}_{\mathcal{T}_{t}},y_{t})}$
$\displaystyle=$
$\displaystyle\int_{\mathbf{z}_{1:T}}\int_{y_{1:T}}\prod_{t=1}^{T}q(y_{t}|\mathbf{x}_{\mathcal{T}_{t}})q(\mathbf{z}_{t}|\mathbf{x}_{\mathcal{T}_{t}},y_{t})\log\frac{p(\mathbf{z}_{1})p(y_{1})}{q(\mathbf{z}_{1}|\mathbf{x}_{\mathcal{T}_{1}},y_{1})q(y_{1}|\mathbf{x}_{\mathcal{T}_{1}})}$
$\displaystyle+\int_{\mathbf{z}_{1:T}}\int_{y_{1:T}}\prod_{t=1}^{T}q(y_{t}|\mathbf{x}_{\mathcal{T}_{t}})q(\mathbf{z}_{t}|\mathbf{x}_{\mathcal{T}_{t}},y_{t})\left[\sum_{t=2}^{T}\log\frac{p(\mathbf{z}_{t}|\mathbf{z}_{t-1},y_{t})p(y_{t}|y_{t-1},\mathbf{z}_{t-1})}{q(y_{t}|\mathbf{x}_{\mathcal{T}_{t}})q(\mathbf{z}_{t}|\mathbf{x}_{\mathcal{T}_{t}},y_{t})}\right]$
$\displaystyle=$
$\displaystyle\int_{\mathbf{z}_{1}}\int_{y_{1}}q(\mathbf{z}_{1}|\mathbf{x}_{\mathcal{T}_{1}},y_{1})q(y_{1}|\mathbf{x}_{\mathcal{T}_{1}})\log\frac{p(\mathbf{z}_{1})p(y_{1})}{q(\mathbf{z}_{1}|\mathbf{x}_{\mathcal{T}_{1}},y_{1})q(y_{1}|\mathbf{x}_{\mathcal{T}_{1}})}$
$\displaystyle+\sum_{t=2}^{T}\int_{\mathbf{z}_{t-1}}\int_{\mathbf{z}_{t}}\int_{y_{t-1}}\int_{y_{t}}q(y_{t}|\mathbf{x}_{\mathcal{T}_{t}})q(y_{t-1}|\mathbf{x}_{\mathcal{T}_{t-1}})q(\mathbf{z}_{t}|\mathbf{x}_{\mathcal{T}_{t}},y_{t})q(\mathbf{z}_{t-1}|\mathbf{x}_{\mathcal{T}_{t-1}},y_{t-1})\log\frac{p(\mathbf{z}_{t}|\mathbf{z}_{t-1},y_{t})}{q(\mathbf{z}_{t}|\mathbf{x}_{\mathcal{T}_{t}},y_{t})}$
$\displaystyle+\sum_{t=2}^{T}\int_{\mathbf{z}_{t-1}}\int_{\mathbf{z}_{t}}\int_{y_{t-1}}\int_{y_{t}}q(y_{t}|\mathbf{x}_{\mathcal{T}_{t}})q(y_{t-1}|\mathbf{x}_{\mathcal{T}_{t-1}})q(\mathbf{z}_{t}|\mathbf{x}_{\mathcal{T}_{t}},y_{t})q(\mathbf{z}_{t-1}|\mathbf{x}_{\mathcal{T}_{t-1}},y_{t-1})\log\frac{p(y_{t}|y_{t-1},\mathbf{z}_{t-1})}{q(y_{t}|\mathbf{x}_{\mathcal{T}_{t}})}.$
(32)

The first term from Eq. A.2 becomes

$\displaystyle\int_{\mathbf{z}_{1}}\int_{y_{1}}q(\mathbf{z}_{1}|\mathbf{x}_{\mathcal{T}_{1}},y_{1})q(y_{1}|\mathbf{x}_{\mathcal{T}_{1}})\log\frac{p(\mathbf{z}_{1})p(y_{1})}{q(\mathbf{z}_{1}|\mathbf{x}_{\mathcal{T}_{1}},y_{1})q(y_{1}|\mathbf{x}_{\mathcal{T}_{1}})}$
$\displaystyle=$
$\displaystyle\int_{y_{1}}q(y_{1}|\mathbf{x}_{\mathcal{T}_{1}})\int_{\mathbf{z}_{1}}q(\mathbf{z}_{1}|\mathbf{x}_{\mathcal{T}_{1}},y_{1})\log\frac{p(\mathbf{z}_{1})}{q(\mathbf{z}_{1}|\mathbf{x}_{\mathcal{T}_{1}},y_{1})}+\int_{y_{1}}q(y_{1}|\mathbf{x}_{\mathcal{T}_{1}})\log\frac{p(y_{1})}{q(y_{1}|\mathbf{x}_{\mathcal{T}_{1}})}\int_{\mathbf{z}_{1}}q(\mathbf{z}_{1}|\mathbf{x}_{\mathcal{T}_{1}},y_{1})$
$\displaystyle=$
$\displaystyle-\mathbb{E}_{q(y_{1}|\mathbf{x}_{\mathcal{T}_{1}})}\left[\text{KL}\left[q(\mathbf{z}_{1}|\mathbf{x}_{\mathcal{T}_{1}},y_{1})||p(\mathbf{z}_{1})\right]\right]-\text{KL}\left[q(y_{1}|\mathbf{x}_{\mathcal{T}_{1}})||p(y_{1})\right].$
(33)

The second term from Eq. A.2 becomes

$\displaystyle\sum_{t=2}^{T}\int_{\mathbf{z}_{t-1}}\int_{\mathbf{z}_{t}}\int_{y_{t-1}}\int_{y_{t}}q(y_{t}|\mathbf{x}_{\mathcal{T}_{t}})q(y_{t-1}|\mathbf{x}_{\mathcal{T}_{t-1}})q(\mathbf{z}_{t}|\mathbf{x}_{\mathcal{T}_{t}},y_{t})q(\mathbf{z}_{t-1}|\mathbf{x}_{\mathcal{T}_{t-1}},y_{t-1})\log\frac{p(\mathbf{z}_{t}|\mathbf{z}_{t-1},y_{t})}{q(\mathbf{z}_{t}|\mathbf{x}_{\mathcal{T}_{t}},y_{t})}$
$\displaystyle=$
$\displaystyle-\sum_{t=2}^{T}\mathbb{E}_{q(y_{t}|\mathbf{x}_{\mathcal{T}_{t}})q(y_{t-1}|\mathbf{x}_{\mathcal{T}_{t-1}})q(\mathbf{z}_{t-1}|\mathbf{x}_{\mathcal{T}_{t-1}},y_{t-1})}\left[\text{KL}\left[q(\mathbf{z}_{t}|\mathbf{x}_{\mathcal{T}_{t}},y_{t})||p(\mathbf{z}_{t}|\mathbf{z}_{t-1},y_{t})\right]\right]$
(34)

and the third term from Eq. A.2 becomes

$\displaystyle\sum_{t=2}^{T}\int_{\mathbf{z}_{t-1}}\int_{y_{t-1}}\int_{y_{t}}q(y_{t}|\mathbf{x}_{\mathcal{T}_{t}})q(y_{t-1}|\mathbf{x}_{\mathcal{T}_{t-1}})q(\mathbf{z}_{t-1}|\mathbf{x}_{\mathcal{T}_{t-1}},y_{t-1})\log\frac{p(y_{t}|y_{t-1},\mathbf{z}_{t-1})}{q(y_{t}|\mathbf{x}_{\mathcal{T}_{t}})}$
$\displaystyle=$
$\displaystyle-\sum_{t=2}^{T}\mathbb{E}_{q(y_{t-1}|\mathbf{x}_{\mathcal{T}_{t-1}})q(\mathbf{z}_{t-1}|\mathbf{x}_{\mathcal{T}_{t-1}},y_{t-1})}\left[\text{KL}\left[q(y_{t}|\mathbf{x}_{\mathcal{T}_{t}})||p(y_{t}|y_{t-1},\mathbf{z}_{t-1})\right]\right].$
(35)

Now we combine Eqs. A.2, A.2, A.2, and A.2 to get the full variational lower
bound:

$\displaystyle\log p(\mathbf{x}_{1:T})\geq$
$\displaystyle\sum_{t=1}^{T}\mathbb{E}_{q(y_{t}|\mathbf{x}_{\mathcal{T}_{t}})q(\mathbf{z}_{t}|\mathbf{x}_{\mathcal{T}_{t}},y_{t})}\left[\log
p(\mathbf{x}_{t}|\mathbf{z}_{t})\right]$
$\displaystyle-\mathbb{E}_{q(y_{1}|\mathbf{x}_{\mathcal{T}_{1}})}\left[\text{KL}\left[q(\mathbf{z}_{1}|\mathbf{x}_{\mathcal{T}_{1}},y_{1})||p(\mathbf{z}_{1})\right]\right]-\text{KL}\left[q(y_{1}|\mathbf{x}_{\mathcal{T}_{1}})||p(y_{1})\right]$
$\displaystyle-\sum_{t=2}^{T}\mathbb{E}_{q(y_{t}|\mathbf{x}_{\mathcal{T}_{t}})q(y_{t-1}|\mathbf{x}_{\mathcal{T}_{t-1}})q(\mathbf{z}_{t-1}|\mathbf{x}_{\mathcal{T}_{t-1}},y_{t-1})}\left[\text{KL}\left[q(\mathbf{z}_{t}|\mathbf{x}_{\mathcal{T}_{t}},y_{t})||p(\mathbf{z}_{t}|\mathbf{z}_{t-1},y_{t})\right]\right]$
$\displaystyle-\sum_{t=2}^{T}\mathbb{E}_{q(y_{t-1}|\mathbf{x}_{\mathcal{T}_{t-1}})q(\mathbf{z}_{t-1}|\mathbf{x}_{\mathcal{T}_{t-1}},y_{t-1})}\left[\text{KL}\left[q(y_{t}|\mathbf{x}_{\mathcal{T}_{t}})||p(y_{t}|y_{t-1},\mathbf{z}_{t-1})\right]\right]$
$\displaystyle=$
$\displaystyle\sum_{t=1}^{T}\sum_{k=1}^{K}\Big{(}q(y_{t}=k|\mathbf{x}_{\mathcal{T}_{t}})\mathbb{E}_{q(\mathbf{z}_{t}|\mathbf{x}_{\mathcal{T}_{t}},y_{t}=k)}[{\log
p(\mathbf{x}_{t}|{\mathbf{z}}_{t}})]\Big{)}$
$\displaystyle-\sum_{k=1}^{K}\Big{(}q(y_{1}=k|\mathbf{x}_{\mathcal{T}_{1}}){\text{KL}\left[q(\mathbf{z}_{1}|\mathbf{x}_{\mathcal{T}_{1}},y_{1}=k)||p(\mathbf{z}_{1})\right]}\Big{)}-\text{KL}\left[q(y_{1}|\mathbf{x}_{\mathcal{T}_{1}})||p(y_{1})\right]$
$\displaystyle-\sum_{t=2}^{T}\sum_{k=1}^{K}\sum_{k^{\prime}=1}^{K}\Big{(}q(y_{t}=k|\mathbf{x}_{\mathcal{T}_{t}})q(y_{t-1}=k^{\prime}|\mathbf{x}_{\mathcal{T}_{t-1}})$
$\displaystyle\quad\quad\mathbb{E}_{q(\mathbf{z}_{t-1}|\mathbf{x}_{\mathcal{T}_{t-1}},y_{t-1}=k^{\prime})}[{\text{KL}\left[{q(\mathbf{z}_{t}|\mathbf{x}_{\mathcal{T}_{t}},y_{t}=k)}||{p(\mathbf{z}_{t}|\mathbf{z}_{t-1},y_{t}=k)}\right]}]\Big{)}$
$\displaystyle-\sum_{t=2}^{T}\sum_{k=1}^{K}\Big{(}q(y_{t-1}=k|\mathbf{x}_{\mathcal{T}_{t-1}})\mathbb{E}_{q(\mathbf{z}_{t-1}|\mathbf{x}_{\mathcal{T}_{t-1}},y_{t-1}=k)}[{\text{KL}\left[q(y_{t}|\mathbf{x}_{\mathcal{T}_{t}})||{p(y_{t}|y_{t-1}=k,\mathbf{z}_{t-1})}\right]}]\Big{)}$
(36) $\displaystyle\approx$
$\displaystyle\sum_{t=1}^{T}\sum_{k=1}^{K}\Big{(}q(y_{t}=k|\mathbf{x}_{\mathcal{T}_{t}}){\log
p(\mathbf{x}_{t}|{\tilde{\mathbf{z}}}_{t}^{k}})\Big{)}$
$\displaystyle-\sum_{k=1}^{K}\Big{(}q(y_{1}=k|\mathbf{x}_{\mathcal{T}_{t}}){\text{KL}\left[q(\mathbf{z}_{1}|\mathbf{x}_{\mathcal{T}_{1}},y_{1}=k)||p(\mathbf{z}_{1})\right]}\Big{)}-\text{KL}\left[q(y_{1}|\mathbf{x}_{\mathcal{T}_{1}})||p(y_{1})\right]$
$\displaystyle-\sum_{t=2}^{T}\sum_{k=1}^{K}\sum_{k^{\prime}=1}^{K}\Big{(}q(y_{t}=k|\mathbf{x}_{\mathcal{T}_{t}})q(y_{t-1}=k^{\prime}|\mathbf{x}_{\mathcal{T}_{t-1}})$
$\displaystyle\quad\quad{\text{KL}\left[{q(\mathbf{z}_{t}|\mathbf{x}_{\mathcal{T}_{t}},y_{t}=k)}||{p(\mathbf{z}_{t}|\tilde{\mathbf{z}}_{t-1}^{k^{\prime}},y_{t}=k)}\right]}\Big{)}$
$\displaystyle-\sum_{t=2}^{T}\sum_{k=1}^{K}\Big{(}q(y_{t-1}=k|\mathbf{x}_{\mathcal{T}_{t-1}}){\text{KL}\left[q(y_{t}|\mathbf{x}_{\mathcal{T}_{t}})||{p(y_{t}|y_{t-1}=k,\tilde{\mathbf{z}}_{t-1}^{k})}\right]}\Big{)}$
(37)

where

$\tilde{\mathbf{z}}_{t}^{k}\sim
q(\mathbf{z}_{t}|\mathbf{x}_{\mathcal{T}_{t}},y_{t}=k)\ \textrm{and}\
\tilde{\mathbf{z}}_{t-1}^{k}\sim
q(\mathbf{z}_{t-1}|\mathbf{x}_{\mathcal{T}_{t-1}},y_{t-1}=k).$

###  A.3 Mix of observed and unobserved labels

The loss function that incorporates both observed and unobserved discrete
labels can be computed in an elegant and straightforward manner. For
notational simplicity we define a new distribution over the discrete
variables, which is equal to the approximate posterior for unlabeled time
points, and equal to a one-hot vector representing the true discrete state for
labeled time points. Let $\mathcal{T}_{\mathcal{L}}$ be the set of all labeled
time points and $\mathcal{T}_{\mathcal{U}}$ be the set of all unlabeled time
points, such that
$\mathcal{T}_{\mathcal{L}}\cup\mathcal{T}_{\mathcal{U}}=\\{1,\ldots,T\\}$ and
$\mathcal{T}_{\mathcal{L}}\cap\mathcal{T}_{\mathcal{U}}=\varnothing$. Then we
define

$\displaystyle
r(y_{t}|\mathbf{x}_{\mathcal{T}_{t}})=\begin{cases}e_{y_{t}}&\text{if
}t\in\mathcal{T}_{\mathcal{L}}\\\
q(y_{t}|\mathbf{x}_{\mathcal{T}_{t}})&\text{if
}t\in\mathcal{T}_{\mathcal{U}}\\\ \end{cases}$

where $e_{y_{t}}$ denotes a vector with a 1 in the $y_{t}^{\text{th}}$
coordinate and 0 elsewhere. We can then compute Eq. 37 using this mixture of
distributions due to the fact that each of the individual terms in the
unsupervised ELBO (Eq. 4.2.2) reduces to the corresponding term in the
supervised ELBO (Eq. 4.2.2) when using one-hot vectors. In addition, we
compute the classification loss on the approximate posterior
$q(y_{t}|\mathbf{x}_{\mathcal{T}_{t}})$ for all observed discrete labels:

$\displaystyle\mathcal{L}_{ss}=$
$\displaystyle\sum_{t=1}^{T}\sum_{k=1}^{K}\Big{(}r(y_{t}=k|\mathbf{x}_{\mathcal{T}_{t}}){\log
p(\mathbf{x}_{t}|{\tilde{\mathbf{z}}}_{t}^{k}})\Big{)}$
$\displaystyle-\sum_{k=1}^{K}\Big{(}r(y_{1}=k|\mathbf{x}_{\mathcal{T}_{t}}){\text{KL}\left[q(\mathbf{z}_{1}|\mathbf{x}_{\mathcal{T}_{1}},y_{1}=k)||p(\mathbf{z}_{1})\right]}\Big{)}-\text{KL}\left[r(y_{1}|\mathbf{x}_{\mathcal{T}_{1}})||p(y_{1})\right]$
$\displaystyle-\sum_{t=2}^{T}\sum_{k=1}^{K}\sum_{k^{\prime}=1}^{K}\Big{(}r(y_{t}=k|\mathbf{x}_{\mathcal{T}_{t}})r(y_{t-1}=k^{\prime}|\mathbf{x}_{\mathcal{T}_{t-1}})$
$\displaystyle\quad\quad{\text{KL}\left[{q(\mathbf{z}_{t}|\mathbf{x}_{\mathcal{T}_{t}},y_{t}=k)}||{p(\mathbf{z}_{t}|\tilde{\mathbf{z}}_{t-1}^{k^{\prime}},y_{t}=k)}\right]}\Big{)}$
$\displaystyle-\sum_{t=2}^{T}\sum_{k=1}^{K}\Big{(}r(y_{t-1}=k|\mathbf{x}_{\mathcal{T}_{t-1}}){\text{KL}\left[r(y_{t}|\mathbf{x}_{\mathcal{T}_{t}})||{p(y_{t}|y_{t-1}=k,\tilde{\mathbf{z}}_{t-1}^{k})}\right]}\Big{)}$
$\displaystyle+\sum_{t\in\mathcal{T}_{\mathcal{L}}}\alpha\log
q(y_{t}|\mathbf{x}_{\mathcal{T}_{t}}).$ (38)

##  Appendix B The International Brain Laboratory consortium members

Larry Abbot1, Luigi Acerbi2, Valeria Aguillon-Rodriguez3, Mandana Ahmadi4,
Jaweria Amjad4, Dora Angelaki5, Jaime Arlandis6, Zoe C Ashwood7, Kush Banga8,
Hailey Barrell9, Hannah M Bayer1, Brandon Benson10, Julius Benson5, Jai
Bhagat8, Dan Birman9, Niccolò Bonacchi6, Kcenia Bougrova6, Julien Boussard1,
Sebastian A Bruijns11, E Kelly Buchanan1, Robert Campbell12, Matteo
Carandini13, Joana A Catarino6, Fanny Cazettes6, Gaelle A Chapuis2, Anne K
Churchland14, Yang Dan15, Felicia Davatolhagh14, Peter Dayan11, Sophie
Denève16, Eric EJ DeWitt6, Ling Liang Dong17, Tatiana Engel7, Michele Fabbri1,
Mayo Faulkner8, Robert Fetcho7, Ila Fiete17, Charles Findling2, Laura Freitas-
Silva6, Surya Ganguli10, Berk Gercek2, Naureen Ghani12, Ivan Gordeliy16, Laura
M Haetzel7, Kenneth D Harris8, Michael Hausser18, Naoki Hiratani4, Sonja
Hofer12, Fei Hu15, Felix Huber2, Julia M Huntenburg11, Cole Hurwitz1, Anup
Khanal14, Christopher S Krasniak3, Sanjukta Krishnagopal4, Michael Krumin8,
Debottam Kundu11, Agnès Landemard13, Christopher Langdon7, Christopher
Langfield1, Inês Laranjeira6, Peter Latham4, Petrina Lau18, Hyun Dong Lee1,
Ari Liu17, Zachary F Mainen6, Amalia Makri-Cottington18, Hernando Martinez-
Vergara12, Brenna McMannon7, Isaiah McRoberts5, Guido T Meijer6, Maxwell
Melin14, Leenoy Meshulam19, Kim Miller9, Nathaniel J Miska12, Catalin
Mitelut1, Zeinab Mohammadi7, Thomas Mrsic-Flogel12, Masayoshi Murakami20,
Jean-Paul Noel5, Kai Nylund9, Farideh Oloomi11, Alejandro Pan-Vazquez7, Liam
Paninski1, Alberto Pezzotta4, Samuel Picard8, Jonathan W Pillow7, Alexandre
Pouget2, Florian Rau6, Cyrille Rossant8, Noam Roth9, Nicholas A Roy7, Kamron
Saniee1, Rylan Schaeffer17, Michael M Schartner6, Yanliang Shi7, Carolina
Soares12, Karolina Z Socha13, Cristian Soitu3, Nicholas A Steinmetz9, Karel
Svoboda21, Marsa Taheri14, Charline Tessereau11, Anne E Urai3, Erdem Varol1,
Miles J Wells8, Steven J West12, Matthew R Whiteway1, Charles Windolf1,
Olivier Winter6, Ilana Witten7, Lauren E Wool8, Zekai Xu4, Han Yu1, Anthony M
Zador3, Yizi Zhang1

1Zuckerman Institute, Columbia University, New York, NY, USA
2Department of Basic Neuroscience, University of Geneva, Geneva, Switzerland
3Cold Spring Harbor Laboratory, Cold Spring Harbor, NY, USA
4Gatsby Computational Neuroscience Unit, University College London, London, UK
5Center for Neural Science, New York University, New York, NY, USA
6Champalimaud Center for the Unknown, Lisboa, Portugal
7Princeton Neuroscience Institute, Princeton University, Princeton, NJ, USA
8Institute of Neurology, University College London, London, UK
9Department of Biological Structure, University of Washington, Seattle, WA,
USA
10Department of Applied Physics, Stanford University, Stanford, CA, USA
11Max Planck Institute for Biological Cybernetics, Tübingen, Germany
12Sainsbury-Wellcome Centre, University College London, London, UK
13Institute of Opthalmology, University College London, London, United Kingdom
14Department of Neurobiology, University of California, Los Angeles, CA, USA
15Department of Molecular and Cell Biology, University of California,
Berkeley, CA, USA
16Département D’études Cognitives, École Normale Supérieure, Paris, France
17Department of Brain and Cognitive Sciences, Massachusetts Institute of
Technology, Cambridge, MA, USA
18Wolfson Institute of Biomedical Research, University College London, London,
United Kingdom
19Center for Computational Neuroscience, University of Washington, Seattle,
WA, USA
20Department of Physiology, University of Yamanashi, Kofu, Yamanashi, Japan
21The Allen Institute for Neural Dynamics, Seattle, Washington, USA