Modeling state-dependent communication between brain regions with switching nonlinear dynamical systems

We thank the reviewer for their time, comments and helpful suggestions. We especially appreciate the reviewer’s positive comments on the clarity and intuitive presentation of the model and approach in our paper, as well as the expressivity of the model in successfully capturing a nonlinear, multiregion latent process.

We comment below on specific weaknesses and suggestions brought up by the reviewer, detailing changes and additional analysis we have included. We hope that we have satisfactorily addressed some of the concerns and weaknesses brought up by the reviewer, and we apologize in advance that we were not able to fix all issues.

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.

2. Clarity of benchmarks table, adding standard deviation of results for benchmarks. We apologize for the lack of clarity and have modified the caption to clarify - the table presents the mean square error, or MSE, of co-smoothed, held out neuron activity predictions. We use MSE since we apply the model to continuously valued calcium imaging traces, for which we use a Gaussian observation model. Regarding the need to include standard deviation of results, we thank the reviewer for this suggestion and apologize that we didn’t have a chance to update the table during the discussion time.

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.

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.

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.

We thank the reviewer for their time and evaluation of our paper, as well as for helpful suggestions and comments. We also appreciate their comment on the expressivity of our model in using deep neural networks to parameterize latent switching dynamics, and on our comparison in the paper with a number of widely used state space models in the neuroscience community.

We comment below on specific points and suggestions brought up by the reviewer. We apologize in advance that we weren't able to address all issues.

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.

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.

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.

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.

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.

We thank the reviewer for their time in reading our paper and providing thoughtful and helpful comments, questions and suggestions.

We especially appreciate the reviewer’s comments regarding the clarity of the intuition behind and mathematical development of our model presented in our paper, and regarding our novel contribution beyond existing methods in quantifying messages between regions. We also appreciate the reviewer’s recognition of our model’s outperformance of many other relevant methods on a number of datasets tested in the paper.

We address the reviewer’s specific comments below, and apologize in advance that we weren’t able to address some of the specific concerns.

1. Clarifying the use case and benefits of the model, improving the discussion section. We thank the reviewer for calling attention to this and regrettably agree that our discussion section was under-developed. We have modified our discussion to reflect the reviewer’s suggestions, focusing on the goals and benefits of our approach and clarifying and justifying our design decisions.

2. Disadvantage of many latent dimensions required in benchmarked piecewise linear models. This is an important question brought up by the reviewer, and we modify our discussion of this point to better address and clarify. Briefly, the main problem with using piece-wise linear dynamics with many latent dimensions to capture low-dimensional nonlinear dynamics is interpretability. Koopman theory shows that it is possible to express any nonlinear dynamical system of finite dimensions as a linear system in a ‘lifted’, infinite dimensional space - so it is generally possible to use linear models of high dimensions to model a low dimensional nonlinear system. However, to visualize the dynamics of these systems in two or 3 dimensions, we have to use some dimensionality reduction method, such as PCA, to project those systems down, and in doing so we lose interpretability, since important features of the dynamics may live in the additional dimensions. A similar problem also exists with high-dimensional nonlinear approaches, such as LFADS. We agree with the reviewer’s insight that there is a trade off between models that can perform summarization, representing core dynamics of a system in fewer dimensions, and models that are more biologically plausible, and use high dimensional dynamics to capture the system. This trade off is actually the key motivation behind our high-dimensional RNN simulation, we regret we did not have a chance to flesh this out further.

3. Uncommented note on inference vs generation with external input in the switching Lotka Volterra simulation. We apologize for our oversight of the visible comment that the reviewer pointed out. Briefly, in the switching Lotka Volterra simulation, we provide an external switch cue to the model as an input, approximating an external stimulus cue or signal from an upstream region. Importantly, these switch times do not depend on the continuous latent state of the model, ie p(z|u,z), not p(z|x,z), similar to an input driven hidden markov model or HMM. Thus, while the model is able to infer the correct switching time without access to this external cue in inference mode, it is unable to generate correct switching times when run in forward generative mode without this cue, since z and x alone do not contain information about the switching times. We note that in the double well simulation, there is no such external switching cue, and the global switch depends directly on the continuous latent state x.

4. Figure formatting. We deeply apologize for the inconsistent formatting, resolution, and captions and text size in some of the images that the reviewer points out, and we have carefully edited and modified all figures to reflect the reviewer’s comments.

