Bubblewrap: Online tiling and real-time flow prediction on neural manifolds

1 Introduction . Systems neuroscience is in the midst of a data explosion . Advances in microscopy [ 1 , 2 ] and probe technology [ 3 , 4 , 5 ] have made it possible to record thousands of neurons simultaneously in behaving 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . animals . At the same time , growing interest in naturalistic behaviors has increased both the volume and complexity of jointly recorded behavioral data . On the neural side , this has resulted in a host of new modeling and analysis approaches that aim to match the complexity of these data , typically using artificial neural network models as proxies for biological neural computation [ 6 , 7 , 8 ] . At the same time , this increase in data volume has resulted in increasing emphasis on methods for dimensionality reduction [ 9 ] and a focus on neural populations in preference to the coding properties of individual neurons [ 10 ] . However , given the complexity of neural dynamics , it remains difficult to anticipate what experimental conditions will be needed to test population hypotheses in post hoc analyses , complicating experimental design and reducing power . Conversely , adaptive experiments , those in which the conditions tested change in response to incoming data , have been used in neuroscience to optimize stimuli for experimental testing [ 11 , 12 , 13 , 14 ] , in closed-loop designs [ 15 , 16 , 17 ] , and even to scale up holographic photostimulation for inferring functional connectivity in large circuits [ 18 ] . Yet , despite their promise , adaptive methods are rarely applied in practice for two reasons : First , although efficient online methods for dimensionality reduction exist [ 19 , 20 , 21 , 22 , 23 ] , these methods do not typically identify stable dimensions to allow low-dimensional representations of data to be compared across time points . That is , when the spectral properties of the data are changing in time , methods like incremental SVD may be projecting the data into an unstable basis , rendering these projections unsuitable as inputs to further modeling . Second , while many predictive models based on the dynamical systems approach exist [ 6 , 24 , 25 , 26 , 27 , 28 , 29 ] , including online approaches [ 30 , 31 , 16 , 32 ] , they typically assume a system with lawful dynamics perturbed by Gaussian noise . However , many neural systems of interest are noise-dominated , with multimodal transition kernels between states . In this work , we are specifically interested in closed loop experiments in which predictions of future neural state are needed in order to time and trigger interventions like optogenetic stimulation or a change in visual stimulus . Thus , our focus is on predictive accuracy , preferably far enough into the future to compensate for feedback latencies . To address these goals , we propose an alternative to the linear systems approach that combines a fast , stable , online dimensionality reduction with a semiparametric tiling of the low-dimensional neural manifold . This tiling introduces a discretization of the neural state space that allows dynamics to be modeled as a Hidden Markov Model defined by a sparse transition graph . The entire model , which we call “ Bubblewrap , ” can be learned online using a simple EM algorithm and handles tilings and graphs of up to thousands of nodes at kiloHertz data rates . Most importantly , this model outperforms methods based on dynamical systems in high-noise regimes when the dynamics are more diffusion-like . Training can be performed at a low , fixed latency ⇡10ms using a GPU , while a cached copy of the model in main memory is capable of predicting upcoming states at < 1ms latency . As a result , Bubblewrap offers a method performant and flexible enough to serve as a neural prediction engine for causal feedback experiments . 2 Stable subspaces from streaming SVD . As detailed above , one of the most pressing issues in online neural modeling is dealing with the increasingly large dimensionality of collected data — hundreds of channels per Neuropixels probe [ 4 , 5 ] , tens of thousands of pixels for calcium imaging . However , as theoretical work has shown [ 33 , 34 ] , true neural dynamics often lie on a low-dimensional manifold , so that population activity can be accurately captured by analyzing only a few variables . Here , we combine two approaches to data reduction : In the first stage , we use sparse random projections to reduce dimensionality from an initial d dimensions ( thousands ) to n ( a few hundred ) [ 35 , 36 ] . By simple scaling , for a fixed budget of N cells in our manifold tiling , we expect density ( and thus predictive accuracy ) to scale as N 1n in dimension n , and so we desire n to be as small as possible . However , by the Johnson-Lindenstrauss Lemma [ 37 , 36 ] , when reducing from d to n dimensions , the distance between vectors u⇤ and v⇤ in the reduced space is related to the distance between their original versions u and v by ( 1 `` ) ku vk2 ku⇤ v⇤k 2 ( 1 + `` ) ku vk2 ( 1 ) with probability 1 if n > O ( log ( 1/ ) / '' 2 ) . Unfortunately , even for `` ⇠ 0.1 ( 10 % relative error ) , the required n may be quite large , making this inappropriate for reducing to the very small numbers of effective dimensions characterizing neural datasets . Algorithm 1 Procrustean SVD ( proSVD ) 1 : Given : Initial data X0 , decay parameter ↵ 2 ( 0 , 1 ] 2 : Initialize : QR Factorization : X0 = Q0R0 3 : 4 : for t = 1 . . . do 5 : Fetch b new columns of data , X+ 6 : C Q > t 1X+ , X ? X+ Qt 1C , Q ? , R ? QR ( X ? ) . Gram-Schmidt 7 : Q̂ [ Qt 1 Q ? ] , R̂ Rt 1 C 0 R ? . QR of augmented data 8 : U , ⌃ , V SVD ( R̂ ) 9 : ⌃ ↵⌃ . Discount old data 10 : M Q > t 1Q̂U1 = [ k⇥k 0k⇥b ] U1 . U1 contains the first k columns of U 11 : Ũ , ⌃̃ , Ṽ SVD ( M ) , T Ũ Ṽ > . Orthogonal Procrustes : minT kQ̂U1T > Qt 1kF 12 : Qt Q̂U1T > . Update left subspace basis 13 : Qv , Rv QR ( V ) , Rt T⌃1Q > v . QR right subspace , update inner block 14 : end for Thus , in the second stage , we reduce from n ⇠ O ( 100 ) to k ⇠ O ( 10 ) dimensions using a streaming singular value decomposition . This method is based on the incremental block update method of [ 20 , 22 ] with an important difference : While the block update method aims to return the top-k SVD at every time point , the directions of the singular vectors may be quite variable during the course of an experiment ( Figure 1d–h ) , which implies an unstable representation of the neural manifold . However , as we show below , the top-k subspace spanned by these vectors stabilizes in seconds on typical neural datasets and remains so throughout the experiment . Therefore , by selecting a stable basis ( instead of the singular vector basis ) for the top-k subspace , we preserve the same information while ensuring a stable representation of the data for subsequent model fitting . More specifically , let xt 2 Rn be a vector of input data after random projections . In our streaming setup , these are processed b samples at a time , with b = 1 reasonable for slower methods like calcium imaging and b = 40 more appropriate for electrophysiological sampling rates of ⇠20kHz . Then , if the data matrix X has dimension n ⇥ T , adding columns over time , the incremental method of [ 20 , 22 ] produces at each time step a factorization X = QRW > , where the columns of the orthogonal matrices Q and W span the left and right top-k singular subspaces , respectively . If the matrix R were diagonal , this would be equivalent to the SVD . In the incremental algorithm , R is augmented at each timestep based on new data to form R̂ , which is block diagonalized via an orthogonal matrix and truncated to the top-k subspace , allowing for an exact reduced-rank SVD ( Appendix A ) . However , as reviewed in [ 20 , 22 ] , since there are multiple choices of basis Q for for the top-k singular subspace , there are likewise multiple choices of block diagonalization for R̂ . In [ 20 , 22 ] , the form of this operation is chosen for computational efficiency . But an equally valid option is to select the orthogonal matrix that minimizes the change in the singular subspace basis Q from one timestep to the next : minkQt Qt 1kF = min T kQ̂U1T > Qt 1kF , ( 2 ) where Q̂ is an augmented basis for the top- ( k + b ) singular subspace , U1 contains the first k left singular vectors of R̂ , and T is an orthogonal matrix ( Appendix A ) . This minimization is known as the Orthogonal Procrustes problem and has a well-known solution [ 38 ] : T = Ũ Ṽ > , where Ũ and Ṽ are the left and right singular vectors , respectively , of M ⌘ Q > t 1Q̂U1 . ( See [ 39 ] for a recent application of similar ideas in brain-computer interfaces ) . This Procrustean SVD ( proSVD ) procedure is summarized in Algorithm 1 . There , lines 1–8 follow [ 20 , 22 ] , while lines 10 and 11 perform the Orthogonal Procrustes procedure . Line 9 serves as a leak term that discounts past data as in [ 40 ] . Figures 1a-c illustrates the performance of the two-stage dimension reduction algorithm for a case of d = 104 randomly generated Gaussian data . While proSVD yields minimal distortion ( due to truncation of the spectrum to k = 6 ) , random projections require k ⇠ O ( 100 ) to achieve the same result ( Figure 1a ) . By contrast , random projections are much faster ( Figure 1b ) . Thus , we can trade off distortion against time by adjusting n , the number of intermediate dimensions . As Figure 1c shows , the optimal tradeoff occurs somewhere around n = 200 for this example . Figures 1d-i show results for neural data from recorded from monkey motor cortex [ 26 ] in a cued reach task . While projection of the data into the basis defined by streaming SVD remains unstable early in data collection ( top ) , the proSVD representation is nearly equivalent to the full offline result after only a few trials ( ⇡15s , middle ) . This is due to the fact that , in all data sets we examined , the top-k SVD subspace was identified extremely quickly ; proSVD simply ensures the choice of a stable basis for that subspace . 3 Bubblewrap : a soft manifold tiling for online modeling . As reviewed above , most neural population modeling approaches are based on the dynamical systems framework , assuming a lawful equation of motion corrupted by noise . However , for animals engaged in task-free natural behavior [ 41 , 42 , 43 ] , trajectories are likely to be sufficiently complex that simple dynamical models fail . For instance , dynamical systems models with Gaussian noise necessarily produce unimodal transition probabilities centered around the mean prediction , while neural trajectories may exhibit multimodal distributions beginning at the same system state . By contrast , we pursue an alternative method that trades some accuracy in estimating instantaneous system state for flexibility in modeling the manifold describing neural activity . Our approach is to produce a soft tiling of the neural manifold in the form of a Gaussian mixture model ( GMM ) , each component of which corresponds to a single tile . We then approximate the transitions between tiles via a Hidden Markov Model ( HMM ) , which allows us to capture multimodal probability flows . As the number of tiles increases , the model produces an increasingly finer-grained description of dynamics that assumes neither an underlying dynamical system nor a particular distribution of noise . More specifically , let xt be the low-dimensional system state and let zt 2 1 . . . N index the tile to which the system is assigned at time t. Then we have for the dynamics p ( zt = j|zt 1 = i ) = Aij p ( xt|zt ) = N ( µzt , ⌃zt ) p ( µj , ⌃j ) = NIW ( µ0j , j , j , ⌫j ) , ( 3 ) where we have assumed Normal-inverse-Wishart priors on the parameters of the Gaussians . Given its exponential family form and the conjugacy of the priors , online expectation maximization updates are available in closed form [ 44 , 45 , 46 ] for each new datum , though we opt , as in [ 45 ] for a gradient-based optimization of an estimate of the evidence lower bound L ( A , µ , ⌃ ) = X ij ( N̂ij ( T ) + ij 1 ) logAij + X j ( Ŝ1j ( T ) + jµ0j ) > ⌃ 1j µj ( 4 ) 1 2 X j tr ( ( j + Ŝ2j ( T ) + jµ0jµ > 0j + ( j + n̂j ( T ) ) µjµ > j ) ⌃ 1 j ) 1 2 X j ( ⌫j + n̂j ( T ) + d+ 2 ) log det⌃j with accumulating ( estimated ) sufficient statistics ↵j ( t ) = X i ↵i ( t 1 ) ij ( t ) N̂ij ( t ) = ( 1 `` t ) N̂ij ( t 1 ) + ↵i ( t 1 ) ij ( t ) ( 5 ) n̂j ( t ) = X i N̂ij ( t ) Ŝ1j ( t ) = ( 1 `` t ) Ŝ1j ( t 1 ) + ↵j ( t ) xt Ŝ2j ( t ) = ( 1 `` t ) Ŝ2j ( t 1 ) + ↵j ( t ) xtx > t where ↵j ( t ) = p ( zt = j|x1 : t ) is the filtered posterior , ij ( t ) is the update matrix from the forward algorithm [ 44 ] , and `` t is a forgetting term that discounts previous data . Note that even for `` = 0 , L is only an estimate of the evidence lower bound because the sufficient statistics are calculated using ↵ ( t ) and not the posterior over all observed data . In setting Normal-Inverse-Wishart priors over the Gaussian mixture components , we take an empirical Bayes approach by setting prior means µ0j to the current estimate of the data center of mass and prior covariance parameters j to N 2 k times the current estimate of the data covariance ( Appendix B ) . For initializing the model we use a small data buffer M ⇠ O ( 10 ) . We chose effective observation numbers ( , ⌫ ) = 10 3 and trained this model to maximize L ( A , µ , ⌃ ) using Adam [ 47 ] , enforcing parameter constraints by replacing them with unconstrained variables aij and lower triangular Lj with positive diagonal : Aij = exp ( aij ) / P j exp ( aij ) , ⌃ 1 j = LjL > j . Finally , in order to prevent the model from becoming stuck in local minima and to encourage more effective tilings , we implemented two additional heuristics as part of Bubblewrap : First , whenever a new observation was highly unlikely to be in any existing mixture component ( log p ( xt|zt ) < ✓n for all zt ) , we teleported a node at this data point by assigning ↵J ( t ) = 1 for an unused index J . For initial learning this results in a “ breadcrumbing ” approach where nodes are placed at the locations of each new observed datum . Second , when the number of active nodes was equal to our total node budget N , we chose to reclaim the node with the lowest value of n̂ ( t ) and zeroed out its existing sufficient statistics before teleporting it to a new location . In practice , these heuristics substantially improved performance , especially early in training ( Appendix D ) . The full algorithm is summarized in Algorithm 2 . Algorithm 2 Bubblewrap 1 : Given : Hyperparameters j , ⌫j , t , forgetting rate `` t , teleport threshold ✓ , step size , initial data buffer M 2 : Initialize with { x1 . . . xM } : µj µ̄ , ⌃j ⌃̄ , aij 1N , ↵j j 1 N . 3 : 4 : for t = 1 . . . do 5 : Observe new data point xt . 6 : if log p ( xt|zt ) < ✓ 8zt then . Teleport 7 : µJ = xt , ↵J ( t ) = 1 for J = argminj n̂j ( t ) 8 : end if 9 : Calculate ij ( t ) via forward filtering [ 44 ] . 10 : Update sufficient statistics via ( 5 ) . . E step 11 : µ̄ P j Ŝ1jP j n̂j , ⌃ P j Ŝ2jP j n̂j µ̄µ̄ T . Global mean and covariance update 12 : ✏j ⇠ N ( 0 , ⌘2 ) , µ0j aµ0j + ( 1 a ) µ̄+ ✏j , j ⌃ N 2 k . Update priors ( Appendix B ) 13 : Perform gradient-based update of L ( A , µ , ⌃ ) ( 4 ) . M step 14 : end for

Summary: This paper presents a method to efficiently characterize high-dimensional neural dynamics using a soft "tiling" of a lower-dimensionality manifold. Specifically, the method is a combination of two algorithms, (i) "Stable Streaming SVD", a dimensionality reduction method that aims to find the top-k-dimensional subspace instead of the top-k bases; (ii) "BubbleWrap", a Gaussian mixture model of the probability density with Hidden Markov Model transitions, that works on the low-dimensional manifold generated in (i). This is a model-free approach, in the sense that it is not based on a model of the dynamical system, and the authors claim that the flexibility of the method makes it outperform other (model-based) approaches in high-noise regimes. The authors also highlight that the proposed inference is fast, possibly allowing online applications.