Drop, Swap, and Generate: A Self-Supervised Approach for Generating Neural Activity

1 Introduction . In the brain , the coordinated actions of groups of neurons are responsible for encoding sensory inputs and movements , as well as all processing and manipulation in between ( 1 ; 2 ; 3 ; 4 ; 5 ) . Understanding what different populations of neurons are doing and how they work together to encode their inputs is a primary goal of neuroscience ( 6 ) . When successful , representations learned from populations of neurons can provide insights into how neural circuits work to encode their inputs and produce behavior , and allow for robust and stable decoding of these correlates . Over the last decade , a number of unsupervised learning approaches have been introduced to build representations of neural population activity without knowledge of specific labels or downstream decoding objectives ( 7 ; 8 ; 9 ; 10 ; 11 ; 12 ; 13 ) . Such methods have provided exciting new insights into the stability of neural responses ( 14 ) , individual differences ( 10 ) , and remapping of neural responses through learning ( 15 ) . However , without labels or additional inputs to guide the network , learning representations that allow different sources of variability to be distinguished is still a major challenge ( 16 ; 17 ) . ∗Contact : rliu361 @ gatech.edu , evadyer @ gatech.edu . Project page : https : //nerdslab.github.io/SwapVAE/ . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Here , we develop a novel unsupervised approach for disentangling neural activity called Swap-VAE . Our approach is loosely inspired by methods used in computer vision that aim to decompose images into their content and style ( 18 ; 19 ; 20 ; 21 ; 22 ) : the representation of the content should give us the abstract “ gist ” of the image ( what it is ) , and the style components are needed to create a realistic image ( or , equivalently , they capture the variation in images with the same content ) . To map this idea onto the decomposition of brain states , we consider the execution of movements and their representation within the the brain ( Figure 1 , Right ) . The content in this case may be knowing where to go ( target location ) and the style would be the exact execution of the movement ( the movement dynamic ) . We ask whether the neural representation of movement can be disentangled in a similar manner . To identify the content within our neural recording , we use a self-supervised approach : we apply a variety of transformations to the input data ( observed firing rates ) which we hypothesize to be contentpreserving , and train the representation to be invariant to these manipulations . These transformed ( or augmented ) views are created by dropping out neurons and jittering samples in time , which intuitively should lead the network to a representation that maintains both temporal consistency and invariance to the specific neurons used to represent the neural state . In addition to this instance-specific alignment loss , we also encourage the network to reconstruct the original inputs using a regularized variational autoencoder ( beta-VAE ) ( 23 ; 24 ) that has access to both the content variables and another set of variables in the model that encodes the style . We show that through combining our proposed selfsupervised alignment loss with a generative model , we can learn representations that disentangle latent factors underlying neural population activity . We apply our method to synthetic data and publicly available non-human primate ( NHP ) reaching datasets from two different individuals ( 25 ) . To quantify how effectively our method disentangles these datasets , we propose several general-purpose measures of representation quality , which characterize the extent to which variation in content and style is isolated by the two different spaces . We show that by using our approach , we can effectively disentangle the behavior and dynamics of movement without any labels . Our model thus strikes a nice balance between view-invariant representation and generation . Our specific contributions are as follows : • In Section 3 , we propose a generative method , Swap-VAE , that can both ( i ) learn a representation of neural activities that reveal meaningful and interpretable latent factors and ( ii ) generate realistic neural activities . • To further encourage disentanglement , we introduce a novel latent space augmentation called BlockSwap ( Section 3.3 ) , where we swap the content variables between two views and ask the network to predict the original view from the content of a different view . • In Section 3.4 , we introduce metrics to quantify the disentanglement of our representations of behavior and apply them to neural datasets from different non-human primates ( Section 4 ) to gain insights into the link between neural activity and behavior . 2 Background and Related Work . 2.1 Variational autoencoders and their application in neural data analysis . Variational auto-encoders ( VAEs ) ( 26 ) are a popular deep generative learning framework used to generate and denoise data . Let x and z denote the data and the latent variables , respectively , where z = qφ ( x ) is the latent representation extracted from x by the encoder qφ . The usual objective of probabilistic generative models is to maximize the log evidence of the observed data maxθ log pθ ( x ) = ∫ pθ ( x|z ) p ( z ) based on the parameterized model pθ ( called the decoder ) . VAEs and their variants instead optimize a tractable lower bound on the original objective , which is also well-known as the evidence lower bound ( ELBO ) in variational Bayesian inference ( 27 ) , log pθ ( x ) ≥ Ez∼qφ ( z|x ) [ pθ ( x|z ) ] − βDKL ( qφ ( z|x ) ||p ( z ) ) = : L VAE θ , φ , ( 1 ) where the encoder qφ is trained to approximate the Bayes posterior pθ ( z|x ) , p ( z ) is the prior over the latent variables , DKL is the Kullback–Leibler divergence , and β ≥ 1 is a trade-off parameter . Specifically , the standard VAE is obtained when setting to β = 1 , while β > 1 corresponds to the beta-VAE ( 23 ) . A larger value of β imposes stronger implicit regularization on the posterior of latent variables to align with the prior , which empirically induces more disentanglement and thus interpretability in the learned representations ( 24 ; 28 ; 29 ) . Recently , a number of generative modeling techniques based upon VAE have been applied to neural data . LFADS and more recent extensions of this model ( 10 ; 9 ) , use a sequential VAE ( 30 ; 31 ) to estimate neural population dynamics and show that this model can faithfully reconstruct single trial firing rates . More recently , pi-VAE ( 7 ) was proposed to learn representations of neural activity using a simple MLP encoder to capture information in each neural state vector ( firing rate of d neurons ) independently . They show that , even with a simplified architecture that treats each sample independently , it is possible to learn an identifiable model that can directly link behaviors to neural responses . In our work , we also use a MLP encoder similar to ( 7 ) , as we do not explicitly incorporate temporal structure into our architecture . Instead , we train the model to disentangle the content ( target ) from the dynamics without regularization from a more complex architecture . 2.2 Instance-specific alignment and self-supervision . Recent self-supervised learning ( SSL ) methods have made impressive advances , now rivaling ( or in some cases surpassing ) supervised methods ( 32 ; 33 ; 34 ; 35 ) . To build representations , these approaches aim to maximize the similarity across multiple augmented “ views ” of the same sample ( positive examples ) ( 36 ; 37 ; 34 ) . Thus , instead of providing labels to regularize the latent representations , we instead supply augmented versions of the same example and align their latent representations . Intuitively , training the network to align the instance-specific views where the semantic content is preserved , will encourage it to learn meaningful representations , which can then be used to solve downstream tasks ( 38 ) . Building on these successes , ( 11 ) recently showed how instance-specific alignment ( using only positive examples ) can be applied to multi-neuron recordings . This work shows that by combining simple augmentations , including : dropout , temporal shift ( selecting nearby points in time ) , and additive noise , with a dual network , it is possible to learn representations that are useful for decoding behavior . In this work , we further show that coupling this type of alignment approach with a generative model , we can build networks that achieve both good approximation quality and disentanglement , all without the need for a dual ( online/target ) network . This greatly simplifies our architecture and optimization approach . 3 Methods . In this section , we introduce our self-supervised approach for generative modeling of neural data . A PyTorch ( 39 ) implementation is provided here : https : //nerdslab.github.io/SwapVAE/ . 3.1 Model for neural datasets . Throughout , we consider collections of d neurons that have been spike sorted and binned to compute an estimate of the firing rate of all neurons at N distinct time points ( bins ) . This results in an input vector si ∈ Rd as an observation at each time point . Let D = { s1 , . . . , sN } denote the neural state vectors generated by this binning process . Let k denote the dimension of the latent space . 3.2 Unifying instance-specific alignment and generative modeling . As we describe in the introduction , our aim is to build a decomposable picture of brain states . To do this , we will leverage the principles of self-supervision to build a view-invariant representation and use this as the building block for our generative modeling approach . Our goal is to learn two functions , an encoder f : Rd → Rk and decoder g : Rk → Rd . Let x1 = t1 ( s ) and x2 = t2 ( s ) denote the views generated after applying two random transformations t1 , t2 ∼ T to a sample s from our dataset D. Let z1 = f ( x1 ) , z2 = f ( x2 ) denote the representations of both views in the network . To decouple the factors of our latent representations , we divide the latent space into two parts , z1 = [ z ( c ) 1 , z ( s ) 1 ] and z2 = [ z ( c ) 2 , z ( s ) 2 ] , with z ( c ) 1 and z ( s ) 1 modeling the behaviour styles and intrinsic neural contents , respectively . x̂1 = g ( z1 ) and x̂2 = g ( z2 ) are the reconstructions of both views obtained after passing them through the decoder . To encourage alignment of the views through the encoder while also solving our generative modeling objective , we propose the following loss : min f , g ∑ i=1,2 Lrec ( xi , g ( zi ) ) ︸ ︷︷ ︸ Reconstruction loss +β ∑ i=1,2 DKL ( z ( s ) i ‖ z ( s ) i , prior ) ︸ ︷︷ ︸ Regularization - style space +α Lalign ( z ( c ) 1 , z ( c ) 2 ) ︸ ︷︷ ︸ Alignment - content space , ( 2 ) where the alignment loss Lalign encourages two views to be close ( here we used a normalized L2-distance ) , α and β are hyperparameters that determine the tradeoff between alignment and reconstruction , and the KL divergence terms measure the deviation between the style latent variables and the prior z ( s ) i , prior which we set to be the isotropic GaussianN ( 0 , I ) . In our experiments on neural datasets , we choose the reconstruction loss Lrec to be the Poisson loss ( 10 ; 40 ) . Further details on the method and our implementation is provided in Appendix A . 3.3 BlockSwap : A novel latent space augmentation for disentanglement To further improve disentanglement in our model , we propose the following novel latent space augmentation , which basically exchanges the content information ( block of variables ) between two augmented views while keeping their style constant . We refer to this latent augmentation as BlockSwap as it holds one part of the representation consistent and then swaps a different subset of latent variables between two augmentations . Specifically , after generating the representations z1 = [ z ( c ) 1 , z ( s ) 1 ] and z2 = [ z ( c ) 2 , z ( s ) 2 ] for each view , we also generate their content-swapped versions z̃1 = [ z ( c ) 2 , z ( s ) 1 ] and z̃2 = [ z ( c ) 1 , z ( s ) 2 ] . To encourage disentanglement , we propose to replace the previous reconstruction term by adding the loss over the swapped representations : Lswaprec = ∑ i=1,2 Lrec ( xi , g ( z̃i ) ) ︸ ︷︷ ︸ Swapped content +Lrec ( xi , g ( zi ) ) ︸ ︷︷ ︸ Original . ( 3 ) When we use this loss , we can also consider removing the alignment term in Equation 2 and simply couple this reconstruction loss with the regularization KL term on the style space ( see Section 4.3 ) . In practice , the reparameterization trick is performed before the swapping of content and style variables .

Summary: This work proposes a method for learning disentangled latent variables to generate neural activity. The proposal consists of (1) a self-supervised VAE that separates the latent spaces into vectors for for 'content', which is preserved across transformations, and 'style', which is altered across transformations; (2) a training augmentation which encourages content vectors to be interchangeable (swappable). The authors then propose a metric to measure the disentanglement between two variables. They illustrate properties of their proposed method using both synthetic experiments that generate data from ground truth latent variables and real neural data from a monkey reaching task, and compare them to other related methods in the literature. For the monkey reaching task, they find that disentanglement between the two latent spaces is higher for the monkey that needs to insert a delay in its response, suggesting that this method could provide insight into how latent variables are encoded under different task conditions.