GeoDynamics: A Geometric State‑Space Neural Network for Understanding Brain Dynamics on Riemannian Manifolds

We sincerely thank the reviewer for their positive and constructive feedback. We are glad that the reviewer recognizes the importance of the problems we address in neuroscience, as well as the novelty of our neural network model based on state-space dynamics on the SPD manifold. We appreciate the encouraging comments on how our method combines tools from different fields and performs intrinsic updates on the manifold. We also thank the reviewer for pointing out areas for improvement. In the following, we respond to each concern and suggestion in detail.

Q1: The presentation and writing need some improvement, specifically, the treatment of the Riemannian geometry of SPD matrices. (see details above)

A1: Thank you for your detailed and insightful comments regarding the use and presentation of Riemannian geometry on the SPD manifold. We have carefully revised the manuscript to address the clarity and terminology issues you highlighted:

• Line 74: We now explicitly specify the manifold as the space of $N \times N$ symmetric positive definite (SPD) matrices, denoted $\text{Sym}^+_N$, a well-studied Riemannian manifold in information geometry and medical imaging.

• Line 80: We have replaced “manifold instance” with “a point on the manifold” to align with standard terminology in differential geometry.

• Section 2.2: We have corrected “Riemannian geometry algebra” to “Riemannian geometry” throughout. In line 120, we now clarify the notation by replacing $\mathcal{M}$ with $\text{Sym}^+_N$ to improve precision.

• Metric choice: Thank you for your insightful comment. We agree that many Riemannian metrics have been proposed for the SPD manifold, including the affine-invariant Riemannian metric (AIRM) [Pennec et al., 2006] and the log-Euclidean metric [Arsigny et al., 2006]. In the revised manuscript (Sec. 2.2), we have included a discussion of these alternatives and clarified our rationale for selecting the Stein metric (also known as Jensen-Bregman log-det divergence).

Compared to AIRM and log-Euclidean, the Stein metric offers several benefits for our application:

○ Like AIRM and log-Euclidean, it respects the intrinsic Riemannian geometry of SPD matrices.

○ It is invariant under congruent transformations, which is important for cross-subject comparisons in functional connectivity.

○ It avoids the computational cost of matrix logarithms and square roots.

To empirically support this choice, we conducted an ablation study on the ADNI dataset, comparing classification accuracy and runtime across three manifold-based metrics:

As shown above, the Stein metric achieves the best trade-off between classification performance and computational efficiency. These findings, combined with its theoretical strengths, support our decision to adopt the Stein metric as the default geometry-aware distance measure in our model. The manuscript has been updated to reflect this comparison and justification.

These clarifications and revisions have been incorporated into Section 2.2 and we believe they improve both the rigor and readability of the presentation. Thank you again for your constructive feedback.

Q2: Implementations details are missing - for example, details on how the Frechet mean is computed

A2: Thank you for your insightful comment. We now provide a detailed explanation of how the weighted Fréchet mean is computed in our implementation:

1. $L$ is initialized as the identity matrix so that $F=L L^{\top}+\varepsilon I$ remains SPD. (the complexity is $O(N^3)$)
2. The algorithm minimizes the weighted sum of squared LogDet distances over the SPD matrices in $X$​ by updating $L$ using an optimizer (Adam). This differentiates the log-determinant functions concerning $F$ while respecting the SPD manifold step typically involves structure. (the complexity is $O(N^3)$)
3. The process stops when the change in loss falls below a given tolerance, ensuring convergence. (the complexity is $O(N^3)$)

In addition, the translation is defined as: $\mathcal{T}(X,g):=g X g^{\top}$. The implementation details are provided below.

1. Reshapes vector $g$ as a row vector;

2. From the row vector, a lower triangular matrix $L$ is constructed by iterating over each row, extracting the required slice of $g$, and padding with zeros; (need interaing N rows, the complexity is $O(N^2)$);

3. A skew-symmetric matrix is formed by subtracting the transpose of the lower triangular matrix: $G_{skew}​= L-L^⊤$;

4. The skew-symmetric matrix is then mapped into a rotation matrix in $\mathbb O(n)$ using the matrix exponential: $G=exp(G_{skew})$; (involves inverting and multiplying operation, so the complexity is $O(N^3)$).

5. The translated SPD matrix is computed by conjugating $X$ with $G$: $\mathcal{T}(X,G):=G X G^{\top}$. (including two matrix multiplications, the complexity is $O(N^3)$).

We will incorporate the implementation details in the final version, thank you so much.

