Across-animal odor decoding by probabilistic manifold alignment

1 Introduction . Collective network dynamics are the foundation of neural computation , from early sensory encoding [ 1 ] , to working memory [ 2 ] , decision making [ 3 ] , or motor control [ 4 ] . Neural population activity often has low-dimensional structure that is qualitatively preserved across sessions [ 3 ] , behavioral states [ 5 ] , and even animals [ 6–8 ] . Aligning neural datasets from multiple recordings and animal subjects into a common latent space could provide a powerful tool for extracting unifying principles of brain computations . Nonetheless , progress is hampered by the lack of robust statistical tools for extracting shared neural population dynamics across datasets . The alignment of low-dimensional neural manifolds is equally important in brain-computer interface ( BCI ) applications , where the assumption of a common low dimensional structure in neural responses is used to compensate for instability in neural recordings and ensure robust performance over time [ 9 , 10 ] . This is not only restricted to motor control ; in the sensory domain , a new generation of chemical detectors exploit the unsurpassed sensitivity of rodent olfactory receptors , but bypass the limitations of training animals to report odor identity by directly decoding it from neural responses [ 11–13 ] . Unfortunately , the applicability of this idea is limited by the need to learn the mapping between neural activity and chemical identity on an animal-by-animal basis , which involves costly data collection [ 14 ] . Statistical tools that make use of data from previous animals to quickly calibrate the decoder in a new animal could dramatically increase the practical use of such technologies . The basic structure of the rodent olfactory bulb ( OB ) is well-understood [ 15 ] : a vast number of different olfactory receptors project to distinct OB subpopulations ( glomeruli ) whose spatial organization is largely preserved across animals [ 16 , 17 ] . Odor identity is encoded in low-dimensional transients of their population dynamics [ 1 , 18 , 13 ] , though how the glomerular topography translates into stimulus-dependent network dynamics is less clear . Nonetheless , given the anatomical and functional similarity of the rodent olfactory bulb , it is likely that the OB response to individual odors is itself 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . somewhat stereotyped across animals . With appropriate alignment , one could construct ‘ universal ’ odor decoders , applicable across animals . Here we propose a novel probabilistic approach to aligning observations across neural recordings taken from different mice so as to extract commonalities in their circuit-level representations of different odors . Our estimation procedure , which uses our ‘ aligned mixture of latent dynamical systems ’ ( amLDS ) graphical model , learns independent mappings of different recordings into a common latent manifold.1 In this space , neural trajectories are similar across animals but distinct across odors which allows for robust odor identification . Moreover , our probabilistic decoder provides uncertainty estimates which can be used for optimal decision making . We used extended numerical simulations to evaluate the properties of the decoder in response to a range of biologically relevant manipulations , such as variations in measurement noise , dimensionality and across-trial variability in the latent dynamics , ( mis ) alignment of electrodes across animals , etc . We then tested our procedure on recordings from the OB of mice in response to several monomolecular odors . When fit to this data , our model confirmed that encoding of odors relies on transient dynamics in a shared low-dimensional manifold and that these dynamics are largely preserved across animals . Taking advantage of this reproducible structure , we were able to obtain remarkable improvements in odor identification performance relative to alternative approaches . 2 Shared odor encoding dynamics across animals . We start by mathematically formalizing the idea of shared stimulus-specific neural dynamics with animal-specific measurements ( Fig . 1A ) . The amLDS graphical model assumes d-dimensional latent variables zt , which evolve according to stochastic linear dynamics with odor-specific parameters ( Fig . 1B ) : zt = Akzt−1 + bk , t + wt , ( 1 ) where t = { 2 . . . T } indexes time within trial ; Ak and bk denote parameters specifying the dynamics in response to odor k , with independent noise , wt ∼ N ( 0 , Qk ) . The prior for the initial conditions z1 is normal , with zero mean and covariance Q0 . The shared odor-specific dynamics are mapped into measured neural responses via an animal-specific linear Gaussian observation model , shared across all odor conditions x ( m ) t = Cmzt + vt , ( 2 ) with i.i.d . noise vt ∼ N ( 0 , Rm ) ; Cm and Rm are observation model parameters specific to animal m. We will use Nm to denote the size of observations for animal m. Note that the animal-specific observation model means that it is not strictly required for the neural measurements to be tightly matched . For instance , different animals may have different numbers of measurement channels ( e.g . due to some electrodes becoming damaged ) . Moreover , the model accounts for variability in the SNR of different electrodes ( across channels and devices ) . Assuming that the measurement process 1Code available : github.com/pedroherrerovidal/amLDS . is consistent in terms of the properties of the electrodes and their alignment on the olfactory bulb surface , one can additionally define a prior that encourages Cs to be similar across animals . To keep the procedure general , we only consider the unconstrained version here . Parameter learning The dataset D combines measurements from M animals , jointly covering all K stimulus conditions . For each animal-stimulus pair , ( m , k ) , we have Im , k trials.2 The goal of the parameter estimation procedure is to derive maximum likelihood estimates of all model parameters , jointly denoted by θ = { A1 : K , b1 : K , Q0 : K , C1 : M , R1 : M } . We use expectation maximization ( EM ) to optimize the parameters , by introducing the variational distribution q ( z ) and using it to lower-bound the marginal likelihood [ 19 ] :3 log P ( D|θ ) = ∑ i , m , k ∫ z log P ( x|θ ) q ( z ) dz = ∑ i , m , k ∫ z q ( z ) ( log P ( z , x|θ ) − log P ( z|x , θ ) ) dz = ∑ i , m , k ∫ z q ( z ) log P ( x , z|θ ) q ( z ) dz− ∑ i , m , k ∫ z q ( z ) log P ( z|x , θ ) q ( z ) dz = L ( q , θ ) + ∑ i , m , k KL ( q ( z ) ||P ( z|x , θ ) ) . Due to the Markovian structure of the dynamics , the joint probability of the latent and observations has a relatively simple joint Gaussian structure . The E-step optimizes L w.r.t . q , by setting q ( z ) = P ( z|x , θ ) ; this is computed by traditional Kalman smoothing using the appropriate trial-specific parameters ( Suppl . S1 ) . The M-step optimizes L w.r.t . to θ . The cost function decomposes into the sum of terms that only depend on the subject-specific observation parameters ( Cm , Rm ) and the condition-specific latent dynamic parameters ( Ak , Qk , Q0 , bk , t ) , respectively , and can be optimized separately . Taking the derivative of the loss and setting it to zero results in parameter updates : bnewk,1 = 1 Ik ∑ i E [ z1 ] ( 3 ) Qnew0 = 1 Ik ∑ i ( E [ z1z1 > ] − E [ z1 ] E [ z1 ] > ) , ( 4 ) where the indicator i covers all trials of the same stimulus type , across all animals Ik = ∑ m Im , k ; expectations are taken under posterior for trial i ( explicit indexing omitted for brevity ) . The same set of trials are used for updating the latent dynamics , and the input drive : Anewk = ( Ik∑ i=1 T∑ t=2 ( E [ ztzt−1 > ] − bk , tE [ zt−1 > ] ) ) ( Ik∑ k=1 T∑ t=2 E [ zt−1zt−1 > ] ) −1 ( 5 ) bnewk , t = 1 Ik Ik∑ i=1 ( E [ zt ] −Anewk E [ zt−1 ] ) . ( 6 ) Finally , the updated noise covariance is computed as Qnewk = 1 Ik ( T − 1 ) Ik∑ i=1 T∑ t=2 ( E [ ztzt > ] −Anewk E [ zt−1zt > ] − E [ ztzt−1 > ] Anewk > + Anewk E [ zt−1zt−1 > ] Anewk > − E [ zt ] bnewk , t > − bnewk , t E [ zt ] > + Anewk E [ zt−1 ] bnewk , t > +bnewk , t E [ zt−1 ] > Anewk > + bnewk , t bnewk , t > ) . ( 7 ) 2Note that it is not necessary that all stimuli are presented to all animals , nor that the number of stimulus presentations is identical across animals and conditions . 3To keep notation simple , we drop the explicit indexing by time , trial , stimulus type , and animal of x and z . The updates for the observation model are the standard Kalman ones , with the distinction that the index i iterates across all trials and stimulus conditions for a particular animal ( total number Im = ∑ k Im , k ) : Cnewm = ( Im∑ i=1 T∑ t=1 xtE [ zt ] ) ( Im∑ i=1 T∑ t=1 E [ ztzt > ] ) −1 ( 8 ) Rnewm = 1 ImT ∑ i , t ( xtxt > −Cnewm E [ zt ] xt > − xtE [ zt ] > Cnewm > + Cnewm E [ ztzt > ] Cnewm > ) ( 9 ) All the required posterior moments are obtained during the E step via Kalman smoothing as E [ zt ] = µt|T , E [ ztzt > ] = Σt|T − µt|Tµt|T > ( see Suppl . S1 for further details ) . As an important practical side note , the quality of the estimated parameter depends critically on a good initialization of θ . We achieve this by fitting a simpler factor analysis model ( FA ) , which ignores all temporal dependencies . This directly yields initial conditions for Cm and Rm . The inputs bk , t are initialized to the average of the corresponding inferred latent variable zt , using all trials from stimulus condition k , while parameters Ak and Q0 : K are initialized heuristically by linear regression using the FA-inferred posterior mean estimates of the latents . We used Bayesian model comparison to determine the dimensionality of the latent space from data ( evaluated on a separate validation set ) . As a more intuitive measure of goodness-of-fit , we also estimated the reconstruction error when predicting unobserved measurement dimensions ( also known as leave-neuron-out error ) , commonly used for assessing the quality of models of neural population dynamics [ 20 , 21 ] . More specifically , we use trials from the validation dataset , infer the posterior over the latent trajectory given data measurements in all electrodes except the j-th by Kalman smoothing P ( z1 : T |x¬j , θ ) , which is used to predict the values of the response of the jth neuron , x̂jt = Cjmµ ( ¬j ) t|T , where Cjm is the jth row of Cm . We report reconstruction error as the squared distance between the reconstructed and measured activity , 〈‖xjt − x̂ j t‖22〉j , i , where ‖ · ‖2 is the L2 norm and 〈·〉j , i is the mean over trials and choices for j . Odor decoding as hierarchical inference At test time , one needs to infer the identity of the stimulus given the observations x1 : T in animal M , and the estimated parameters θ , obtained by Bayes rule , P ( k|x1 : T , θ ) ∝ P ( x1 : T |k , θ ) P ( k ) . This requires estimating the marginal likelihood for each possible stimulus condition , which can be computed iteratively during the Kalman filtering procedure . More specifically , using the chain rule and then the Markov structure of the stimulus-specific LDS , we have P ( x1 : T |k , θ ) = T∏ t=1 P ( xt|x1 : t−1 , k , θ ) = T∏ t=1 ∫ P ( xt , zt|x1 : t−1 , k , θ ) dzt . Taking the logarithm , rearranging the terms and taking advantage of the fact that the posterior marginals are normal with parameters obtained by Kalman filtering ( see Suppl . Info . S1 ) , P ( zt|x1 : t−1 , k , θ ) = N ( µt|t−1 , Σt|t−1 ) , yields : log P ( x1 : T |k , θ ) = T∑ t=1 log ∫ P ( xt|zt , x1 : t−1 , θk ) P ( zt|x1 : t−1 , k , θ ) dzt = T∑ t=1 logN ( xt ; Cmµt|t−1 , CmΣt|t−1Cm > + Rm ) Lastly , here we assume a uniform prior over possible stimuli , to match the statistics of our OB dataset . Nonetheless , for real world applications the prior should reflect natural stimulus statistics .

Summary: This submission is a very well done application of the probabilistic inference to decoding of odour identity from neural recordings in mice. Interestingly, the latent code is assumed to have temporal dynamics shared across animals. The main benefits (high accuracy, data savings,  robustness to incomplete data) are shown in simulation and neural recordings. The method outperforms the existing (non-probabilistic) neural alignment procedures.