Venue: NeurIPS 2021 Conference

Reviews:

final_score: rating: 6: Marginally above the acceptance threshold
review_content: summary: This paper details a dimensionality reduction and dynamical modeling technique that partitions the neural variability into behaviorally-relevant and behaviorally-irrelevant factors. This idea was introduced and explored in preferential subspace identification (PSID), a reformulation of SSID (subspace identification), which models the latent dynamics as evolving linearly. This paper extends that to the nonlinear regime, using recurrent neural networks to model the evolution of the latent variables. The authors use the LFADS framework to model the temporal evolution, while partitioning the state into behaviorally-relevant and behaviorally-irrelevant factors, and not allowing for the behaviorally-irrelevant state to influence the behaviorally-relevant state.
main_review: The authors may have severely misrepresented PSID. From Equation (4) of the Methods in Sani et al. (2021), PSID actually seems to allow the behaviorally-irrelevant states as having an effect on the behaviorally relevant states, and not in fact a block diagonal transition matrix A as the authors have mentioned (Question: is this why their implementation of PSID fails?). TNDM does not allow this influence to occur, which is a significant departure from the assumptions in PSID, and in fact, the assumptions for TNDM require a strong motivation. The assumption that the behaviorally-irrelevant dynamics are completely disjoint from the behaviorally-relevant dynamics implies that there are two completely independent processes occurring in the neural activity that require a separate set of dynamics to explain. On the contrary, papers such as Sussillo et al. (2015) and Russo et al. (2018) show that if one trains a network to produce behavioral outputs in the context of motor control, the network states can actually recapitulate neural activity, traditionally thought of as behaviorally-irrelevant, but that are in fact necessary to produce behavior. Moreover, the authors mention that PSID is hindered since it cannot incorporate lags between the neural activity and the behavior. However, the state z is a latent representation that can in fact accumulate past neural activity to lead to current behavior, and thus can incorporate lags between the two representations.

The authors find that the reconstruction of the firing rates using TNDM is dependent almost purely on the number of total factors that are used, regardless of how many of them are constrained to be behaviorally-relevant (Figure 2C). Please discuss. It would be interesting to include in Figure 2A,C the behavioral and neural predictions for # of behaviorally-irrelevant factors = 0.

Please explain why TNDM achieves a lower firing rate RMSE than LFADS with the same number of factors (Figure 2C) for # total factors > 4.

Please show the post-hoc behavior reconstruction using the LFADS factors in Figure 3 using the 4 factors. The behavior R^2 was high, so the reconstructions should be visually comparable to TNDM.

Please discuss and / or show the post-hoc partition of LFADS factors into behaviorally-relevant and behaviorally-irrelevant by decoding behavior from the factors (as done for PSID in this paper).

The paper would benefit from a better discussion of the structure and role of behaviorally-relevant vs. behaviorally-irrelevant dynamics.

The authors refer to the behaviorally-irrelevant factors and networks as simply ‘irrelevant’. I would encourage them to use the full term ‘behaviorally-irrelevant’ instead.

Originality: The paper is an extension of LFADS using ideas in PSID. 
Quality: The paper could lead to interesting insights into the structure of behavioral vs. neural activity. However, in its current form, it is unclear what it is adding to the realm of different dimensionality reduction models of neural activity.
Clarity: The paper is written very clearly, and the figures are of high quality.
Significance: While the paper explores an interesting idea, it would be much improved with proper motivation of their model structure. 
limitations_and_societal_impact: Yes
review_title: null

final_score: rating: 7: Good paper, accept
review_content: summary: In this paper authors provide a novel architecture for modelling neural dynamics with the aim of separating dynamics that are related to behaviour from the ones that are not. Using amortized inference, observed neural activities are used to set the initial state of two RNNs, one for relevant dynamics and one for irrelevant dynamics, which in turn predict behavioural and neural data. The results show that compared to baseline models the model is indeed able to predict behavioural responses, while maintaining prediction ability for the neural data.
main_review: The paper is clearly written and the model, contributions, and the results are well explained. The model itself is also well designed for the purpose and potentially provides important contributions to the literature. 

I have two concerns as explained below.

I am rather confused about why neural dynamics at time t are modelled to affect future behaviours at c > t directly (in addition to the behaviour at time t). Normally, I would assume that current neural dynamics at time t influence neural activity at time c > t, which in turn affect behaviour at time c. But the current setup implies that neural activities at time t had to directly affect behaviour at time c without mediation of neural responses at time c. It seems to imply that the RNN network was unable to capture the whole neural dynamics over time, which seems a serious limitation. 

From the results, it's hard to evaluate whether irrelevant RNN is contributing at all to the prediction of neural activities; i.e., it could be the case that the whole neural dynamics are explained by the relevant RNN; is there any evidence supporting the role of relevan RNN in predicting neural activities? Similarly, how much is the contribution of the irrelevant RNN in predicting neural activities? Is there any evidence that not all the neural dynamics is explained by the irrelevant RNN?



limitations_and_societal_impact: They are adequately discussed.
review_title: null

final_score: rating: 6: Marginally above the acceptance threshold
review_content: summary: In this work, the authors propose a state-space model and inference algorithm that aims to recover the latent nonlinear dynamics from sequential data with particular interests in neuroscience. The proposed model inherits the latent factor analysis via dynamical systems (LFADS) and incorporates the advantage of disentangling behavior relevant/irrelevant dynamics in the latent space as preferential subspace identification (PSID) does. It overcomes the limitation of linear dynamics and more faithfully captures the behavior. The proposed method was tested on the data from a monkey reaching task.
main_review: Originality: The proposed model uses the same structure as LFADS with extra conditional Gaussian observation model for behavior variables of the similar idea as that of PSID. The objective function modifies the ELBO with terms that regularize the entanglement and behavior likelihood. Previous work are discussed and compared. The contributions are clearly stated.