Q3: The computational complexity is mentioned only in Appendix A.11 in the context of limitations without details. From a practical standpoint, this is an important issue that should be better discussed. Specifically, what is the computational complexity of the algorithm and how does it scale with the size of the FC matrices?

A3: Thank you for this constructive comment. If the FC matrices are $N\times N$, and both the input $X$ and the convolution kernel $K$ are $N\times N$, then each output element requires an inner product over $N^2$ elements. Since the convolution is computed at each spatial location—amounting to $N^2$ output elements—the overall time complexity for one time step is $O(N^4)$. When considering $T$ time points, the total complexity becomes: $O(T \cdot N^4)$. nn.DataParallel in PyTorch can achieve parallel computing, which distributes the computation across multiple GPUs. We have included the inference time in Table 5 and will add this analysis in the final version.

Q4: Some limitations are mentioned in the appendix (A.11) but I find it quite hidden, and some of the discussion could be moved to the main text

A4: Thank you for your helpful suggestion. It should be an easy fix by highlighting the existing write-ups. In the revised version, we will move a concise summary of the key limitations to the main text (end of Section 4.3), including aspects such as computational cost, assumptions about manifold structure, and interpretability challenges. We will retain the detailed discussion in the appendix for completeness but make cross-references more prominent to ensure these important considerations are accessible to readers.

Thanks again for your feedback, if the reviewer has additional questions, we would be glad to provide further clarifications.

We sincerely appreciate the reviewers' constructive feedback and thoughtful critiques. Below, we provide point-by-point responses to address each concern and outline the improvements made in the revised manuscript.

Q1: Please clarify the novelty of your approach to the part papers on dynamic brain functional connectivity.

A1 Thank you for pointing this out. We fully agree that modeling brain functional connectivity as a trajectory of SPD matrices has been explored in prior work, such as Dai et al. (2020). Our contribution builds upon this foundation but introduces a novel integration of geometric deep learning and state space modeling (SSM) within the Riemannian manifold framework.

Specifically, our approach differs in three important ways (Line 71-94 in the original manuscript):

• Structural novelty: We formulate dynamic FC modeling as a learning task on the Riemannian manifold of SPD matrices, treating each FC matrix as a point on the manifold and encoding whole-brain interactions more holistically than Euclidean approximations.

• Behavioral modeling via SSM: Unlike prior approaches that either lack temporal modeling or rely on RNNs, we adopt a continuous-time state space model (SSM) to capture latent neural dynamics. This dual formulation of state and observation equations enables a principled mapping between brain states and observed neural signals over time.

• Geometric–neural integration: We replace traditional Euclidean algebra in SSMs with Riemannian geometry, enhancing the model's ability to handle the curvature and structure of SPD space. This is paired with a geometric-adaptive attention mechanism that improves efficiency and generalization (Sec. 3).

While Dai et al. focused on geometry-aware trajectory analysis (a valuable contribution), our work adds a temporal modeling layer via SSM and neural parametrization to learn both latent dynamics and manifold structure in an end-to-end fashion. To our knowledge, our work is the first framework to unify these components within a deep geometric setting for functional brain analysis.

We have validated our method not only on multiple brain datasets (e.g., ADNI, OASIS, HCP, PPMI, ABIDE), but also on human action recognition tasks, demonstrating generalizability across domains. We will make more clearer in the final version.

Q2: The computational cost of performing analysis on SPDM manifold is expensive. The matrix computations grow much faster than the size of the matrices. The method will not scale as well if the number of nodes in the brain increases.

A2: Thank you for highlighting this important point. We acknowledge that operations on the SPD manifold—especially matrix logarithms and exponentials—can become computationally expensive as the number of brain regions increases (we have discussed this limitation (Appendix A.11-Page 22) in the original manuscript). This is a general challenge for manifold-based methods. To address this, we currently leverage parallel computing and batched matrix operations to significantly reduce the computational burden during training and inference. Additionally, since many SPD operations are independent across time steps or subjects, they are naturally suited to GPU acceleration and distributed computing.

Looking ahead, we plan to explore more scalable approximations, such as low-rank factorizations, tangent-space projections, or spectral truncation of covariance matrices, which can maintain performance while further improving efficiency. We appreciate the reviewer’s comment and will clarify these scalability considerations in the revised manuscript.

Q3: Even though the empirical results are impressive, this formulation is hard to analyze theoretically due to nonlinearity of the underlying manifold.

A3: Thank you for raising this insightful point. We agree that theoretical analysis becomes more complex when working with nonlinear Riemannian manifolds such as the space of SPD matrices. This is a known limitation of many geometry-based deep models.

