Paper: One-hot Generalized Linear Model for Switching Brain State Discovery (ICLR-2024)

Reviewer comment (Writing):
While the overall presentation of the paper is commendable, there is room for improvement in Sections 2 and 3. These sections could benefit from more intuitive and lucid explanations accompanying the mathematical equations, making it more accessible for readers.

Author response:
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.

Reviewer comment (Novelty):
I believe the prior work by Glaser et al. [1] deserves acknowledgment. It might also be valuable to include it in the comparative models, given that their focus on cluster (population) structures aligns with the theme of underlying structures.

Author response:
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.

Reviewer comment (Theory):
Going forward, are the authors considering other factors that might influence interactions? For instance, within an E-I balanced network, given identical anatomical structures and brain states, interactions could vary based on the stage and phase of short-term synaptic depression. This suggests that intrinsic governing features could arise when adding more biological constraints or features. I'd appreciate the authors' perspective on this.

Author response:
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.

Reviewer comment (Novelty):
In the method, the authors first describe OHG and then introduce GHG. They are both variants of HMM-GLM but OHG outperforms GHG. Thus, the order seems unreasonable. What’s more, the experimental results showed that GHG was unable to achieve this paper’s goal. Then what’s the value of GHG?

Author response:
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.   &nbsp;&nbsp;&nbsp;&nbsp;     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".

Reviewer comment (Evaluation):
As shown in Table 2, the results of different numbers of states were similar to that of one-state GLM. It can be explained that global static connection patterns dominate functional interactions in all states as mentioned in the manuscript. Then was the state division biologically reasonable? Perhaps only the features of the global prior were extracted or there was only one state.

Author response:
1. **One-state GLM vs OHG**: As the reviewer points out, the absolute difference between the one-state GLM and the OHG is not large. Nevertheless, we believe that the results of the multi-state OHG are significant for the following reasons. First, the difference in the log-likelihood of the OHG model compared to the one-state GLM is statistically clearly significant, as the standard error in the log-likelihood for the 5-state OHG is ~0.03 while the OHG has an increased log-likelihood of ~0.64. Furthermore, the latent states predicted by the multi-state OHG have a clear interpretation: as shown in Fig. 4, the states correspond to different behavioral stages as the animal performs the task. The multi-state OHG is thus capable of detecting intra-trial changes in functional connectivity that meaningfully reflect behavior. The fact that other multi-state models perform worse than the one-state GLM indicates that the way in which the prior is imposed is critical to obtaining reliable state prediction, i.e., a simple Gaussian prior is not enough and the one-hot component is necessary. So the multiple states are necessary, but only provide useful information with structural prior as in OHG.

Reviewer comment (Experiments):
The experiments fixed the generative hyperparameters and claimed that this set was noninformative priors and insensitive to different datasets. Is there any support for this declaration?

Author response:
2. **Noninformative priors**:
    * We agree with the reviewer that a more principled way of choosing hyperparameters of the OHG via cross-validation would be desirable. In the current treatment, we manually chose a flat prior (i.e., large variance/heavy tail) which reflects our prior knowledge of the parameters. We did this because there was no principled manner in which to select a prior distribution for weights, biases, and other model parameters. Following the reasoning laid by Zeller [1], we selected a noninformative, flat prior which expresses our lack of prior information about the parameters. This approach is akin to selecting a small regularization coefficient $\lambda$ in a ridge or lasso regression.
    * To provide further justification for this hyperparameter choice, we performed additional experiments comparing the performance of OHG under different prior variances of  $\tilde w_{s,n\gets n'}$. The larger the $1/\sigma_w$, the stronger the prior. The following shows that the average log-likelihood on 10 different datasets decreases as the prior becomes stronger. This implies that it would be better to set a conservative weak prior to allow the model to learn parameters more reliant on data, rather than a strong prior assumption.

    | $\propto \frac{1}{\sigma_w}$ | 0 (absolute noninformative prior) | 0.001 | 0.01 | 0.1 | 1 | 10 | 100 |
    | - | - | - | - | - | - | - | - |
    | average log-likelihood | 2.26 | 2.03 | 1.81 | 1.37 | 1.34 | 0.97 | 0.17 |
   * We believe that this result provides some evidence that our choice of hyperparameter is reasonable. However, as noted above, a more rigorous treatment with cross-validation would be desirable. We thank the reviewer for the suggestion and in the interest of preserving the scope of the paper, leave this for a future study.

[1] Zellner, Arnold (1971). "Prior Distributions to Represent 'Knowing Little'". An Introduction to Bayesian Inference in Econometrics. New York: John Wiley & Sons. pp. 41–53.

Reviewer comment (Experiments):
The synthetic study seems quite limited to the type of data the model is designed for (a single global state). It is not clear how well it will work if the neurons are organized into groups with their own state dynamics (which evolve largely independently) and only rarely communicate. The topology of the network could be loose connections between tightly interconnected subnetworks.

Author response:
* **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.

Reviewer comment (Theory):
A principled approach for the selection of the number of states is not discussed. At one point the paper mentions that the log-likelihood is higher with additional states although these states are rare: "there are many sessions with rarely occupied states, and the distinction between states becomes subtle". This seems to be a flaw in the modeling if someone does not know how many true states. Should the reader be suggested to look at the distribution of states to decide? Perhaps a model selection criterion is needed.

Author response:
* **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.

Reviewer comment (Evaluation):
An analysis of the decoding of task information from the latent state would help understand, in real-world tasks, the utility of the state estimate.

Author response:
* **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.

Reviewer comment (Experiments):
How scalable is the model and/or the algorithm? New recording technology including optical calcium imaging can record from hundreds to close to thousands of neurons. The number of neurons in the synthetic study could be ramped up to assess this.

Author response:
* **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.

Reviewer comment (Experiments):
How would the number of states be selected in practice?

Author response:
* **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.

Reviewer comment (Experiments):
How scalable is the model in terms of subpopulations with their own dynamics?

Author response:
* **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.