{"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)