To mitigate this challenge, our framework is carefully designed to preserve theoretical tractability where possible. Specifically:

• We replace Euclidean algebra in state space models with Riemannian counterparts in a principled way, allowing certain analogs of stability and controllability theory to still be studied in the manifold setting (see Section 3).

• The structure of the SSM provides an interpretable framework, with latent dynamics governed by well-defined ODEs and output mappings, which serve as a bridge between complex geometry and system theory.

• For practical robustness, we regularize manifold operations using exponential map truncation and intrinsic norms, which helps constrain model behavior and improve empirical generalization.

We also note that the field is actively developing new theoretical tools for analyzing learning dynamics on manifolds (e.g., Riemannian optimization and tangent space approximations). We believe our work contributes to this ongoing effort by demonstrating a successful application in the neuroscience domain, and future work will aim to formalize the theoretical underpinnings further.

We will clarify this discussion in the revised manuscript.

Q4: I don't see any note on the limitations of the model/method proposed here.

A4: Thank you for your important observation. Due to space constraints in the main text, we have included a detailed discussion of the model’s limitations in Appendix A.11 (Page 22) of the original manuscript. This section outlines key limitations, such as the computational overhead introduced by Riemannian operations and challenges related to noisy or rank-deficient input data. We will make this cross-reference more prominent in the revised version to ensure readers are aware of these important considerations. Thank you so much.

We thank the reviewers for their constructive feedback and positive assessment of our manuscript. Below is a brief summary of our work, followed by our responses to the key concerns raised.

Q1: Are those improvement really big enough? (the Accuracy increase from 80.4 to 81.2 in ADNI (+1%), from 89.26 to 89.6 in OASIS)

A1: Thank you for your thoughtful comment. While the observed accuracy gain may appear modest (e.g., 80.4% → 81.2% on ADNI), even small improvements can be meaningful in clinical or population-scale applications, especially when models are already operating near performance ceilings. To assess the robustness of this improvement, we conducted a paired t-test over the 10-fold cross-validation results. For ADNI, the improvement was statistically significant (p<0.05), suggesting the gain is unlikely due to random variation. For OASIS, however, the improvement is not signifiant at the level p<0.05. We also note important differences between the datasets. ADNI is more heterogeneous, spanning broader disease stages, acquisition protocols, and demographics—conditions under which modest gains are more likely to emerge and generalize. In contrast, OASIS is more homogeneous with higher baseline performance, making further improvement inherently more challenging.

We have added this statistical analysis (Table 1 '*' denotes the significant improvement ($p<0.05$)) to the revised version and will include it in the final manuscript.

Q2: What are the hyper-parameters of GeoDynamics?

A2: The hyper-parameters are as follows: Learning rate: $5\times10^{-5}$, Weight decay=0, Batch size=16, Epochs=300, Hidden dim=number of brain region, network layer=2, window size=15. The full configuration is summarized in Table 5 of the original manuscript for reproducibility. We have conducted an ablation study in terms of difference window size (Appendix A.9 in the original manuscript). We will ensure this information remains clearly accessible in the final version.

Q3: Are those values (Acc, Pre, and F1) shown in previous models copied from original papers or the authors reproduced previous models by themself?

A3: Thank you for your question. All results reported for the baseline models—including accuracy, precision, and F1-score—were reproduced by us using publicly available implementations or reimplementations based on the descriptions in the original papers. To ensure a fair comparison, we applied the same 10-fold cross-validation splits, consistent preprocessing, and identical evaluation metrics across all models. This uniform setup ensures that performance differences reflect the intrinsic modeling capabilities rather than external confounding factors.

We would like to clarify that the complete benchmarking code was already uploaded at the time of the initial submission to the anonymous GitHub repository (name "Comparison methods" folder) linked in the original manuscript. We respectfully encourage the reviewers to revisit the provided link for verification purposes. This inclusion complies with the rebuttal policy, as no new materials have been introduced. To further enhance transparency and reproducibility, we will make this information more prominent in the final version. Thank you so much.

Q4: L44 and Figure 1b legend: Does the "black solid box" indicate Euclidean space or Manifold space?

A4: In Fig.1b, the "black solid box" denotes the vanilla SSMs, which operate in Euclidean space. We agree that this distinction could be made more explicit, and we will revise the figure legend in the final version to clearly indicate that the “black solid box” corresponds to Euclidean space.

Q5: L251: So it is a 1 out 8 classification task?

