One-hot Generalized Linear Model for Switching Brain State Discovery

Dear Reviewer G7kx,

Thank you very much for your time and valuable comments on our paper. We appreciate your recognition of the strengths of our paper. Hopefully, the following responses could resolve most of your concerns and answer your questions.

Weaknesses

1. Following the reviewer's suggestion, we have reorganized the Methods section to improve readability for readers. We have changed the order of the presentation to introduce the models in order of increasing complexity. Furthermore, we have added a subsection detailing the relation between GHG and OHG. We hope these changes have made the presentation more intuitive for readers. Please check the current latest revision.
2. We agree that the work by Glaser et al. (2020) should be mentioned in the paper.
   • The citations regarding SLDS, rSLDS, and mp-SLDS are included in the introduction section.
   • Relationships to SLDS-based methods: Instead of considering state switching on neural connectivities as our HMM-GLM-based models, SLDS, rSLDS, and mp-rSLDS consider state switching on underlying linear dynamics governing the observed spike trains. Besides, mp-rSLDS also considers incorporating prior information (such as anatomy) into the linear mappings in different states. However, their shared priors are known hyperparameters, but the shared global priors in our GHG and OHG are learnable.
   • We have finished the comparison of our OHG with SLDS-based methods shown in Fig. 16 and Appendix A.7 in our latest revision. Specifically, the inferred states from SLDS-based methods exhibit fast-switching phenomena, which hinders their interpretability. This result might imply that state switching over neural connectivities could be an important assumption of an effective component that should be taken into account when dealing with spike train data collected from multi-stage experiments.

Questions

Future Directions: We greatly appreciate the reviewer's perspective shown through their questions. Indeed, placing a biologically motivated constraint on the structural connectivity with the OGH improved inference of both state and functional connectivity; our results thus indicate that including relevant biological constraints can lead to a better outcome. On the other hand, adding too many biological details can lead to overparameterization of the model, so each added degree of complexity should be carefully chosen and evaluated. We agree that one promising extension would be to add effects of short-term synaptic potentiation and depression, which is biologically well-characterized but not included in our models. Adding constraints to control the E-I balance could also be beneficial. We leave these promising topics to a future study. Another extension would be to add different assumptions about the network structure, for example, the inclusion of multiple regions with dense intra-regional connectivity and sparse inter-regional connectivity. This direction was suggested by Reviewer hZX1 and is also explored in the aforementioned study by Glaser et al. Please check Fig. 9 in Appendix 4 in the current revision for the result regarding subpopulations of neurons.

We hope that this response addresses the majority of your concerns and questions. Your feedback is valuable, so please don't hesitate to provide additional comments or ask further questions. Thank you again for your time and valuable feedback!

Dear Reviewer zShZ,

Thank you very much for your time and valuable comments on our paper. We highly appreciate your recognition of the strengths of our paper. Hopefully, the following responses could resolve most of your concerns and answer your questions.

Weaknesses

1. Mentioning GHG and OHG, and the value of GHG: Thanks for your valuable suggestion. We have modified our abstract, method, result, and conclusion parts to include more discussion on GHG. Specifically:
   • Abstract: We include the names of the two frameworks (two variants of HMM-GLM): GHG and OHG explicitly.
   • Method: We add an extra comparison between OHG and GHG. Compared with OHG, GHG serves as an intermediate model with a straightforward prior directly on the weights representing shared global connectivity but without the one-hot decomposition.
   • Experiment results: In the experimental evaluation section, GHG is not bad, but still worse than OHG, validating that the one-hot decomposition component in OHG does account for the better performance achieved by OHG.
   • Conclusion: We add new sentences mentioning the role of GHG. I.e., GHG serves as an intermediate model with shared prior directly on weight matrices without (one-hot) strength-connection decomposition, confirming that such a decomposition is critical for the success of multi-state inference.
