Drop, Swap, and Generate: A Self-Supervised Approach for Generating Neural Activity
Authors: Ran Liu
Mehdi Azabou
Max Dabagia
Chi-Heng Lin
Mohammad Gheshlaghi Azar
Keith B Hengen
Michal Valko
Eva L Dyer
Venue: NeurIPS 2021 Oral (2021)
Keywords: Generative modeling
Self-supervised learning
Neural decoding
Neural population activity
Representation learning

Abstract: Meaningful and simplified representations of neural activity can yield insights into how and what information is being processed within a neural circuit. However, without labels, finding representations that reveal the link between the brain and behavior can be challenging. Here, we introduce a novel unsupervised approach for learning disentangled representations of neural activity called Swap-VAE. Our approach combines a generative modeling framework with an instance-specific alignment loss that tries to maximize the representational similarity between transformed views of the input (brain state). 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. Through evaluations on both synthetic data and neural recordings from hundreds of neurons in different primate brains, we show that it is possible to build representations that disentangle neural datasets along relevant latent dimensions linked to behavior.

Drop, Swap, and Generate: A Self-Supervised Approach for Generating Neural Activity
Ran Liu" Georgia Tech
Mehdi Azabou Georgia Tech
Max Dabagia Georgia Tech
Chi-Heng Lin Georgia Tech
Mohammad Gheshlaghi Azar DeepMindd
Keith B. Hengen Washington Univ. in St. Louis
Michal Valko
Eva L. Dyer"
DeepMind
Georgia Tech
Abstract
Meaningful and simplified representations of neural activity can yield insights into how and what information is being processed within a neural circuit. However, without labels, finding representations that reveal the link between the brain and behavior can be challenging. Here, we introduce novel unsupervised approach for learning disentangled representations of neural activity called Swap-VAL VAE. Our approach combines generative modeling framework with an instance-specific alignment loss that tries maximize the representational similarity between trans- formed views of the input (brain state). These transformed (or augmented) views are created by dropping out neurons and jittering samples in time, which intu- itively should lead the network to representation that maintains both temporal consistency and invariance to the specific neurons used to represent the neural state. Through evaluations on both synthetic data and neural recordings from hundreds of neurons in different primate brains, we show that is possible to build representations that disentangle neural datasets along relevant latent dimensions linked to behavior.
Introduction
In the brain, the coordinated actions of groups of neurons are responsible for encoding sensory inputs and movements, well as all processing and manipulation in between 2 3 4 (). Understanding what different populations of neurons are doing and how they work together to encode their inputs is 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 TO II 12 13). Such methods have provided exciting new insights into the stability of neural responses (1), individual differences (10), and remapping neural responses through learning (55). However, without labels or additional inputs to guide the network, learning representations that allow different sources of variability to be distinguished still a major challenge I7).
"Contact: rliu361 @ edu, evadyer @ateeh..aaeee...... Project page: http:///aechihhgghhubb.osoaaaaa
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