A5: Yes, this is an 8-class classification task. We selected 1,081 subjects from the HCP-WM dataset, where the working memory paradigm included both 2-back and 0-back conditions across four stimulus categories: bodies, places, faces, and tools, along with fixation periods. The eight specific task events are: 2bk-body, 0bk-face, 2bk-tool, 0bk-body, 0bk-place, 2bk-face, 0bk-tool, and 2bk-place. A detailed description of the dataset and task design is provided in Appendix A.6 (Table 3, Page 18) of the original manuscript.

Q6: Where is the Discussion of this paper? Section 4.3 is not the real Discussion.

A6: Thank you for pointing this out. We have revised Section 4.3 to more accurately reflect its focus on Model Validation. The main discussion, including limitations, future directions, and broader implications, is provided in Appendix A.10-A12 (Page 22 -23) of the original manuscript. We will clarify this organization in the final version to improve readability and structure.

Q7: Format

A7: Sorry about that. The overlapping of line numbers (e.g., 112+113, 177+178, 197+198, 233+234) results from our use of \vspace commands to condense vertical spacing. We will correct this and restore proper spacing in the final version of the manuscript.

Thank you for your time and consideration. We are committed to incorporating all suggestions and improvements into the final version of the manuscript.

We appreciate your recognition of our method's theoretical grounding and geometric formulation. Below we respond to each of your concerns and will incorporate all suggestions into the final version.

Q1: Core concepts such as system input, hidden state, and output are not clearly mapped to the neuroscience application (e.g., is the FC the input or the hidden state?)

