Parsing neural dynamics with infinite recurrent switching linear dynamical systems
Authors: Victor Geadah
International Brain Laboratory
Jonathan W. Pillow
Venue: ICLR 2024 poster (2024)
Keywords: Markov switching processes
Neural data analysis
State-space models
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
TL;DR: We extend switching state-space models with a PDE prior that induces a state geometry, and showcase the model's abilities for uncovering trial-varying structure in neural data

Abstract: Unsupervised methods for dimensionality reduction of neural activity and behavior have provided unprecedented insights into the underpinnings of neural information processing. One popular approach involves the recurrent switching linear dynamical system (rSLDS) model, which describes the latent dynamics of neural spike train data using discrete switches between a finite number of low-dimensional linear dynamical systems. However, a few properties of rSLDS model limit its deployability on trial-varying data, such as a fixed number of states over trials, and no latent structure or organization of states. Here we overcome these limitations by endowing the rSLDS model with a semi-Markov discrete state process, with latent geometry, that captures key properties of stochastic processes over partitions with flexible state cardinality. We leverage partial differential equations (PDE) theory to derive an efficient, semi-parametric formulation for dynamical sufficient statistics to the discrete states. This process, combined with switching dynamics, defines our infinite recurrent switching linear dynamical system (irSLDS) model class. We first validate and demonstrate the capabilities of our model on synthetic data. Next, we turn to the analysis of mice electrophysiological data during decision-making, and uncover strong non-stationary processes underlying both within-trial and trial-averaged neural activity.