2. Order of introducing GHG and OHG: We thank the reviewer for pointing out this issue. The value of GHG is not in that it outperforms the OHG, but lies in its role as an intermediate model. The GHG is the simplest way to add a prior to the state-dependent weight matrices $\boldsymbol W_s$. The OHG has an additional layer of complexity in that it decomposes the weight matrices into a product of a discrete connection matrix and a strength matrix, allowing the model to specifically constrain the structural connectivity. The value of including the GHG in our study is thus not that the GHG outperforms the OHG, but rather that it allows us to examine whether the structure-strength decomposition of the OHG is necessary, or whether a simple Gaussian prior over the weight matrix is enough. The fact that the OHG significantly outperforms the GHG in both state prediction and inference of the underlying functional connectivity shows that the structural decomposition method provides additional benefits that a simple prior cannot.      However, we understand that the presentation of the Methods may be confusing in that the final method introduced is an intermediate model. Accommodating the reviewer's comment, we rearranged the Methods to introduce the models in increasing order of complexity: The GHG is now introduced before the OHG, which is introduced last. Furthermore, in order to clarify the role of GHG in our study we added a subsection in the Methods titled "The relationship between GHG and OHG".

• Scalable to the number of neurons:
  • The time complexity of the learning in one epoch is $\mathcal O(N^2TS)$ for all HMM-GLM-based methods (including HG, GHG, and OHG), where $N$ is the number of neurons, T is the number of time bins, and $S$ is the number of states. Therefore, the time increases quadratically to the number of neurons $N$ since there will be $\frac{N(N-1)}{2}$ edges in each state. However, since the parameter set of OHG is larger than GHG and larger than naive HG, the constant coefficient of OHG before the complexity term $N^2TS$ is larger than GHG and larger than naive HG. To understand the scalability of our algorithm to large numbers of neurons, we are running experiments on this and will keep you posted when the result can be accessed in the Appendix in our next revision.
  • From the synthetic dataset, we observe that the state accuracies of all HMM-GLM-based methods start to drop when there are 40-50 neurons. To this scale, there will be more than $S\times [1600,2500]$ edges (connectivities) that need to be learned in the network, which introduces challenges to determining the correct state switches. Under such situations, the effectiveness of both the E-step and the M-step in the algorithm would be mutually influenced by each other. However, this does not mean that the HMM-GLM-based methods cannot be applied to real-world datasets. It might still be worth trying HMM-GLM-based methods even if the data consists of a large number of neurons, since Fig. 17 in Appendix A.8 shows that the per-neuron log-likelihood does not drop when increasing the number of neurons. This means the HMM-GLM-based models (especially OHG) is still to predict firing rates effectively for spike train data reconstruction. We would like to view this as an important future direction.
• Inference time: Once the model parameter is learned, the inference should be very quick. Specifically, the time complexity to run the forward-backward algorithm to infer the state sequence is $\mathcal O(2ST)$. On the synthetic dataset with 5 states and 20 neurons, the running time of the forward-backward algorithm on a sequence of length 5000 time bins is about 1 second (on one core of the Intel Xeon Gold 6226 "Cascade Lake" @ 2.7Ghz CPU). Once the model parameters are learned from previous spike trains, we can rapidly infer the state sequence for the current spike train, and then use it to control the stimulus for the next trial.

We hope that this response addresses the majority of your concerns and questions. Your feedback is valuable, so please don't hesitate to provide additional comments or ask further questions. Thank you again for your time and valuable feedback!

Dear Reviewer hZX1,

Thank you very much for your time and valuable comments on our paper. We highly appreciate your recognition of the strengths of our paper. Hopefully, the following responses could resolve most of your concerns and answer your questions.

• Experiments on subpopulation configurations: Although all HMM-GLMs in this paper including our newly proposed GHG and OHG only work for one global prior, it is still possible to apply the model to multi-region data. For example, $C$ groups of neurons $N_1,N_2,\dots,N_C$ evolve largely independently and only rarely communicate. Further, assume each group of neurons has its own global structure and its own number of states $S_1,S_2,\dots,S_C$. Then, we can reduce such a problem into a model with one single globally shared structure $\operatorname{diag}(\boldsymbol A_{1,0},\boldsymbol A_{2,0},\dots,\boldsymbol A_{C,0})$. Then, there will be $S = S_1\times S_2\times\dots\times S_C$ states for the whole groups of neurons in total. For better illustration, we run a simple example on $C=3$ subpopulations (groups) and each group has 5 neurons and 2 states. Therefore, we need a model of 15 neurons and 8 states to accommodate this configuration. The results are shown in Fig. 14 in Appendix 6 in the current revision. Although this is a feasible solution to using our HMM-GLMs framework to tackle subpopulations problem, we think it is better and more reasonable to directly develop a new model for subpopulations problem, and this could be one of our great future directions.
• Number of states:
  • In real-world scenarios, we don't know the true number of true states or there is even no definitive concept of absolute true states. Cross-validation is certainly a rigorous way to determine the number of states. However, increasing the number of states may increase the validation log-likelihood, as the OHG result in Tab. 2. Too many states may hurt interpretability. Therefore, there should be a trade-off between the validation log-likelihood and the interpretability. It might be better to select a turning point where the likelihood stops its rapid growth, and in the meantime, the number of states is not too large to be interpreted.
  • For example, for $S>5$ in the PFC-6 dataset, the performance of OHG does not increase significantly and becomes flat (Fig. 10} in Appendix 4.3. Although more numbers of states are assumed, only around four effective states are learned and the interpretations are all similar (Fig. 11 in Appendix 4.3) to that of the 4 states explained in the main content.
  • Generally speaking, the problem of determining a suitable number of states exists in most of the multi-state models, and this could be one of our important future directions.