5. Discerning the number of brain regions. We thank the reviewer for this insightful question, as in some cases and species brain regions and their functional roles may be less well defined. Because our model constrains the emissions process so that each latent dimension uniquely and ‘privately’ maps to particular observation dimensions, we have to pre-specify which observations come from each region. It is in principle possible to additionally learn these mappings, we add this as a future direction in our discussion. To address this question partially, however, we run a 16 region model on the widefield dataset, for which we had collapsed the original 16 regions across hemispheres into 8 regions in order to simplify the presentation of results. Table 2 in the Appendix shows 16 regions provides only a marginal benefit over 8 on 1 step generated co-smoothing. This is a more challenging target than typical co-smoothing, which maps form inferred latents to observations, by requiring the model to co-smooth through a single timestep of generated dynamics.

We thank the reviewer for their time and attention in reading our paper, for their very detailed and insightful comments and for their helpful suggestions and specific recommendations in improving our paper. We especially appreciate the reviewer’s kind remarks on the technical strength and novelty of our method in generalizing existing algorithms, characterizing it as promising and a tour-de-force. We also appreciate the comments regarding the quality of our simulation experiments and benchmarks, modeling assumptions, and the clarity of presentation of the introduction and method.

Below, we address the comments the reviewer has made. We apologize that we weren’t able to directly address all of the concerns in the given time.

1. Small typos, ‘rough edges’, and additional explanation of figures. 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.

2. Hyperparameter selection and tradeoffs of interpretability vs performance 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.

3. Robustness of the model’s scientific results across hyperparameter choices. 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.

4. Scientific interpretation and explanation of application to neural datasets. 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.

Many thanks to the authors for their detailed responses to my questions and comments, and their acknowledgement of the current limitations and the lack of time to address them fully in revision.

EDIT: I realize now that I would effectively be arguing for the paper's rejection given my current score, and I certainly don't feel confident enough to do so. Therefore, given my above comments, noting especially the authors' revision that improved low-level readability issues and acknowledgement of limitations, I updated my score from 5 to 6.

We thank the reviewer deeply for their time and attention in reading our paper, and for their very helpful comments. We also appreciate their characterization of our work as novel and for noting that our model extends and generalizes earlier models and performs well on benchmarks. Below we respond to specific weaknesses pointed out by the reviewer, and detail changes to the paper and new analysis that addresses these.

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.

2. Interpretability of transformer inference model in MR-SDS vs RNN-based generative models. The reviewer rightly points out that while we mention in the paper that high-dimensional RNN generative models can suffer from a lack of interpretable dynamics, our model uses a transformer, which is also hard (and maybe harder) to interpret. We would like to emphasize however that we constrain our use of the transformer to the inference step alone, and use simple feed forward or MLP networks to parameterize the nonlinear dynamics functions in the generative model. Thus, while the transformer inference network is indeed complex to interpret, our model does not require interpreting it to analyze the dynamics of the system. On the other hand, the generative model we use is easy to interpret because it produces low-dimensional dynamics gradient fields which are designed for easy visualization in 2 dimensions. This property contrasts with high-dimensional RNN models, (such as LFADS) in which the generative dynamics are produced in a very high-dimensional space, often in the hundreds or higher, thus requiring dimensionality reduction methods such as PCA to visualize dynamics, and then only partially (as important aspects of dynamics may exist beyond the first 2-3 PCs). We have updated our discussion section to make this issue clearer.

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.

4. MR-SDS embeds a richer representation of coupled high-dimensional RNNs vs PCA. We apologize for the lack of clarity on this point, made worse by the unclear figure and reference. This is a comparison of the explained variance in the simulated observations reconstructed under PCA vs MR-SDS embeddings. The author rightly points out that PCA dimensions have an orthogonality constraint that complicates their comparison with the latent dimensions of our model. We add to this, that PCA embeds data linearly, an additional constraint, while our model embeds data nonlinearly. In this sense, it is an ‘unfair’ comparison, as PCA is clearly a less expressive method. Our point in this comparison is simply to compare a practical method used by many researchers in the field to visualize and understand the dynamics of high-dimensional network models (such as LFADS), with our method, which directly produces low-dimensional dynamics fields that explain the data. With access to the full true network dynamics (which a researcher would not have in the case of true data, only in a simulation), a PCA embedding of the full, ground-truth dynamics still captures less variance in the observations than the low-dimensional dynamics inferred by MR-SDS, and thus would offer a less complete description of the data and dynamics. This makes the point that MR-SDS does a good job of modeling a high-dimensional system, despite model mismatch.