{"title": "Point process latent variable models of larval zebrafish behavior", "book": "Advances in Neural Information Processing Systems", "page_first": 10919, "page_last": 10930, "abstract": "A fundamental goal of systems neuroscience is to understand how neural activity gives rise to natural behavior.  In order to achieve this goal, we must first build comprehensive models that offer quantitative descriptions of behavior.  We develop a new class of probabilistic models to tackle this challenge in the study of larval zebrafish, an important model organism for neuroscience.  Larval zebrafish locomote via sequences of punctate swim bouts--brief flicks of the tail--which are naturally modeled as a marked point process.  However, these sequences of swim bouts belie a set of discrete and continuous internal states, latent variables that are not captured by standard point process models.  We incorporate these variables as latent marks of a point process and explore various models for their dynamics.  To infer the latent variables and fit the parameters of this model, we develop an amortized variational inference algorithm that targets the collapsed posterior distribution, analytically marginalizing out the discrete latent variables.  With a dataset of over 120,000 swim bouts, we show that our models reveal interpretable discrete classes of swim bouts and continuous internal states like hunger that modulate their dynamics.  These models are a major step toward understanding the natural behavioral program of the larval zebrafish and, ultimately, its neural underpinnings.", "full_text": "Point process latent variable models of\n\nlarval zebra\ufb01sh behavior\n\nAnuj Sharma\n\nColumbia University\n\nRobert E. Johnson\nHarvard University\n\nFlorian Engert\n\nHarvard University\n\nScott W. Linderman\u2217\nColumbia University\n\nAbstract\n\nA fundamental goal of systems neuroscience is to understand how neural activity\ngives rise to natural behavior. In order to achieve this goal, we must \ufb01rst build\ncomprehensive models that offer quantitative descriptions of behavior. We develop\na new class of probabilistic models to tackle this challenge in the study of larval ze-\nbra\ufb01sh, an important model organism for neuroscience. Larval zebra\ufb01sh locomote\nvia sequences of punctate swim bouts\u2014brief \ufb02icks of the tail\u2014which are naturally\nmodeled as a marked point process. However, these sequences of swim bouts belie\na set of discrete and continuous internal states, latent variables that are not captured\nby standard point process models. We incorporate these variables as latent marks of\na point process and explore various models for their dynamics. To infer the latent\nvariables and \ufb01t the parameters of this model, we develop an amortized variational\ninference algorithm that targets the collapsed posterior distribution, analytically\nmarginalizing out the discrete latent variables. With a dataset of over 120,000 swim\nbouts, we show that our models reveal interpretable discrete classes of swim bouts\nand continuous internal states like hunger that modulate their dynamics. These\nmodels are a major step toward understanding the natural behavioral program of\nthe larval zebra\ufb01sh and, ultimately, its neural underpinnings.\n\n1\n\nIntroduction\n\nComputational neuroscience\u2014the study of how neural circuits transform sensory inputs into be-\nhavioral outputs\u2014is intimately coupled with computational ethology\u2014the quantitative analysis of\nbehavior [1, 2]. In order to understand the computations of the nervous system, we must \ufb01rst have\na rigorous description of the behavior it produces. To that end, comprehensive, quantitative, and\ninterpretable models of behavior are of fundamental importance to the study of the brain.\nFor many organisms, overt behaviors manifest as a sequence of discrete and nearly-instantaneous\nevents unfolding over time, often with some associated measurements, or marks. Multiple times a\nsecond, our eyes saccade in a quick, jerking motion to \ufb01xate on a new point in our \ufb01eld of view [3].\nSome electric \ufb01sh emit pulsatile discharges to navigate, detect objects, and communicate [4]. In this\npaper we study larval zebra\ufb01sh, a model organism for neuroscience. They swim with brief tail \ufb02icks,\nor bouts, that propel them forward, reorient them, and enable them to pursue and capture prey [5, 6].\nImportantly, larval zebra\ufb01sh offer exciting opportunities: if we can better quantify their behavioral\npatterns, we can use whole brain functional imaging technologies to search for correlates of these\npatterns in the neural activity dynamics of behaving \ufb01sh [7\u201311].\nFigure 1 illustrates our experimental setup for collecting behavioral data of freely swimming larval\nzebra\ufb01sh [12]. Each \ufb01sh swims in a large (30cm) tank for 40 minutes while feeding on paramecia and\nis recruited to the center to initiate each observational trial (a.). Each trial consists of a sequence of\nup to 350 swim bouts (b.) and we recorded over 120,000 bouts from 130 \ufb01sh over about 1000 trials.\n\n\u2217Corresponding author: scott.linderman@columbia.edu.\n\n32nd Conference on Neural Information Processing Systems (NeurIPS 2018), Montr\u00b4eal, Canada.\n\n\fFigure 1: Overview of our experimental setup for studying zebra\ufb01sh behavior over multiple time-scales. a. We\ncollected many trials of larval zebra\ufb01sh freely swimming in a large tank with paramecia, the \ufb01sh\u2019s prey. b. Each\ntrial consists of a sequence of punctuated swim bouts separated by longer periods of rest. c. Most swim bouts\nlast less than 200ms, nearly instantaneous for our modeling purposes. d. As the \ufb01sh swims, we track it with an\noverhead camera and record high-resolution video at 60fps. In each video frame, we identify the \ufb01sh\u2019s 2 eye\nangles and the change in its 20 tail tangent angles over consecutive frames to describe its posture. For each\nbout, we use ten frames starting with movement onset, giving us a 20D representation of the eyes and 180D\nrepresentation of the tail. We then use PCA to reduce the tail representation to the same dimension as the eye\nangles and use the resulting 40D representations as the marks in our point process latent variable model.\nBouts are nearly instantaneous events, most lasting under 200ms (c.). As the \ufb01sh swims, we track it\nwith a moving overhead camera and collect high-resolution video of its postural dynamics (d.). We\nuse eye angles and the change in tail shape through ten frames starting with movement onset as a\nhigh-dimensional quanti\ufb01cation of each bout.\nWe aim to answer two scienti\ufb01c questions with this dataset. First, what dynamics govern how swim\nbouts are sequenced together over time? Second, how are these dynamics modulated by internal\nstates like hunger? We develop a new class of probabilistic models to address these questions.\nLarval zebra\ufb01sh behavior is naturally viewed as a marked point process, a stochastic process that\ngenerates sets of events in time with corresponding observations, or marks. Here, each bout is a\ntime-stamped event marked with a corresponding vector of tail postures and eye angles. Marked\npoint processes offer a probabilistic framework for modeling the rate at which the observed events\noccur. However, our scienti\ufb01c questions pertain to discrete and continuous states that are not directly\nobservable. This motivates the new point process latent variable models (PPLVM) we introduce in\nSection 3, which blend deep state space models and marked point processes. This work builds upon\nand extends many existing models, as we discuss in Section 2 and Section 5. Section 4 develops an\namortized variational inference algorithm for inferring the latent states and \ufb01tting the parameters of\nthe PPLVM. Sections 6 and 7 present our results from applying our methods to synthetic and real data.\n\n2 Background\n\nWe start by introducing the key modeling ingredients that underlie our model.\n\nPoint processes and renewal processes. Point processes are stochastic processes that generate\ndiscrete sets of events in time and space. In our case, each swim bout is characterized by a time-\nstamp tn and a corresponding mark yn, here a vector of eye and tail angles. Generally, point processes\nare characterized by a rate function, which implies a probability density on sets of events [13].\nUnfortunately, evaluating this density requires integrating the rate function, which is intractable for\nall but the simplest models. However, when the events admit a natural ordering\u2014for example, when\nevents can be sorted in time\u2014we can use a renewal process (RP) instead. Renewal processes specify\na distribution on the intervals in (cid:44) tn+1 \u2212 tn between consecutive events, and the joint probability\nof sets of intervals is typically easy to compute. For example, gamma renewal processes (GRP) treat\neach interval as an independent gamma random variable so that the joint distribution factorizes over\nintervals. When the intervals are independent exponential random variables, we recover the standard\nPoisson process (PP). By changing the interval distribution or introducing dependencies between\nintervals, we develop point processes with more complex structure yet still tractable distributions.\nMoreover, renewal processes are easily extended to handle sets of marked events by specifying a\nconditional distribution over marks given the intervals.\n\n2\n\n30o2 secswimboutsfish tank200 ms4 cm27 strial 2trial 8trial 11a.b.c.d.1 mm\fDeep generative models.\nIn practice, it can be dif\ufb01cult to model distributions over high dimensional\nmarks. Recent advances in deep generative modeling [14\u201316] offer new means to tackle this challenge\nwith neural networks. For example, deep latent Gaussian models use neural networks to capture\nnonlinear mappings between low dimensional latent variables and observed data. In this way, simple\npriors on latent variables give rise to complex conditional distributions over data. However, learning\nthe neural network weights is far from trivial because the marginal log probability of the data, or\nevidence, is intractable. Instead, we resort to approximate methods like variational expectation-\nmaximization, which maximize a more tractable evidence lower bound (ELBO). Two advances\nmake this practical: recognition networks, which model the variational approximation as a learnable\nfunction of the data, again implemented as a neural network; and the reparameterization trick, which\nallows for lower variance estimates of the gradients of the ELBO for stochastic gradient ascent. These\nideas will be key to articulating and \ufb01tting our models of zebra\ufb01sh behavior.\n\nState space models. State space models capture dependencies between latent variables over time.\nDeep generative models offer a very \ufb02exible approach to modeling dependencies, but we can often\nmake more restrictive assumptions about the nature of the temporal dynamics. In doing so, we hope\nto recover more interpretable latent structure. For example, we believe that zebra\ufb01sh behavior is\ngoverned by discrete and continuous latent variables that evolve over time; these are naturally captured\nby hidden Markov models (HMM) [17] and Gaussian processes (GP) [18]. HMMs model sequences of\ndiscrete latent states with Markovian dynamics, and when the discrete states govern a distribution over\nintervals of an RP, we obtain Markov renewal processes (MRP). GPs are nonparametric models for\nrandom functions x(t) with covariance structure determined by a kernel K(t, t(cid:48)). Under a GP model,\nthe set of function evaluations x1:N at times t1:N is jointly Gaussian distributed with covariance\nmatrix C, where Cn,n(cid:48) = K(tn, tn(cid:48)). Given the kernel function, it is straightforward to compute the\nGaussian predictive density p(xn+1 | x1:n, t1:n+1) and its predictive covariance Cn+1|1:n. With the\npredictive distribution, we can simulate the function forward in time at asynchronous time stamps.\n\n3 Mixed Discrete and Continuous Point Process Latent Variable Models\n\nWe propose a class of point process latent variable models that blend renewal processes, deep\ngenerative models, and state space models to build a model for sets of marked events in time. The\nkey idea is to view the latent variables as unobserved elements of the events\u2019 marks. Each event has\nan observed time stamp tn and mark yn; rather than modeling the time stamps directly, we model\nthe intervals in (cid:44) tn+1 \u2212 tn; n = 1, . . . , N. (Technically, we model t1, i1:N\u22121, and the probability\nthat iN > T \u2212 tN .) We augment these marks with three latent variables: a continuous latent state xn,\na discrete state zn, and an embedding of the high dimensional mark hn. We use state space models\nto link these latent variables across sequences of events, and deep generative models to relate the\nembedding to the observed mark. In modeling larval zebra\ufb01sh behavior, we expect these latent\nvariables to capture continuous internal states, like hunger, discrete states, like the type of swim bout,\nand low dimensional properties of the bout kinematics. There are many ways to relate these latent\nvariables. We motivate one model and discuss other special cases.\n\n3.1 Gaussian process modulated Markov renewal process\n\nOur choice of conditional distributions is guided by three desiderata: we desire \ufb02exibility in the\naspects of the model about which we are less certain, we want to express prior knowledge when\nit is available, and we want to build models that admit ef\ufb01cient inference algorithms. To that end,\nwe propose a semi-parametric point process latent variable model that we call the Gaussian process\nmodulated Markov renewal process (GPM-MRP).\nThe \ufb01rst component of the GPM-MRP is a deep latent Gaussian model of the high-dimensional marks.\nWe assume that each bout\u2019s observed eye and tail angles yn re\ufb02ect a low-dimensional continuous\nlatent embedding hn \u2208 RH. This embedding is transformed through a neural network, which outputs\nthe mean and diagonal variance of a distribution over the observed mark yn \u223c N (\u00b5\u03b8(hn), \u03a3\u03b8(hn)).\nWe expect this latent embedding to act as a low-dimensional summary of the bout\u2019s most salient\nattributes, and hence, conditioned on hn, yn is assumed to be independent of all other variables.\nBased on past ethological studies of larval zebra\ufb01sh [5\u20137], we believe that swim bouts can be catego-\nrized into discrete types, and that these types are correlated over time. Intuitively, a bout\u2019s discrete\n\n3\n\n\fFigure 2: Generative models and recognition network. Left: The full generative model relates discrete and\ncontinuous latent states to the low-dimensional mark embeddings and the observed inter-bout intervals and\nmarks. The continuous states follow a Gaussian process, so the preceding values and past intervals are necessary\nto predict the next continuous state. These dependencies are shown in light gray. The mapping from embeddings\nto observed bout kinematics is implemented via a neural network, as indicated by the square-tipped arrows.\nMiddle: Since the discrete latent states are connected in a Markov chain, we can ef\ufb01ciently sum over them via\nmessage passing to obtain a collapsed generative model. Marginalization yields a purely continuous, densely\nconnected latent variable model. Right: We infer the continuous latent variables via a recognition network with\na bidirectional LSTM. The LSTM states (blue squares) are read out at only a subset of points (here, two middle\nbouts), which then determine the other continuous states.\n\ntype determines the distribution over its attributes hn and subsequent intervals in. We formalize this\nintuition by introducing a discrete state zn \u2208 {1, . . . , B}, which determines the conditional mean and\ncovariance of a Gaussian prior on the embedding hn and contributes to a generalized linear model for\nthe following interval in. To capture the temporal correlation of these types, we include a Markovian\ndependency between zn and zn+1.\nWhile MRPs are able to model the evolution of discrete states over time, their assumption of stationary\ntransition distributions is overly restrictive for our application, as we expect zebra\ufb01sh to vary their\ntransition probabilities over time. To model non-stationarities in both the discrete transitions and\ninterval distributions, we introduce a scalar-valued continuous latent state xn that modulates the\ntransition probabilities and interval distributions. In the context of modeling zebra\ufb01sh behavior, we\nexpect these continuous states to capture slowly varying internal states like hunger, which are not\ndirectly observable but manifest in different patterns of swim bouts and intervals. At the same time,\nwe do not have strong prior beliefs about the dynamics of these states, except that they are smoothly\nvarying with a relatively long time constant. We capture these intuitions with a zero-mean Gaussian\nprocess prior on the continuous states, x(t) \u223c GP(K(t, t(cid:48))), with a squared exponential kernel.\nConditioned on xn = x(tn), we model the discrete transition probabilities with a generalized linear\nmodel, \u03c0\u03b8(zn\u22121, xn) = softmax(Wxxn + Pzn\u22121) where \u03b8 consists of Wx \u2208 RB and Pzn\u22121 \u2208 RB.\nThe matrix formed by stacking the row vectors {P T\nb=1 can be seen as a baseline log (unnormalized)\ntransition matrix, which is modulated by the continuous states xn. Similarly, we model the non-\nstationary interval distributions as gamma random variables parameterized by generalized linear\nmodels a\u03b8(xn, zn) and b\u03b8(xn, zn) with exponential link functions.\nIn sum, we sample the GPM-MRP by iteratively drawing from the following conditional distributions,\n\nb }B\n\nxn | {xn(cid:48), in(cid:48)}n(cid:48)<n \u223c N (mn|1:n\u22121, Cn|1:n\u22121),\n\nzn | xn, zn\u22121 \u223c \u03c0\u03b8(xn, zn\u22121),\nhn | zn \u223c N (\u00b5zn , \u03a3zn ),\nyn | hn \u223c N (\u00b5\u03b8(hn), \u03a3\u03b8(hn)).\n\nin | xn, zn \u223c Ga(a\u03b8(xn, zn), b\u03b8(xn, zn)),\n\n(GP predictive distribution)\n(Discrete transition probability)\n(Gaussian mixture of latent embeddings)\n(Gamma gen. linear model of intervals)\n(Deep generative model of marks)\n\n4\n\nobserved mark(eye and tail movement)observed intervallatent mark embeddingdiscretelatent statecontinuouslatent stateparametersFull Generative ModelCollapsed Generative ModelBidirectional LSTMRecognition Network...latentobserveddependencyneural net dep.cliqueLSTM state\u03b8xnznhninynn=1n=N...n=1n=N...n=1n=N\fHere, \u03b8 denotes the set of parameters that we must learn: the parameters of the generalized linear\nmodels of discrete transition probabilities and interval densities, the means and covariances of the\nlatent embeddings, and the weights of the neural network observation model. We treat the GP\nhyperparameters as \ufb01xed. Figure 2 (left) shows the complete graphical model.\n\n3.2 Special cases and extensions\n\nBy restricting the form of these dependencies we obtain many well-known models as special cases.\nBy removing the discrete and continuous states, we recover standard renewal processes. Given\nonly a continuous latent state, we can model piecewise constant conditional intensity functions and\napproximate log Gaussian Cox processes [19]. With only a discrete state, we recover the standard\nMarkov renewal process and, in discrete time, a hidden Markov model.\nThe GPM-MRP is only one of many possibilities for mixed discrete and continuous point process latent\nvariable models, and there are many clear extensions. For example, it is straightforward to allow the\nmarks to depend on both the discrete and the continuous states. Likewise, the interval model can also\nbe readily extended to more complex history dependence via an autoregressive process (AR) that\nconsiders in\u2212p, . . . , in\u22121. The continuous latent states could be multidimensional rather than scalar.\nFinally, the discrete transition probabilities can be extended to include semi-Markovian dependencies\nas well; i.e. to depend not only on the preceding discrete state, but also how long that state has been\nused. However, as we will show in the next section, it is critical that the discrete dependencies remain\ntractable so that we can ef\ufb01ciently compute the marginal distribution by summing them out.\n\n4\n\nInference\n\nNs\n\n1 , . . . , a(s)\n\nOur data consists of a set of S sequences of marked events. To simplify notation, let bold vari-\nables as (cid:44) a(s)\ndenote the values of variable a in sequence s of length Ns. Given a set of\nsuch sequences, we aim to estimate the global model parameters \u03b8 and infer a posterior distribution\nover the latent variables for each sequence p\u03b8(xs, zs, hs | is, ys). Computing this posterior and its\nreparameterization gradients is complicated by the presence of both continuous and discrete latent\nvariables. To handle these hybrid states, we develop an amortized variational inference algorithm that\ntargets the posterior of the collapsed distribution, analytically marginalizing out the discrete latent\nvariables,\n\np\u03b8({xs, hs, is, ys}S\n\ns=1) =\n\np\u03b8(xs, zs, hs, is, ys).\n\n(1)\n\nS(cid:89)\n\n(cid:88)\n\ns=1\n\nzs\n\nS(cid:88)\n\n(cid:104)\n\nThe key to this approach is that the discrete variables zs are connected in a Markov chain. Thus,\nfor any values of xs, hs, is, and ys, we can compute the marginal densities in (1) in O(Ns) time\nusing standard message passing algorithms, just as in an HMM [17]. Summing over the discrete states\nyields the densely connected but purely continuous generative model shown in Figure 2 (middle).\nWe approximate the intractable posterior distribution of xs and hs with a variational approximation\nq\u03c6(xs, hs) \u2248 p\u03b8(xs, hs | is, ys). We seek parameters \u03c6 that minimize the Kullback-Leibler diver-\ngence between the approximate and true posterior and the parameters \u03b8 that maximize the likelihood\nof the data. We \ufb01nd both simultaneously by optimizing the ELBO,\n\nL(\u03c6, \u03b8) =\n\nE\nq\u03c6(xs,hs)\n\nlog p\u03b8(xs, hs, is, ys) \u2212 log q\u03c6(xs, hs)\n\ns=1\n\nWe optimize this lower bound with stochastic gradient ascent using mini-batches of sequences.\nComputing gradients of the ELBO requires back-propagating through the HMM message passing\nroutine. See Appendix A for details on this routine and its gradients.\nOnce we have obtained an approximate posterior over the continuous latent variables, we reintroduce\nthe discrete states zs and compute their posterior. For a given con\ufb01guration of xs, hs and \u03b8, the\nconditional distribution p\u03b8(zs | xs, hs, is, ys) admits ef\ufb01cient algorithms for a variety of queries. We\ncan compute its mode, its marginal distributions, and draw samples from it, all using similar message\npassing algorithms. Thus, by optimizing the variational bound, we obtain the desired approximate\nposterior over discrete and continuous variables q(zs, xs, hs) = p\u03b8(zs | xs, hs, is, ys) q\u03c6(xs, hs).\n\n5\n\n(cid:105) \u2264 log p\u03b8({is, ys}S\n\ns=1).\n\n\fRecognition networks To accelerate inference, we also learn a recognition network that maps\na sequence is, ys to a set of variational parameters over the distribution of xs, hs [14, 15]. Our\nnetwork, shown in Figure 2 (right), assumes the approximate posterior factorizes as,\n\nq\u03c6(xs, hs; is, ys) =\n\nq\u03c6(xs; is, hs)\n\n.\n\n(2)\n\n(cid:104) Ns(cid:89)\n\nn=1\n\n(cid:124)\n\n(cid:105)\n\n(cid:125)\n\n(cid:124)\n\n(cid:123)(cid:122)\n\nq\u03c6(h(s)\n\nn ; y(s)\nn )\nfeed-forward\n\n(cid:123)(cid:122)\n\n(cid:125)\n\nbidirectional RNN\n\nn ; y(s)\n\nThe \ufb01rst term, q\u03c6(h(s)\nn ), is parameterized by a feed-forward neural network. Given a single\nbout\u2019s vector of eye and tail angles, this network outputs a mean and covariance of a Gaussian over\nn | is, hs) depends on observations\nthe inferred latent embedding. Since the true posterior p\u03b8(x(s)\nboth before and after the n-th event [20\u201322], we use a bidirectional recurrent neural network for\nq\u03c6(xs; is, hs). Given an input sequence, the network outputs a mean and covariance of xs. Complete\ndetails are in Appendix B.\nSparse GP inference The Gaussian process prior on xs imposes a substantial computational burden:\nevaluating the ELBO requires inverting the GP covariance matrix C, which is O(N 3\ns ) complexity.\nTo overcome this computational bottleneck, we use a sparse approximation to the full GP [23,\n24], computing the inverse covariance matrix at a subset ts,u \u2282 ts of Ns,u \u201cinducing\u201d points,\nwhere Ns,u (cid:28) Ns. (These are not technically inducing points as de\ufb01ned in Snelson and Ghahramani\n[24] since they are \ufb01xed, not learned.) For instance, in our experiments with zebra\ufb01sh behavior, we\ntake every 20th point in a sequence to be in this subset. For this sparse GP setup, the variational\nmodel only decodes the RNN hidden state at each point in ts,u. For a particular con\ufb01guration of xs,u\nat these events, the continuous states at the times of all other events follow deterministically.\n\n5 Related Work\n\nWe build upon a great deal of existing work on point processes, state space models, and approximate\nBayesian inference. These classes of methods have had signi\ufb01cant impact in computational neuro-\nscience [21, 25\u201334]. Of particular interest is the work of Cunningham et al. [27, 28], which develops\nGaussian process models of the underlying intensities of renewal processes and inference algorithms\nvia discretization of the underlying continuous intensity. The class of Gaussian process-modulated\npoint processes are well-studied in statistics and machine learning more generally. Prime among\nthese is the log Gaussian Cox process, which models the log intensity of a Poisson process as a\nGaussian process [19]. Several sampling and variational inference schemes have been proposed\nfor these types of models [35\u201339]. Most closely related to our work, Rao and Teh [36] propose a\nGaussian process-modulated renewal process and an accompanying uniformization-based sampling\nprocedure for inference of the latent continuous state. While these classes of models offer reasonable\napproaches for our scienti\ufb01c problem, they do not model co-evolving discrete and continuous latent\nstructure over time or incorporate deep generative models of marked data.\nA more recent body of work has combined deep generative models and state space models and\ndeveloped new inference methods for these deep, dynamic models. Particularly, advances in structured\nvariational inference provide us with methods for ef\ufb01cient inference in a variety of deep state space\nmodels [20\u201322]. The speci\ufb01c challenges of modeling mixed discrete and continuous states has also\ngarnered interest [40]. While our work draws upon these recent advances, we emphasize that our\nwork focuses on point process observations, which pose unique modeling and inference challenges.\nFinally, others have used neural networks for modeling point process data [41\u201343]. However, these\nmodels typically do not incorporate latent states in the dynamics. Moreover, in fully-general recurrent\nneural network models like these, it is more challenging to incorporate explicit prior knowledge about\nthe type and dynamics of latent variables. We make use of recurrent neural networks in our amortized\nvariational inference procedure, but their purpose is to accelerate scalable Bayesian inference in a\nstructured and interpretable probabilistic model.\n\n6 Synthetic Validation\n\nWe test our models and inference algorithm on synthetic data and ensure that we can accurately\nrecover the true underlying discrete and continuous latent structure from noisy marked point process\n\n6\n\n\fFigure 3: Synthetic data validation. We simulate a continuous latent state x(t) from a Gaussian process and\nevaluate it on a \ufb01nely spaced grid. These continuous states modulate the transition probabilities of an underlying\nset of discrete states, which in turn determine the likelihood of the observed time stamps and marks (not shown).\nThe base transition probabilities are shown at the top, along with the weights with which the continuous states\nbias them. Our amortized variational inference algorithm accurately recovers the true underlying continuous\nstates as well as the transition probabilities at each point in time. Two time points are shown here as examples.\n\nobservations. We simulate a synthetic dataset consisting of S = 1000 sequences, each of which\ncontains Ns = 300 events and shares the same global parameters \u03b8. We use Str = 750 of these\nsequences for training and save 250 for evaluation. We \ufb01x the true number of discrete states\nto B = 3, and we simulate H = 2 dimensional embeddings. For simplicity, we start by treating these\nembeddings as directly observable and focus on learning the discrete and continuous latent states of\nthe model. We learn the model parameters by maximizing a lower bound on the marginal likelihood,\nas described above, using a subset of size Ns,u = 15 for the sparse GP approximation.\nFigure 3 shows an example of true and inferred continuous latent states from one sequence\nin our synthetic dataset. The true latent states (evaluated on a \ufb01ne grid) are shown in black,\nand the inferred mean and 95% posterior credible intervals are shown in blue and light blue,\nrespectively. These are deterministic given a sample from the inferred posterior at the sub-\nset of points ts,u. The continuous latent states determine the transition probabilities at each\npoint in time by modulating the base transition matrix with a linear set of weights, as shown\nin the top left. Since x and z are de\ufb01ned up to a linear transformation and permutation, re-\nspectively, we solve for the optimal transformations to align the true and predicted latent states.\nWe see that the learned weights accurately recover the true underlying transition probabilities.\n\nStandard models like gamma renewal processes\nand Markov renewal processess can only ap-\nproximate the effects of the mixed discrete and\ncontinuous latent variables. This is evident in\nthe decreased log likelihoods on held-out test\ndata, as we show in Table 1. As the number of\ntraining sequences increases, the discrepancy in\ntest performance increases.\n\n# TRAIN SEQ.\n\nStr = 10\nStr = 50\nStr = 100\nStr = 250\n\nGPM-MRP\n-239.12\n-230.96\n-226.68\n-226.50\n\nMRP\n\nGRP\n\n-248.34\n-244.76\n-244.15\n-245.19\n\n-359.45\n-349.51\n-353.42\n-353.95\n\nTable 1: Test marginal likelihood on synthetic data for\nincreasing numbers of training sequences, Str.\n\n7 Experimental Results on Large-Scale Larval Zebra\ufb01sh Behavior Data\n\nFinally, we use these point process latent variable models to study latent states of larval zebra\ufb01sh\nbehavior. Figure 1 provides an overview of our experimental setup. Each \ufb01sh is observed one at a\ntime while swimming freely in a large tank, preying on paramecia. As described in Section 1, we\ntrack the \ufb01sh and record a 20 dimensional representation of the eyes and 180 dimensions for the\ntails in each bout. To place these features on the same footing, we \ufb01rst reduce the tail features to 20\n\n7\n\n020040060080010001200time1.51.00.50.00.51.0x(t)trueinferredinducing pointscredible intervalzt+1ztbaseweightsinferredtrueinferredtrue-404+x(t)\u00d7=\fFigure 4: Inferred discrete states of zebra\ufb01sh behavior, their baseline transition probabilities, and their following\ninterval distributions. Discrete states can be understood in terms of their corresponding eye angle: positive\nangles indicate hunting, negative indicate exploration. The states follow characteristic transition patterns and\nintervals. We use cross-validation to select B = 8 discrete states and H = 10 embedding dimensions, a model\nwith high test likelihood and interpretable results.\n\ndimensions using PCA, giving us a D = 40 dimensional mark for each bout. Each of 130 \ufb01sh were\nobserved in trials over a 40 minute period, resulting in over 120,000 swim bouts. We use 105 \ufb01sh for\ntraining and 25 for model comparison. We \ufb01t the model with 50 epochs of stochastic gradient ascent.\n\nin, yn\nin, yn\nin, yn\nin, yn\n\nin, yn, zn\nin, yn, hn\n\nGPM-MRP\n\nin, yn, zn, hn\n\nin, yn, zn, hn, xn\n\nMethod\n\nVariables\n\nPP\nGRP\nAR1\nAR10\nMRP\nGRP+\nMRP+\n\nTest LL\n-57.24\n-57.06\n-50.01\n-49.75\n-41.15\n-24.88\n-23.12\n-22.61\n\nFigure 4 shows the inferred bout types for a ran-\ndomly chosen 1000 bouts. We see that they cluster\ninto B = 8 groups in our latent space. In the top left\npanel, we display two dimensions of our H = 10\ndimensional latent space. Upon inspection, we \ufb01nd\nthat these two dimensions correlate with two key\ncharacteristics of the eye-angles over the course of\na bout. Particularly, we compute the change in eye\nangle between the \ufb01rst and last frames of a bout,\nas well as the mean eye angle over the 10 frames.\nPer these features, we \ufb01nd that the inferred clus-\nters correspond to known bout types related to head\nmovement (grey), exploratory locomotion (pink, orange, red, and crimson), J-turns (yellow) that\nsignal the entrance to a hunt [5], pursuits (green), and hunt-ends (blue). These bouts follow an\ninterpretable transition matrix that suggests \ufb01sh alternate between exploration and pursuing prey, and\nthe transition between these two modes is gated by J-turns and hunt-ends. Moreover, each bout type\nentails a characteristic distribution over the following interbout interval. We chose the number of\nstates and embedding dimension based on the held-out ELBO and inspection. We found that with more\nthan B = 8 states and H = 10 dimensions, the gains in held-out likelihood diminished. Moreover,\nthe inferred clusters appear to further subdivide the explore bouts without fundamentally changing\nthe transition or interval distributions, suggesting that these re\ufb01nements are less meaningful.\nMRPs could identify these latent types of bouts, but they cannot easily capture the in\ufb02uence of internal\nstates like hunger. In this experiment, 57 of the \ufb01sh were starved for 2-4 hours prior to entering\n\nTable 2: Test log likelihood of zebra\ufb01sh data in units\nof nats/bout.\n\n8\n\n\fFigure 5: Hunger modulates transition probabilities into each bout and the intervals following them. Top: Our\ncontinuous latent states capture a bias in transition probabilities over time. On average, we see that starved \ufb01sh\n(dark lines) up-regulate hunting related bouts (J-turns and pursuits). Bottom: Starved \ufb01sh also swim more often\nregardless of bout type, as indicated by the decreased interbout intervals. After 40 minutes in tank preying on\nparamecia, fed and starved \ufb01sh equalize.\n\nthe tank. Our GPM-MRP \ufb01nds that these \ufb01sh change their transition probabilities and intervals as a\nfunction of how long they have been in the tank.\nFigure 5 shows the effect of the continuous state on log transition probabilities (top) and interbout-\nintervals (bottom), averaged over all \ufb01sh in the fed or starved groups in \ufb01ve minute intervals. Starved\n\ufb01sh up-regulate the probability of entering a hunt and pursuing prey, whereas fed \ufb01sh show increased\nprobability of ending hunts. Across all bout types, starved \ufb01sh show shorter expected inter-bout\nintervals. In sum, starved \ufb01sh swim more often and are more likely to engage in hunts, as we might\nexpect. Table 2 shows that the GPM-MRP is not only interpretable, it also outperforms existing models\nin predicting held-out data. See Appendix B for further details on the baseline comparisons.\n\n8 Discussion\n\nThe principal output of an animal\u2019s nervous system is a sequence of actions selected from its\nbehavioral repertoire. Understanding the set of possible actions [44] and the ways in which they\nare \ufb02exibly and adaptively combined is critical to constraining our understanding of how even the\nsmallest animal brains function in the natural world. The larval zebra\ufb01sh is studied in thousands\nof labs worldwide [45] and its behavior is unique among model organisms in that it is naturally\nsegmented into punctuated bouts. This simple behavioral structure lends itself well to be modeled as\na marked point process.\nWe develop new PPLVMs and show how a hidden continuous internal variable like hunger can\nmodulate both action selection and timing. Our models blend co-evolving discrete and continuous\nlatent states to generate marked point process observations. We show how one member of this class,\nthe GPM-MRP, is able to capture meaningful dynamics in a large-scale dataset of zebra\ufb01sh behavior.\nWhile the models we develop are able to uncover meaningful latent structure, there are several\npotential areas for improvement. For instance, our discrete state dynamics are limited by our ability\nto analytically marginalize them out, but semi-Markovian models [46] are a natural extension. In\naddition, while we build on prior work on point processes, we have only explored PPLVMs within the\ncontext of temporal observations. We leave an examination of blending point processes, state space\nmodels, and deep generative models in the spatiotemporal domain to future work.\nAs models of behavior grow to incorporate multiple internal variables (e.g. stress, arousal, attention,\nfear), interpretable models will be necessary to understand how unobserved variables interact to yield\nnatural behavioral sequences. Such models will aid in generating hypotheses about how the brain\nimplements behavioral algorithms that are modulated by latent internal states. For example, we \ufb01nd\nthat increased hunger promotes shorter wait times between actions. This knowledge may be used\nin conjunction with whole-brain imaging studies to identify neural populations which regulate the\nprecise timing of action initiation in both health and disease.\n\n9\n\n0.10.00.1transition biashead mvmt.explore1explore2explore3explore4j-turnpursuithunt-end0-55-1010-1515-2020-2525-3030-3535-40time (min)0.20.40.60.8interval (s)0-55-1010-1515-2020-2525-3030-3535-40time (min)0-55-1010-1515-2020-2525-3030-3535-40time (min)0-55-1010-1515-2020-2525-3030-3535-40time (min)0-55-1010-1515-2020-2525-3030-3535-40time (min)0-55-1010-1515-2020-2525-3030-3535-40time (min)0-55-1010-1515-2020-2525-3030-3535-40time (min)0-55-1010-1515-2020-2525-3030-3535-40time (min)fedstarved\fAcknowledgements. The authors thank John Cunningham and Liam Paninski for helpful advice and\nfeedback. SWL thanks the Simons Foundation for their support (SCGB-418011). FE received funding from the\nNational Institutes of Health\u2019s Brain Initiative U19NS104653, R24NS086601 and R43OD024879, as well as\nSimons Foundation grants (SCGB-542973 and 325207).\n\nReferences\n[1] Andr\u00b4e EX Brown and Benjamin de Bivort. Ethology as a physical science. Nature Physics,\n\n2018.\n\n[2] Gordon J Berman. Measuring behavior across scales. BMC biology, 16(1):23, 2018.\n[3] Eileen Kowler. Eye movements: The past 25 years. Vision research, 51(13):1457\u20131483, 2011.\n[4] Ann Kennedy, Greg Wayne, Patrick Kaifosh, Karina Alvi\u02dcna, LF Abbott, and Nathaniel B\nSawtell. A temporal basis for predicting the sensory consequences of motor commands in an\nelectric \ufb01sh. Nature neuroscience, 17(3):416, 2014.\n\n[5] Allan V Kalueff, Michael Gebhardt, Adam Michael Stewart, Jonathan M Cachat, Mallorie\nBrimmer, Jonathan S Chawla, Cassandra Craddock, Evan J Kyzar, Andrew Roth, Samuel\nLandsman, et al. Towards a comprehensive catalog of zebra\ufb01sh behavior 1.0 and beyond.\nZebra\ufb01sh, 10(1):70\u201386, 2013.\n\n[6] Jo\u02dcao C Marques, Simone Lackner, Rita F\u00b4elix, and Michael B Orger. Structure of the zebra\ufb01sh\nlocomotor repertoire revealed with unsupervised behavioral clustering. Current Biology, 2018.\n[7] Timothy W Dunn, Yu Mu, Sujatha Narayan, Owen Randlett, Eva A Naumann, Chao-Tsung\nYang, Alexander F Schier, Jeremy Freeman, Florian Engert, and Misha B Ahrens. Brain-wide\nmapping of neural activity controlling zebra\ufb01sh exploratory locomotion. Elife, 5:e12741, 2016.\n[8] Misha B Ahrens, Jennifer M Li, Michael B Orger, Drew N Robson, Alexander F Schier, Florian\nEngert, and Ruben Portugues. Brain-wide neuronal dynamics during motor adaptation in\nzebra\ufb01sh. Nature, 485(7399):471, 2012.\n\n[9] Lin Cong, Zeguan Wang, Yuming Chai, Wei Hang, Chunfeng Shang, Wenbin Yang, Lu Bai,\nJiulin Du, Kai Wang, and Quan Wen. Rapid whole brain imaging of neural activity in freely\nbehaving larval zebra\ufb01sh (Danio rerio). eLife, 6, 2017.\n\n[10] Dal Hyung Kim, Jungsoo Kim, Jo\u02dcao C Marques, Abhinav Grama, David GC Hildebrand,\nWenchao Gu, Jennifer M Li, and Drew N Robson. Pan-neuronal calcium imaging with cellular\nresolution in freely swimming zebra\ufb01sh. Nature methods, 14(11):1107, 2017.\n\n[11] Martin Haesemeyer, Drew N Robson, Jennifer M Li, Alexander F Schier, and Florian Engert. A\nbrain-wide circuit model of heat-evoked swimming behavior in larval zebra\ufb01sh. Neuron, 98(4):\n817\u2013831, 2018.\n\n[12] Robert E. Johnson*, Scott W. Linderman*, Thomas Panier, Carole Wee, Erin Song, Kristian\nHerrera, Andrew C. Miller, and Florian Engert. Revealing multiple timescales of structure in\nlarval zebra\ufb01sh behavior. Computational and Systems Neuroscience (Cosyne) Abstracts, 2018.\n\n[13] John F. C. Kingman. Poisson processes, volume 3. Clarendon Press, 1992.\n[14] Diederik P Kingma and Max Welling. Auto-encoding variational Bayes. In Proceedings of the\n\nInternational Conference on Learning Representations (ICLR), 2014.\n\n[15] Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation\nand approximate inference in deep generative models. In Proceedings of the 31st International\nConference on Machine Learning, pages 1278\u20131286, 2014.\n\n[16] Rajesh Ranganath, Linpeng Tang, Laurent Charlin, and David Blei. Deep exponential families.\nIn Proceedings of the International Conference on Arti\ufb01cial Intelligence and Statistics, pages\n762\u2013771, 2015.\n\n[17] Lawrence R Rabiner and Biing-Hwang Juang. An introduction to hidden Markov models. IEEE\n\nASSP Magazine, 3(1):4\u201316, 1986.\n\n[18] Carl Edward Rasmussen and Christopher KI Williams. Gaussian processes for machine learning.\n\nThe MIT Press, Cambridge, MA, USA, 2006.\n\n[19] Jesper M\u00f8ller, Anne Randi Syversveen, and Rasmus Plenge Waagepetersen. Log Gaussian Cox\n\nprocesses. Scandinavian Journal of Statistics, 25(3):451\u2013482, 1998.\n\n10\n\n\f[20] Rahul G Krishnan, Uri Shalit, and David Sontag. Structured inference networks for nonlinear\n\nstate space models. In AAAI, pages 2101\u20132109, 2017.\n\n[21] Yuanjun Gao, Evan W Archer, Liam Paninski, and John P Cunningham. Linear dynamical\nneural population models through nonlinear embeddings. In Advances in Neural Information\nProcessing Systems, pages 163\u2013171, 2016.\n\n[22] Matthew Johnson, David K Duvenaud, Alex Wiltschko, Ryan P Adams, and Sandeep R Datta.\nComposing graphical models with neural networks for structured representations and fast\ninference. In Advances in Neural Information Processing Systems, pages 2946\u20132954, 2016.\n\n[23] Lehel Csat\u00b4o and Manfred Opper. Sparse on-line gaussian processes. Neural computation, 14\n\n(3):641\u2013668, 2002.\n\n[24] Edward Snelson and Zoubin Ghahramani. Sparse Gaussian processes using pseudo-inputs. In\n\nAdvances in Neural Information Processing Systems, pages 1257\u20131264, 2006.\n\n[25] Anne C Smith and Emery N Brown. Estimating a state-space model from point process\n\nobservations. Neural Computation, 15(5):965\u2013991, 2003.\n\n[26] Wilson Truccolo, Uri T. Eden, Matthew R. Fellows, John P. Donoghue, and Emery N. Brown. A\npoint process framework for relating neural spiking activity to spiking history, neural ensemble,\nand extrinsic covariate effects. Journal of Neurophysiology, 93(2):1074\u20131089, 2005.\n\n[27] John P Cunningham, Byron M Yu, Maneesh Sahani, and Krishna V Shenoy. Inferring neural\n\ufb01ring rates from spike trains using Gaussian processes. Advances in Neural Information\nProcessing Systems, pages 329\u2013336, 2007.\n\n[28] John P Cunningham, Krishna V Shenoy, and Maneesh Sahani. Fast Gaussian process methods\nfor point process intensity estimation. In Proceedings of the 25th International Conference on\nMachine Learning, pages 192\u2013199. ACM, 2008.\n\n[29] Jonathan W Pillow, Jonathon Shlens, Liam Paninski, Alexander Sher, Alan M Litke,\nEJ Chichilnisky, and Eero P Simoncelli. Spatio-temporal correlations and visual signalling in a\ncomplete neuronal population. Nature, 454(7207):995, 2008.\n\n[30] Liam Paninski, Yashar Ahmadian, Daniel Gil Ferreira, Shinsuke Koyama, Kamiar Rahnama\nRad, Michael Vidne, Joshua Vogelstein, and Wei Wu. A new look at state-space models for\nneural data. Journal of computational neuroscience, 29(1-2):107\u2013126, 2010.\n\n[31] Jakob H Macke, Lars Buesing, John P Cunningham, M Yu Byron, Krishna V Shenoy, and\nManeesh Sahani. Empirical models of spiking in neural populations. In Advances in Neural\nInformation Processing Systems, pages 1350\u20131358, 2011.\n\n[32] Scott Linderman, Ryan P Adams, and Jonathan W Pillow. Bayesian latent structure discovery\nfrom multi-neuron recordings. In Advances in Neural Information Processing Systems, pages\n2002\u20132010, 2016.\n\n[33] Yuan Zhao and Il Memming Park. Variational latent Gaussian process for recovering single-trial\n\ndynamics from population spike trains. Neural Computation, 29(5):1293\u20131316, 2017.\n\n[34] L Paninski and JP Cunningham. Neural data science: accelerating the experiment-analysis-\n\ntheory cycle in large-scale neuroscience. Current opinion in neurobiology, 50:232, 2018.\n\n[35] Ryan Prescott Adams, Iain Murray, and David JC MacKay. Tractable nonparametric Bayesian\ninference in Poisson processes with Gaussian process intensities. In Proceedings of the 26th\nAnnual International Conference on Machine Learning, pages 9\u201316. ACM, 2009.\n\n[36] Vinayak Rao and Yee Whye Teh. Gaussian process modulated renewal processes. In Advances\n\nin Neural Information Processing Systems, pages 2474\u20132482, 2011.\n\n[37] Chris Lloyd, Tom Gunter, Michael Osborne, and Stephen Roberts. Variational inference for\nIn International Conference on Machine\n\nGaussian process modulated Poisson processes.\nLearning, pages 1814\u20131822, 2015.\n\n[38] Yves-Laurent Kom Samo and Stephen Roberts. Scalable nonparametric Bayesian inference on\npoint processes with Gaussian processes. In International Conference on Machine Learning,\npages 2227\u20132236, 2015.\n\n[39] Maneesh Sahani, Gergo Bohner, and Arne Meyer. Score-matching estimators for continuous-\ntime point-process regression models. In Machine Learning for Signal Processing (MLSP),\n2016 IEEE 26th International Workshop on, pages 1\u20135. IEEE, 2016.\n\n11\n\n\f[40] Scott W. Linderman*, Matthew J. Johnson*, Andrew C. Miller, Ryan P. Adams, David M. Blei,\nand Liam Paninski. Bayesian learning and inference in recurrent switching linear dynamical\nsystems. In Proceedings of the 20th International Conference on Arti\ufb01cial Intelligence and\nStatistics (AISTATS), 2017.\n\n[41] Nan Du, Hanjun Dai, Rakshit Trivedi, Utkarsh Upadhyay, Manuel Gomez-Rodriguez, and\nLe Song. Recurrent marked temporal point processes: Embedding event history to vector. In\nProceedings of the 22nd ACM SIGKDD International Conference on Knowledge Discovery and\nData Mining, pages 1555\u20131564. ACM, 2016.\n\n[42] Hongyuan Mei and Jason M Eisner. The neural Hawkes process: A neurally self-modulating\nmultivariate point process. In Advances in Neural Information Processing Systems, pages\n6757\u20136767, 2017.\n\n[43] Shuai Xiao, Mehrdad Farajtabar, Xiaojing Ye, Junchi Yan, Le Song, and Hongyuan Zha. Wasser-\nstein learning of deep generative point process models. In Advances in Neural Information\nProcessing Systems, pages 3247\u20133257, 2017.\n\n[44] Michael B Orger and Gonzalo G de Polavieja. Zebra\ufb01sh behavior: opportunities and challenges.\n\nAnnual review of neuroscience, 40:125\u2013147, 2017.\n\n[45] Katie Lidster, Gareth D Readman, Mark J Prescott, and Stewart F Owen. International survey\non the use and welfare of zebra\ufb01sh Danio rerio in research. Journal of \ufb01sh biology, 90(5):\n1891\u20131905, 2017.\n\n[46] Kevin P Murphy. Hidden semi-Markov models (HSMMs). Technical report, MIT, 2002.\n\n12\n\n\f", "award": [], "sourceid": 7992, "authors": [{"given_name": "Anuj", "family_name": "Sharma", "institution": "Columbia University"}, {"given_name": "Robert", "family_name": "Johnson", "institution": "Harvard University"}, {"given_name": "Florian", "family_name": "Engert", "institution": "Harvard University"}, {"given_name": "Scott", "family_name": "Linderman", "institution": "Columbia University"}]}