• Decoding task information from inferred latent:
  • Thanks and we think this is a good suggestion. To further confirm our interpretation regarding the inferred states from OHG, we use a logistic regression model to decode the correctness of each trial from the inferred states. Specifically, each trial is viewed as a data point in logistic regression. The input variable is the one-hot representation of the inferred hidden states of size number of states $\times$ number of time bins. We use the one-hot representation since the state in each time bin is a categorical variable. The output is a binary variable, representing the correctness of a trial. We use 2/3 trials to train a logistic regression model and test on the remaining 1/3 trials. Fig. 9 in Appendix 4 shows that logistic regression fitted to the inferred states from OHG obtains the highest decoding accuracy. This means the inferred states from OHG do include enough information related to the correctness of each trial. From the interpretation in the main content, we mention that the animal will enter state 4 if it is a correct trial because of obtaining the reward at the correct target location. This is consistent with the positive coefficients of state 4 at the end period of trials, indicated by the black square in Fig. 9.
  • To further confirm the irreplaceable role of the inferred latent in predicting trials' correctness, we train a multilayer perceptron (MLP) neural network with one hidden layer of size 100 (we tried different MLP configurations and selected the best one) to predict the correctness of each trial directly using the neural spike train as the input. The test accuracy is only 0.72, which is significantly worse than OHG. This implies that OHG plays a very important role in summarizing the task information from the neural spike train, similar to the irreplaceable role of the convolutional layer in CNN.
  • Please check the details and results in Appendix 4 in the current revision.

Thanks for your further questions.

1. Given the subpopulations problem is not the primary goal of our models, we agree with you that the relative performance might not be as good as the example in the main content, as you pointed out the spurious/noisy weights should not have appeared at the off-diagonal blocks. As we mentioned before, we first reduced the subpopulations problem and then used our OHG to solve it. Therefore, OHG doesn't know that the off-diagonal blocks should be all 0s. There are two possible solutions:
   • Increase the data size, so that OHG could learn better from the more sufficient dataset that the off-diagonal blocks could be 0s. We doubled the data size and the results became slightly better (but not visually significant). Those off-diagonal blocks are still not completely silenced.
   • If we know the configurations of the subpopulations, we can certainly use a regularization term to suppress the off-diagonal blocks, or even silence them in a hard style (e.g., through a mask). Then, we can obtain weight matrices with less noisy elements on those off-diagonal blocks (Fig. 15 in Appendix A.6).
2. Thanks. Please check the updated Fig. 10 in Appendix A.4.3. The lowest range of BIC $S\in\\{3,4,5\\}$ roughly matches the turning point of the log-likelihood plot.
3. Adding a prior to strengthen the diagonal of the transition matrix is an intuitive way of suppressing the fast state switches in HG. This is equivalent to applying an L2 regularization term on the off-diagonal elements in the transition matrix. For some applications, it might be helpful. On the PFC-6 dataset, however, Fig. 18 in Appendix A.9 shows that when the number of state switches is suppressed, the inferred states are still meaningless. Particularly, most switches happen within a short duration, which looks like nothing but noisy state switches. Only one major state governs the whole trial. Besides, Fig. 8 in Appendix A.4.1 shows that both GHG and OHG don't have such a fast switches phenomenon. This might further imply that a global structural prior plays an important role in HMM-GLM-based models.