Published as a conference paper at ICLR 2024
PARSING NEURAL DYNAMICS WITH INFINITE RECUR- RENT SWITCHING LINEAR DYNAMICAL SYSTEMS
Victor Geadah', The International Brain Laboratory, Jonathan W. Pillow
Program in Applied and Computational Mathematics, Princeton University.
Princeton Neuroscience Institute, Princeton University. (victor. geadah, pillow) @princeton. edu, anfrodamtterrneiiiomlbbaaaaaaaa
ABSTRACT
Unsupervised methods for dimensionality reduction of neural activity and behavior have provided unprecedented insights into the underpinnings of neural information processing, One popular approach involves the recurrent switching linear dynami- cal system (rSLDS) model, which describes the latent dynamics of neural spike train data using discrete switches between finite number of lww-dimesiiona linear dynamical systems. However, a few properties of rSLDS model limit its deployability on trial varying data, such as fixed number of states over trials, and no latent structure nor organization of states. Here we overcome these limitations by endowing the rSLDS model with a semi-Markov discrete state process, with latent geometry, that captures key properties of stochastic processes partitions with flexible state cardinality. We leverage partial differential equations (PDE) theory to derive efficient, semi- par-maareettric formulation for dynamical sufficient statistics to the discrete states. This process, combined with switching dynamics, defines our infinite recurrent switching linear dynamical system (irSLDS) model class. We first validate and demonstrate the capabilities of our model on syn- thetic data. Next, we turn to the analysis of mice electrophysiological data during decision- decision-making, and uncover strong non-stationary processes underlying both within-trial and trial-averaged neural activity.
INTRODUCTION
Parsing behavioral and neural dynamics into simpler alternating states is providing unprecedented insights into the nature of computation in the brain (Wiltschko et al. 2015 Calhoun et al. 2019 Ashwood et al. 2022 Bolkan et al. 2022). Such data, conventionally collected over multiple trials and spanning considerable lengths of time, is often related to latent processes that exhibit complex dynamics across multiple timescales (Cowley et al.) [220 [oyy et al. 2021 Nassar et al. 2019 Yu et al. 2009). Learning is one important example of such non-stationariy, but other changes may be present, such fatigue (Marcora al. 2009) arousal (Schriver et al. (2018). However, inferring those latent processes from data, either within or across trials, remains challenging.
The state-space models that perform this segmentation into discrete states commonly revolve around the hidden Markov model (HMM). This model posits an underlying, hidden, discrete Markov chain, with different observation model for each state giving rise the data. Each discrete state can also be extended to capture a different dynamical regime, governing the dynamics of some continuous-space stochastic process. The resulting switching state-space models capture activity that alternate between discrete set of dynamical regimes, and have proved useful in the modeling of complex nonlinear activity (Fox et al., 2010, Smith al., (02T). A significant step in capturing online dependencies in such models is through recurrence (Linderman et al.) 2016 Zoltowski et al.) 2020), where the continuous dynamical variables govern switches between the discrete states. Such models have been providing powerful insights into the dime-dependence of those processes (e.g. Glaser al. (2020)). but suffer from training challenges due to the fixed cardinality of the discrete states and a lack of geometry over states.
Indeed, traditional HMMs, and models built upon them such as switching state-space models, use a fixed number of discrete states. This fixed cardinality is typically determined by cross-validation,

---

Published as a conference paper at ICLR 2024
procedure that can prove computationally expensive. Furthermore, it does not encourage the use of fewer discrete states on any subset of trials or sessions, so that they can be interpreted and determined needed. To overcome these limitations, multiple avenues have considered expending HMMs with flexible state cardinality (Beal et al. 200T Fox et al.) 2011). These models revolve around the Hierarchical Dirichlet Process (HIP), which places flexible non-parametric prior over the transition probabilities between the (theoretically infinite) states. Unfortunately, while the HDP and its generalizations can capture state dependence and persistence (Teh et al. 2006 Fox al., 2008), in their standard formulation they do not allow for recurrence and other dependencies between the discrete states allocations.
Furthermore, HMMs do not have any priori state geometry. In particular, the discrete states are permutation invariant, such that the model is equivalent under any permutation in the label ordering of the discrete states (see schematic in Appendix Fig. 4 B This induces non-identifiabiiity and limits the interpretability of the inferred latent discrete states. While recent work has focused on variability
in discrete dynamical regimes (Nassar et al. 2019: Linderman et al. 2019 Costacurta et al. 2022) yielding impressive results in terms of flexibility, most of these can be cast as "local" variations ; they consider perturbations parametrized variations of the parameters associated with each discrete states. To our understanding, none specifically tackle the geometric nature of the discrete states space.
In this work, we use partial differential equations (PDE) theory to develop an alternate prior over the discrete state transitions in recurrent switching state-space models. Our formulation allows us to (1) capture important properties of the stochastic processes for flexible state cardinality, while supporting the use of recurrent connections, (2) induce geometry over the discrete states, and (3) induce eemi-Markovian discrete states dynamics capture more intricate temporal dependencies. After reviewing relevant background in Section 2 we present in Section our model, the infinite recurrent switching linear dynamical system (irSLDS), and provide details on generative formalism and inference procedure. In Section $ we first validate the model on synthetic data and showcase its properties, before turning to mice dlectro-hyysooloy data from the International Brain Laboratory (IBL). The PDE prior is defined by less parameters than the traditional HMM-based transitions model, yet we show that it maintains or even outperform the traditional model while offering an interpretable structure that is amenable uncovering trial-varying structure.
2 BACKGROUND
We review in this section the key models considered in this work. We consider time-stamped data {(fns yn) yn))²== with time steps t, and data yn R
2.1 SWITCHING LINEAR DYNAMICAL SYSTEMS
Hidden Markov Models (HMM) posit an underlying discrete-state Markov chain with states E {1...., K}. and conditionally independent observations p(æn|2z)) Auto-regressive HMMs (AR- HMM) extend this framework by endowing the observations on E R with linear dynamics, dependent on the discrete state 2n The generative model reads as
Zn+1 (124+1/2::71, x.,)
In+1 N (A(zn) xn + (Zu)) Exe] E,,)
(2)
at time n E (O...., Th. with different set of dynamics ( E DxD and bias (=x)) E R for each discrete state 2n, and covariance Σ, The switching process in (eq. [T] is described by transition matrix, written in full generality, but is usually taken to be time- (n) and observation (ær,)
independent such that P(zx+1 = žn = = Auj for some transition matrix A.
In this work, we are interested in leveraging those models to study possibly high-dimensional data yn E RM,, M 2 D. such as spike train data. We thus turn to Switching Linear Dynamical Systems (SLDS) models, which model the AR-HMM observations Rn (low-dimensional) representation of the data, and model the data V1.T as conditionally independent linear Gaussian emissions
Yn N' (CIn + C, Eny
(3)
for n E T). with decoding weights C E MMx c. E R M and emission covariance Σ,- Finally, inputs un e.g. the time steps Ln. may be linearly encoded in the continuous dynamics (eq. 2
2

---

Published as a conference paper at ICLR 2024
or directly along with the emissions in (eq. 33 Following Fox et al. (2010), refer to AR-HMM and SLDS models as instances of Markov switching processes, or generally switching state-spact models.
A vital augmentation of Markov switching processes is to allow observations internal states guide the switches in discrete states. Linderman et al. (2016) introduced such a 2i - 2x+11 dependency (blue arrows in Fig. (E) in the SLDS model class, coining recurrence. Now in the recurrent SLDS (rSLDS), in place of the general equation in eq. T the continuous states En guide the switching through
Zn+1 2ns Un R In + r
for recurrent encoding matrices R( E RRXD and biases r.( E R.KK This is the central model we consider in this work-sce Fig. TE for the graphical model.
2.2 DISTANCE DEPENDENT CHINESE RESTAURANT PROCESS
The switching process in (eq. 1 sets the discrete dynamical modes of the continuous dynamics. Modeling and inferring trajectories under this process amounts to the problem of clustering the dynamics in discrete set of dynamical modes. As stated in the introduction, we wish to infer the number states K from the data, contrary the formulation above. Dirichlet Processes (DP) are a classical tool in this case, they provide an infinite random measure over cluster partitions (Ferguson 1973 Antoniak 1974). They can be described by the Chinese Restaurant Process (CRP, form of Pólya urn process, see primer in Appendix $A) which captures the self-reinforcing and non-parametric structure of the prior. key limitation is that points under the CRP are exchangeable (Blei & Frazier 2011) (see schematic in Appendix Fig. ) underr any permutation of the ordering, the probability of a configuration is the same. The purpose of recurrence is to actively control the dynamical mode transition, which fundamentally breaks exchangcability.
We seek a stochastic process over discrete state partitions that will allow us to capture recurrent (i.c. online) dependencies on the data latent states. To this end, turn to the distance dependent CRP (dist-CRP) from Blei & Frazier (2011). This variant the CRP breaks exchangeability by associating each time step another, and then perform clustering based on pairwise assignments. Specifically, at time step n, this process assigns time step E
[v)] with ei E [r]] following
{ [(Dij: 3) i aa
plG = -DD,j,,,a
(4)
else
with distance matrix Dij. decay function f(:; and decay parameter 3 > 0. As we consider time-stamped data, we consider the time difference D,j = t, t, for i 2 j. and enforce that no step is assigned with future steps by setting (D 0 i j. Finally,
the entire history C.n = [C.m-1, } of pairwise assignments to perform clustering: if c = Cj. then and j are in the same cluster zę. Applying this the history Cans get Z.n, thereby setting Zn. The graphical model for this process is illustrated in Fig. TA
3 INFINITE RECURRENT SWITCHING LINEAR DYNAMICAL SYSTEMS
Our technical contribution consists in using key properties of the dist-CRP to guide the switching process in the iSLDS in manner that supports efficient inference and generation. Section $3.1 presents the fallouts of a naive combination, providing details on how conventional methods relying on Pólya-gamma augmentation yield intractable Bayesian inference in this case. Section $3.2 motivates our alternative approach to modeling based on sufficient statistics and partial differential equations (PDEs), and highlights key non non-exchanngable and geometrical properties of our model. Finally, Section 3.3 provides implementation details for both generation and inference.
3.1 PÓLYA-GAMMA AUGMENTATION YIELDS INTRACTABLE INFERENCE
Pólya-gamma gamma augmentation (Possnn et al. 2013 Linderman et al. 2015) is powerful augmentation
strategy that handles non-Gaussian factors by introducing additional latent variables to obtain joint Gaussian likelihoods. It was used in the original formulation of the rSLDS model by Linderman
et al. (016)) to allow tractable inference with message passing. We show in Appendix CC that
3

---

Published as a conference paper at ICLR 2024
A Distance Dependent CRP
B
Time evolution
w(i
wij-
Space smoothness prior
wi
Recurrent Switching LDS
Infinite Recurrent Switching LDS
Discretize
Dicrrete state
Sufficient
4
20
TE Discrete states
30
400
Continuous states
Discrete
Date
Figure 1: Modeling details. (A) Graphical model for the distance dependent CRP (dist-CRP) in its conventional form. (B) The generative process of the dist-CRP can be summarized with a sufficient statistic w(j,t) t) for state j at time showcasing a typical choice-driven exponential decay. (C) We induce geometry over the state space j E J, and add a spatial smoothness prior. (top) Each time a state is chosen (red bar), its probability is increased for future steps and so is the probability of nearby states. (D) Modeling the influence function w(j. t) as a solution to the heat equation. (E) Graphical model for the recurrent switching linear dynamical system (rSLDS). The recurrence is highlighted by the blue arrows. (F) The infinite rSLDS model combines the heat-equation generalization of the
dist-CRP for the discrete states with the rSLDS switching linear dynamics emission process. (G) The irSLDS discrete state prior better supports geometric analyses (top) and a variable number of occupied states (bottom).
for appropriate choices of decay function f. we can also leverage Pólya-gamma augmentation to similarly handle the non-Gaussian factors emerging from recurrence. However, the resulting Gaussian augmentation grows linearly for each step n, making it computationally inefficient. Luckily as we'll below, can use sufficient statistics and recurrent dynamics circumvent this problem.
3.2 PDE PRIOR FOR EFFICIENT PARAMETRIZATION AND NON-EXCHANGEABILIT
To remedy the challenges arising out of a naive combination of the original dist-CRP and SLDS models, we rely on the following sufficient statistic express the cluster allocations
[wwj [wij if j in history else
pz = jic.n a, B) = p(2, j|w(, (n), a, 3)
(5)
a
where w(j, 1) denotes the influence function state j at time t I - w(j()
Σ
J(I (1-4;;81 4;
(1:A5A)))
with distance function f(:: 3) and decay parameter 8 > 0. Using this sufficient statistic makes the entire process Markovian. Setting f(x; 3) exp(-3z) to be exponentially decreasing function, notice that we can rewrite w(j. t) above as a solution to the continuous time ODE
Øt ((j -ww(j, (=.=)) VjEJ
with càdlág trajectories w(j in time t. Each time t, discrete state 2n j is chosen, its influence function is bumped by x >> 0, increasing the weight of this state for future time steps (see Fig. B This has similar effect the "stickiness" parameter in the sticky HDP-HMM (Fox et al.! 2008).
4

---

Published as a conference paper at ICLR 2024
In practice, need to evolve in parallel the influence functions for each state j. For any fixed t, we make the modeling choice of relating these states to one another by using a smoothness prior (see Fig. [TC. Inspired by the temporal dynamics in (eq. ( model the sufficient statistic to evolve according to a partial differential equation (PDE), the heat equation (see Fig. 1
8
((j.t) t) ? 0j²
w(i.t)
Ōt
with diffusion coefficient > and with initial conditions for 2 fm
w w(,, (,) lim w(+,t)+ +k.., w 0) =0
rft,
over the state space j J. While many alternate priors could offer form of spatial smoothness with similar temporal dynamics, we choose the heat-equation for several reasons:
1. The temporal dynamics encompass the dist-CRP. Indeed, by standard separation of variables argument, the temporal component of solution w to (eq. 7 exhibits the exponential decay desired in (eq. 6 for some appropriate 3 0.
2. The spatial dynamics offer an intuitive notion of diffusion of probabilities. Each time a discrete state j is chosen, "nearby" states also become more probable the next time step (see Fig. [ top).
3. The heat-equation as a differential operator functions possesses numerous linearity qualities, which makes general numerical behavior and approximations favorable.
We stress the meaning of "nearby": in models based around an HMM, the states do not have an inherent geometry In comparison, here discrete states can be interpreted as points on a continuous state space J. We can restrain this space to discrete number of states, and we discuss how to do so in practice in the next subsection.
3.3 IMPLEMENTATION AND INFERENCE IN INFINITE RECURRENT SLDS
We proceed by describing some of the implementation implications of the PDE prior presented in the previous section. We assume from now on segularly-sample data in time, of difference Åt, and drop the inclusion of the time data fy only consider the emissions data yn. The code for this work builds on the SSM package (Linderman et al. 2020), and we present in Appendix C.2 the relevant modules.
PDE prior implementation with finite difference methods We use finite difference methods to implement sufficient statistic w solving the heat-equation prior. Let [wa+il5+1 := w(j + Aj, (n 1)At) be our discrete approximation. We a forward difference of the time derivative and central difference approximation to the second order spatial partial derivative to obtain the solution
Wn+1 Uwn, = tridiag(6, 1 1-2,
23,3) E K
general requirement for stability, and let Aj be adjusted accordingly given (model parameter) and At (data parameter). Inputs can be added to drive the system, including (1) the desired =1,,= adding self-reinforcemmnt to the system, and (2) the past internal states encoded by a module of weight E RR and bias r E R for recurrence. The dimension K of w,, R' is fixed hyperparameter. Importantly, however, the heat equation dynamics capture properties of the dist-CRP, such that K rather acts as an upper bound on the parsimonious number of states used per trajectory. In all, the dynamics of w, follow
where -
At We impose AL (Aj)²
Wn++1 (Wi, Fn, Zn = i} Uw,, + xl, + Ræ,, +
with parameters of diffusion and decay x B > 0 and self renforrinfore E R. The j-th entry of the vector Wx+1 in (eq. { defines, up to necessary row-normalization factor, the ij-th entry of our transition matrix [W... n+ Tij- Dynamical modes 2n are then sampled according to (eq. 5). with final 0. Now few scalar parameters {a, Y, x)) govern the discrete state on process, contrasting with the K x K matrix the rSLDS. In all eqs. ( with the switching continuous dynamics in (eq. 2 and emissions in (eq. 3) together define the infinite recurrent SLDS (irSLDS, see graphical model in Fig. [.... We refer to Appendix $D.3 for a study of the behavior of the model scale K, the influence of {a, 7, ss.. and the values used for the latter.
random transition parameter
'One can consider spectral clustering from the graph Laplacian defined from the transition matrix. This however remains transformation, and the underlying states originally have geometry.
5

---

Published as a conference paper at ICLR 2024
Inference with variational Laplace-EM EMM The added sufficient statistic wn, in (eq. is determin- istically determined given 2n-1 and In-1. thus adds no component to the posterior inference. We follow prior work and perform inference in this model using variational Laplace-En, EM from Zoltowski et al. (022)). As an overview, we posit the same structured mean-field approximation q(>1:T ) (::)) to the states posterior. For our continuous states posterior, the most likely path L:TT (Paninski et al. 2209). Given a continuous states trajectory under this posterior, the transition matrices WI.T can be computed to define our model joint p(y1:T, 21:T, ,1.T,31.T) .I:T). From there, the discrete state posterior approximation is found by locally maximizing the expected model joint under the continuous state posterior. Generally, conditioned on samples T1:7, the factor graph for our z posterior is equivalent to that of an heterogeneous HMM. Common tools can then be used, such as importantly the Viterbi algorithm to obtain the most likely discrete state sequence $1:T conditioned on T1.7, which we will consider below.
Laplace approximation around
Discrete state geometry The support of W,, is taken to be the set (A,, 2Aj...., KA, KAjJ. which lives on the latent geometry I used in our PDE formalism. In this work we consider one possible way to leverage this geometry: to define intermediate states. Indeed, the continuity of allows for interpolation between states. From any distribution continuous interpolation thereof, one can define different statistics of interest such as the mean or the mode. Such statistics can take continuous values in J, and are ill-defined for traditional HMM-based models due to permutation invariance. In particular, given the most likely sequence Ê1: and our posterior q(z Zn, 2n) with interpolation function over J, we will consider the interpolated sequence 11. as the interpolated posterior modes
discrete states 2n E {1...., }), or
En = arg max (zn = z 8n-1, 1n) n E {I....., TT, 20 = arg ((20) SEJ
Parallelizing transition matrix dynamics with seans In both generation and inference, the core computational difference in terms of complexity between the irSLDS and the (r)SLDS models is that the transition matrix W,, for the discrete states 2n now possesses its own recurrent dynamics (from eq. (8) see Appendix eq. (13). While the (r)SLDS models also have time-varying transition matrices, they do not have such recurrent term. To illustrate the difference, considering eq. (8) for row of the transition matrix Wn, the equivalent notation for the rSLDS would be
Wn+11
{I 2n [wo] + Rx +
Computing the on transition matrices in the rSLDS thus amounts to the cost of encoding the continu- ous state sequence T1:r, which can be distributed over different processors. In comparison, we are forced to compute sequentially the transition matrices statistics Wa, naively taking (T3 opera- tions in addition to the sequential encoding of æ1.7. Fortunately, notice that the dynamics (Appendix eq. [3) are linear. In this case, presented by Blelloch| (190),, computing the transition matrices can cast as a scan operation (see /Smith et al. (2023) for a recent application of this concept). efficiently parallelize our transition matrix dynamics and reduce the computation
With this, to (T L + log L)) operations over L processors, resulting in similar time complexity as the rSLDS model. We refer to Appendix C.2 for code and more details our use of the parallel scan.
4 EXPERIMENTS
We train models by maximizing the Evidence Lower Bound (ELBO) using variational Laplace-EM from Zoltowski et al. (2020). We compare model performance using the marginal log likelihood (LL)
logr(ūr.) = f p(ünr a Z (..T)((#T.) F1.r ddrr)dd
which is the log-probability of held-out test data ÿ1.T under a given model, where test data arises from 4:1 train-to-test split of the full dataset (see details in Appendix $B). The required integral is high-dimensional and intractable, and we thus resort to sequential Monte Carlo (SMC), also known as particle filtering, compute it (Del Moral et al. 2006 Kantas et al. 2009
4.1 VALIDATION ON THE SYNTHETIC NASCAR TASK
We consider the synthetic NASCAR task as used by Linderman et al. (016). In the original task, the true underlying model is rSLDS with K states, dissecting the 2-D state space in four

---

Published as a conference paper at ICLR 2024
Extended NASCAR
C
Continuous states x, posterior mean
NASCAR
new dynamics
NASCAR
Extended NASCAR
State space dynamics
B Extended dynamics
Altemate
or 2 times
D Samplex: NASCAR
Extended NASCAR
Sample
Inferred discrete
Original
Extended dynemics
ISLDS
ILDD
6:29
H00 9.
10
1065
Disonetee state
iSLDS (ours)
InSLDS
new nynamis
SA
500
3000
Time step
Figure 2: Synthetic NASCAR experiments. (A-B) True flow field for the NASCAR task in A. The extended NASCAR task alternates between the flow field in A and the extended flow field in B. (C) Mean continuous trajectories 11.7 under the variational Laplace posterior, per task. (D-E) Sample continuous latent states Fı.r trajectories per model (NASCAR in D, extended NASCAR in E). (F) Most likely discrete state trajectories 21:T for the rSLDS (top) and irSLDS (ours, bottom) for one data trajectory in the extended NASCAR task (color scheme matches E), along with the interpolated sequence hI. (eq. 9 in black for the irSLDS.
Table 1: Performance metrics on NASCAR experiments. We report the marginal ligeliiel on test data (Test LL), evaluated with SMC. Test data-sets consist of T = 200 excluded steps in NASCAR, and excluded trajectory of = 1000 steps in extended NASCAR. Column values for K denote value set for the model. We also report the dynamical meea--squre error (Dyn. MSE) between the sample æ state-space flow and true flow field. See Appendix Table 2 for dynamical MSE error bars and number of parameters.
NASCAR
Extended NASCAR
Test (t)
Dyn. MSE (4)
Test LL (1) Dyn. MSE (L)
Model
K =4 4
K K=6 8
K 44
K =
K=6
SLDS
1331.35
1592.81
0.930
305.046
7353.22
316.75
rSLDS
1628.57
1613.53
0.010
0.011
8015.93
0.35
irSLDS (ours)
1656.10
1611.45
0.015
0.011
8219.41
0.24
quadrants with rotational and bias dynamics (see Fig. 2 The resulting continuous latent states resemble car tracks on a NASCAR driving track (see Fig. I--f for corresponding mean posterior estimates under variational Laplace-EM), and the data is obtained from high-dimensional (N 10) Gaussian emissions of the æ, dynamics. We compare the isSLDS against the rSLDS and SLDSS models, presenting results over 5 random initialization seeds.
First, we found both our model, the irSLDS, and the true generative model, the rSLDS, to be similarly able to generate accurate samples of the dynamics (see Fig 2). In terms of performance time steps of the trajectory (T = 200), both models attained similar accuracy (see Tab. I.. The irSLDS attained the higher mean on the true number of states K = 4, and both models performed similarly for the over-complete K = 8. The SLDS however did not perform as well, as to be expected- this NASCAR task introduced test the inclusion of recurrence in the rSLDS,
held-out
7

---

Published as a conference paper at ICLR 2024
compared to the standard SLDS. While the SLDS model attained relatively high test LL (see Tab. it struggled to accurately depict the true dynamics (see Fig. 2D-left). We refer to this generating process accuracy as dynamical accuracy, which quantify by reporting the mean squared error (MSE) between the learned dynamics and the true, available, dynamics. We minimize the MSE over reflections and rotations of the learned dynamics field. The SLDS model showed poor dynamical accuracy (see Tab. [T. while the irSLDS matched the accuracy of the true rSLDS. In all, this validates the incorporating of recurrence into our model, and furthermore provides evidence that the PDE prior can mimick the performance and generative capabilities of the fully-trained HMM prior.
The original NASCAR task has a periodic temporal structure (a property unchanged by permutation invariance) with a constant number of required states K bbtboth stationary properties. Before turning to our target experimental data, we assess the handling of non-stationary challenges by considering an extended NASCAR task. In this new task, the dynamics alternate between the original dynamics (Fig. [ and a set of new extended dynamics (Fig. 23) where two additional bias dynamics are introduced. This effectively changes the true number of discrete states from K = 4 + 6 during this extended dynamics block. We consider E {2 blocks of alternating dynamics, each of same expected block lengths T/ B. In this new task, we found the irSLDS to outperform the rSLDS, in turn outperform the SLDS (Tab. D Looking into the learned solutions for the rSLDS and isSLDS models, we observe that the sample trajectories (Fig. E in the irSLDS more closely match the posterior mean from data. Furthermore, the irSLDS model correctly identifies two new states after the switch (see Fig. -bobottom). Interestingly, the posterior ends up attributing these two states to be in between previous states, such that the sequence of discrete states exhibits this traversing trajectory. In contrast, the rSLDS model mistakenly maintains only active states after the switch to the alternate dynamics ( Fig. -ttpp) Together, this provides evidence for the performance of the irSLDS in con-stationary tasks and an example on how we can leverage the discrete state geometry of the irSLDS
4.2 THE IRSLDS UNCOVERS TRIAL-VARYING STRUCTURE IN NEURAL DYNAMICS REGIMES
Next, we turn our attention to an electrophysiolggiaal dataset from the International Brain Laboratory (IBL), recorded in mice during sensory- motor seeision-making task. In this task, mice reported the location of a sinusoidal grating stimulus by turning a wheel either left or right, with task difficulty controlled by stimulus constrast (Fig. 3 (Laboratory et al. 2021). We consider Neuropixels probe recordings from the "Brainwide map" data release (Laboratory [002]]. For our analyses, we projected spike train data onto the top principal components to obtain firing rates, making it amenable to analysis by state-space models with Gaussian emissions (Fig. B-CC (further methodological and data details can be found in Appendix $B including firing-rate and continuous latent dimensions). As the irSLDS directly expands on the rSL.DS we focus on comparing these two models in this section. We picked K = 8 discrete states see Appendix Tab. for test marginal LL values for K E {2, 4, 8J)..
First, we found that the models uncovered discrete latent states that switched at task-relevant event times. Indeed, we did not provide the models with task event times such as the stimulus onset ("Stim On")) or reward, or behavioral measurements such as the movement onset time ("First movement"). Given that a task event occurred at time step t, we plot in Fig. 3D the estimated probability of a switch time-step + for lag I (z(switch+++/evene the baseline probability p( ((switche), obtained by trial shuffling. The models learned task switches significantly different from chance, capturing a switch in discrete states either preceding or following a relevant task event. We thus posit that the statistical models capture relevant dynamics for the tasks.
)) for the irSLDS. We compare this against
Second, we found that the irSLDS uncovered differences in the discrete state distribution over trials. In Fig. 3 we plot the number of active states used in the most likely 21:7 trajectories under the variational posterior. We fit spline function to this data for both models, determined from the minimal p-value of the F-test for a comparison against a constant function (computed over various degrees of smoothness and degree). While both models differed significantly from a constant function (p < 0.001), the irSLDS uncovered larger fluctuations in the cumber-of-discrett www-diireeestttte-ppeeee sitee-stttte-e-rriiid curve. Fluctuations were much smaller for the rSLDS model, and the spline provided a poorer fit (It2 of 0.16 v.s. 0.46 for irSLDS). This indicates that discrete state distribution over trials fluctuated more dramatically and systematically under the irSLDS model, with substantially more discrete states used in the middle of the session than at the beginning or end. The irSLDS attained slightly higher log- likelihood on held-out test trials (see Tab. 33. so in all this indicates that those fluctuations are important, and might be missed by traditional HMM-basser priors.
8

---

Published as a conference paper at ICLR 2024
Spike train data
Friing RC.
Nouron ic.
Pre procossing
IBL., 2022)
IBL Task
Time from smoc start
LikE
uis Lika
10
In al trme steps
50
Trial-level modeling
Session-level modeling
D
Stimulus On
First Movement
Reward
FSLDS
Switch probability
Contitenal
Discrete 14
Lag
E
ISLDS
irSLDS (ours)
irSLDS
0.16
R²
Interp.2
Number of octive
Discrete state
200
400
232
400
600
Trial number
Trial number
Trial number
Figure 3: Results on neural electro-physiology data from the IBI. (A) Experimental paradigm in the IBL task, taken from Laboratory (2022). See text for further details. (B) Example spike train data, plotting spike times per neuron, for one trial in one session. (C) The trial in B. pre-processed with kernel smoothing and PCA, along with relevant trial events. (D) Switching probability as a function of the time-step difference (lag 1) for two key trial events, either conditioning on the event happening lag ("conditional"", black), or without conditioning ("baseline", gray). (E) Number of actives states occupied by most likely discrete states trajectories per trial. Curve fits, either constant of spline, are plotted against the number of states per trial in gray. The R²2 for the spline fit is indicated top right. (F) Most likely discrete states 21:Neaate trajectories underlying per-trial average neural responses.
To further investigate the task-induces neural activity, we turn from trial-level modeling to session- level modeling. Specifically, we focus on the average firing rate y/: E R averaging from 200ms before the stimulus onset to the reward, for each trial n E [Nrrint]] Fig. 3 shows the most likely discrete state sequences &1:Nerat. under each model. Both models achieve similar training ELBO (see Appendix Tab. (3). We observe that the irSLDS showcases a distinctive slow geometric drift in the discrete states occupied, as observed from and further solidified by the trace of the interpolated (defined in eq. eq. Finally, we note how the interpolation also indicates the presence of uncertainty between states and 4 the end of the session. Such conclusions cannot be drawn in the (r)SLDS, as any permutation of the discrete states yields an equivalent visualization of the discrete state process.
5 CONCLUSION
In this work, we extend recurrent switching state-space models with an input-driven heat-equation prior over the dynamics of the transition matrix. This results in a semi-Markov discrete state process that capture the behavior of stochastic processes on partitions for variable number of states per trajectory, as well as inducing a continuous geometry on the discrete states. We show that while this process is defined by less parameters than the traditional HMM models, it still matches or outperforms the original model while providing insights into varying processes data. We showcase the model on synthetic task. before turning to clectrophysiological data from the IBL. The IBL hosts extensive datasets covering multiple repeated- to brain-wide wide recordings (Laboratory et al. 2022 Laboratory| 2022), and the modeling presented here provides grounds for further investigation into the vire-narying processes underlying neural data.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGMENTS
The authors would like to thank Rich Pang and David Zoltowski for helpful discussions and code suggestions early in the project. VG was supported by doctoral scholarships from the Natural Sciences and Engineering Research Council of Canada (NSERC) [PGSD3--57885--222 and the Fonds de Recherche du Québec Nature technologies (FRQNT) [B2X 297667]. JWP was supported by grants from the Simons Foundation (SCGB AWD543027), the NIH (R01EY017366), the NIH BRAIN initiative (N$104899 and (40))))56404-044, and the CAREER award (IIS-1 150186).
REFERENCES
Charles E. Antoniak. Mixtures of dirichlet processes with applications to bayesian nonparametric problems. Ann. Statist,, 2(6), 1974. doi: 10.1214/aos/1 176342871. URL https://doi
org ooc/:00.114422225522111113313333
Zoe C.
Ashwood, Nicholas A. Roy, Iris R. Stone, Anne E. Urai Anne K. Churchland, Alexandre Pouget, and Jonathan W. Pillow and. Mice alternate between discrete strategies during per- ceptual decision-making, Nature Neuroscience, 25(2):201-212 February 2022. doi: 10.1038/
141593-021-01007-z.
Matthew Beal, Zoubin Ghahramani, and Carl Rasmussen. The Infinite Hidden Markov Model. In Dietterich, S. Becker, and Z. Ghahramani
in Neural Information Processing Sys-
tems, volume 14. MIT Press, 2001. URL http::tt////p/oeidiiiii
wtwpa//nnnnn......eeeeeeeeeee.ea http:pp.ccppppppe Paper.pdl
200.
D. M. Blei and P. I. Frazier. Distance Dependent Chinese Restaurant Processes. Journal of Ma- chine Learning Research, 12(74):2461-2488 2011. URL
bleil h
Guy E. Blelloch. Prefix sums and their applications. Technical Report CMU-CS-90-190, School of Computer Science, Carnegie Mellon University, November 1990.
Scott S. Bolkan, Iris R. Stone, Lucas Pinto, Zoe C. Ashwood, Jorge M. Iravedra Garcia, Alison L. Herman, Priyanka Singh, Akhil Bandi, Julia Cox, Christopher A. Zimmerman, Jounhong Ryan Cho, Ben Engelhard, Jonathan W. Pillow, and Ilana B. Witten. Opponent control of behavior by dorsomedial striatal pathways depends on task demands and internal state. Nat Neurosci, 25 (3):345-357, mar 2022. doi: 022 01021-9
URL [RLltttss////oioorgg
103882Fs4
Adam J. Calhoun, Jonathan W. Pillow, and Mala Murthy. Unsupervised identification of the internal states that shape natural behavior. Nature Neuroscience, 22(12):204--244) November 2019. doi: 10. 10.1038//41593-01--533 841593-019-0533-
URL https:://ol..rg//00...
Julia C Costacurta, Lea Duncker, Blue Sheffer, Winthrop Gillis, Caleb Weinreb, Jeffrey Evan Markowitz, Sandeep R. Datta, Alex Williams, and Scott Linderman. Distinguishing discrete and continuous behavioral variability using warped autoregressive HMMs. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022. URL
net/ffrunt d (Kj]66iiijii
Benjamin R. Cowley, Adam C. Snyder, Katerina Acar, Ryan C. Williamson, Byron M. Yu, and Matthew A. Smith. Slow drift of neural activity signature of impulsivity in macaque visual and prefrontal cortex. Neuron, 108(3):551-56) 1575115567,e0 e8, November 2020. doi: 10. wel.1066..euronn222222222. 200000000000000000012. URL htt.thttp:///o..o. .rg...rgg/. /00 101
neuron 2020.07 07.021
Pierre Del Moral, Amaud Doucet, and Ajay Jasra. Sequential monte carlo samplers. Journal of the Royal Statistical Society: Series B (Statistical Methodology). 68(3):411 1--46, 2006.
Thomas S. Ferguson. A bayesian analysis of some nonparametric problems. Ann. Statist,, 1(2), mar 1973. doi: 10, 10. /0114///////1342336666 URL /20_
2F 17177644360
10

---

Published as a conference paper at ICLR 2024
E. Fox, E. B. Sudderth, M. Jordan, and A. S. Willsky. Bayesian nonparametric inference of switching dynamic linear models. IEEE Transactions on Signal Processing, 59(4):1569 1585, 4 2011. doi: 10.1 10.1100/s.2221021102222.
/ 11w.11..111110000 1tt..:orr....111122............. 2102755
Emily B. Fox, Erik B. Sudderth, Michael I. Jordan, and Alan S. Willsky. An hdp-hmm for systems with state persistence. In Proceedings of the 25th International Conference on Machine Learning, ICML '08, pp. 312-319, New York. NY, USA. 2008. Association for Computing Machinery. ISBN 9781605582054, doi: 10.1 10.1145/1391156. 1390196. httphttps:///oo.oooooooo htt.ddoooogg/..00 1390156. 1390196
Emily B. Fox, Erik B. Sudderth, Michael I. Jordan, and Alan S. Willsky. Bayesian nonparametric methods for learning markov switching processes. IEEE Signal Processing Magazine, 27(6):43-54, 2010. doi: 10.1 1000/M5P2200..9779999.
Joshua Glaser, Matthew Whiteway. John P Cunningham, Liam Paninski, and Scott Linder- man. Recurrent Switching Dynamical Systems Models for Multiple Interacting Neural Pop- ulations. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin (eds.), Ad- vances in Neural Information Processing Systems, volume 33, pp. 14867-14878. Curran Asso- ciates, Inc., 2020. URL URRJuttpss.pprrceednnsss.. neurpss.......ssss............ ccppaper.ppppppeeepepe fi files/paper 2220/fileeaalffff/33 7ha hat.a//SSS./22///eodooeeeeeeeeea 200 Paper.
Nicholas Kantas, Arnau Doucet, Sumeetpal Sindhu Singh, and Jan Marian Maciejowski. overview of sequential monte carlo methods for parameter estimation in general state-space models. IFAC Proceedings Volumes, 42(10))774-785, 2009.
International Brain Laboratory. 2022. URL 0402000022222220111133
Data release Brainwide map Q4 2022, 11
wtt.o///rrticceepppeeeeeee prepriint Ceeerintttattttt Data release_
Brai nwi ____map___04___0
International Brain Laboratory, Valeria Rourigon-Rd Dora Angelaki, Hannah Bayer, Niccolò Bonacchi, Matteo Carandini, Fanny Cazettes, Gaelle Chapuis, Anne K Churchland, Yang Dan, et al. Standardized and reproducible measurement of decision-making in mice. Elife, 10:c63711, 2021.
International Brain Laboratory, Kush Banga, Julius Benson, Niccolò Bonacchi, Sebastian A Bruijns, Rob Campbell, Gaëlle A Chapuis, Anne K Churchland, M Felicia Davatolhagh, Hyun Dong Lee, Mayo Faulkner, Fei Hu, Julia Hunterberg, Anup Khanal, Christopher Krasniak, Guido T Meijer, Nathaniel Miska, Zeinab Mohammadi, Jean-Paul Noel, Liam Paninski, Alejandro Pan- Vazuuzz Noam Roth. Michael Schartner, Karolina Socha, Nicholas A Steinmetz, Karel Svoboda. Marsa Taheri, Anne Urai, Miles Wells, Steven West, Matthew Whiteway, Olivier Winter, and Ilana Witten. Reproducibility of vivo electrophy measurements in mice. May 2022. URL https:://di..rgg /10.
10. http///..111111100 1101/2022 05. 09. 491042
S. W. Linderman, M. Johnson, and R. P. Adams. Dependent multinomial models made easy: Stick-breaking with the polya-gamma augmentation. In Corinna Cortes, Neil D. Lawrence, Daniel D. Lee, Masashi Sugiyama, and Roman Garnett (eds.), Ad- vances in Neural Information Processing Systems 28: Annual Conference on Neural Infor- mation Processing Systems 2015, December 7-12, 2015, Montreal, Quebec, Canada, pp. 3456-3466 2015. URL https: http:e/ee/pccceeeeeieeee nwwrips pttp:/ppppppppp///........ 0 07a4e20a7bbbeebe 1 ottaaaaee.....22222ebeb-e--aa...
Abst .htll
S. W. Linderman, C. Miller, P. Adams, D. Blei, Paninski, and M. J. Johnson. Recurrent switching linear dynamical systems, 2016.
Scott Linderman, Annika Nichols, David Blei, Manuel Zimmer, and Liam Paninski. Hierarchical recurrent state space models reveal discrete and continuous dynamics of neural activity in c. elegans. ", April 2019. URL (RLLittts:///a// ort/://///oooooo/.//ooooo/
Scott Linderman, Benjamin Antin, David Zoltowski, and Joshua Glaser. Ssm: Bayesian learning and inference for state space models, 2020. htll.ttts:.// /ttp:////ttttttttt com/ inderman lab/ smm

---

Published as a conference paper at ICLR 2024
Samuele M. Marcora, Walter Staiano, and Victoria Manning. Mental fatigue impairs physical performance in humans. Journal of Applied Physiology,
March 2009. doi:
URL httpas.//o...org/1..11112.
91324. 2008
Josue Nassar, Scott Linderman, Monica Bugallo, and Il Memming Park. Tree-structured recur- rent switching linear dynamical systems for multi-scale modeling. International Confer- ence on Learning Representations, 2019. URL http:ss///penreeee...eeee nerm:/
Hk HkzRQhR9YX
Liam Paninski, Yashar Ahmadian, Daniel Gil Ferreira, Shinsuke Koyama, Kamiar Rahnama Rad, Michael Vidne, Joshua Vogelstein, and Wei Wu. A look state-space models for neural data. Journal of Computational Neuroscience, 29(1-2)11071266, August 2009. doi: 10.1007/
s10627-009-0179-x URL htt.https:.. //dp:p:://do.... org/ / 0 1007 002000000000000----
Nicholas G. Polson, James G. Scott, and Jesse Windle. Bayesian inference for logistic models using polya- gamma latent variables. Journal of the American
Associaaion,,,, 108(504):
2013. ISSN 01621459. URL http: htTE::://ww..sstt.... wttp://ttggt....l////o/oo
Nicholas A. Roy, Ji Hyun Bak, Athena Akrami, Carlos D. Brody, and Jonathan W. Pillow. Extracting the dynamics of behavior in sensory decision-making experiments. Neuron, 109(4):597-660.c6
February 2021. doi: 10. www..i.eeroo..2222...... neuron 2020 12. 004
Brian Schriver, Svetlana Bagdasarov, and Qi Wang. Pupil-linked arousal modulates behavior in rats performing whisker deflection direction discrimination task. Journal of Neurophysiology, 120(4):165--666) October 2018. doi: 10.1152/jn.00290.2018 URL/aattp:///o..ooro 1152/ jn. 00290. 2018
Hideaki Shimazaki and Shigeru Shinomoto. A method for selecting the bin size of time histogram. Neurul Computation, 150) June 2007. doi: 10. 22.1166/ecco.207..11 1503. URL
https::/doi.com .or..orgg//o. 10..0.112//eeee 2007.19.6.... 1503
Jimmy Smith, Scott Linderman, and David Sussillo. Reverse engineering recurrent neural networks with Jacobian switching linear dynamical systems. In Advances in Neural Information Processing Systems, volume 34, pp. 16700-16713. Curran Asso- ciates, Inc., 2021. URL httpp::://prccccciinngi/oo/////t neureeeidiinnnnnniieeee neurips........ cc http:/pppe/222///aae/ 887755115c1eee11222611161----3-3 Abstraat .ataattttttlll
Jimmy T.H. Smith, Andrew Warrington, and Scott Linderman. Simplified state space layers for sequence modeling. In The Eleventh International Conference on Learning Representations, 2023.
net hwt.tnettoornnoee forum ? 8Hw 3H
Yee Whye Teh, Michael Jordan, Matthew
Beal, and David M Blei. Hierarchical dirichlet processes. Journal of the American Statistical Association, 101(476):1566 1581, dec 2006. doi: 10. 14118//00624440000000033333
Alexander B. Wiltschko, Matthew J. Johnson, Giuliano Iurilli, Ralph E. Peterson, Jesse M. Katon, Stan L. Pashkovski, Victoria Abraira, Ryan Adams, and Sandeep Robert Datta. Mapping sub-second structure in mouse behavior. Newron, 88(6): December 2015. doi: 10. 1016/ j j.neuron. 2015. 1.031. /ttptssss://dii wwt.:doooorr.g../ 1016/ www.on...oonmmmmmm 2015 11.
Byron M. Yu, John Cunningham, Gopal Santhanam, Stephen Ryu, Krishna V. Shenoy, and Maneesh Sahani. Gaussian-process factor analysis for low-dimensional single-trial analysis of neural population activity. Journal of Neurophysiology, 102(1):614 -355,,,,,,6661, July 2009. doi: 10. (1/52/nn.99411 .2008. URL httpt:tt///e/t / orw. http:/10..11////0
90941.
D. M. Zoltowski, J. W. Pillow, and S. W. Linderman. Unifying and generalizing models of neural dynamics during decision decision-making, 2020. URL http:t////aavvvoo////////t 2050 04571
12

---

Published as a conference paper at ICLR 2024
Broader Impact The modeling here presented is targeted towards enhancing our understanding of neural computation. Analysis of electrophysiologien data can have long-term implications for the treatment and understanding of medical treatment and neurological disorders. However, these considerations are far removed for the preliminary analysis and modeling presented, and we foresee immediate societal consequences of this work.
Limitations We highlight a few limitations. First, in practice for our experiments, our implemen- tation of the scan operation for parallel computing, as presented in section $3.3] does not support the automatic differentiation used for our variational inference. We thus only use the scan partially- we refer to section C.2 for more details on the manner. Second, we present early results on ephys data, for limited number of sessions, probe locations, animals, etc. We do not aim to make general claims about computation in neural systems, but rather highlight the properties of our models in uncovering time-dependent processes. Expanding to the aforementioned data modalities is grounds for rich future work.
external clusters. For reference, training on single session for the ephys IBL data might take up 12 hours for a single model, with some multiprocessing. Training single model for the synthetic NASCAR data can be accomplished on the order of 30 minutes.
Compute All experiments were
A BACKGROUND
A.1 CHINESE RESTAURANT PROCESS
Dirichlet Processes (DP) mixture models provide powerful random measure over clusterings. We refer to Ferguson 1973): Antoniak (1974) for a peeasur-thhoreee treatment of DPs, and Tch et al. (2006) for machine-learning overview. They can be alternatively defined through the Chinese Restaurant Process (CRP), a process akin to the Pólya Um process. The analogy goes as follows upon entering restaurant, a customer selects sit table with probability proportional the number of people already sat at that table. With some fixed rate o, they may decide to start new table. Put otherwise, for a new customer i, its table allocation 2; follows
T2 ax
if <
plz = klzz. a) x
(10)
k +1
with ng the size of cluster E [K] In this work, one should think of tables as clusters or dynamical modes, and the customers time-steps. It can be easily seen that this process induces joint probability over cluster assignments that is invariant to the order of customers entering. We call this property exchangeability. This enforces strong and limited prior on distributions of partitions that can arise from this model.
B TRAINING DETAILS
B. NASCAR TASK
Data generation We generate NASCAR track trajectories by instantiating the model as described, and generating sample trajectory running for T = 1000 steps. We then split the first 800 steps as train data, and keep the last 200 steps as test data (4:1 split). We repeat this procedure for random seeds. For all models, we use K 4. D = 2 (æ, dimension) and M = 10 (observation y/ dimension).
B.2 IBL TASK
Pre- Pre-processing Given spike train data from given period, we first perform the following pre- processing steps to obtain the firing rates over the period
1. Kernel smoothing against Gaussian kernel obtain firing rate graphs. We use standard- deviation of 100ms if the period a whole trial, and 30ms if the period is the Open-Loop period.
13

---

Published as a conference paper at ICLR 2024
A
Exchangeability
Permutation invariance
Discrete states
Discrete NJ U
equiprobable sequences
equivalent visualizations
Figure 4: Schematic representation of (A) exchangeability. notably in the Chinese Restaurant Process, and (B) permutation invariance, notably in HMMs and HMM-based models.
2. The firing rate graphs are evaluated over an evenly space grid over the period interval. The bin size for the grid is determined following Algorithm from Shimazaki & Shinomoto (2007), taking the minimal value, after removing outliers, over trials. We define outliers as values falling outside the [Q1 1.5 * IQR, Q3 1.5 IQR] quartile interval. The chosen bin size is on the order of 40ms for whole-trial period, and 10ms for the Open-loop period.
After the firing rates are obtained, the next step in processing depends on the nature of the data analyzed, and involves using PCA. If consider trial-level modeling, then perform PCA project the firing rates onto the top PCs (M = 50 PCs, over 90% variance explained). If we consider session-level modeling, we first average the firing rates from 200ms before the stimulus onset to the reward, then perform PCA to project the per-trial activity onto the top PCs from the concatenated trials (M 10 PCs, over 85% variance explained).
Finally set our internal states to be of dimension D = 4 for per-trial neural activity, and D = 2 for per-trial-averraeed activity.
Ephys data details We use Neuropixels recordings tied to the Brainwide data release Labo- ratory (2022). The results this paper pertain to the first available session, with session id We refer the reader to the data release website to obtain more information on the specific probe: https://viis http://ii...tttrraaiiaaaaaaaaaaa org/app
ae878 7b1 1011-2223 2e2--956-b022-6660-----
Data we use 4:1 train to test split for the per-trial spike train analysis. 20% of the trials are randomly selected as test trials- these trials span whole session.
C MODELING
C.I PÓLYA-GAMMA AUGMENTATION FOR GIBBS SAMPLING INFINITE RECURRENT SLDS
We provide details inference challenges in infinite recurrent SLDS model naively combining the dist-CRP and the SLDS. Following previous methodologies, we could use Gibbs sampling and leverage message passing perform inference. would require handle the conditional density
p(11:N CCIN. ZI,N, (VINN, tinh) x II (w Fnn 2n sn)bb(e Cns t.n ( Un)
n=1
where (an,,, +1 1 Ca++1 is the potential from the continuous recurrent dynamics, and (Tn, Un) the evidence potentials. The recurrent connections introduce the dependencies captured (5 Ce+1). 1), which adds significant challenges for inference. Without it, in the standard SLDS, the potentials are all Gaussian, allowing for analytical integration.
Following Linderman et al. (2016), can leverage Pólya-gamma augmentation (Polson et al. 2013 Linderman et al.) (2115) to deal with these non- non-Gaussian factors. The key is that instead of performing categorical choice over pre-determined set of K dynamical modes, perform association, categorical choice, with a previous time-step j E {1,...,t).
t). Because the conceptual categorical similarities in updating, we find similarities in inference methodology with the rSLDS. This is where
14

---

Published as a conference paper at ICLR 2024
the choice of decay function f comes in, in enforcing that link. This non- non- -au-Saas factor is
1, Gn. (.)) - p( km) a II
[In - l3i B(x 1) 1)
(e [(n=
j=1
n-1 II jml
eluals
a [ex=n
[+ e.
for Mn E RR-1 [un]i = 3(xxx-1) (In tz). We can leverage the following integral quantity
(1+ = 2- eN 1se - pro (wib,, Ojdw
(e")o
b
b b> = =
2
to introduce auxiliary variables {wr}s=1 such that the conditional density p(cn+1 2n fm+1, becomes Gaussian
(FF, Gx+ 1; x a N (un, 22 Ki, (2,²²)
(11)
where 2, diag(w):- and [k..]; [[,, = j]. Krn E TR.. With this augmentation, the required
potentials are Gaussian and the integral can be calculated analytically. We refer the reader to Blei &
Frazier (2011) for subsequent details how one handles the messages mm--n+1 (Gx+1)-
While we do obtain an analytical expression for the augmented potentials in (eq, I1 such potentials multivariate Gaussians of size n for each time-step computational complexity, both in handling the potentials and in any required marginalization.
[I,...,NN, This adds significant
C.2 IMPLEMENTATION OF THE IRSLDS
In this appendix subsection, we present code modules and functions relevant to the modeling.
Parallel scan We cast the linear dynamics of the transition matrices W, as efficiently parallelize the computation. The SSM package relies on Numba and aut auroggad,, the latter not having readily implemented associate scan operation. To this end, implemented in Numba and Numpy the scan operation, using numba for the parallization (Listing [ Unfortunately, this scan operation cannot leverage the automatic differentiation from autograd. We found however that including the scan in the computation does not significantly change the learned solutions in the iSSLDS parameters. Thus in practice, we use the scan to do faster hyperparameter search or gradient-free free computation (e.g. generating samples, evaluating likelihoods), but rely on the full sequential version for the final training.
operation to
(IL
oopartyo (aLt, 9_1):
geen Binazy operetoz for parallei sean of lineez socurssaco.
Axgas
18 [, dit tuple containing EL1 bud at position. 1, (K, 3
tuple containing and position
Reburnst
naw olement Louut, blas_out
10
11
tiy b_1 a
12
op.natnul HLI, U_)), oponatiull (b (_)))
13
Dpunba. 1it (parcllol-Teue Iorsocb9-Truul
15
def apply_ bles_elene)as)::
16
cean
17
Apply perallel For Tength Trear
Form
18
19
xtO: biia_elementa[0;
20
xit] x[t-1] blann x(t]
Aros:
23
Ue Recurrent rix, upe ndarray
(,)
24
tiaaseiementa Addil bass elenents, operdaceay
[T-1,8]
initiali
Becurnes
27
X: State dynaniva, ap.niarsay
28
ne*
(lenent)
dhshage[3]
31
32
Initalizet lon
np, op.erpty(11 XI, atypennp, tloatt66)
g (Kt, rp-xeros, 1ike (biss elenerts(0))
15

---

Published as a conference paper at ICLR 2024
36
Use nunte prance for parallization for in mumaz.prange (T)
37
38
biaasy_oparator (q_4, (_4) X[4) 2_411]
40
socurn X[1:]
Listing I: Numba implementation of associative scan
Transition module We implement the irSLDS within the SSM package from Linderman et al. (2020). This can be accomplished by using the same SLDS parent class the (r)SLDS mod- els, and specifying new discrete state process. We implement this discrete state process as ssm transtanstttt ons Transi tion module.
clase Weatiecurreet (ransitnnooonooonoonniiiiioo (Craneitionn):
wat. probability dist to sample x.lt Exon. XLS tollows the heat equations wut Vin t/1 eatistying D.t Agaama D.xx
set init Iselt, K, D, M-C, gamm--1.0, kappa-0.4, scen=Falsel: super (lgertl..intt init 08, D, и)
Paraneters linking past
self. bets
12
selt.scan ccan
13
selt._.et__eet betail sclt._set__V FDM foz heat equation
15
kepp. kappa
16
17
Poranctors linking coscrvations distrizution
selt.Ve npr.randn
Inputs ensoding, per state
19
selt.Rs ngr.randn
Previous obeervations encoding, per etate
20
selt. apt.randn
21
_=tt__otat delta_t 1.0
23
ascune reqular enissions.
24
delta_x np. 148 colt.ganna delta deita_t (delsa_x 21/14 seim.
deternine the delta ouch that END is stable general
25
stability
selt.beta iscit. daita_t) (delts_x
27
28
daf _eat_]izo111:
29
solt.3 scit.betaenp. 11-2eseif.beta eap. diag tnp.ones (aelC.813 K13
(se15.K-1), +)
31
self. htt..bbtaappp..ee diae tap, ones (se1f. ke-1)
32
it selt.mod:
33
salf.Ur3,-2 self.beta
solf.U[-i,0] selt.beta
35
36
def get_ourficiontsstatisccc (solf, data, input, mask, tas, K_atata-Rene, SCAN-Falser:
37
39
Initiallzation
40
v2
ln Nonel Iail
1/self.X np.ones (120l2.K, self.8)!
43
44
45
Jynanics
SCAN:
47
a Use winaun to parallellze over Line the blases uncoding
48
kappa_u ap.einsun 1h 1k-i1k self.U, self axisel) biss_tlementa np. repeat Dapppa_DDpp.neee
ap. repeat fnp, ainsun(" tnyak-oekf, inpal, self.Va.ry no.neaaxia, ale self.s, axia=l) ap. acnsum (' dk-ekk, data, s61t.85.7) Ey op.navaxis, 13 poir.8, exis+1)
51
52
initial condition lon initial blas
biss_olenenta cp.comeatenate IIV novaxis, i, f2, bias_elaneate), axis=01
55
56
4 Uso the scan oparation to parallelize the linear dynamics over time F.U,,
zelucii
59
clae:
60
Use seguent fall dynanics conputation
ap.arange (T-
ut ap. (vt, self.n.r)
FD8
63
wt wt i sel f.l.fkaapepng pag.kea.n...eeeeeee. (pely
self twi owinformeeets inpais
64
wi ugedut Cinpul [c], aull.Va.1t op.dol (data[l), self.Bs. (7) ve.epgend op.erray sebucn
67
68
logut ranelt ion va Cuge wwwtfiiiettttttt.. isciiii (data, input, mask, Lag, vistalemmmsataee SCAN-swIC. normal ii sed log Ps Toggumesg (ws, exie=2, keepdins-Trues if swlf wahhoo0 reweighted. _og__e al-salt.a Log_F_ out reve gc.
date, ingut, teg, xustatteNone];
70
71
72
logsurexp inwweighheeeeeee ig lag Pe, axia=2, keeptinseTruel
16

---

Published as a conference paper at ICLR 2024
15
76
return nommalaand
ouE
77
Listing 2: Basic implementation of the transition matrix dynamics with heat-equation prior
C.3 TRANSITION MATRIX DYNAMICS
We explicit here the dynamics of the transition matrix Wn, E RR The equation for the sufficient statistic w, dynamies in 8 which rewrite here for reference,
Wx+1 (w,, Lny 2n = i} Wwn, + xll; + Ræn + T,
refers to row i of the transition matrix Wn+1, so that
[W+tkks = [wa+1| (æn: 2n = i}li
(12)
We note that all vectors are columns vectors.
To get the dynamics of the transition matrix itself, must expand out the computation for each row. Specifically, given W we obtain
(UIW.)..)
n + T
W W++1
+
+
(U[WaI,,)
k1 k
RRo + T
[Wil, U
+ kld + (Ræn + +r) @
[W.] UT
W. kld + (Ræ, + +r) 81
(13)
as the dynamics of the transition matrix. One can find these dynamics in our implementation of the lines 62-64. Unless otherwise noted, we use 0.4, and 1.0 for
transition module Listing 2 defining U.
17

---

Published as a conference paper at ICLR 2024
D RESULTS
D.1 ACCURACY METRICS
All models trained for 100 iterations, for every task. We only report the "mean" test log- likelihood, i.. the log of the mean of the test likelihoods.
Table 2: Performance metrics on NASCAR experiments--xpandee We report the marginal log- likelihood on test data, evaluated with Sequential Monte Carlo. Test data-sets consist of 200 excluded steps NASCAR, and excluded trajectory T = 1000 steps in extended NASCAR. We also report the dynamical meean-quaree emror (MSE) between the sample state-space flow and true flow field. We first reflect and rotate the sample trajectories to best align them with the flow field.
Model
K
parameters
Test log-likelihood (t)
Dynamical MSE (4)
(nats)
Mean
IQR
NASCAR
SLDS
92
1331.35
0.930
0.247, 0.629
8
176
1592.81
305.046
0.058, 0.759
rSLDS
100
1628.57
0.010
0.005,0.011.
192
1613.53
0.011
0.010, 0.013
irSLDS
88
1656.10
0.015
0.007, 0.015
136
1611.45
0.011
0.006, 0.014
Extended NASCAR
SLDS
6
130
7353.22
316.75
[0.29, 5.46
ILLS
6
142
8015.93
0.35
0.04, 0.48
iSLLD
112
8219.41
0.24
00.10, 0.27
Table 3: Performance metrics on IBL experiments: Marginal Log- Likelihood of held-out test data, evaluated with Sequential Monte Carlo. Higher is better.
Per trial
Per session
Test ligliikelhood (f)
ELBO (t) K
Model
K =
K
K=
rSLDS
759
-1437.519
-1399.030
-6.67 x 103
isSLDS
1446.692
-1461.777
-1392.508
-6.69 10
D.2 EMPIRICAL RUNNING TIMES
Table Empirical running times, in seconds per EM iteration, reported for the standard NASCAR task and trial-level IBL. The experiments were on an AMD EPYC 7H12 64-Core Processor, on 10 cores. No GPU was used.
(s/iter) Model
NASCAR K = 4
Per trial IBL K
K =
K = 8
SLDS
0.13 0.02
6.02 0.06
12.9 0.2
8.7 0.3
iSLDS
0.18 + 0.03
7.32
12.3 0.6
18.2 + 0.3
irSLDS
3.21 0.71
2.94
5..4
70.0 5.3
irSLDS (w/ scan)
0.54 + 0.08
11.73
14.0 0.8
18.4 + 1.0
D.3 HYPER-PARAMER INFLUENCE
Scaling the number of states K We include results in Fig. 5 on the impact of scaling K across models. We consider the standard NASCAR task, and evaluate model performance on = 200 held-out time steps, averaged random initialization seeds, and report mean and standard deviation.
18

---

Published as a conference paper at ICLR 2024
1000
Test log-likelhhood
-1000
-2000
model
-3000
SLD5
ISLDS
44000
HSLDS
Number of states K
Figure 5: Impact of scaling K on model performance. Test log-likelihood NASCAR task as function the hyperparameter K. See text for details.
35.0
15.0
33.5
39.0
30.2
2.5
3.0
5.0
2.5
2.5
0:0
150
305
0/
37
100
139
202
300
= 16
1600
1600
1540.3
1385.1
1553.9
1429.1
1562.8
1595.0
1599.1
1550
1550
1523.6
1379.7
kagpa
kapps
1599.7
1563.9
1441.0
1424.8
1508.3
1507.2 2
1503.2
1529.3
1440
1400
1.0
2.0
1.0
2.0
Figure 6: Impact of the hyperparameters {7: k} on the irSLDS model. We report test log- on NASCAR task function of the hyperparameters, and plot example discrete state trajectories. See text for details.
We trained and tested models using K {4, 12, 16) states, thus considering the over-parametrized regime over the true K = states for the NASCAR task. We observe that the irSLLDS maintains high performance across seeds, matched by and slightly over-performing the true rSLDD model class. The SLDS model class exhibits more variability. Note that in contrast to Table I train/test datasets for each seed.
we generated new
Heat equation parameters in the irSLDS We consider the impact performance of the self- reinforcement parameter k E IR and the parameter of decay ? > O in the heat equation prior of eq. equation in $3.3 For all experiments and results reported throughout the paper, we chose y 1.0 and k 0.4, we let At 1.0, and set Aj / 2v2 (for stability). We report in Fig. 6 test log-likelihood for the standard NASCAR task, on one random seed, over grid of hyper-parameter {y, x) combinations. In terms of performance, we observe that the specific choice of hyper parameters has minimal effect for low number of states K = 4. The impact becomes more evident when increased K 16, as a consequence of the more fine grained discrete state space support for the heat equation PDE. We plot prior sample trajectories, with a 0., for two exemplar combination of parameters to illustrate their influence.
19


Reviews:
[{"review_id": "ZWmLuo564Y", "reviewer": "Reviewer_ots3", "summary": "This work presents an addition to recurrent switching state-space models by using PDEs to develop an input-driven prior that induces state geometry over the HMM while still retaining its recurrent capability. The proposed infinite recurrent switching linear dynamical system (irSLDS) allows for more expressible and flexible state cardinality over a fixed number of discrete states while yielding fewer parameters compared to previous HMM models but greater performance on two time-varying datasets: NASCAR and the “Brainwide map” dataset from the International Brain Laboratory. This is achieved by differentiating the influence functions of each state using the heat equation as a spatial smoothness prior. The forward difference of the time derivative and central difference approximation to the second order spatial partial derivative are utilized to retrieve an infinite prior for the dynamics of the influence functions and sampled to compose the overall discrete state process.", "questions": "Smith et al. is cited in reference to the use of a scan operation to compute transition probabilities. However, S5 assumes a diagonal state matrix to efficiently compute the linear recurrence which takes a different form than the time-varying transition matrices in this work. How can this scan method scale efficiently?\nWhy does the heat equation prior directly enable discovery of greater fluctuations/switching in the data? Can you attribute the performance gains to this specific infinite prior?\nWhy not compare to other linear state space models and generalize outside of rSLDS and irSLDS that can allow for input-dependent transitions and recurrent relations without the explicit use of the influence functions as a geometric prior?", "limitations": "", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "The paper's originality derives from improvements it makes in the class of AR-HMM and SLDS models by removing exchangeability without removing input recurrence. This enables greater efficiency with expressibility. Overall, this work nicely integrates the heat-equation generalization and recurrent switching state space models. They present overall improvements on two different experiments with greater efficiency of parameters. Visuals are strong, particularly ones pertaining to the flow fields and the switching states of the rSLDS and irSLDS.", "weaknesses": "Most of the comparisons are done between rSLDS and irSLDS, yet the results are not convincing and substantially novel in the experiments conducted. If the authors are adamant about just comparing these two models, then it would be nice to see comparisons at greater scales instead of just testing 4-8 states for NASCAR and the IBL experiments. The results at just these scales do not appear novel.\nFurthermore, both experiments show that the irSLDS models uncovered switching at task-relevant states, yet it isn't clear why adding the heat equation prior itself is the best choice in discovering this switching. The authors discuss their reasons for using the heat equations but don't explain further into infinite priors that possibly could have achieved the same if not more.\nThe efficient of the parallel scan is not clear. Other works that utilize scans to efficiently compute the transition probabilities such as S4D (Gu et al., 2022), DSS (Gupta et al., 2022), S5 (Smith et al., 2023) are done on diagonalized transition matrices. The connection to the parallelization of the linear dynamics in Equation 8 is unclear since the transition matrix in this work is not diagonalized, so this work does not reap the same computational benefits that they claim, and their algorithm cannot scale effectively since it incurs cubic cost.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "InrpCh4oLC", "reviewer": "Reviewer_u3HM", "summary": "The authors proposed a novel infinite recurrent switching linear dynamical system (irSLDS) model, which combines partial differential equations (PDE) theory, semi-parametric formulation, and switching dynamical models. The authors validate and demonstrate the capabilities of their model on simulated data, and compare their model with existing popular methods. Finally, they use the model to explore mouse electrophysiological data during decision-making and uncover strong non-stationary processes underlying both within-trial and trial-averaged neural activity.", "questions": "See 'Weaknesses'.", "limitations": "", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The authors introduced an influence function of the states in switching linear dynamical systems, which controls the switches between discrete states by adding time-dependent parameters to constrain the duration of a state, and a space smoothness prior to the influence function.", "weaknesses": "How is the smoothness prior determined? In the paper, the authors state “Each time a discrete state j is chosen, “nearby” states also become more probable at the next time step”, it may be better to show the comparison with no smoothness or other smoothness formulations. \n\nThe prior based on dist-CRP is sensitive to time, what if there is a state with a long duration in the real dynamics? Or what if the dynamics are influenced by time a lot, such as in a Lorenz attractor?\n\nHow was the upper bound of K in irSLDS determined? Since the method is called infinite rSLDS, can the K be very large?\n\nClarity:\nIn 4.2 paragraph3, “we fit a spline to this curve a ...” seems to have a typo.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "X44yYSQwtb", "reviewer": "Reviewer_uvcW", "summary": "The authors propose a method for fitting dynamical models onto neural data that estimates the latent state of the system/organism.  A commonly used approach in the recent years is to fit a recurrent switching linear model, that approximates the latent dynamics with a finite set of low dimensional linear systems, where the system state is assumed to switch at discrete points between these discrete states. However this formalism requires stationary dynamics across all trials and does not impose any geometry on the discrete latent state space. Here the authors propose an extension to the recurrent switching linear systems, by considering that the evolution between the discreet states follow a Markov process itself (a distance-dependent Chinese restaurant process). To perform inference on the proposed model, they introduce a time-dependent sufficient statistic for the discrete state, the influence function, that is assumed to follow the heat (PDE) equation. The advantage of this statistic is that it is deterministically determined given the latent discrete state and the observations, and therefore does not arise in the posterior inference. Thus the authors can apply existing variational Laplace-EM proposed for these type of models like in Zoltowski et al. (2020). \n\nThey validate their approach on a synthetic dataset devised to imitate o stationary discrete latent state dynamics, and on an neural electrophysiology dataset from IBL with neural recordings from mice performing  a decision-making task.\n\nOverall the paper is very well written and is a sound extetion to the existing literature of switching linear dynamical systems.", "questions": "- In Figure 3 F lower,  irSLDS identifies the state 2, although it does predict (as I understand from the absence of orange dots) that the discrete state of the system was never at that state through the trials. Isn’t that strange? Shouldn’t the framework have estimated 7 states in total for this dataset?\n\n- In the NASCAR experiment, can you provide the same plot as in Figure 2F that shows the sequence of the identified states also for the rSLDS?\n\n- What are the limitations for considering Poisson instead of Gaussian observation process for the irSLDS framework?\n\n\nMinor:\n- Conclusion has various typos, i.e. first sentence of conclusion: “a”-> “an”, missing ’s’-es.\n- I understand that in the caption of Figure 1 you do not want to write out the full abbreviated terms, but you can at least provide a hyperlink to their mention in the main text for the reader to follow.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "- They endow the switching linear dynamical system framework with a latent discrete-state Markov structure that allows the continuous state of the system to guide the discrete state switches.\n- The proposed formalism induces a geometry in the discrete latent state between switches. \n- The paper is very well written providing a very well constructed introduction and build-up to their proposed extention.", "weaknesses": "- The resulting computations for this model are considerably more expensive compared to the previous approaches recurrent switching linear system framework, however this might be the tradeoff of dealing with non-stationarity.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "LFkHvyspxs", "reviewer": "Reviewer_Rx1k", "summary": "This paper builds upon the recurrent switching linear dynamical system (rSLDS) in the context of modeling latent dynamics of neural spike train data. The infinite recurrent SLDS (irSLDS) is proposed to overcome the fixed number of discrete states in the rSLDS, which can limit its capabilities on trial-varying data that may require a varying number of states. Prior methods have proposed using Dirichlet processes to infer the number of states from the data, often using a Chinese Restaurant Process (CRP). But this is limited in the rSLDS case, since the CRP has exchangeability assumptions which the state-dependence of the rSLDS breaks. The proposed method introduces a sufficient statistic to express the cluster allocations of a distance-dependent Chinese Restaurant Process (dist-CRP) which can be used to address the exchangeability issue and determine the discrete state switching probabilities. This sufficient statistic is implemented following a heat-equation combined with finite difference methods, resulting in linear recurrences. The proposed irSLDS is compared with SLDS and rSLDS on a synthetic NASCAR experiment and compared to the rSLDS on Neuropixel data recorded from mice.", "questions": "For major points, see weaknesses above. \n\nAdditional questions/comments:\n\n- In Table 1, can you make it more clear what the true number of states in the system are and what size $K$ is used by the models? I believe the number currently listed is the true number of states in the system, but it is unclear what $K$ was chosen for the models.\n\n- Unless I am missing something, I believe the last paragraph of Section 3.3 has an error regarding the computational complexity of the sequential vs parallel versions of the algorithm. It is stated that the sequential scan requires $\\mathcal{O}(K^3T)$ ops, but shouldn't this be  $\\mathcal{O}(K^2T)$ ops? This is because it requires $T$ sequential matrix-vector multiplications (compared to the cubic cost of matrix-matrix multiplications of the parallel version of the scan). Thus the choice of parallel vs sequential scan would be dependent on the state size and sequence length, since the cubic cost could grow quite large for larger state sizes.\n\n- This paper https://arxiv.org/abs/2111.01256 seems relevant to infinite switching states and could be cited as related work", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- For the most part, the paper is clear and presented well\n\n- The method is motivated well. In the scientific community, there is a strong desire for models that are both expressive and interpretable that can be used to gain insights into data, e.g. neural spike trains.  While rSLDS has been popular in this regard, addressing its weakness of having a fixed number of states is of interest to the community\n\n- The idea to characterize the distance-dependent CRP in the context of SLDS with a heat equation appears (to my knowledge) to be novel and is an interesting idea that seems to lead to a relatively straightforward implementation (the basic implementation of the code in the appendix is also appreciated).\n\n- The considered experiments appear to show that the irSLDS can have a performance/interpretability edge in practice and can potentially lead to insights that rSLDS cannot", "weaknesses": "- The paper would benefit from providing a bit more context/intuition behind the Polya-gamma approach in Section 3.1 and its weaknesses. Just reading this section currently does not make it as clear why the approach in section 3.2 is necessary. Section C.1 in the appendix is helpful. Perhaps a bit more details about the Gaussians with time-dependent sizes could be added to the main paper to make this point more clear?\n\n- The experimental results are the biggest weakness of the paper. It is not currently obvious that a practitioner should clearly reach for irSLDS as the preferred method.\n  -  Are there other datasets (synthetic or real) that could be added to make the benefits of irSLDS more clear? \n   - Can the extended NASCAR example be made more extreme (e.g. more states) to make the differences between irSLDS and rSDLS more apparent? It is not clear that the differences in test log-likelihood are that drastic (and there do not appear to be underlying states of rSDLS displayed to compare to)\n   - Why not also show rSLDS continuous states in Figure 2.E? I would expect these to not look good (unlike in Figure 2.D where rSLDS and irSLDS look similar) and think highlighting this would be helpful\n\n- It is claimed in Section 3.3 that irSLDS has a similar time complexity to rSLDS. It would be nice to see some kind of empirical runtime comparison to support this. How does this change when choosing different upper bounds of $K$?\n\n- Related to the 2 previous points above, it seems 2 potentially simple baselines would be an rSLDS with a large number of states or alternatively performing a grid search with rSLDS over different state sizes. On the other hand, given the flexibility of irSLDS, I would think it should be able to achieve better performance than these 2 baselines in one-shot. Instead, it seems the same amount of grid search was done for each method over the number of states. Alternatively, if either of these two baselines still takes less time/cost than running irSLDS (this is where the empirical runtime comparison comes in), then it seems perhaps rSLDS would be preferable for a practitioner in many cases? I think empirical evidence to resolve these doubts, both in terms of the empirical runtime under different scenarios and the time/cost per methods to achieve strong results, would strengthen the paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "dZfUrYWlNW", "reviewer": "Reviewer_k1CY", "summary": "This work extends the rSLDS. It embeds the latent modes (discrete states) in Euclidean space with PDE, and endows the former nominal states with geometry. The proposed method was validated on synthetic data and real neural recordings. The results show that the proposed method better captures the change of dynamics on nonstationary DS.", "questions": "In the extended NASCAR example\n- Why the trajectory is not smooth? The bias is too large?\n- Why not show MSE for extended NASCAR in Tab 1.\n\nIt would be more convincing to see the discrete state inference from continuous recording for real nonstationary system.", "limitations": "", "rating": 6, "confidence": 5, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "This work shows good performance for nonstationary systems.\nThe extension brings up new interpretations on the discrete states.", "weaknesses": "rSLDS has the intuition of partitioning the latent space into regions that are governed by different linear DS individually, and the resultant nonlinear DS globally. It's not very suitable to do comparison on a nonstationary system.\nThe goodness of fit numbers do not show irSLDS significantly outperform rSLDS. It seems better state inference does not help much with explaining the data.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)