Quality: There're more have to be done to support the claims.
- Only one empirical result on the reaching task was shown. It isn't adequate to evaluate the strength and weakness of the method. Synthetic data and other tasks are needed.
- The work is focusing on the reaching task overly. TNDM seems to be tightly tailored for the task. These particular settings prohibits the method from generalizing to other task and weaken its easiness for use. E.g. the choice of initial conditions is needed. It is also mentioned that a random initialization might lead to weird dynamics; The causal relationship is not always true for behavior and neural activity, considering feedback scenarios. 
- "to learn low-dimensional dynamics" is ambiguous. In TNDM, the dynamical system is modeled by a high-D RNN as well as LFADS. Only the latent factor $z$ are low-D. 
- The power of SSM like LFADS is not only to recover the latent trajectories but also to learn the dynamical system (velocity/phase portrait/simulated trajectories). This should be also evaluated. 
- It is expected that TNDM outperforms LFADS on behavior encode at small number factors, because of imperfect encoding of behavior variables and other internal dynamics. 
- It is also important to evaluate the disentanglement, e.g. the decoding performance of irrelevant dynamics.
- PSID used Gaussian filtered spikes. This made the data and likelihood are both different from that of TNDM. Since TNDM is not limited to Poisson for spikes. It's more fair to fit TNDM to the same data and likelihood as PSID. 

Clarity: The submission is well organized. The writing can be improved.
- The authors spent a lot for the details of LFADS and PSID but not say much about how PSID separates dynamics which is a key point to the work. 
- In Eq(2), the encoder network is only a function of the spike $x$ without behavior $y$? and why not?
- It is unclear that how the initial conditions $\Phi_r$ and $\Phi_i$ are parametrized in Eq(2).
- Where are $\mathbf{g}_{0r}$ and $\mathbf{g}_{0i}$ defined? These bold symbols have different indexing order from their plain counterparts. Are they vectors?
- How is the mean of the sample cross-correlations between $\mathbf{g}_{0r}$ and $\mathbf{g}_{0i}$ computed?

Significance: The results show  that TNDM can capture, from neural recordings, the latent trajectories that faithfully reconstructs the neural activity and behavior. It could help to separate the latent factors into behavior-related and internal dynamics. The demonstration could be improved though.

Minor comments:
- Why to trim the trials to minimum? Varying length trials should not be a prolem.
- Why the covariance of behavior is fixed at the identity matrix?
limitations_and_societal_impact: The authors discussed the limitations and broader impact.
review_title: null

final_score: rating: 5: Marginally below the acceptance threshold
review_content: summary: The authors introduce a new sequential VAE that simultaneously models behavioral dynamics and neural activity for experimental neuroscience data. The model includes partitioned latents which map to either neural activity and behavior, or just behavior.  The model uses RNNs to describe the evolution of each of these latent states (similar to LFADS), and the model is trained by optimizing an objective function which contains separate variational distributions for behaviorally relevant and irrelevant latents, as well as a penalty to prevent entanglement of the latent states. The authors test their model on neural data recorded from Monkey V1 during a reaching task.
main_review: I think this work is an interesting application and extension of sequential VAEs for neuroscience. The writing is clear, the model is novel, and I believe the work would be of interest to the neuroscience community. My primary concern is that this paper, in its current form, may not be appropriate for NeurIPS.  My reason for this is that there is not a lot of mathematical or theoretical development -- specifically, the objective function is weakly motivated, and there is minimal to no attention given to inference. I elaborate on each of these pieces below.

Weak motivation of an objective function: The introduction of the ELBO in (3) feels like it naturally follows from the previous explanation and description of the model.  However, immediately after introducing this, the authors say that the variational distributions can become entangled and thus introduce the cross-correlation penalty (4). Though they show some qualitative comparisons with and without this penalty in the appendix, the motivation feels ad hoc and quickly thrown in. As this is a fundamentally new objective, I think elaboration is warranted.

I'm curious if the authors could bake this type of regularization into the model itself, as opposed to just modifying the objective function and exploring the qualitative changes. Perhaps one could impose a prior on the latent state(s)? 

Additionally, the authors introduce a parameter balancing the behavioral and neural objectives (4) with minimal justification or motivation.

Lack of attention given to inference: The current inference section is mostly devoted to specifying the objective function, not actually how inference is done. I am not familiar with inference procedure in sequential VAEs. Does it use the reparameterization trick? Is there anything specific in this model that can be done to perform inference quickly or efficiently?

I do think that overall this work is interesting and well-done.  However, given that the work is a somewhat natural combination pieces of existing models, additional attention should be devoted to expanding on partitioned seq VAEs in general, some more fleshing out of inference, or better theoretical justification/exploration of the objective function, especially for NeurIPS. 

A few other smaller concerns:

The authors fail to mention statespace models that more have flexible than linear dynamics - namely GPFA. See 

Zhao, Yuan, and Il Memming Park. "Variational latent gaussian process for recovering single-trial dynamics from population spike trains." Neural computation 29.5 (2017): 1293-1316.

Keeley, Stephen, et al. "Efficient non-conjugate Gaussian process factor models for spike count data using polynomial approximations." International Conference on Machine Learning. PMLR, 2020

The authors talk about the importance of time-lags in the model, but there are no results supporting this. 
limitations_and_societal_impact: None
review_title: null

Meta-review: This paper proposes a new sequential VAE that simultaneously models behavioral and neural time series. The reviewers agree that it provides meaningful contribution to neural data analysis. The discussion among the reviewers and the authors was very constructive and productive, and the overall writing of the final version should reflect the nuances raised through the process and results from the additional experiments.
Decision: Accept (Poster)