A1: Sorry for the unclear description. The FC matrices $X(t)$ are treated as the system input in our model. The hidden state $S(t)$ represents the internal geometric representation on the manifold (learned from $X(t)$ that summarizes historical FC dynamics, and the system output $Y(t)$ corresponds to the predicted FC state used for downstream inference tasks (e.g., brain state identification). We have updated Sec. 3.1-Overview to clearly define the mapping between the neuroscience data and the model variables. These changes will be reflected in the revised manuscript to improve clarity and align the mathematical formulation with the neuroscience context.

Q2: Equations (e.g., Equation 3) are difficult to parse and are not explained with sufficient care

A2: Eq. 3 is derived as the manifold-based extension of Eq. 1. The key distinction lies in the underlying geometry of the system:

• Eq. 1 corresponds to a vanilla SSM defined in Euclidean space, where system evolution and observations are governed by standard linear operations.
• Eq. 3 reformulates the SSM on the Riemannian manifold of symmetric positive definite (SPD) matrices, where FC matrices are naturally embedded.

In this manifold setting, standard operations (e.g., averaging, updating states) must be replaced with geometrically consistent operations, including:

• the Fréchet mean $\mathcal{F}$, which generalizes the concept of averaging to Riemannian spaces, and
• the translation operator $\mathcal{T}$, which allows us to evolve system states and outputs across the manifold.

We will revise Sec. 3.1 to walk readers through these concepts more carefully and to explain each term in Eq. 3 with explicit reference to their role in modeling the time-varying FC matrices on the SPD manifold.

Q3: Basic elements of the SSM formalism (e.g., matrices A, B, C, D) are not introduced for readers unfamiliar with control theory.

A3: We agree that clarity around the role of A,B,C,D is important, especially for readers unfamiliar with control theory. Fortunately, it should be an easy fix in the final version by highlighting the existing write-ups. As noted in Sec. 2.1 (line 110-117), we have already introduced the standard linear state-space model and defined each matrix. In our proposed model, these parameters retain similar functional roles but operate on SPD manifolds through geometry-aware operations (e.g., Fréchet mean and manifold translation), as described in Sec. 3.1 and Eq. (3). We will make this connection between the classical and manifold-based formulations more explicit in the revised version.

Q4: About group action definitions

A4: Thank you for these insightful comments. We have unified and clarified our definitions in Section 2.2 as follows:

• 1. Group action on SPD matrices: $g \cdot X := gXg^\top, \quad g \in \mathbb{G},; X \in \mathcal M.$ Here, $\mathbb{G}$ denotes the general linear group of full‑rank $N\times N$ matrices.
• 2. Restriction to orthogonal subgroup: We then restrict to the orthogonal subgroup $\mathbb{O}\subset\mathbb{G}$, consisting of all $N\times N$ orthogonal matrices that preserve the Stein metric.
• 3. Translation operator: $\mathcal T: \mathcal M \times \mathbb{O} \to \mathcal M, \quad \mathcal T(X,g) = gXg^\top.$ In this definition, the first argument $X$ may be any SPD matrix (e.g., the output of the Fréchet mean), and the second argument $g$ is an orthogonal matrix.
• 4. Application in Eq. (3): $S^{(k)} = \mathcal T\bigl(\mathcal F(\cdots),,\widetilde A\bigr),$ which denotes the action of the orthogonal matrix $\widetilde A \in \mathbb{O}$ on the SPD output of $\mathcal F$. Since $\mathcal F(\cdot)$ always returns an SPD matrix, $\mathcal T(\mathcal F(\cdot),\widetilde A)$ is well‑defined.

We have updated Lines 126–133 accordingly to remove any redundant or unnecessary repetitions. Thanks.

Q5: Equation 3 confusion: What is the set over which you are computing the Fréchet mean in Equation 3? S^(k-1) appears to be a single SPD matrix, not a set. What is Ostensibly it is a set of weights (scalars).

A5: Thank you for this insightful question. We appreciate the opportunity to clarify Eq. 3, particularly the structure and interpretation of the weighted Fréchet mean. In Eq. 3 the Fréchet mean is not taken over a singleton {$S^{(k-1)}$}, but over a sliding window of past states (and inputs) with their associated weights. Concretely, letting the window length be $\tau$, we compute

$$S^{(k)} = \mathcal{T} \left( \mathcal{F}_1(\cdot), \, \mathcal{F}_2(\cdot) \right)$$

$\mathcal{F}_1(\cdot) = \mathcal{F}\left( \left\{ S_j \right\}_{j=k-\tau}^{k-1}, \left\{ A_j \right\}_{j=k-\tau}^{k-1} \right)$

$\mathcal{F}_2(\cdot) = \mathcal{F}\left( \left\{ X_j \right\}_{j=k-\tau}^{k}, \left\{ B_j \right\}_{j=k-\tau}^{k} \right)$

The full expression is too long to display clearly in a single line, so we break it down into individual components for clarity:

so that the set {$S_j$} $_{j={k-\tau}}^{k-1}$ contains the $\tau$ most‑recent system–state SPD matrices (window size), with transition matrix $\{A_j\}$,

the set {$X_j$} $_{j={k-\tau}}^{k}$ contains the same window of inputs plus the current $X^k$. We have updated Eq. 3 and the related text to make this explicit. We have also conducted an ablation study in terms of difference window size (Appendix A.9 in the original manuscript). We will make it more clear in the final version, thanks.

Q6: What are the guarantees on the existence and uniqueness of the Fréchet mean under your assumptions? Do you enforce closeness on the manifold during training?

A6: Under the affine‑invariant metric, SPD(n) is a Hadamard manifold (complete, simply‑connected, non‑positive curvature) [Afsari, B. (2011). Riemannian $𝐿^𝑝$ center of mass: existence, uniqueness, and convexity. Proceedings of the American Mathematical Society]. Hence, there exists a unique Fréchet mean given a set of FC matrices. We initialize each translation matrix $g$ close to the identity (so $g \thickapprox I$) and train with small learning rates using our manifold‑aware optimizers (M‑Adam) (Table 5 in the original manuscript). This guarantees that each update exp(Δ) moves only a short geodesic distance on $\mathbb O(n)$, keeping all points for the Fréchet mean within a convex neighborhood.

Q7: Could the model be evaluated more thoroughly on neuroscience-relevant tasks?

A7: We appreciate this constructive suggestion and agree that the core manuscript should remain squarely focused on modeling brain dynamics. To address this, we will: (1) Move the HAR experiment to a concise “Method Validation” section in the Supplement, so that the main text is entirely devoted to neuroscience applications. (2) Enrich our neuroscience evaluation by adding three EEG‑based benchmarks, such as Emotion recognition on the SEED [1], DEAP [2], and HCI [3] datasets, where our geometry‑aware SSM outperforms baseline methods (as below). We are committed to include the following EEG results in the final version.

[1] Zheng, W. L., et al (2015). Investigating critical frequency bands and channels for EEG-based emotion recognition with deep neural networks. IEEE TAMD.

[2] Koelstra, S., et al. (2011). Deap: A database for emotion analysis; using physiological signals. IEEE TAC.

[3] Soleymani, et al (2011). A multimodal database for affect recognition and implicit tagging. IEEE TAC.

Q8: Terminological inconsistency: Terms like "system state" vs. "hidden state" are used interchangeably, it would help to stick to one of them.

A8: Thank you for the suggestion. We have standardized the terminology throughout the manuscript by replacing all occurrences of both “system state” and “hidden state” with the single term “system state”. This change has been applied in Sections 2, 3, and all figure captions to ensure consistency.

Thanks again for your feedback.