Here, we develop novel unsupervised approach for disentangling neural activity called Swap-Vas 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 [ 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 be disentangled in similar manner.
Right). The content in this case may be knowing where
To identify the content within our neural recording, we use self- self-supereieed approach: we apply a variety of transformations to the input data (observed firing rates) which we hypothesize to be content- preserving, 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. addition to this instance-specific alignment loss, we also encourage the network to reconstruct the original inputs using regularized variational autoencoder (beta- VAEE) (23 244 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 self- supervised 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-purposs measures of representation quality, which char- acterize 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 nice balance between view-invariant representation and generation.
Our specific contributions are as follows:
In Section 3 we propose a generative method, Swap-VA that can both (i) learn repre- sentation of neural activities that reveal meaningful and interpretable latent factors and (ii)
generate realistic neural activities.
To further encourage disentanglement, we introduce 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 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 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) (6) 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 = 4o (x) is the latent representation extracted from x by the encoder go- The usual objective of probabilistic generative models is to maximize the log evidence of the observed data maxg log (x) = f po (x(2))((() based on the parameterized model pe (called the decoder). VAEs and their variants instead optimize a tractable lower bound the original objective, which is also well- as the evidence lower bound (ELBO) in variational Bayesian inference (27).
log (x) 2 (x (x(x)) - 3D (10)=((s(2))))llll (z)x)]l((()) =: CYAE
where the encoder ds is trained to approximate the Bayes posterior pp (2(x), p(z) is the prior over the latent variables, DKL. is the Kullback-Leibler divergence, and 3 > 1 is trade-off parameter. Specifically, the standard VAE obtained when setting to = 1, while 3 > corresponds to the beta- VAB (23). A larger value of B imposes stronger implicit regularization on the posterior of

---

Content
Style
Mi walloae Nu
DROP
Enc
Dec
directio
SWAP
GENERATE
Denamic
Dec
Overview of approach. On the left, show illustrative figure of the motor cortex reaching datasets, where rhesus macaques are trained to reach one of eight targets while their brain state (activity of many neurons at an instance in time) is recorded. Different views of a brain state are generated by dropping out neurons and exchanging samples that are close in time. The encoder (Enc) of the model extracts and partitions the latent variables of the two views into content space and a style space. An instance-specific loss is applied the content representations of augmented brain states encourage alignment, while neural activity reconstructed through a decoder (Dec) using hoth parts of the latent space. To further enhance alignment in the content space, we introduce BlockSwap: the content variables of the representations of the two augmented views are swapped before being passed through the decoder. To the right, we show how the behavior of the rhesus macaques movements can be disentangled. where reach target be decomposed into the direction of the movement (blue) and its undcrlying dynamies (yellow).
Figure 1:
latent variables to align with the prior, which empirically induces more disentanglement and thus interpretability 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 (0. use 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 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 ( 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 more complex architecture.
2.2 Instance-specifie alignment and self-supervision
Recent self-super learning (SSL) methods have made impressive advances, now rivaling (or in some cases surpassing) supervised methods (32) 33 34 33). 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 representa- tions, we instead supply augmented versions of the same example and align their latent representations. Intuitively, training the network to align the instance-specife 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, 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 dual (online/target) network. This greatly simplifies our architecture and optimization approach.
recently showed how instance-specific alignment (using only
3

---

3
Methods
In this section, we introduce self-supervised approach for generative modeling of neural data. A PyTorch (39) implementation is provided here: https::/nerddsaat https////eedalaggithhbb... So/SsapVAE/
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 N distinct time points (bins). This results in input vector E IRd as an observation at each time point. Let 'D = {81,... Sn denote the neural state vectors generated by this binning process. Let denote the dimension of the latent space.
3.2 Unifying instance-specific alignment and generative modeling
As we describe in the introduction, our aim to build a decomposable picture of brain states. To do this, we will leverage the principles of self-supervision build view-invariant representation and use this as the building block for our generative modeling approach.
Our goal is to learn two functions, an encoder / Rd + R and decoder 9 IR* + Rd. Let X1 = (s) and X2 fz(s) denote the views generated after applying two random transformations f1, t2 ~ T to sample from our dataset D. Let 2 = f(x1) ), Z2 = f(x2) denote the representations of both views the network. To decouple the factors of our latent representations, we divide the
latent space into two parts, 2 [z] 2 and 22 = [z (*) with 2(() and z modeling the behaviour styles and intrinsic neural contents, respectively. x g(z)) ) and X2 g(za) 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:9
Σ
Cree (x;;, +B Σ DEL (z)))
2 2ipplor to Calign
(2)
Reconstruction loss
Regularization style space
Alignment content space
where the alignment loss Lalign encourages two views to be close (here we used a normalized L22-ditancee, a 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, zi,prior which we set be the isotropic Gaussian NN 1). In our experiments on neural datasets, choose the reconstruction loss Lree to be the Poisson loss (I0 40). Further details the
method and our implementation 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 holds one part of the representation consistent and then swaps different subset of latent variables between two augmentations. Specifically, after generating the representations
(c) and Z2
(c)
2
[z ziz for each view, we also generate their sontent-swapped versions
2i 22 1 and Z2 = [z] [e) To encourage disentanglement, we propose to replace the
previous reconstruction term by adding the loss over the swapped representations:
sswap
(xec g(z))) Cree (xe.
1=12
Swapped content
Original
When we use this loss, we can also consider removing the alignment term in Equation and simply couple this reconstruction loss with the regularization KI. term on the style space (see Section (43). In practice, the reparameterization trick is performed before the swapping of content and style variables.

---

3,4 Representational quality and disentanglement metrics
Determining whether representation captures the underlying factors of interest is, in general, very challenging (16) 41 (42). In neuroscience this is certainly true. We will consider two main measures of representation quality and disentanglement to guide our investigation.
Multi-task disentanglement score. To confirm that our method effectively disentangles the repre- need to measure the extent which latent
sented behavior of neural signals from the dynamics, variables respond to one or the other with specificity. Concretely, let z denote a k-dimensional latent representation, and let ye: V/ be discrete values that label the reaching direction and dynamics that correspond to z, respectively. For each latent variable Z, (1 S i S k),, we assess the degree to which it is aligned with the content of the behavior (ye) the temporal structure (ya). Computing the covariance score for a particular latent variable Z; consists of three steps:
1. Compute the variance when changing Ve with fixed Y. and average over values of Va-
2. Compute the variance when changing y. with fixed Ve. and average over values of ye-
3. Compute the absolute difference of the two variances. This is the score for Z,.
Intuitively, if the score is large, then z; changes more dramatically in response to one parameter than the other, so displays specificity, while if it low then the amount that z, changes is nearly the same. Averaging across all latent variables after normalization gives a final score, which provides a measure of how disentangled the entire representation is.
Linear readout from representation layer. To further quantify the representation quality and stability of representations in downstream decoding, frequently in self-supervised learning approaches (4), 34), In particular, we will train the model on training dataset, freeze the weights in the network, and then train linear layer decode the reach directions from the output of the encoder. However, because we are also interested in disentanglement, will consider the prediction of two different class labels from either the full, content, style factors in the network. When decoding either reach directions or temporal structure Ya. we will retrain the linear weights but keep the representation fixed.
linear readout strategy employed
As in (II, we quantify reach decoding with two scores, acc and delta-ace, for the linear decoding accuracy reach direction. Consider the reaching task as regression over circle with total of I discrete labels, we count the decoded angle that falls within [(: - 1)T/l, /1, (2i + 1)T/(l] as the correct
classification in acc, and that falls within [2i 1.5)m/ /L, (2i + 1.5)x/( as the correct classification in delta-acc. The two scores both provide a measure of the representation quality in terms of the
precision of reach direction decoding. The time decoding accuracy is computed similarly, but in this case, the goal is to predict how far into each reach each sample is.
Experiments
The usefulness of this method depends on its ability decompose neural data into meaningful latent factors. To quantify this, we devised experiments to reveal three desirable properties:
1. Performance on downstream classification tasks (linear separability of the representations).
2. Faithfulness of the latent space structure to the ground-truth structure of the task.
3. Separability of the content and style components across the respective latent subspaces.
4.1 Synthetic experiments
We first trained and evaluated our model on synthetic dataset that was designed to resemble our neural datasets. The data are generated and the experiments are designed following the approach used 43
Synthetic reaching dataset. We generated latent variables from 2-dimensional mixture of four Gaussians, where the ith component Gaussian has mean (5sin u, 5 cos u,) and variance (0.6
0.3| s3l, 0.3) sin uzl) and u; is uniformly sampled from the interval [ix= ((+1)x= For each sequence, we randomly sampled / = data points within each cluster and order them in clockwise manner to create consistent dynamics in the latent space (as shown in Figure 2 on the left). The formed sequences were fed to a RealNVP network (44) to generate 100-dimensional Poisson observations of the firing rates. The generated synthetic firing rates are then shuffled and
{0. 1, 2, 3}.
5

---

Ground truth
Our model
beta-
Supervised
Componentz
Compenant
-25
-20
-25
6.0
2.5
3.0
Component
Component:
Component
Figure 2: Synthetic Experiments. On the left, we show the ground truth latent space and their dynamics, from which the firing rate is generated in clockwise manner. To the right, we show the results of our model, a Veta-VAE, and supervised model. Our model recovers both the discrete classes and the sequential structure present within this synthetic dataset.
split into 80% training set and 20% test set. The ground truth latent distributions and the generated results are shown in Figure 2 Ideally, we want the model to capture and disentangle both the discrete groups/clusters ye and the dynamics of the sequence (which are implicitly encoded via the Re network).
Results on synthetic datasets. To compare the representational power of our approach, applied our model, the bet,-VAE, and a supervised decoder (trained to predict ye) to these synthetic datasets. All models have the same backbone with latent space of 32-dim, which in our model is split into a content space of 16-dim and a style space of 16-dim. Each model is trained on the training set for 100,000 iterations, and optimized using Adam with learning rate of 0.0005 (further details on model selection and hyperparameter optimization in Appendix AA.. When we examined the representations formed by each of these models, we found that Swap- VAE was very effective at both preserving the sequence dynamics (as highlighted by the connection between class centroids) as well as separating the different target classes (Figure 2.. From the figure, we can see that while the VAE successfully separated different clusters, the ordering of clusters is not aligned with the true dynamics (the black line), while the elongated distribution formed by our model more accurately reflected the true distribution of each component, regardless of the noise. Using the metrics described in Section (3.4), we further confirmed that our model provides good disentanglement, producing a multi-task disentanglement score of 0.93. The corresponding score for the beta- VAE and supervised model were 0.46 and 0.12, respectively.
4.2 Experiments on neural datasets
After testing the model on synthetic datasets, we applied our model to datasets collected from the primary motor cortex of two non-human primates (rhesus macaques) performing reaching task (25). Neural recordings from these same individuals have been used in recent studies of deep representation learning (1),, interpretable generative modeling 77, and adversarial domain adaptation (45).
Motor cortex reaching datasets. We use reaching a simplified laboratory task to test hypothesis that the what and how of movements could be disentangled. We consider spike sorted datasets from two rhesus macaques, Chewie and Mihi, both trained perform a reaching task towards one of eight different directions after a cue. The reaching task has two different settings: Chewie performs the reaching task immediately after seeing the target the screen (no waiting), while Mihi performs the reaching task with waiting period of time (between 500-1500 ms) after receiving an auditory "Go" cue. While carrying out these movement tasks, neural activities in primary motor cortex (MI) were recorded from both individuals, In these examples, the activity of population of roughly one hundred single neurons was binned into 100ms intervals to generate approximately 1.3k data points per dataset. For each target direction, there are multiple trials/repeats. For each trial, the first 9 binned time points are selected for temporal decoding. For each individual, two days of neural recordings are considered, where different groups of neurons are recorded on different days.
Experimental setup. In our model, we apply combination of two augmentations: (i) spatial augmentations, where neurons are randomly dropped (masked with zeros) from the input (with p 0.6), and (ii) temporal augmentations that select nearby point in time (randomly in window,

---

(A)
Content space
Style space
Beta VAE
Supereised
-20
31722107
-0
(8)
-25
(D)
12
BVDs.
100
-00
-39
(C)
20
36
-9 30
Reaching
Time
-20
Mihi
-22 Mhh
Chowe:1
Figure 3: Disentangling neural representations of movement the primate motor cortex. (A), we show the representations formed by our model's content and style space when compared with the beta-VAE and a supervised network trained to decode reach direction. All of the visualizations are obtained after embedding the representations into 2D using iSNE. Below, decompose the content and style space further by averaging over all trials towards a specific reach and visualizing their trial-averaged trajectory for Chewie (Day in (B) for
Mihi (Day 2) in (C). (D), we compare decoding accuracies over both reach direction and Chewie-1 for our Full, Content, and Style spaces, the benchmark models, and supervised decoders trained on either task. In (E), we show the results of our disentanglement score on all four reaching datasets. In this case. we compare the disentanglement over all of our latent space (Full) with the Veta-VAE and supervised model trained reach direction.
+5 samples from target sample) positive example. All models have 128-dim latent space, which is split for our model between into 64-dim blocks for the style and content space. All models trained using one Nvidia Titan RTX GPU for 200 epochs using the Adam optimizer with learning rate of 0.0005 (further details can be found in Appendix A).. With as the total number of neurons, all generative models have an encoder and symmetric decoder, where the encoder has three fully connected layers with size [d, 128, 128), batch normalization, and ReLU activations. All discriminative models have encoder of 4 linear layers with size [d. 128, 128, 128), which was determined to be preferable after extensive hyperparameter optimization. our experiments, we split the dataset into 80% for train and 20% for test.
Investigating disentanglement in neural representations of movement. After training our model, we examined the latent space structure by applying (SNE (46) the Full space (considering both Content and Style jointly). as well as the Content and Style spaces individually. Chewie-1 and Mihi-2 are visualized in Figure which admitted the highest reach decoding accuracy and most disentangled representations, respectively; visualizations of the remaining datasets can be found in Appendix C.I where similar trends hold. When compared with a beta-VAE and Supervised decoder trained on the reach direction task, we observe that the Content space in our model and the beta-VAs VAE have similar overall structure, with our model providing further separability and preservation of the task structure (circular positioning of targets, which is not intentionally encouraged by the loss function). The Style space provides a good embedding direction has been collapse but time is nicely organized. These results suggested that our model is good at separating semantic structure without any labely while also preserving the overall structure of the behavior.
entire dataset along an axis where reach

---

To understand how much information our model extracts relevant to the two different downstream tasks, we explored the decoding accuracies in our reach and temporal decoding tasks on Chewie- ! (see Figure P, Appendix (C.)) We examined the Full, Content, and Style spaces for our model on both tasks, and compared with the benchmark models and two separate supervised models trained on two tasks as the upper bounds. These measures provided further evidence of disentanglement our Content spaces provide good decoding accuracies on reach decoding, while the Style space has little predictive power over reach direction (as anticipated). The reverse true for the temporal decoding in the Content space. These results are promising indicators that disentanglement indeed possible with our approach and that these decoding measures capture what we observe in visualization.
We next measured the multi-task disentanglement scores across all four datasets. When examining the separability of our latent space across the two individuals (see Figure BEE, we found that the disentanglement (i.e., separation between the reach direction and the dynamics of the movement) quantified by the multi-task covariance scores for Chewie is on average lower than for Mihi in both cases; this observation may be interesting given that Mihi needs to wait before making a reach and has to delay their movement at the beginning. While the results of this analysis need to be studied further, this result provides initial evidence that our unsupervised method for disentanglement provides a useful lens into the distinction between the neural representation of different movement tasks.
Quality of representations as measured through linear readouts. Next, we conducted a com- prehensive evaluation of the decoding of the reach direction (Table [] and the temporal dynamics (Table 2). In this case, we compared our model with a supervised decoder trained on either task (Sup-Reach Reach for reach decoding, Sup-Time for temporal decoding), a supervised (pi- VAE (7)), and an unsupervised (beta- VAE (33 24)) generative approach for disentanglement (pi- VAE (7) and 3-VAE (23) 24)). and two self-supervised methods for general representation learning (B YOL (47)
and MYOW ( These models provide a comprehensive collection of both supervised and unsupervised competitors for our tasks of interest.
The decoding accuracy analysis yielded the following interesting observations. First, we find that our model performs on par with or better than all the benchmark models on the reach direction decoding task. Second, we find that our model performs comparably with beta-VAB and the supervised model (Sup-Time) the time decoding task, while other sulf-super models fail capture the temporal structure that they have built invariance to. These results indicate that both Swap- VAE and MYOW, two self-supervised approaches, outperform supervised decoders on reaching direction classification, due to the strong regularization from their alignment losses.
The power of our model shines in its ability to extract both the reaching target as well as the temporal structure from the data, As we show in both Fi and Table using a SSL-bbased alignment loss YOOOL, MYOW, and Swap-VAL VAE's 's Content space) with temporal jitter augmentations builds temporal invariance, which makes it challenging to decode temporal structure of the reaches. On the other hand, Swap- VAE introduces an additional set of latent variables (Style space) that carry additional information and make it possible to achieve high accuracy in both reach and temporal decoding.
Testing the generative quality of the model. A key component of Swap-VAE VAE is the integration of SSL with generative modeling framework. Thus, it is important to gauge the quality of the generated data. As shown in Figure f the direct reconstruction of the neuron firing rate is realistic in terms of both the flass-oondiiooee firing rates, and the dynamics of individual neurons. When we analyzed the RMSE of the fitted rate for all neurons in our model against those from a VAE, we found our model has lower error and could reconstruct data more faithfully than the VAE. Our results suggest that Swap- VAE provides a good denoised estimate of neural activity that captures information related to the behavior (rather than factors extraneous the movement).
To further demonstrate that our generated neuron activities are useful for downstream tasks, we use our trained model to generate new samples and mix them with original training samples when training a supervised classifier. All supervised models are trained for 400 epochs, with the same model settings as mentioned in the experimental setup. We tested one dataset from each individual (Chewie-l,
'Chewie-1 I yields the highest reach decoding accuracy of all datasets, which motivated this choice. We include the reported numbers for reach decoding, and also reproduced MYOW for temporal decoding (see reproduced results in Appendix B.I. With changes to the network architecture and set of augmentation operations, MYOW recently reported even higher accuracies on reach decoding.
8

---

Table Accuracy %) for reach direction classification on neural dalasets. See Appendix deviation of the benchmark models.
standard
Sup-Reach
pi-VAE
VAta-VAE
BYOL
MYOW(I
Ours
acc 61.59
66.30
64.34
62.12
70.41
72.81 1
Chewie-1
ò-acc
77.58
82.93
80.83
81.27
86.24
85.04(+0..94)
Chewie-2
acc
69.71
61.33
60.24
57.25
60.95
68.97(13,41)
ò-acc
78.18
73.63
80.09
76.89
81.36
83.66(+1.81))
Mihi-1
acc
62.86
62.63
58.11
60.03
70.48
64.26(110889
ò-acc
79.10
79.20
75.98
78.82
83.24
82.18(+:1.55)
acc 60.72
62.70
60.23
59.94
64.35
66.12(2:0.8)
Mihi-2
ò-acc
74.02
76.89
77.89
78.10
80.58
82.74
Table 2: Accuracy (in %) for temporal decoding classification on neural datasets. See Appendix standard deviation of the benchmark models.
for the
Sup-Time
Time
pi VAE
beta- VAE
BYOL
MYOW
Ours
Chewie-1 53.80
49.56
48.24
42.30
22.45
56.09(+1.80)
Chewie-2
55.58
54.45
63.76
39.69
31.96
62.54
Mihi-1 55.93
52.66
54.10
43.46
21.88
57.77(#3.30)
Mihi-2 58.75
48.03
47.89
38.16
22.09
49.65(+2.01)
Mihi-2) and computed the improvement in accuracy as we increase the number of generated samples included in the training set (50%, 100%, 200%). In all cases, we found some improvement in the aceuracy of the model, with roughly 5% and 3.5% gains over the supervised baseline in Chewie and Mihi, respectively. We note that the supervised models trained with generated samples still does not surpass the Swap-VAL VAE (details Appendix D)
(D)
poth 24
35
VAE
(B)
2
30
2
Swap-VAE
Swap-VAE
Reach directions
(F)
[
pen
9%
Chewie
3.5% Whhi
ong
om
Generate virtual samplesi
Sapervised
50%
100%
200%
50%
100%
200%
Figure 4: Testing Swap- VAE's ability to reconstruct and generate new neural activity. Reconstruction of
firing rates from example neurons over (A) and across different reaching directions in (B), further validate the reconstruction accuracy of our model by comparing the RMSE obtained with our model (x-axis) vs. the RMSE for the VAE (y-axis) for different reaching directions (C) and over time (D). (F) shows sketch of how we generate virtual samples (green) and them train supervised classifier (orange). (F) highlights improvement classification accuracy we increase the amount generated data fed into supervised decoder (left, Chewie-l; right, Mihi-2.)
4.3
Model ablations: Testing our BlockSwap augmentation
We studied different variants proposed loss functions in Equations 2 and (3) and how different
data augmentation operations impact decoding accuracy (Table 3). The three variants of the loss are the following: (i) when only removing the alignment term Eq. 2 (no 1.2). (ii) when only removing the BlockSwap augmentation (No-Swap), and (ii) when removing the alignment term and the original

---

Table 3: Model ablations. Accuracy (in %) of different variants of our proposed model when tested on Chewie-1.
L2
no-swap
swap-only
vanilla VAE
S-Aug 66.77
T-Aug
Ours
acc
71.63
63.17
68.36
63.79 79.10
71.15
72.81 85.04
delta-acc
85.62
83.20
83.02
81.03
83.34
reconstruction term (Swap-only). We also report the results obtained with a vanilla VAE, while beta-VAE results are in Table Ablations of different dimensions of the content and style space are included in Appendix
In our experiments, find that almost all of the variants of our decomposed loss functions performs better than a vanilla VAE or a beta-VAE. Adding the BlockSwap loss term improves performance overall, with our highest decoding accuracies being obtained with this model. When use the content swapping technique, we can remove the 1.2 alignment loss with minimal change in performance, but in general, including this alignment terms provides additional parameter give more flexible control. This shows that our proposed model is stable and that our proposed swapping loss provides a strong boost in performance.
We tested the spatial-only (S-Aug) and temporal-only (T-Aug) conditions separately in Table 3 In this case,
that they all perform reasonably well, although they are both worse than our final model where we combine both spatial and temporal augmentations. As we know, the selection of the data augmentations critical for the performance of representation learning model (32) 48 (49). Our model needs even fewer data augmentation operations than MYOW to achieve a good performance, highlighting the power of our approach.
5
Conclusion
This paper introduces new self-supervised approach, Swap-VAE, for generative modeling of neural activity. Our proposed method leverages self-supervised alignment to decompose neural activity into different latent subspaces, with the goal of provide insights into the relationship between neural activity and animal bchaviour.
While our model provides good balance between both reach and temporal decoding using simple augmentations (dropout, temporal jitter), we note that, just like other self-supervised learning methods, the performance of our method relies on our choice of augmentations, which might be dependent on the downstream task of interest. Exploring how different augmentations affect performance on different downstream tasks may yield insights into how information is coded in the brain, which will be an exciting line of future research.
In our current formulation, we only consider dropout and localized temporal shifts as augmentations of the neural activities. However, other works like (II) use a nearest-neighhor approach to link brain states that are temporally nonlocal or may span different trials. Through combining our approach with this nonlocal view mining strategy, we may be able to build even further invariance into our model's content space. Combining our SSL-backed approach with a sequential encoder (I0) or transformer (50) another exciting line of future research that can be used to model latent structure over longer timescales.
Beyond demonstrating the effectiveness of our method, our analysis of neural activity patterns facilitated by Swap-VAL VA had other intriguing outcomes. There is an active debate in contemporary neuroscience about the role of preparation in the neural activity associated with movement 521 53).
with some evidence suggesting that preparation makes the neural activity driving movement more stable and robust. As we point out in Section 4.2, Chewie and Mihi were instructed to make reaches differently, with Mihi required to wait for an auditory cue before performing each movement. We found that the disentanglement in Mihi was more pronounced than Chewie, both in terms of decoding accuracy across content and style spaces and multi-task covariance scores. It is possible that this due to a weakness of our method, which could perhaps be resolved by tuning hyperparameters, increasing the size of the network, using some new augmentations. However, it might be possible that preparation leads more disentangled neural activity and that this discernible in the evaluation of the multi-task disentanglement score, which should be explored carefully across more individuals.
10

---

Acknowledgements
This project supported by NIH award 1R01EB029852-01 (RI,, MA, CHL, KBH, ELD), CKOINS118442-0 (KBH), NSF awards IIS-1755871 (ELD) and IIS-2039741 (ELD), a NSF GRFP (MD), as well as generous gifts from the Alfred Sloan Foundation (RL, ELD) and the McKnight Foundation (RL, MA, CHL, ELD).
References
[1] V. Mante, D. Sussillo, K. Shenoy, and W. Newsome, "Context-dependent computation by recurrent dynamics in prefrontal cortex," Nature, vol. 503, no. 7474, pp. 78-84, 2013.
[2] R. Yuste, "From the neuron doctrine to neural networks." Nature Reviews Neuroscience, vol. 16, no. pp. 487-497, 2015.
[3] S. Fusi, E. K. Miller, and M. Rigotti, "Why neurons mix: high dimensionality for higher cognition," Current Opinion in Neurobiology, vol. 37, pp. 66-74, 2016.
[4] Gao, E. Trautmann, B. M. Yu, G. Santhanam, S. Ryu, K. Shenoy, and S. Ganguli, "A theory of multineuronal dimensionality, dynamics and measurement," bioRxiv doi: 101/214262.
[5] H. Eichenbaum, "Barlow versus hebb: When is it time to abandon the notion of feature detectors and adopt the cell assembly as the unit of cognition?." Neuroscience Letters, vol. 680, pp. 88-93, 2018.
[6] S. Saxena and J. P. Cunningham, "Towards the neural population doctrine," Current Opinion in Neurobiology, vol. 55, pp. 103-111, 2019.
[7] D. Zhou and X.-X. Wei, "Learning identifiable and interpretable latent models of high- dimensional neural activity using pi-vae," arXiv preprint arXiv:2011.04799 2020.
[8] L. Y. Prince, S. Bakhtiari, C. J. Gillon, and B. A. Richards, "Parallel inference of hierar- chical latent dynamics in two-photon calcium imaging of neuronal populations," bioRxiv Not:2021.03.05.44444
2021.
[9] M. R. Keshtkaran, A. R. Sedler, R. H. Chowdhury, R. Tandon, D. Basrai, S. L. Nguyen, H. Sohn, M. Jazayeri, L. E. Miller, and C. Pandarinath, "A large-scale neural network training framework for generalized estimation single-trial population dynamics," bioRxiv doi: 100.10..110/2221.11133.222
[10] C. Pandarinath, D. J. O'Shea, J. Collins, R. Jozefowicz, S. D. Stavisky, J. C. Kao, E. M. Trautmann, M. Kaufman, S. Ryu, L. R. Hochberg, al., "Inferring single-trial neural population dynamics using sequential auto-encoders," Nature Methods, vol. 15, no. 10, pp. 805- 815, 2018
[11] M. Azabou, M. G. Azar, R. Liu, C.-H. Lin, E. C. Johnson, Bhaskaran-Nair M. Dabagia, K. B. Hengen, W. Gray-Roncal, M. Valko, and E. Dyer, "Mine your own view: Self-supervised learning through across-sample prediction," arXiv preprint arXiv:2102.10100 10106 2021.
[12] G. "Deep random splines for point process intensity estimation of neural population data," arXiv preprint arXiv: (903v. 026100 2019.
M. Perkins, K. E. Schroeder, M. M. Churchland, and J. P. Cunningham,
[13] M. Y. Byron, J. P. Cunningham, G. Santhanam, S. I. Ryu, K. V. Shenoy, and M. Sahani, cauussian-proeess factor analysis for low-dimensional single-trial analysis of neural population activity," in Advances in Neural Information Processing Systems, pp. 1881-1888 2009.
[14] J. A. Gallego, M. G. Perich, R. H. Chowdhury, S. A. Solla, and L. E. Miller, "Long-term stability of cortical population dynamics underlying consistent behavior," Nature Neuroscience, vol. 23, no. 2, pp. 260 160-270,220 2020.
[15] M. D. Golub, P. T. Sadtler, E. R. Oby, K. M. Quick, S. I. Ryu, E. C. Tyler-Kabara, A. P. Batista, M. Chase, and B. M. Yu, "Learning by neural reassociation," Nature Neuroscience, vol. 21,
no. 4, p. 607. 2018.

---

[16] F. Locatello, S. Bauer, M. Lucic, G. Raetsch, S. Gelly, B. Schölkopf, and O. Bachem, "Chal- lenging common assumptions the unsupervised learning of disentangled representations"" in International Conference on Machine Learning, pp. 41 14244,24, PMLR, 2019.
[17] I. Higgins, D. Amos, D. Pfau, S. Racaniere, L. Matthey, D. Rezende, and A. Lerchner, "Towards a definition of disentangled representations." arXiv preprini arXiv: 2018.
[18] M. Mathieu, J. Zhao, P. Sprechmann, A. Ramesh, and Y. LeCun, "Disentangling factors of variation in deep representations using adversarial training," arXiv preprint arXiv:1611.0338. 2016.
[19] R. Zhang. Tang, Y. Li, J. Guo, Y. Zhang, Li, and Yan, "Style separation and synthesis via generative adversarial networks," in Proceedings of the 26th ACM International Conference on Multimedia, pp. 183-191, 2018.
[20] X. Huang and S. Belongie, "Arbitrary style transfer in real-time with adaptive instance normaliza- tion," in Proceedings of the IEEE International Conference on Computer Vision, pp. 1501-1510, 2017.
[21] H. Zhang and K. Dana, "Multi-style generative network for real-time transfer," in Proceedings of the European Conference on Computer Vision (ECCV) Workshops, pp. 0-0, 2018.
[22] T. Karras, S. Laine, and T. Aila, "A style-based generator architecture for generative adversarial networks," " in Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 1401-4410, 2019.
[23] I. Higgins, L. Matthey, A. Pal, C. Burgess, Glorot, M. Botvinick, S. Mohamed, and A. Ler- chner, "beta-vae: Learning basic visual concepts with a constrained variational framework," 2016.
[24] C. P. Burgess, I. Higgins, A. Pal, L. Matthey, N. Watters, G. Desjardins, and A. Lerchner, "Understanding disentangling B-vae," arXiv preprint arXiv:1804.03599 2018.
[25] E. L. Dyer, M. G. Azar, M. G. Perich, H. L. Fernandes, S. Naufel, L. E. Miller, and K. P. Körding, "A cryptography-basee based approach for movement decoding," Nature Biomedical Engineering, vol. 1, no. 12, pp. 967-976, 2017.
[26] D. P. Kingma and M. Welling, "Auto-encoding variational bayes," arXiv preprint 1311v:1312.6114, 2013.
[27] D. J. MacKay and D. J. Mac Kay, Information Theory, Inference and Learning Algorithums. Cambridge university press, 2003.
[28] E. Mathieu, T. Rainforth, N. Siddharth, and Y. W. Teh, "Disentangling disentanglement in variational autoencoders," in International Conference Machine Learning, pp. 4412-4412, PMLR, 2019.
[29] R. Т. Chen, X. Li, R. Grosse, and D. Duvenaud, "Isolating sources of disentanglement in variational autoencoders," arXiv preprint arXiv: 140400420 2018
[30] J. Chung, K. Kastner, I.. Dinh, K. Goel, A. Courville, and Y. Bengio, "A recurrent latent variable model for sequential data," arXiv preprint arXiv:/506.022/6 2015.
[31] O. Fabius and J. R. Van Amersfoort, "Variational recurrent auto-encoders," arXiv preprint 1rXiv:1412.6587, 2014.
[32] T. Chen, S. Kornblith, M. Norouzi, and G. Hinton, "A simple framework for contrastive learning of visual representations," in International Conference Machine Learning, pp. 1697-1607, PMLR, 2020.
[33] K. He, H. Fan, Y. Wu, S. Xie, and R. Girshick, "Momentum contrast for unsupervised visual representation learning," " in Proceedings of the IEEE/CVE Conference on Computer Vision and Pattern Recognition, pp. 9729-9733 2020.
12

---

[34] Z. D. Guo, B. A. Pires, B. Piot, J.-B. Grill, F. Altché, R. Munos, and M. G. Azar, "Boot- strap predit-pe representations for multitask reinforcement learning," " in International Conference on Machine Learning, pp. 3875-3886, PMI.R 2020.
[35] X. Chen and K. He, "Exploring simple siamese representation learning"" arXiv preprint arXiv:2011. 10566, 2020.
[36] I. Misra and L. v. d. Maaten, "Self-supervised learning of pretext-invariant representations" "
in Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition.
pp. 6707-6717, 2020.
[37] M. Caron, I. Misra, J. Mairal, P. Goyal, P. Bojanowski, and A. Joulin, "Unsupervised learning of visual features by contrasting cluster assignments," arXiv preprint arXiv:2006. 2006008820..,888,, 2020.
[38] M. Caron, P. Bojanowski, A. Joulin, and M. Douze, "Deep clustering for unsupervised learning of visual features," in Proceedings of the European Conference on Computer Vision (ECCV), pp. 132-149, 2018.
[39] A. Paszke, S. Gross, F. Massa, A. Lerer, J. Bradbury, G. Chanan, T. Killeen, Z. Lin, N. Gimelshein, L. Antiga, et al., "Pytorch: An imperative style, high- sigg--erformance deep learning library," arXiv preprint arXiv:1912.01770
2019.
[40] Y. Gao, E. Archer, L. Paninski, and J. P. Cunningham, "Linear dynamical neural population models through nonlinear embeddings." arXiv preprint arXiv:
2016.
[41] S. Van Steenkiste, F. Locatello, J. Schmidhuber, and O. Bachem, "Are disentangled representa- tions helpful for abstract visual reasoning?," arXiv preprint arXiv:1905. /2506, 2019.
[42] Locatello, G. Abbati, Rainforth, S. Bauer, B. Schölkopf, and O. Bachem, "On the fairness of disentangled representations." arXiv preprint arXiv: 1905. /3662, 2019.
[43] I. Khemakhem, D. Kingma, R. Monti, and A. Hyvarinen, "Variational autoencoders and nonlinear ica: A unifying framework," in International Conference on Artificial Intelligence and Statistics, pp. 2207-2217, PMLR, 2020.
[44] L. Dinh, J. Sohl-Dickstein, and Bengio, "Density estimation using real nvp." arXiv preprint arXiv:
2016.
[45] A. Farshchian, J. A. Gallego, J. P. Cohen, Y. Bengio, L. E. Miller, and S. A. Solla, "Adversarial domain adaptation for stable brain-machine interfaces," arXiv preprint arXiv: 18100000, 2018.
[46] L. Van der Maaten and G. Hinton, "Visualizing data using t-sne.," " Journal of Machine Learning Research, vol. no. 11, 2008.
[47] J.-B. Grill, F. Strub, F. Altché, C. Tallec, P. H. Richemond, E. Buchatskaya, C. Doersch, B. A. Pires, Z. D. Guo, M. G. Azar, et al., "Bootstrap your own latent: A new approach to self-supervised learning," arXiv preprint arXiv:2006.0771 2020.
[48] H. Lee, S. J. Hwang, and J. Shin, "Rethinking data augmentation: Self-supervision and self- distillation," 2019.
[49] Y. Tian, C. Sun, B. Poole, D. Krishnan, C. Schmid, and P. Isola, "What makes for good views for contrastive learning." arXiv preprint 2rXiv:2005.10243 2020.
[50] J. Ye and C. Pandarinath, "Representation learning for neural population activity with neural data transformers," bioRxiv doi: 100000000000000001111666 14.6699695 2021.
[51] M. T. Kaufman, M. M. Churchland, S. I. Ryu, and K. V. Shenoy, "Cortical activity in the null space: permitting preparation without movement," Nature Neuroscience, vol. 17, no. 3, pp. 440-448. 2014.
[52] K. Svoboda and N. Li, "Neural mechanisms of movement planning: motor cortex and beyond," Current Opinion in Neurobiology, vol. 49, pp. 33-41, 2018
[53] A. A. Russo, S. R. Bittner, S. M. Perkins, J. S. Seely, B. M. London, A. H. Lara, A. Miri, N. J. Marshall, A. Kohn, T. M. Jessell, et al., "Motor cortex embeds muscle-like commands in an untangled population response," " Neuron, vol. 97, no. 4, pp. 953-966, 2018.
13


Reviews:
[{"review_id": "ttccCxEozJz", "reviewer": "Reviewer_sJZd", "summary": "The study proposed an unsupervised (self-supervised) learning strategy to model neural activity in terms of meaningful latent factors. The authors apply the proposed method to both synthetic and actual neural datasets and show that their approach can both successfully disentangle neural activity along meaningful behaviorally-associated latent factors, and also generate realistic patterns of neural activity.  A significant methodological advance is that the disentanglement is achieved in the learned latent space in an entirely unsupervised manner, without requiring any attribute labeling. While some evaluation decisions could be criticized, the results (and success of the approach) are overall convincing. ", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "This is an interesting paper where the methods (associated loss functions) are clearly stated and explained, and the described analyses seems well-executed. I have some comments about the lack of thorough analysis, which potentially prevent drawing strong conclusions about the benefits of different model components in driving the learned representations. Although, that said, I think the paper does pack a lot in 8 pages, so the lack therein could be attributed to too much content. I am not sure that the proposed approach extends previous methods for generative modeling (and disentanglement) significantly. Although coupled with the specific application (modeling ‘of neural activity’), the combination of methods employed (contrastive learning + generative modeling) in this study could be considered novel. The application of the proposed method for neural datasets is well-motivated; specifically, the augmentations and other decisions seem like good choices for neural representations. Visualization and other quantitative analyses on the synthetic datasets clearly illustrate the merits of the proposed approach and the results on the real neural datasets are promising for future applications geared towards discovering the structure in dynamic patterns of neural activity. ", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The paper is overall interesting, well-written and makes a valuable contribution. I do, however, have some comments for the authors to consider (which in my mind, are potential limitations of the study):  \n-\tComparison of the proposed unsupervised method with the supervised baseline is not suggestive because of the absence of augmentations in the supervised baseline. The authors should consider reporting performance on the decoding task when the supervised method employed data augmentations as well. \n-\tFor completeness, the authors should also report how the hyperparameters for the linear decoder were determined. Ideally, I would’ve liked to see error bars for decoding accuracies as well (maybe by bootstrapping training set for the decoder?) \n-\tIn future, the authors could also consider replacing the acc metric for decoding with better evaluation metrics for circular data, like circular correlation. This would treat the reach direction as a continuous variable (which it is) rather than as a discrete unordered variable (which it theoretically isn’t).  \n-\tThe authors consider swapping only the block of variables belonging to the `content’ group. What would happen if the reconstruction term in the BlockSwap method swapped both the content and style of the augmented views? Does swapping only the content block necessarily facilitate disentanglement? If the BlockSwap was essential, does the proposed method required knowing the number of latent factors in advance. The authors could discuss these aspects in their conclusion/discussion. \n-\tWhen comparing the proposed SwapVAE against vanilla VAE, the authors should also consider reporting other metrics more commonly employed in VAE evaluation (likelihood etc.) and not just the reconstruction error (which can be trivially minimized). This is important, since the authors mention ‘generating realistic neural activity’ as a significant contribution of their paper. \n-\tThe authors should also consider defining content and style more broadly as it relates to their specific neural application (e.g., as in Gabbay &Hosehn (2018)) where style is instance-specific(?) and content includes information that can be transferred among groups. More specifically, since their model is not sequential and does not capture the temporal dynamic structure, what do they really mean by ‘style’ represents the ‘movement dynamic’?\n", "review_text": "", "extra_scores": {"ethical_concerns": "None", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4-5 hours"}, "extra_text": {}}, {"review_id": "aCF34I38GX8", "reviewer": "Reviewer_r6nv", "summary": "The authors propose a self-supervised approach to learning representations of neural population activity in motor cortex. The approach is based on disentangling slowly varying action classes (content) from temporal dynamics of the action (style). They show that their approach facilitates decoding of eight action classes (reach directions) and compare with previous methods based on autoencoding (beta-VAE) or contrastive learning (MYOW). Overall the approach seems sound and promising. Unfortunately, it is difficult to verify the claims of the paper, since the analyses are very selective and not rigorous. \n\nWhile I am generally supportive of the paper, its current issues let me lean towards rejection unless the problems will be resolved so the claims can be assessed. My suggested improvements could potentially be done during the author response phase and I'm happy to revise my score if the authors provide solid evidence for their claims.\n\n### Update after discussion period\nThe authors have thoroughly addressed my concerns and performed additional experiments. I am confident that after the promised revisions the paper will be very worthwhile publishing at NeurIPS and changed my score to 7 accordingly.", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "## Strengths\n+ Simple and novel self-supervised representation learning method\n+ Potentially improves decoding of task-relevant variables\n\n\n## Weaknesses\n- Not clear if claims are fully supported by the data\n- Possible selective reporting/cherrypicking of results\n- Writing not always clear\n\n\n## Detailed comments on weaknesses\n\n### Selective reporting\n\nUnfortunately, the data analysis is not done very rigorously (no statistics at all) and on subsets of the data. Also, I can't help but getting the impression that some results that don't fit as nicely into the story as others are conveniently moved to the supplement. \n\n(a) The authors perform the analyses shown in Fig. 3D and Table 2 only on the Chewie-1 dataset. Given the substantial variation across datasets shown in Fig. S2, Table 1 and Table S1, this is a problem – especially since it's not stated when and why this dataset was selected. \n\nAction items: All analysis that are used to draw conclusions should be done on all datasets and a measure of uncertainty (e.g. standard error of the mean/confidence intervals across datasets) should be reported, so the reader can assess to what extent the reported results can be attributed to randomness in the data vs. the method proposed.\n\n(b) The reader would get a better impression of the reliability of the method (or the variability across datasets) if the examples in Fig. 3B+C were not cherry-picked to show the circular arrangement of the latent space, but rather show the datasets of the first days for both monkeys. This is also worrisome as the authors make a claim about the method identifying the circular arrangement in the latent space earlier in Fig. 2 and it is not clear to me what aspect of the loss function/training procedure would encourage such behavior (see below for a comment on this point).\n\nAction item: show more representative data in Fig. 3 B+C.\n\n(c) In the section on decoding, the authors show only Table 1, the accuracy for reach direction, to make the point that their method is superior to earlier approaches. However, Table S1, which shows that the method does not lead to an improvement in decoding time, is relayed to the appendix without even stating the results in the main text. Moreover, Table S2 is lacking the stronger baselines, so presumably the proposed method would look even worse in this respect if this table was complete.\n\nAction items: Table S1 should be complete with all baselines, moved into the main text and the paper's claims be adjusted to be consistent with the full set of results. If space is a problem, the authors could summarize reach direction and time decoding in one table by showing only the averages across datasets in the main paper and providing the full results per dataset in the appendix.\n\n\n### Further comments (action required)\n\n- Line 311ff. The difference between (i) and (ii) is not clear to me. Could you clarify what removing \"the original reconstruction term\" means? If there is no reconstruction term, how is the decoder learning anything? Or is it meant such that (i) includes both reconstruction terms from Eq. (3) while (ii) reconstructs only the content-swapped sample?\n\n- Multi-task disentanglement score (line 158ff.). This section needs a substantial revision. I did not understand what's happening here, despite several reading attempts. The variables y_c and y_s are not defined anywhere, it's not clear to me how they relate to z_i and what variances are computed in lines 166–168.\n\n- Fig. 3D: bars should start at zero, not at 0.3\n\n- Fig. 4D: The y axis is labeled \"VAE\". Is this the same beta-VAE as everywhere else or a vanilla VAE (beta=1)? Given the selective analyis at other places, I'd like to make sure that this figure is not using a strawman as baseline.\n\n\n### Minor comments (action recommended, but not required)\n\n- The meaning of content vs. style could be discussed a bit better. I am not sure if the \"style\" is a well-defined concept here. The authors seem to equate it with temporal dynamics, but the temporal dynamics are clearly also encoded in the content space (as can be seen quite well in Fig. S2). So while the content space enables good decoding of the reach direction, time can be decoded from both subspaces. In that sense, the disentangling isn't very clear. \n\n- Evaluation metric for decoding: Why turning it into classification if you treat it as regression? Instead of applying two different sets of boundaries for acc and delta-acc, why not just report mean squared angular error as a metric of reconstruction accuracy?\n\n- Fig. 2: I understand that the four classes lie on a circle, but it's not clear to me why this arrangement should be retained in the latent space. The authors seem to think of this as important/desirable and their method retains the arrangement while the beta-VAE doesn't. However, it's not clear to me what aspect of the loss function/training procedure would encourage such behavior and whether it's not just a random effect of initialization. If I'm not mistaken, if you pick four locations at random, it should be 50:50 whether they maintain the circular arrangement or not.\n\n- Line 296 \"rather than noise\": I am not sure if noise is the right term here. Presumably a substantial fraction of neuronal variability is due to variability in the exact arm trajectory across trials, which isn't really noise.\n\n- Line 297ff. I am not sure what's the point of this paragraph. As the authors state themselves at the end, if you want to decode, you're better off just using the SwapVAE latent than training a supervised classifier from generated samples. I think this paragraph can be removed to make space for more important points and showing some results that are currently hidden in the appendix (see above).\n\n- Line 267ff. This paragraph seems like pure speculation. Given the amount of variability across datasets and the fact that it's just a single animal, this results could be due to many things, the delay period being only one of them. I think the authors would do themselves a favor by not jumping to a premature conclusion here. So again, this paragraph could be cut and replaced by more important material.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Limitations could be addressed more openly, e.g. the limits on disentanglement w.r.t. decoding time.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}, {"review_id": "_jAk6hNaQPl", "reviewer": "Reviewer_hEnc", "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.", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The description of the method and results are quite clear. However, I have questions about how well the proposed method compares to prior work. I will note that I am not very familiar with the related methods and do not have a good judgment on the significance of this contribution compared to that work. But I can see how a method for separating neural representations into 'content' and temporal dynamics (style) is clearly useful for neuroscience. On those terms I think the contribution is well worth accepting this paper. I also have some questions about the latent 'style' vector for the neural data.\n\n- The decoding results (Fig 3D) seem to indicate that the supervised model does quite well at capturing elapsed time, performing better than the latent style space. From my understanding this is the case even though the supervised model is only trained to predict the reach locations. Can the authors comment on why this occurs?\n    - This seems related to the fact that the content space also does reasonably well at decoding time. Both pieces of evidence suggest that disentanglement with the proposed approach is incomplete. This seems like a limitation of the current work -- do the authors have a proposal for how to address this?\n- Table 1 shows accuracy in classifying reach direction for the proposed method and related approaches. It is not clear whether the differences are significant, especially because there are no reported error bars or statistical tests. (I do not see them in the Appendix either.) This could be easily done by looking at the accuracy over different validation folds, or over random seeds.\n- Can the authors comment on how the 'style' latent space which captures temporal dynamics might be affected by temporal noise in neural activity? This method seems like it could be useful to create generative models of data where the sources of temporal variability are less well-understood, but the implications here are not clear to me.\n\nMinor comments\n- The monkeys are introduced by initials in the text, but referred to by their full names later. I would recommend standardizing this to avoid confusion when reading the figures (e.g. either write Monkey C everywhere, or Chewie everywhere).\n- Fig 2: please include descriptions of the solid black and dotted lines in the legend. \n- Table 1: would be useful to highlight/bold the best model(s) for each dataset\n- Fig 4A/B: there are no x-axis tick labels\n- Fig 4C/D: axes are not well-labeled to distinguish the two plots from each other. Could write \"reach RMSE\" and \"time RMSE\" on the plots", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "My questions about possible limitations are included in the main review above.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "2"}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper present a method for learning disentangled latent variables inherent in brain activity, and illustrate their methods with synthetic and real monkey data from a reaching task. After the initial response from the reviewers, the authors responded to their concerns extensively, including running additional experiments and addressing the concerns of the reviewers. Given that all the reviewers expressed that their concerns were addressed by the responses, I recommend acceptance and ask the authors to modify the paper to include these changes and address all the comments from the reviewers.

Decision: Accept (Oral)