Paper: Modeling state-dependent communication between brain regions with switching nonlinear dynamical systems (ICLR-2024)

Reviewer comment (Writing):
The -2th line of page 3, $f_z^{kk}$ should be $f_{kk}^z$.

Author response:
**1. Main weakness: typos, incomplete legends and incorrect references affecting clarity of the paper.** We deeply apologize to the reviewer for the typos, incomplete legends, and incorrect figure references pointed out by the reviewer as the main weakness of our paper. We understand that these oversights deeply affect the clarity and accessibility of our work, as well as the overall impression of the reader of its quality. We also realize these presentation errors were a practical barrier to understanding and evaluating our experiments, and we again deeply apologize for this. **We have carefully edited and fixed** all typos and other mistakes, and we hope the paper is now clear, and thus can be better assessed.

Reviewer comment (Presentation):
Fig. 1 is a summary of the whole generative model, but it is informal to not mention it in the main text.

Author response:
**1. Main weakness: typos, incomplete legends and incorrect references affecting clarity of the paper.** We deeply apologize to the reviewer for the typos, incomplete legends, and incorrect figure references pointed out by the reviewer as the main weakness of our paper. We understand that these oversights deeply affect the clarity and accessibility of our work, as well as the overall impression of the reader of its quality. We also realize these presentation errors were a practical barrier to understanding and evaluating our experiments, and we again deeply apologize for this. **We have carefully edited and fixed** all typos and other mistakes, and we hope the paper is now clear, and thus can be better assessed.

Reviewer comment (Presentation):
Page 6, 2nd line should be Figure 2(c).

Author response:
**1. Main weakness: typos, incomplete legends and incorrect references affecting clarity of the paper.** We deeply apologize to the reviewer for the typos, incomplete legends, and incorrect figure references pointed out by the reviewer as the main weakness of our paper. We understand that these oversights deeply affect the clarity and accessibility of our work, as well as the overall impression of the reader of its quality. We also realize these presentation errors were a practical barrier to understanding and evaluating our experiments, and we again deeply apologize for this. **We have carefully edited and fixed** all typos and other mistakes, and we hope the paper is now clear, and thus can be better assessed.

Reviewer comment (Presentation):
Page 6, 3-5 lines and Figure 2 — Where is the $R$ and $u$ in Figure 2?

Author response:
**1. Main weakness: typos, incomplete legends and incorrect references affecting clarity of the paper.** We deeply apologize to the reviewer for the typos, incomplete legends, and incorrect figure references pointed out by the reviewer as the main weakness of our paper. We understand that these oversights deeply affect the clarity and accessibility of our work, as well as the overall impression of the reader of its quality. We also realize these presentation errors were a practical barrier to understanding and evaluating our experiments, and we again deeply apologize for this. **We have carefully edited and fixed** all typos and other mistakes, and we hope the paper is now clear, and thus can be better assessed.

Reviewer comment (Presentation):
In Figure 2(b), what is true and what is the dashed curve? There is no clear legend.

Author response:
**1. Main weakness: typos, incomplete legends and incorrect references affecting clarity of the paper.** We deeply apologize to the reviewer for the typos, incomplete legends, and incorrect figure references pointed out by the reviewer as the main weakness of our paper. We understand that these oversights deeply affect the clarity and accessibility of our work, as well as the overall impression of the reader of its quality. We also realize these presentation errors were a practical barrier to understanding and evaluating our experiments, and we again deeply apologize for this. **We have carefully edited and fixed** all typos and other mistakes, and we hope the paper is now clear, and thus can be better assessed.

Reviewer comment (Presentation):
There are other mismatches in Figure 2 and etc.

Author response:
**1. Main weakness: typos, incomplete legends and incorrect references affecting clarity of the paper.** We deeply apologize to the reviewer for the typos, incomplete legends, and incorrect figure references pointed out by the reviewer as the main weakness of our paper. We understand that these oversights deeply affect the clarity and accessibility of our work, as well as the overall impression of the reader of its quality. We also realize these presentation errors were a practical barrier to understanding and evaluating our experiments, and we again deeply apologize for this. **We have carefully edited and fixed** all typos and other mistakes, and we hope the paper is now clear, and thus can be better assessed.

Reviewer comment (Writing):
There is no "." in the last sentence of section 4.1.

Author response:
**1. Main weakness: typos, incomplete legends and incorrect references affecting clarity of the paper.** We deeply apologize to the reviewer for the typos, incomplete legends, and incorrect figure references pointed out by the reviewer as the main weakness of our paper. We understand that these oversights deeply affect the clarity and accessibility of our work, as well as the overall impression of the reader of its quality. We also realize these presentation errors were a practical barrier to understanding and evaluating our experiments, and we again deeply apologize for this. **We have carefully edited and fixed** all typos and other mistakes, and we hope the paper is now clear, and thus can be better assessed.

