SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Connections with Switching Linear Dynamical Systems With this foundation, it is natural to hypothesize that a switching low-rank tensor autoregression like SALT could approximate a switching LDS. There are two ways this intuition could fail: first, if the dynamics in a discrete state of an SLDS are unstable, then Proposition 1 would not hold; second, after a discrete state transition in an SLDS, it may take some time before the dynamics reach stationarity. We empirically test how well SALT approximates an SLDS in Section 5 and find that, across a variety of datasets, SALT obtains commensurate performance with considerably simpler inference and estimation algorithms.
Low-rank tensor decompositions of time-invariant autoregressive models Similar to this work, Wang et al. [2021] also modeled the transition matrices as a third-order tensor  $\mathcal{A} \in \mathbb{R}^{N \times N \times L}$  where the  $\mathbf{A}_{:l}$  is the  $l$ -th dynamics matrix. They then constrained the tensor to be low-rank via a Tucker decomposition, as defined in (1). However, unlike SALT, their model was time-invariant, did not have an ARHMM structure, or, make connections to the LDS and SLDS, as in Proposition 1.
Low-rank tensor decompositions of time-varying autoregressive models Low-rank tensor-based approaches have also been used to model time-varying AR processes [Harris et al., 2021, Zhang et al., 2021]. Harris et al. [2021] introduced TVART, which first splits the data into  $T$  contiguous fixed-length segments, each with its own AR-1 process. TVART can be thought of as defining a  $T\times N\times N$  ARHMM dynamics tensor and progressing through discrete states at fixed time points. This tensor is parameterized using the CP decomposition and optimized using an alternating least squares algorithm, with additional penalties such that the dynamics of adjacent windows are similar. By contrast, SALT automatically segments, rather than windows, the time-series into learned and re-usable discrete states.
Zhang et al. [2021] constructed a Bayesian model of higher-order AR matrices that can vary over time. First,  $H$  VAR dynamics tensors are specified, parameterized as third-order tensors with a rank-1 CP decomposition. The dynamics at a given time are then defined as a weighted sum of the tensors, where the weights have a prior density specified by an Ising model. Finally, inference over the weights is performed using MCMC. This method can be interpreted as a factorial ARHMM, hence offering substantial modeling flexibility, but sacrificing computational tractability when  $H$  is large.
Figure 2: SALT approximates LDS: Data simulated from an LDS for which  $n = 1$  and  $m = 3$  (see Proposition 1). (A-B): Average mean squared error of the autoregressive tensor corresponding to the LDS simulation and the log-likelihood of test data, as a function of SALT rank. According to Proposition 1, to model the LDS Tucker-SALT and CP-SALT require 7 and 10 ranks respectively (indicated by vertical dashed lines). Note the parameter error increases above the predicted threshold as a result of overfitting. (C-D): Mean squared error of the learned autoregressive tensor and log-likelihood of test data as a function of training data.
Low-rank tensor decompositions of neural networks Low-rank tensor decomposition methods have also been used to make neural networks more parameter efficient. Novikov et al. [2015] used the tensor-train decomposition [Oseledets, 2011] on the dense weight matrices of the fully-connected layers to reduce the number of parameters. Yu et al. [2017] and Qiu et al. [2021] applied the tensortrain decomposition to the weight tensors for polynomial interactions between the hidden states of recurrent neural networks (RNNs) to efficiently capture high-order temporal dependencies. Unlike switching models with linear dynamics, recurrent neural networks have dynamics that are hard to interpret, their state estimates are not probabilistic, and they do not provide experimentally useful data segmentations.
Linear dynamical systems and low-rank linear recurrent neural networks Valente et al. [2022] recently examined the relationship between LDSs and low-rank linear RNNs. They provide the conditions under which low-rank linear RNNs can exactly model the first-order autoregressive distributions of LDSs, and derive the transformation to convert between model classes under those conditions. This result has close parallels to Proposition 1. Under the conditions identified by Valente et al. [2022], the approximation in Proposition 1 becomes exact with just one lag term. However, when those conditions are not satisfied, we show that one still recovers an LDS approximation with a bounded error that decays exponentially in the number of lag terms.
We now empirically validate SALT by first validating the theoretical claims made in Section 3, and then apply SALT to two synthetic examples to compare SALT to existing methods. We conclude by using SALT to analyze real mouse behavioral recordings and  $C$  elegans neural recordings.
5.1 SALT Faithfully Approximates LDS
To test the theoretical result that SALT can closely approximate a linear dynamical system, we fit SALT models to data sampled from an LDS. The LDS has  $S = 7$  dimensional latent states with random rotational dynamics, where  $\Gamma$  has  $n = 1$  real eigenvalue and  $m = 3$  pairs of complex eigenvalues, and  $N = 20$  observations with a random emission matrix.
For Figure 2, we trained CP-SALT and Tucker-SALT with  $L = 50$  lags and varying ranks. We first analyzed how well SALT reconstructed the parameters of the autoregressive dynamics tensor. As predicted by Proposition 1, Figure 2A shows that the mean squared errors between the SALT tensor and the autoregressive tensor corresponding to the simulated LDS are the lowest when the ranks of CP-SALT and Tucker-SALT are  $n + 3m = 10$  and  $n + 2m = 7$  respectively. We then computed log-likelihoods on 5,000 timesteps of held-out test data (Figure 2B). Interestingly, the predictive performance of both CP-SALT and Tucker-SALT reach the likelihood of the ground truth LDS model with rank  $n + 2m = 7$ , suggesting that sometimes smaller tensors than suggested by Proposition 1 may still be able to provide good approximations to the data. We also show in Figures 2C and 2D that, as predicted, SALT models require much less data to fit than ARHMMs. We show extended empirical results and discussion on Proposition 1 in Appendix D.1.
Figure 3: SALT reconstructs simulated SLDS data and Lorenz attractor: (Top row) Observation generated from a low-dimensional trajectory. (A) shows ten observations generated from a recurrent "NASCAR" SLDS trajectory Linderman et al. [2017]. (B) 20-dimensional observations generated from a Lorenz attractor (5 observed dimensions are shown). (Middle and bottom rows): filtered observations and inferred low-dimensional trajectories from SLDS and SALT models. Colors indicate discrete state for ground truth (if available) and fitted models. SLDS and SALT find comparable filtered trajectories and observations. It is important to note that the latent spaces in both SLDS and SALT are only identifiable up to a linear transformation. We therefore align the latent trajectories for ease of comparison. This latent structure is reliably found by both SALT and SLDS.
5.2 Synthetic Switching LDS Examples
Proposition 1 quantifies the convergence properties of low-rank tensor regressions when approximating stable LDSs. Next we tested how well SALT can approximate the more expressive switching LDSs. We first applied SALT to data generated from a recurrent SLDS [Linderman et al., 2017], where the two-dimensional ground truth latent trajectory resembles a  $\mathrm{NASCAR}^{\text{®}}$  track (Figure 3A). SALT accurately reconstructed the ground truth filtered trajectories and discrete state segmentation, and yielded very similar results to an SLDS model. We also tested the ability of SALT to model nonlinear dynamics – specifically, a Lorenz attractor – which SLDSs are capable of modeling. Again, SALT accurately reconstructed ground truth latents and observations, and closely matched SLDS segmentations. These results suggest that SALT models provide a good alternative to SLDS models. Finally, in Appendix D.3, we used SLDS-generated data to compare SALT and TVART [Harris et al., 2021], another tensor-based method for modeling autoregressive processes, and find that SALT more accurately reconstructed autoregressive dynamics tensors than TVART.
5.3 Modeling Mouse Behavior
Next we considered a video segmentation problem commonly faced in the field of computational neuroethology [Datta et al., 2019]. Wiltschko et al. [2015] collected videos of mice freely behaving in a circular open field. They projected the video data onto the top 10 principal components (Figure 4A) and used an ARHMM to segment the PCA time series into distinct behavioral states. Here, we compared ARHMMs and CP-SALT with data from three mice. We used the first 35,949 timesteps of each recording, which were collected at  $30\mathrm{Hz}$  resolution. We used  $H = 50$  discrete states and fitted ARHMMs and CP-SALT models with varying lags and ranks.
The likelihood on a held-out validation set shows that the ARHMM overfitted quickly as the number of lags increased, while CP-SALT was more robust to overfitting (Figure 4B). We compared log-likelihoods of the best model (evaluated on the validation set) on a separate held-out test set and found that CP-SALT consistently outperformed ARHMM across mice (Figure 4C).
We also investigated the quality of SALT segmentations of the behavioral data (Appendix E.3). We found that the PCA trajectories upon transition into a discrete SALT state were highly stereotyped, suggesting that SALT segments the data into consistent behavioral states. Furthermore, CP-SALT used fewer discrete states than the ARHMM, suggesting that the ARHMM may have oversegmented and that CP-SALT offers a more parsimonious description of the data.
Figure 4: CP-SALT consistently outperforms ARHMM on mouse behavior videos and segments data into distinct behavioral syllables: (A) An example frame from the MoSeq dataset. The models were trained on the top 10 principal components of the video frames from three mice. (B) CP-SALT and ARHMM trained with different ranks and lags. Mean and standard deviation across five seeds evaluated on a validation set are shown. CP-SALT parameterization prevents overfitting for larger lags. (C) Test log-likelihood, averaged across 5 model fits, computed from the best ARHMM and CP-SALT hyperparameters in (B). CP-SALT outperforms ARHMM across all three mice.
5.4 Modeling C. elegans Neural Data
Finally, we analyzed neural recordings of an immobilized C. elegans worm from Kato et al. [2015]. SLDS have previously been used to capture the time-varying low-dimensional dynamics of the neural activity [Linderman et al., 2019, Glaser et al., 2020]. We compared SLDS, ARHMM, and CP-SALT with 18 minutes of neural traces (recorded at  $3\mathrm{Hz}$ ;  $\sim 3200$  timesteps) from one worm, in which 48 neurons were confidently identified. The dataset also contains 7 manually identified state labels based on the neural activity.
We used  $H = 7$  discrete states and fitted SLDSs, ARHMMs, and CP-SALT with varying lags and ranks (or continuous latent dimensions for SLDSs). Following Linderman et al. [2019], we searched for sets of hyperparameters that achieve  $\sim 90\%$  explained variance on a held-out test dataset (see Appendix F for more details). For ARHMMs and CP-SALT, we chose a larger lag ( $L = 9$ , equivalent to 3 seconds) to examine the long-timescale correlations among the neurons.
We find that SALT can perform as well as SLDSs and ARHMMs in terms of held-out explained variance ratio (a metric used by previous work [Linderman et al., 2019]). As expected, we find that CP-SALT can achieve these results with far fewer parameters than ARHMMs, and with a parameter count closer to SLDS than ARHMM (as more continuous latent states were required in an SLDS to achieve  $\sim 90\%$  explained variance; see Appendix F). Figure 5A shows that SALT, SLDS and ARHMM produce similar segmentations to the given labels, as evidenced by the confusion matrix having high entries on the leading diagonal (Figure 5B and Appendix F).
Figure 5C shows the one-dimensional autoregressive filters learned by CP-SALT, defined as  $\sum_{i=1}^{D_1}\sum_{j=1}^{D_2}\sum_{k=1}^{D_3}g_{ijk}^{(h)}u_{pi}^{(h)}v_{qj}^{(h)}\mathbf{w}_{:k}^{(h)}$  for neurons  $p$  and  $q$ . We see that neurons believed to be involved in particular behavioral states have high weights in the filter (e.g., SMDV during the "Ventral Turn" state and SMDD during the "Dorsal Turn" state [Linderman et al., 2019, Kato et al., 2015, Gray et al., 2005, Kaplan et al., 2020, Yeon et al., 2018]). This highlights how switching autoregressive models can reveal state-dependent functional interactions between neurons (or observed states more generally). In Appendix F, we show the autoregressive filters learned by an ARHMM, an SLDS, and a generalized linear model (GLM), a method commonly used to model inter-neuronal interactions [Pillow et al., 2008]. Interestingly, the GLM does not find many strong functional interactions between neurons, likely because it is averaging over many unique discrete states. In addition to its advantages in parameter efficiency and estimation, SALT thus provides a novel method for finding changing functional interactions across neurons at multiple timescales.
We introduce switching autoregressive low-rank tensor (SALT) models: a novel model class that parameterizes the autoregressive tensors of an ARHMM with a low-rank factorization. This constraint
C
Figure 5: CP-SALT provides good segmentations of C. elegans neural data, and inferred low-rank tensors give insights into temporal dependencies among neurons in each discrete state: (A) Example data with manually generated labels (Given), as well as segmentations generated by SALT, SLDS and ARHMM models. Learned states are colored based on the permutation of states that best matches given labels. All methods produce comparable segmentations, with high agreement with the given labels. (B) Confusion matrix of SALT-generated labels. (C) One-dimensional autoregressive filters learned in two states by SALT (identified as ventral and dorsal turns). Colors indicate the area under curve (red is positive; blue is negative). The first four rows are neurons known to mediate ventral turns, while the last two rows mediate dorsal turns [Kato et al., 2015, Gray et al., 2005, Yeon et al., 2018]. These known behavior-tuned neurons generally have larger magnitude autoregressive filters. Interestingly, AVFL and AVFR also have large filters for dorsal turns. These neurons do not have a well-known function. However, they are associated with motor neurons, and so may simultaneously activate due to factors that co-occur with turning. This highlights how SALT may be used for proposing novel relationships in systems.
allows SALT to model time-series data with fewer parameters than ARHMMs and with simpler estimation procedures than SLDSs. We also make theoretical connections between low-rank tensor regressions and LDSs. We then demonstrate, with both synthetic and real datasets, that SALT offers both efficiency and interpretability, striking an advantageous balance between the ARHMM and SLDS. Moreover, SALT offers an enhanced ability to investigate the interactions across observations, such as neurons, across different timescales in a data-efficient manner.
However, SALT is not without limitations. Foremost, SALT cannot readily handle missing observations, or share information between multiple time series with variable observation dimensions. "Hierarchical SALT" is an interesting extension, where information is shared across time series, but the factors of individual time series are allowed to vary. Furthermore, SALT could be extended to handle non-Gaussian data. For example, neural spike trains are often modeled with Poisson likelihoods instead of SALT's Gaussian noise model. In this case, the E-step would still be exact, but the M-step would no longer have closed-form coordinate updates. Despite these limitations, SALT offers simple, effective and complementary means of modeling and inference methodology for complex, time-varying dynamical systems.
Ethical Concerns We note that there are no new ethical concerns as a result of SALT.
This work was supported by grants from the Simons Collaboration on the Global Brain (SCGB 697092), the NIH (U19NS113201, R01NS113119, R01NS130789, and K99NS119787), the Sloan Foundation, and the Stanford Center for Human and Artificial Intelligence. We thank Liam Paninski and the anonymous reviewers for their constructive feedback on the paper. We also thank the members of the Linderman Lab for their support and feedback throughout the project.
Sandeep Robert Datta, David J Anderson, Kristin Branson, Pietro Perona, and Andrew Leifer. Computational neuroethology: A call to action. Neuron, 104(1):11-24, 2019.
Alexander B Wiltschko, Matthew J Johnson, Giuliano Iurilli, Ralph E Peterson, Jesse M Katon, Stan L Pashkovski, Victoria E Abraira, Ryan P Adams, and Sandeep Robert Datta. Mapping sub-second structure in mouse behavior. Neuron, 88(6):1121-1135, 2015.
Julia Costacurta, Lea Duncker, Blue Sheffer, Winthrop Gillis, Caleb Weinreb, Jeffrey Markowitz, Sandeep R Datta, Alex Williams, and Scott Linderman. Distinguishing discrete and continuous behavioral variability using warped autoregressive HMMs. Advances in Neural Information Processing Systems, 35:23838-23850, 2022.
Aram Giahi Saravani, Kiefer J Forseth, Nitin Tandon, and Xaq Pitkow. Dynamic brain interactions during picture naming. eNeuro, 6(4), 2019.
Stefano Recanatesi, Ulises Pereira-Obilinovic, Masayoshi Murakami, Zachary Mainen, and Luca Mazzucato. Metastable attractors explain the variable timing of stable behavioral action sequences. Neuron, 110(1):139–153, 2022.
Yariv Ephraim, David Malah, and B-H Juang. On the application of hidden Markov models for enhancing noisy speech. IEEE Transactions on Acoustics, Speech, and Signal Processing, 37(12): 1846-1856, 1989.
Zoubin Ghahramani and Geoffrey E Hinton. Variational learning for switching state-space models. Neural Computation, 12(4):831-864, 2000.
Biljana Petreska, Byron M Yu, John P Cunningham, Gopal Santhanam, Stephen Ryu, Krishna V Shenoy, and Maneesh Sahani. Dynamical segmentation of single trials from population neural data. Advances in Neural Information Processing Systems, 24, 2011.
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. bioRxiv, page 621540, 2019.
Joshua Glaser, Matthew Whiteway, John P Cunningham, Liam Paninski, and Scott Linderman. Recurrent switching dynamical systems models for multiple interacting neural populations. Advances in Neural Information Processing Systems, 33:14867-14878, 2020.
Aditya Nair, Tomomi Karigo, Bin Yang, Surya Ganguli, Mark J Schnitzer, Scott W Linderman, David J Anderson, and Ann Kennedy. An approximate line attractor in the hypothalamus encodes an aggressive state. Cell, 186(1):178-193, 2023.
David Barber. Expectation correction for smoothed inference in switching linear dynamical systems. Journal of Machine Learning Research, 7(11), 2006.
Emily Beth Fox. Bayesian nonparametric learning of complex dynamical phenomena. PhD thesis, Massachusetts Institute of Technology, 2009.
Kevin Murphy and Stuart Russell. Rao-Blackwellised particle filtering for dynamic Bayesian networks. In *Sequential Monte Carlo methods in practice*, pages 499–515. Springer, 2001.
Scott Linderman, Matthew Johnson, Andrew Miller, Ryan Adams, David Blei, and Liam Paninski. Bayesian learning and inference in recurrent switching linear dynamical systems. In Artificial Intelligence and Statistics, pages 914-922. PMLR, 2017.
Jihye Yeon, Jinmahn Kim, Do-Young Kim, Hyunmin Kim, Jungha Kim, Eun Jo Du, KyeongJin Kang, Hyun-Ho Lim, Daewon Moon, and Kyuhyung Kim. A sensory-motor neuron type mediates proprioceptive coordination of steering in C. elegans via two TRPC channels. PLoS biology, 16(6): e2004929, 2018.
Jonathan W Pillow, Jonathon Shlens, Liam Paninski, Alexander Sher, Alan M Litke, EJ Chichilnisky, and Eero P Simoncelli. Spatio-temporal correlations and visual signalling in a complete neuronal population. Nature, 454(7207):995-999, 2008.
Tohru Katayama. Subspace Methods for System Identification, volume 1. Springer, 2005.
Martin Chalfie, John E Sulston, John G White, Eileen Southgate, J Nicol Thomson, and Sydney Brenner. The neural circuit for touch sensitivity in Caenorhabditis elegans. Journal of Neuroscience, 5 (4):956-964, 1985.
Beverly J Piggott, Jie Liu, Zhaoyang Feng, Seth A Wescott, and XZ Shawn Xu. The neural circuits and synaptic mechanisms underlying motor initiation in C. elegans. Cell, 147(4):922-933, 2011.
Amirreza Farnoosh, Bahar Azari, and Sarah Ostadabbas. Deep switching auto-regressive factorization: Application to time series forecasting. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pages 7394-7403, 2021.
Andreas S Weigend and Neil A Gershenfeld. Time series prediction: Forecasting the future and understanding the past. Santa Fe Institute Studies in the Sciences of Complexity, 1994.
Supplementary Materials for: Switching Autoregressive Low-rank Tensor Models
- Appendix A: SALT Optimization via Tensor Regression.
- Appendix B: SALT Approximates a (Switching) Linear Dynamical System.
- Appendix C: Single-subspace SALT.
- Appendix D: Synthetic Data Experiments.
- Appendix E: Modeling Mouse Behavior.
- Appendix F: Modeling C. elegans Neural Data.
- Appendix G: Deep Switching Autoregressive Model.
A SALT Optimization via Tensor Regression
Let  $\mathbf{y}_t \in \mathbb{R}^{N_1}$  be the  $t$ -th outputs and  $\mathbf{X}_t \in \mathbb{R}^{N_2 \times N_3}$  be the  $t$ -th inputs. The regression weights are a tensor  $\mathcal{A} \in \mathbb{R}^{N_1 \times N_2 \times N_3}$ , which we model via a Tucker decomposition,
$$
\mathcal {A} = \sum_ {i = 1} ^ {D _ {1}} \sum_ {j = 1} ^ {D _ {2}} \sum_ {k = 1} ^ {D _ {3}} g _ {i j k} \mathbf {u}: i \circ \mathbf {v}: j \circ \mathbf {w}: k, \tag {11}
$$
where  $\mathbf{u}_i, \mathbf{v}_j$ , and  $\mathbf{w}_k$  are columns of the factor matrices  $\mathbf{U} \in \mathbb{R}^{N_1 \times D_1}$ ,  $\mathbf{V} \in \mathbb{R}^{N_2 \times D_2}$ , and  $\mathbf{W} \in \mathbb{R}^{N_3 \times D_3}$ , respectively, and  $g_{ijk}$  are entries in the core tensor  $\mathbf{G} \in \mathbb{R}^{D_1 \times D_2 \times D_3}$ .
We define  $\times_{j,k}$  to be a tensor-matrix product over the  $j^{th}$  and  $k^{th}$  slices of the tensor. For example, given a three-way tensor  $\mathcal{A} \in \mathbb{R}^{D_1 \times D_2 \times D_3}$  and a matrix  $\mathbf{X} \in \mathbb{R}^{D_2 \times D_3}$ ,  $\mathcal{A} \times_{2,3} \mathbf{X} = \sum_{j=1}^{D_2} \sum_{k=1}^{D_3} \mathbf{a}_{jk} x_{jk}$ . This operation is depicted in Figure 6.
Consider the linear model,  $\mathbf{y}_t\sim \mathcal{N}(\mathcal{A}\times_{2,3}\mathbf{X}_t,\mathbf{Q})$  where  $\mathcal{A}\times_{2,3}\mathbf{X}_t$  is defined using the Tucker decomposition of  $\mathcal{A}$  as,
$$
\begin{array}{l} \mathcal {A} \times_ {2, 3} \mathbf {X} _ {t} = \mathcal {A} _ {(1)} \operatorname {v e c} \left(\mathbf {X} _ {t}\right) (12) \\ = \mathbf {U} \mathcal {G} _ {(1)} \left(\mathbf {V} ^ {\top} \otimes \mathbf {W} ^ {\top}\right) \operatorname {v e c} \left(\mathbf {X} _ {t}\right) (13) \\ = \mathbf {U} \mathcal {G} _ {(1)} \operatorname {v e c} \left(\mathbf {V} ^ {\top} \mathbf {X} _ {t} \mathbf {W}\right) (14) \\ \end{array}
$$
where  $\mathcal{A}_{(1)}\in \mathbb{R}^{N_1\times N_2N_3}$  and  $\mathfrak{G}_{(1)}\in \mathbb{R}^{D_1\times D_2D_3}$  are mode-1 matricizations of the corresponding tensors. Note that these equations assume that matricization and vectorization are performed in row-major order, as in Python but opposite to what is typically used in Wikipedia articles.
Equation (14) can be written in multiple ways, and these equivalent forms will be useful for deriving the updates below. We have,