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: The authors present an online method for estimating low-dimensional neural dynamics that scales well to large observation dimensions and high sampling rates. The first step of their method is an online version of SVD, designed so that it converges to a stable subspace with very little training data. The second step of their method is an online version of an HMM with a large number of states that tile the low-dimensional space provided by SVD. They demonstrate their method on both simulated and real data, demonstrating both its speed and accuracy.