Reviewer comment (Experiments):
Forward Euler is not appropriate for generating the data. Although the different curves in Figure 2(b) are not understood, the numerical error amplifies significantly because of forward Euler.

Author response:
**Regarding question 4**, this is an important question, we apologize this point was left unclear and that we forgot to mention it in our response. As the reviewer suggests, the forward Euler is problematic because it will drift from the true solutions as a function of the step size, and this was indeed contributing to the issues with the earlier simulation. When generating the new, longer simulation **based on the reviewer’s feedback, we simulated the system with a much smaller step size**, dt=0.001, generating 75,000 steps for the numerical integration for stability, and then **downsampled the resulting timeseries**, with one sample for every 500 steps -- thus producing a final simulated time series of 150 timesteps for each trial. We found this improved the results in addition to the new objective, and as can be **seen in Figure 2**, we are able to recover the latent trajectories of the system nearly perfectly. We thank the reviewer again for this earlier comment, and **we have updated the text to clarify this issue**. Finally, we note that while this is still also a simplistic approach to integration of these equations, we decided not to resort to more sophisticated approaches, as the Lotka-Volterra equations themselves were not our object of study, and our intention was simply showing recovery of overall dynamics and communication in this simulation (as a model of oscillatory communication) -- and therefore may not require yet higher precision. We hope this is clear and satisfactory and thank the reviewer again for this helpful earlier question.

Reviewer comment (Experiments):
Only one state switch is allowed in experiment 4.1 and it has to be within 40-60. This generating process is too simple and not complicated enough to validate the model.

Author response:
**3. Single discrete switch in Lotka-Volterra simulation constrained in time is too simple to validate the model; clarity in the simulation figures.** We thank the reviewer for this insightful comment on the limitations of this simulation. **We have added results of an updated version of the simulation** which uses multiple switches , between 1 and 3 switches picked randomly on each trial. We **have also changed** the potential switching window to be uniformly sampled between timepoints 10-140 in a 150 timepoint long simulation. We present this new result **in Figure 2**. Additionally, we apologize for any confusion in the legend and traces in the original presentation of these results and **have updated these figures** to have clearer traces and legends.

Reviewer comment (Presentation):
Overall, there are too many errors in this paper. Most of the experiments cannot be understood due to wrong text-figure assignments, no legend, etc.

Author response:
**1. Main weakness: typos, incomplete legends and incorrect references affecting clarity of the paper.** We deeply apologize to the reviewer for the typos, incomplete legends, and incorrect figure references pointed out by the reviewer as the main weakness of our paper. We understand that these oversights deeply affect the clarity and accessibility of our work, as well as the overall impression of the reader of its quality. We also realize these presentation errors were a practical barrier to understanding and evaluating our experiments, and we again deeply apologize for this. **We have carefully edited and fixed** all typos and other mistakes, and we hope the paper is now clear, and thus can be better assessed.

Reviewer comment (Experiments):
What about other methods for section 4.1 and 4.2? Are other methods in Table 1 not applicable to experiment 4.1 and 4.2? For example, the generative model of the Lokta-Volterra model is nonlinear. Authors show that MR-SDS is able to capture the nonlinear evolution of the latent process $x_t$. But for SLDS, for example, $x_{t+1} \sim \mathcal N(Ax_t,\Sigma)$. It seems like $Ax_t$ is a linear process, but since there is a Gaussian distribution, it is actually a nonlinear model. Regardless of how “linear/nonlinear” is defined, why not test whether SLDS can learn the Lokta-Volterra latent?

Author response:
**4. Additional comparisons with other models on simulation experiments.** We thank the reviewer for pointing out this shortcoming in our presentation. **We include additional comparisons** on simulations in sections 4.1, 4.2, in *Figures 2 and 3* respectively. We compare with models the reviewers mentions, such as mr-LDS, mr-SLDS, and mr-rSLDS, visualize the latents recovered for each model, as well as quantify the relative performance.

Reviewer comment (Evaluation):
Is Table 1 the held-out likelihood? Where is the variance? Are the likelihood results of MR-SDS significantly better than others? Is it likelihood or log-likelihood?

