Venue: NeurIPS 2021 Conference

Reviews:

final_score: rating: 7: Good paper, accept
review_content: 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. 
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. 
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):  
-	Comparison 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. 
-	For 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?) 
-	In 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).  
-	The 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. 
-	When 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. 
-	The 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’?
review_title: null

final_score: rating: 7: Good paper, accept
review_content: 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.
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.

- 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?
    - 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?
- 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.
- 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.

Minor comments
- 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).
- Fig 2: please include descriptions of the solid black and dotted lines in the legend. 
- Table 1: would be useful to highlight/bold the best model(s) for each dataset
- Fig 4A/B: there are no x-axis tick labels
- 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
limitations_and_societal_impact: My questions about possible limitations are included in the main review above.
review_title: null

final_score: rating: 7: Good paper, accept
review_content: 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. 

While 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.

### Update after discussion period
The 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.
main_review: ## Strengths
+ Simple and novel self-supervised representation learning method
+ Potentially improves decoding of task-relevant variables


## Weaknesses
- Not clear if claims are fully supported by the data
- Possible selective reporting/cherrypicking of results
- Writing not always clear


## Detailed comments on weaknesses

### Selective reporting

Unfortunately, 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. 

(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. 

Action 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.

(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).

Action item: show more representative data in Fig. 3 B+C.

(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.

Action 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.


### Further comments (action required)

- 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?

- 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.

- Fig. 3D: bars should start at zero, not at 0.3

- 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.


### Minor comments (action recommended, but not required)

- 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. 

- 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?

- 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.

- 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.

- 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).

- 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.
limitations_and_societal_impact: Limitations could be addressed more openly, e.g. the limits on disentanglement w.r.t. decoding time.
review_title: null

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)