Learning Time-Varying Multi-Region Brain Communications via Scalable Markovian Gaussian Processes

Dear Reviewer sX24,

Thank you for the suggestion! We hope these improvements clarified your concerns, and that they can be taken into account when deciding the final score.

The additional results are: (https://anonymous.4open.science/r/rebuttal-figures-for-ICML-2025-42E6/rebuttal_figures.pdf), including Rebuttal-Figures 1&4.

Claim anything causal from the delay?

• We thank the reviewer for this thoughtful comment. We agree that delay or directional estimates from observational data alone do not prove causality. We do not make causal claims here; instead, we use these estimates to form hypotheses about across-region information flow. We acknowledge that only direct manipulations can establish true causal effects. Our model’s value is in highlighting putative, time-varying interactions, which can help pinpoint the circuits and time points most relevant for further causal testing. We will elaborate on this limitation in future revisions.

More complex sensory stimulus

• Thanks for this suggestion. We agree that more complex stimuli offer exciting possibilities, and we see this as an important direction for future work.

Parallel scan access to O(T) GPUs.

• Parallel scan algorithm relies on having enough parallel threads rather than literally requiring $T$ separate GPUs. We will clarify this in limitation part.

Improve the diagonal covariance model to a full covariance model?

• In FA model, the latents approximates a full kernel GP, so the marginal covariance $\mathbf{C}Cov(\boldsymbol{x})\mathbf{C}^\top+\mathbf{R}$ is full.

Where are you computing the complete LL?

• In the E-step, we use a Kalman filter/smoother to compute the posterior of the latents and compute the expected complete LL. In the M-step, we maximize this expectation with respect to model parameters using gradient descent, as part of the EM framework rather than directly optimizing the marginal LL. We will provide a pseudo algorithm in future revisions.

Approximation gap between SSM and GP.

• Although our model is linear in state evolution, it can capture the nonlinear covariance structure of GP with a sufficient order $P$ in Eq. 5. In practice, choosing an appropriate $P$ results in performance that closely matches that of a GP.

How do you deal with Poisson data?

• We preprocess using Gaussian smoothing, following prior work setting (Li et al., 2024). We acknowledge this is not ideal and developing a Poisson Kalmen Parallel scan is a key direction for future work.

Why not comparing against SLDS?

• SLDS is not designed to capture temporal delays between regions. But we also show log-likelihood comparisons with multi-region SLDS in Rebuttal-Figure 4.

What enforces orthogonality of the within vs. across variables?

• We do not claim strict orthogonality; instead, the FA model encourages decorrelation between variables through a block-diagonal observation matrix that assigns latents exclusively to shared or independent activity. Across-region variables use separate kernels, encouraging different delay characteristics; however, if the data indicates similar delays, the latents can reflect that similarity.

Do you also perform cross-validation on the number of latents?

• Yes, we perform cross-validation on all hyperparameters, including the number of latents, in Figure 3(C). Also, in Rebuttal Figure 1(D), we test cases with larger and smaller latent numbers than true value. Incorrect latent settings result in larger variance across runs and lower test log-likelihood. We will address this issue in future revisions.

How are the shaded errorbars computed? Degrees of freedom?

• The shaded error bars represent the variance of the learned delay across runs. We initialize the projection matrix $\mathbf{C}$ using CCA, which yields stable fits across random seeds. Our model incorporates a shared kernel length scale over time, which reduces the model’s degrees of freedom by determining the temporal dynamics and constraining the evolution of delays.

Can you include state space visualization?

• We visualized our model’s communication subspace on both neural datasets in Figure 2(A) and Appendix E.

Performance vs number of regions, latents, and length?

• We thank the reviewer for this insightful suggestion. As shown in Rebuttal Figure 1(A-C), the model demonstrates stable performance across different conditions.

The interplay between the kernel and the delay?

• The estimated delay is an effective parameter that captures the time shift between signals, as defined by the chosen kernel. While different kernels may yield varying delay values, the directional information they convey remains consistent. The sign and relative ordering of delays are stable across reasonable kernel choices, as shown in prior work (Li et al., 2024), where a different kernel produced the same directional insights on the two regions data.

Again, we thank the reviewer for insightful suggestions, which improved the quality of our paper!

Dear Reviewer 9Ksf,

Thank you for the constructive comments! We have made several improvement according to your questions and comments. We hope these sufficiently clarified your concerns, and that they can be taken into account when deciding the final review score.

The additional results are provided (https://anonymous.4open.science/r/rebuttal-figures-for-ICML-2025-42E6/rebuttal_figures.pdf), including:

• A comparison of our model with existing methods on the synthetic dataset (Rebuttal-Figure 4(C)).
• The model performance vs. the number of regions, latents, and length on the synthetic dataset (Rebuttal-Figure 1(A-C)).
• Decoding visual stimulus-related information from the learned latents (Rebuttal-Figure 3).

The aforementioned references should be added (Williams et al. NeurIPS 2021, Safaie et al. Nature 2023, Balzani et al. ICLR 2023).

• Williams et al. (NeurIPS 2021) introduce a metric framework for comparing static neural representations, Safaie et al. (Nature 2023) show preserved latent dynamics across subjects using alignment, and Balzani et al. (ICLR 2023) propose a task-aligned model capturing within- and across-region neural variability. In contrast, our work focuses on modeling delay-based brain communication across regions, aiming to capture dynamic, time-varying delays that mediate across-region interactions. This enables insights into the directionality and strength of functional coupling.

• Therefore, we respectfully disagree that our work should be directly compared to alignment-based approaches such as Williams et al. (NeurIPS 2021) and Safaie et al. (Nature 2023), as our modeling goals differ. On the other hand, we greatly appreciate the novel perspective introduced by Balzani et al. (ICLR 2023) on task alignment within communication subspaces. We will include a discussion of their work in our revised manuscript to clarify how it complements and contrasts with our approach.

They minimally evaluate the proposed model in a synthetic dataset, but the authors fail to show comparative performance with other models in this data.

• We thank the reviewer for raising this concern. In Rebuttal-Figure 4(C), we present the test observation log-likelihoods on the time-varying synthetic dataset used in Sec. 4.1. Our model outperforms existing multi-region modeling methods, demonstrating its effectiveness in capturing time-varying multi-region communications.

How does the model behave with respect to the latent dimensionality, or number of areas, or length of recording?.

• We thank the reviewer for this valuable suggestion. In Rebuttal-Figure 1(A-C), we evaluate our model using MSE and Pearson's correlation coefficient (CC) between estimated and ground truth delays under varying numbers of regions, latent dimensions, and lengths. The model shows stable performance across conditions. While MSE increases due to amplitude variability with more regions, CC remains stable, indicating reliable recovery of temporal delay patterns. As shown in Rebuttal-Figure 2, this amplitude variability is not a practical concern because the temporal patterns, which are crucial for understanding brain communication, are well preserved.

The method uses vanilla FA to estimate the latent dimensionality, ..., why not use likelihood estimates from the model itself?

• We follow the strategy from prior work (Gokcen et al., 2022), see Sec. 4.3, using FA to estimate total latent dimensionality and then performing a grid search with model log-likelihood to determine across- and within-region dimensions. This approach balances accuracy with computational efficiency, as a full grid search of latent dimensionality is time-consuming.

Axes ticks and labels are missing in multiple figures.

• We apologize for the oversight and thank the reviewer for pointing it out. We will add axis ticks and labels in the revised version.

Possible extensions would be to include comparisons of the latent representations across models, or decoding performance with respect to relevant neural computation variables, such as visual stimulus class.

• Thank you for the suggestion. We have visualized our model’s latent representations on both neural datasets in Figure 2(A) and Appendix E. In the revised version, we will include latent representations from other models for comparison. Additionally, Rebuttal-Figure 3 demonstrates that decoding from our latent variables shows stronger task-related information (e.g., grating orientations) in the communication subspaces compared to the within-region subspaces. We also view the discovery of further task-relevant information in these subspaces as a promising direction for future work.

Different dimensionality across different latent spaces?

• Yes, we allow the number of within-region latents to vary over regions.

Again, we thank the reviewer for the insightful suggestions, which improved the quality of our paper. If you have further questions, we are happy to discuss them!

Dear Reviewer MsHr,

Thank you for the encouraging feedback and suggestion! We have made several clarifications according to your questions and comments. We hope these sufficiently clarified your concerns, and that they can be taken into account when deciding the final review score.

The additional results are provided (https://anonymous.4open.science/r/rebuttal-figures-for-ICML-2025-42E6/rebuttal_figures.pdf), including:

• Delay estimation on the synthetic dataset with five regions (Rebuttal-Figure 2).

It would be interesting to see how well the model reconstructs true delays in the setting where there are more than two regions though.

• We thank the reviewer for this insightful suggestion. In Rebuttal-Figure 2, we present the estimated time-varying delays from synthetic data with five regions. Rebuttal-Figure 1(A-C) further shows model performance in terms of MSE and Pearson’s correlation coefficient (CC) between the estimated and ground truth delays as the number of regions increases. While MSE rises due to greater amplitude variability, CC remains stable, indicating reliable recovery of the temporal delay patterns. As shown in Rebuttal-Figure 2, this variability in amplitude is not a practical concern, as the temporal patterns, which are crucial for understanding brain communication, are well recovered. We will add these results in the revised version.

To what degree is the inference of delays vs. communications an ill-posed problem? Is there a degeneracy where the true message could have static delays in terms of the transmission channel but it looks like a time-varying delay due to the time-varying nature of the communications?

• We thank the reviewer for raising this question. Separating the effects of time-varying delays from changes in the message content is inherently challenging. In principle, the inference problem is ill-posed because the same observed data could be explained either by static delays combined with time-varying messages or by genuinely time-varying delays. In other words, there is a potential degeneracy: even if the true communication channel has static delays, the variability in the messages may cause the inference process to interpret these as time-varying delays.

• To mitigate this issue, our model incorporates a regularization strategy by using a shared length scale parameter $l$ across all time points. As described in Sec. 2.2, the state transition matrix $\mathbf{\hat{A}}_t$ is uniquely determined by the delay $\theta_t$ and the length scale $l$. By sharing $l$ across time, we constrain the temporal evolution of the delay parameters, enforcing similar temporal dynamics and reducing the risk of misattributing variability in the messages to changes in delays. Furthermore, our synthetic experiments with known ground truth demonstrate that the model reliably recovers the true dynamic delays, indicating that the inferred time-varying delays reflect genuine changes in the communication pathways rather than artifacts arising from the message content.

• However, we acknowledge that further research could explore additional constraints or independent measures to better address this problem, but our current results suggest that the degeneracy is minimal in the context of our application.

Again, we thank the reviewer for the comments, constructive feedback, and insightful suggestions, which improved the quality of our paper. If you have further questions, we are happy to discuss them!

Dear Reviewer Wfpo,

Thank you for the encouraging feedback and practical suggestions! We have made several clarifications according to your questions and comments. Hopefully these will resolve most of your concerns, and that they can be taken into account when deciding the final review score.

Figure 3(B) is challenging to understand. I suggest presenting these results in a different format, potentially via a matrix where each element represents a signed pairwise connection.

• We thank the reviewer for pointing this out. Figure 3(B) is intended to highlight the time-varying nature of temporal delays across five brain regions. It visualizes the learned temporal delays from Figure 3(A) at time points $t=3$ and $t=50$. In the figure, the orange arrows indicate the direction of communication: a positive delay denotes forward communication, while a negative delay indicates feedback communication. Additionally, the length of each edge represents the absolute value of the corresponding temporal delay.

• To provide a clearer and more quantitative view, we will add a matrix alongside the network diagram to explicitly show both the magnitude and sign of the learned temporal delays.

Can the authors provide additional details on the formulation and fitting of the time-varying delay? It appears to be one of the crucial innovations. I'm curious on how flexible this model is and whether it presents any issues in fitting. The paper states that the resulting $\mathbf{\hat{A}}_t$ are still constrained at each time point via a shared length scale. However, it is not obvious to me how strongly that constrains these parameters.

• At each time step $t$, we construct a time-specific Markovian GP conditioned on the MOSE kernel corresponding to that time step. The state transition matrix $\mathbf{\hat{A}}_t$ is uniquely determined by the delay $\theta_t$ and length scale $l$ as mentioned in Sec. 2.2. While the delay $\theta_t$ is allowed to change over time, the length scale $l$, which is shared across all time steps, limits the freedom of $\mathbf{\hat{A}}_t$. This shared length scale fully determines the temporal dynamics of the latent variables and thus constrains the evolution of dynamics from $\mathbf{\hat{A}}_t$ to follow the same smoothness. In practice, this allows the model to flexibly capture time-varying delays while maintaining similar temporal dynamics across time.

• The blue and red latent variables shown in Figure 1(A) illustrate this behavior: they share overall temporal dynamics but exhibit different delays at each time point. Empirical results in Sec. 4.1 further confirm that this shared length scale leads to robust and smooth delay estimates without fitting issues if a reasonable initialization is used, e.g., setting a relatively large length scale to encourage initial smoothness. This initialization aligns with standard practice in GP regression, where initializing with a larger length scale helps guide optimization during early iterations.

Additionally, Appendix B does not appear to describe the time-varying dynamics, as $\mathbf{A}$ is not indexed by in the EM equations. Is there a different objective for the time-varying formulation? If so, it would be very helpful and important to put that in the appendix as well.

• We apologize for the confusion regarding Appendix B. In our time-varying extension, the underlying EM objective remains the same: maximizing the expected complete-data log-likelihood, but the state-space model is augmented to include time-specific transition matrices $\mathbf{\hat{A}}_t$ and process noise covariances $\mathbf{\hat{Q}}_t$. The derivations for the EM updates are analogous to those in Appendix B, with the key difference that the E-step now involves running a vectorized Kalman filter and smoother over the entire time sequence. As a result, the shapes of $\mathbf{\hat{A}}$ and $\mathbf{\hat{Q}}$ become $(T, NP, NP)$. We will revise Appendix B to explicitly incorporate these modifications and clearly describe the time-varying dynamics.

Again, we thank the reviewer for the comments, constructive feedback, and insightful suggestions, which significantly improved the quality of our paper. If you have further questions, we are happy to discuss them!