Author response:
**Regarding weakness 3** we apologize if this remained unclear in our responses. Table 1 shows the mean square error (MSE) on a particular measure of held out prediction, co-smoothing, popularized in the recent Neural Data Benchmarks to compare across neural data models (https://arxiv.org/abs/2109.04463). In co-smoothing, a subset of neurons (here 25%), are dropped out on held out test trials, and then predicted by the model using the remaining neurons in those trials. Thus, both trials and neurons are held out. This is a challenging target because it requires the model to learn to represent relationships amongst the neurons in its latent space so it can predict unseen neurons on unseen trials, as opposed to simply compressing and reproducing held out entire test trials. In Table 2 in the Appendix, we show results on a sweep of model hyperparameters for a related, even more challenging target, which we refer to as 1 step generative co-smoothing. In this form of co-smoothing, the model cannot generate predicted held out neurons directly from inferred latents, and has to pass the inferred latents through one timestep of the generative dynamics before passing through emissions, ie. $E_{q(x|y)} [p(y_t|f(x_{t-1}))]$. Thus, this target requires the generative dynamics to be accurate in addition to the latent inference, and this evaluation target is behind our updated training objective in Equation 14. **We have updated the captions for Tables 1 and 2 to make this clearer**, we also refer the reviewer to **section A.5 and Figure 9 in our updated Appendix** which gives more detail on co-smoothing and 1-step generative co-smoothing.


Finally, we realize MSEs can be hard to compare and interpret, our choice of this metric was due to use of a Gaussian observation model for calcium imaging data, and the simplicity of computing MSE on saved predictions without loading the models - we will plan to include log likelihood instead as a more intuitive measure and apologize for not getting a chance to do this during the discussion period. (We also plan to include variance of each result as the reviewer suggested). To address the question of significance simply for the moment, **we added a plot (Figure 7 in the updated Appendix)** showing a Gaussian fit to the LDS model scores **along with the score and p-value for our model, which is highly significant** (mesoscope=2.209e-99, widefield=5.345e-4).

Reviewer comment (Evaluation):
There is no significant benefit of the proposed model over others, in terms of both prediction on future data and latent estimation.

Author response:
**5. Benefits of our model for future prediction and latent estimation over other models.**
We thank the reviewer for calling attention to the fact that the relative benefit of our model over others was not sufficiently clear in the paper, despite achieving better performance on the presented benchmarks. We consider two aspects of this shortcoming, (a) the relative quantitative performance, and (b) the relative qualitative fits and scientific interpretation. We do 3 things to address both of these aspects. First, **we have introduced a modified objective in training the network**, which forces coupling between the generative dynamics and emissions networks during training (equation 14 in the updated paper) -- **resulting in higher accuracy of recovered latents, dynamics and messages** in our simulation experiments, as shown in **the updated Figures 2 and 3**.

Reviewer comment (Experiments):
The experiment part needs clarification. Section 4.2 simulates a multi-region decision-making task, but there is no comparison of MR-SDS with other methods such as LDS, SLDS, and rSLDS. Besides, the author mentions their method, MR-SDS, is able to learn the nonlinear dynamics in this complex case, but there needs to be a figure to visualize such latent dynamics. Similarly, section 4.3 simulates a multi-region system with three interacting and high-dimensional RNNs, but there needs to be a figure to support their claim.

Author response:
**2. Clarification of experiments section, additional comparisons with other models on some of the simulations.** We thank the reviewer for pointing out this shortcoming in our presentation. First, **we have moved the key results into the main text** for section 4.2, to help support and visualize the result on this experiment. Second, **we include additional comparisons** on these simulations with other models the reviewers mentions, such as mr-LDS, mr-SLDS, and mr-rSLDS, in **Figures 2 and 3**. We visualize the latents recovered by mr-LDS and mr-rSLDS for the Lotka simulation, and for mr-rSLDS (the most advanced of these models) on the Double Well simulation for qualitative and quantitative assessment.

Reviewer comment (Evaluation):
The comparisons with LDS, SLDS, and rSLDS are not enough. This paper only compares their prediction/reconstruction performance on calcium imaging data. It would be better to show the proposed method, MR-SDS, could capture more meaningful latent dynamics than other commonly used models, which may be more important than prediction performances in terms of neuroscience research. For example, could you provide some cases where MR-SDS captures some nonlinearity in the neural recordings while rSLDS cannot?

Author response:
**4. Qualitative advantage of learned MR-SDS latents over rSLDS.** The reviewer makes a very important point - improved prediction performance is an important measure but may not be key to producing improved scientific interpretation, and thus it is important to demonstrate a qualitative advantage in latents learned by MR-SDS for neuroscience research. We agree that more work needs to be done in this direction and apologize that we were unable to provide additional analysis in the discussion time.

Reviewer comment (Evaluation):
No standard deviation (variance) in Table 1.

Author response:
**3. Standard deviation of results for benchmarks.** We thank the reviewer for this suggestion and apologize that we were unable to update the table in the discussion time.

Reviewer comment (Writing):
Many typos and wrong figure index/reference. E.g., in section 4.1 and the second line of page 6, it should be "Figure 2(c) shows ....". In section 4.1, there is no reference to Figure 2(d), etc.

Author response:
**1. Fixing typos and incorrect figure reference.** We deeply apologize for the typos and incorrect references pointed out by the reviewer, and we understand that these deeply affect the clarity and accessibility of our presentation, as well as the overall impression of the reader of the work. **We have carefully edited and fixed** all typos and other mistakes, and we hope the paper is now clear.

Reviewer comment (Experiments):
Are there any insights as to why you choose Transformer instead of some simpler and faster deep neural networks like RNNs?

Author response:
**5. Choice of Transformer encoder vs RNN.** This is an important question - in developing the MR-SDS model we indeed tested an RNN based encoder (a stochastic, bi-drectional RNN, or Bi-RNN specifically, similar in spirit to the one used in LFADS), which did not perform as well as the Transformer, but we did not include those results in the submitted paper due to space constraints. **We have added an additional result in Table 1 showing performance of the Bi-RNN encoder** compared to the transformer. We note that because we constrain our use of the Transformer (or RNN) to the encoder alone, but use simple feed forward or MLP networks for the generative dynamics, we are not directly concerned with the complexity or interpretability of the encoder network, and thus can benefit from the high trainability and expressivity of the Transformer without complicating the generative model and resulting scientific interpretation. We have added additional detail on this last point to the discussion section.

Reviewer comment (Experiments):
The method is powerful by design, leveraging multiple nonlinear DS per region. However, the flexibility comes with issues, most prominently the choice of hyperparameter values, such as the dimensionality of the system and, critically, the number of switching states. In the simulation experiments, the model was given the ground-truth number of states, but it is unclear how the inferred systems would look when starting with a different number (e.g., for the switching Lotka-Volterra system). This is further an issue with the real neural data, where one does not have access to the ground-truth, and it is arguable whether such a number exists.

Author response:
We thank the reviewer for calling attention to this important issue. As the reviewer mentions, the expressivity and flexibility of a model with switching nonlinear dynamics and nonlinear emissions requires well-motivated choices of hyperparameters, specifically the choice of latent dimensions and number of discrete states. For the number of dimensions, we purposely choose 2 dimensions to represent each latent space across all simulations and datasets, reflecting the perspective that we desire to learn dynamics that are easy to visualize, ie can be visualized in 2d. This is so that we can avoid using a dimensionality reduction method such as PCA to only partially visualize and understand the dynamics of the model post-hoc, after it has been fit. However, it is important to understand the performance consequences of this design choice - therefore we **include in Table 2 in the Appendix** a sweep of hyperparameters, showing the single step generated co-smoothing performance across 2-5 dimensional latent models fit to the neural data, as well as for 1-3 states. This benchmark is a more challenging one, as it requires co-smoothing using a single time step of generated dynamics under the generative model (vs going directly from inferred latents to observations). These results show that while adding additional dimensions above 2 to the latent space improves performance, the 2 dimensional model compares well (and is easier to visualize and interpret). Similarly, adding an additional discrete latent state improves performance, but is comparable.

Reviewer comment (Experiments):
There is a major question of whether the existing results are robust to different choices of the number of switching states, latent dimensions, etc., especially in the simulation experiments when using an incorrect number, as well as how results change when applied to experimental data. For a practicing neuroscientist using such a method, one would expect robust and consistent results independent of hyperparameter choices.

Author response:
This is an important point, and we regret that we weren't able to address it fully. First, we note that hyperparameter choices in many models, including ‘simple’ ones, can lead to different scientific interpretations - for example, fitting a clustering model such as a hidden markov model (HMM) or K-means to data with different numbers of states can result in ‘lumping’ or ‘splitting’ of eg different neural activity or behavioral patterns, with potentially different scientific conclusions. There is a similar issue with the number of discrete states in our model, and **we have updated our discussion**, as per the reviewer’s suggestion. Second, **we include a simple analysis of robustness** of inferred message representations, by showing the across trial mean message norm profiles for the mesoscope data across models in the parameter sweep (Figure 6 in the Appendix), showing these look similar across models. We realize that more thought and analysis should go into assessing scientific robustness, we focus here on messages since we believe they are a key target for downstream scientific analysis.

Reviewer comment (Evaluation):
The inferred systems from the real data (Figures 3 & 4) do not appear compelling. In Figure 3B, the flow fields in a state all look like linear 1D flows. It is unclear if this is meaningful, if this could have been combined with the first state, or what happens if one uses a larger number of states—are all the states after state 1 simple linear-like flows? Interpretation of these patterns remains uncertain.

Author response:
We thank the reviewer for their thoughtful and nuanced comments on the matter of scientific interpretation of the neural datasets using the model. We agree that we have more work to do in this respect, and regret that we couldn’t include further analysis which is a work in progress. A small comment on the dynamics gradient fields plotted in Figures 4 and 5, those fields are only the local component of the dynamics acting on each of the regions, whereas the full induced dynamics field is a sum of the additional terms. To compute these, it is possible to average the empirical incoming gradients from all other regions on each point on the grid, we didn’t have a chance to include this regrettably - but the induced fields are likely much less linear.

Reviewer comment (Presentation):
The figures could use more explanation. For example, what are the main occupied states? Why is the same trajectory plotted in both states but with different portions highlighted in blue in Figure 3B? How should one interpret Figures 3C and 4D? Various plots lack labeled axes (e.g., Figures 3D, 4B/C), and if they represent time, is it in milliseconds?

Author response:
We thank the reviewer for calling attention to specific small typos and places where the paper’s clarity could be improved. We regrettably agree that our results section was lacking in the clarity and polish of the introduction and model presentation, and **we have carefully edited and updated the paper** to reflect the reviewer’s suggestions.

Reviewer comment (Presentation):
There are small typos and figure-related issues that should be fixed to improve the overall quality of the paper. Examples include: on page 6 line 2, “Figure 2a” should be “2c”; the line directly above Figure 2 is a runaway LaTeX comment; on page 8 section 4.6, “Figure 5 shows a macro…” should be “Figure 3”; supplemental text refers to “Figure AX” while the numbers continue from the main figures; Figure 3D is missing a caption label; and clarification is needed on what the dark brown color represents.

Author response:
We thank the reviewer for calling attention to specific small typos and places where the paper’s clarity could be improved. We regrettably agree that our results section was lacking in the clarity and polish of the introduction and model presentation, and **we have carefully edited and updated the paper** to reflect the reviewer’s suggestions.

Reviewer comment (Evaluation):
In Figure 2, what is rho? Additionally, the reconstruction of states and the reconstruction of states’ communication are actually not good, thus, it is hard to say the model reconstructs the latent states successfully.

Author response:
**3. Accuracy of dynamics and cross-region message reconstruction in Lotka Volterra simulation and interpretability of the model.** We thank the reviewer for calling attention to the fact that recovered dynamics gradient fields and communications profiles were imperfect despite high performance of the model, harming interpretability of the model. To address the issue of generative dynamics and messages accuracy, **we have introduced a modified objective** in training the network, which forces coupling between the generative dynamics and emissions networks during training (**equation 14** in the updated paper) -- **resulting in more accurate recovered dynamics and messages** as seen in our **updated Figure 2**. We also apologize for the omission in the original figure - rho is a correlation coefficient computed between true and inferred messages. We now cite the updated correlation coefficient in the figure caption.

Reviewer comment (Reproducibility):
Could the authors detail the parameters, such as alpha, beta, and so on, in section 4.1.

Author response:
**1. Clarity could be improved.** We thank the reviewer for pointing out specific proofreading errors in the paper, as well as issues with the figures such missing legends in subfigures. **We have carefully edited the paper** to remove all errors, correct omissions, and ensure clarity of the figures. We have also **added the specific parameters** used for the simulation parameters (alpha, beta, gamma, delta) inline in section 4.1 as the reviewer suggests. Finally, as per the reviewer’s suggestion, we **have moved core result figures** to the main text following the paragraphs for section 4.2, with additional figures in the appendix. We hope the reviewer finds these improvements satisfactory.

Reviewer comment (Writing):
Clarity is not as high as it should be. The authors would need to proofread the submission carefully and generate the figures professionally.

Author response:
**1. Clarity could be improved.** We thank the reviewer for pointing out specific proofreading errors in the paper, as well as issues with the figures such missing legends in subfigures. **We have carefully edited the paper** to remove all errors, correct omissions, and ensure clarity of the figures. We have also **added the specific parameters** used for the simulation parameters (alpha, beta, gamma, delta) inline in section 4.1 as the reviewer suggests. Finally, as per the reviewer’s suggestion, we **have moved core result figures** to the main text following the paragraphs for section 4.2, with additional figures in the appendix. We hope the reviewer finds these improvements satisfactory.