BrainOOD: Out-of-distribution Generalizable Brain Network Analysis

[Q6. Conventional ML Methods as Baselines.] Including traditional machine learning models as baselines is a good suggestion. We conduct experiments with SVM and LR by following the setting in [1], where these ML methods take the flattened upper-triangle connectivity matrix as vector input, instead of using the brain network. We update the results in Table 3 of our revision, in which our proposed BrainOOD still achieves the best performance.

[Q8. How other parameters affect the performance.] Thank you for your comment. For the temperature hyperparameter in Eq. (10), we just set it as $\tau=1.0$ by following the default setting in the previous work [3]. This setting works consistently well, allowing us to use these default values without further tuning when adapting to new datasets. For the number of sampling $k$, we conduct additional experiments by tuning it from {1, 3, 5, 10, 20}. We added the results and discussion in Appendix E.5 of our revision.

[1] Data-driven network neuroscience: On data collection and benchmark. NeurIPS 2023

[2] Braingnn: Interpretable brain graph neural network for fmri analysis. MIA 2021

[3] How interpretable are interpretable graph neural networks? ICML 2024

[W1, Q2. Experiments for ADNI with 5 Classes] Thank you for pointing this out. To maintain consistency with prior studies, we initially adopted the 6-class setting. However, we acknowledge that distinguishing MCI, EMCI, and LMCI introduces complexity. In response to your suggestion, we conducted experiments using a 5-class setting (CN, SMC, EMCI, LMCI, AD) by merging MCI with LMCI. As shown in the following table, our proposed BrainOOD consistently achieved the best results on both ID and OOD sets under this revised setting. We have included the detailed results and discussion in Appendix E.1 of our revision.

[Q1. Thresholding Detail] The adjacency matrix was constructed by thresholding the connectivity matrix to retain the top 20% of values based on correlation magnitude. To ensure compatibility with GNNs, which do not handle negative edges effectively, we retained only positive connections. This approach follows the dataset paper [1] and is a common practice in GNN-based brain network analysis [2].

[Q3. Reconstruction Loss for Feature Selection.] The reconstruction loss plays a crucial role in guiding the model to focus on informative features. By minimizing the discrepancy between the original and reconstructed features, the loss encourages the model to mitigate the limited information recovery (Theorem 4.1), and hence to retain only the most significant features while filtering out noise. This mechanism ensures that the extracted features align with the functional relevance of the brain network.

[Q4. In-Domain Testing.] In-domain (ID) testing was conducted using data from sites that the model has seen in the training set. The ID performance was evaluated using a separate held-out portion of the training data to ensure that the results reflect the model's ability to generalize to data from known distributions.

[Q5, Q7. ID and OOD checkpoints.] Thank you for highlighting potential ambiguities. Both Tables 2 and 3 follow the 10-fold CV setting. We revised the caption of Table 3 to avoid misunderstanding.

• Table 2: ID and OOD accuracies are reported using separate models selected via the validation set for ID and OOD data, respectively.
• Table 3: Overall accuracy combines ID and OOD results through a weighted average for a fair comparison with non-OOD methods.
  In Figure 3, the ID and OOD checkpoints refer to the models selected for ID and OOD evaluation. The visualized edge scores (e.g., max value ~0.08 in Fig. 3) represent $W_{Mask}W_{Mask}^T$ without sigmoid from Eq. (4). Post-sigmoid, even edges with scores as low as 0.08 exhibit strong relative importance. Regarding the subgraph, edge sampling is performed $k$ times based on learned probabilities for independent predictions. This dynamic sampling approach makes it difficult to define the exact number of edges in the subgraph.

[Q1. Ablation Study on Feature Selector and Structure Extractor.] Thank you for the insightful suggestion. To evaluate the contributions of the feature selector and structure extractor, we performed additional experiments as shown in the following table, where “feat” represents using $\mathbf{X}$ or $\mathbf{X}^\prime$ as the feature matrix, while “adj” represents using $\mathbf{A}$ or $\mathbf{A}^\prime$ as the adjacency matrix. The best results are achieved when both modules are used together, confirming the complementary nature of the feature selector and structure extractor. We can also observe that the adjacency matrix generated by the structure extractor contributes more to the improvement of the OOD set. We have included these results and discussions in Section 5.4 of our revision.

[Q2.1. Advantages of HPGNN with Multiple Layers] The advantage of HPGNN with multiple layers lies in its ability to aggregate information across increasingly distant neighborhoods. Unlike directly multiplying the graph Laplacian matrix, which assumes linear propagation of features, multiple layers in HPGNN enable non-linear transformations at each hop. This allows the model to better capture complex, hierarchical patterns and deviations from local structures, which are particularly relevant in brain networks where functional and structural connectivity often span multiple scales.

[Q2.2. Retaining $\mathbf{A}$ During Feature Selection.] Thank you for your question. The output of the feature selector, the masked feature matrix $\mathbf{X}'$, serves as the input to the structure extractor to generate $\mathbf{A}'$. This sequential processing is intentional. We prioritize feature selection first because it allows the structure extractor to work with cleaner and more relevant features, which is essential for learning a high-quality structure matrix. The decision to prioritize structure extraction after feature selection is supported by our ablation study, which demonstrates that the structure plays a more significant role in the final performance. By providing clearer features, we ensure that the structure extractor can generate a more accurate and robust $\mathbf{A}'$, ultimately enhancing the overall model performance.

[W1. Technical Contribution.]

We need to clarify that the technical contributions of this work lie in both the methodology level in addressing several unique challenges of BrainOOD, as well as the benchmarking level in illustrating the OOD generalization challenge in brain network analysis.

From the benchmark level, while multi-site datasets like ABIDE and ADNI are commonly used, they are not inherently designed for the evaluation of OOD generalization in brain network analysis. Our proposed benchmark introduces a structured OOD scenario by partitioning these datasets based on site-specific differences. This setup systematically evaluates the robustness of models against unseen site distributions, a novel challenge that existing benchmarks do not explicitly address. Furthermore, our results demonstrate that current methods often struggle with such OOD scenarios, validating the need for a specialized benchmark to drive future research in this domain.

From the methodology level,

• BrainOOD indeed raises several unique challenges for OOD generalization and interpretability, which are the noisy inputs and the consistency requirements, respectively.
• To tackle the noises in the brain networks, we propose a novel feature selection strategy as well as a noise filtering strategy to better extract the meaningful information in the brain network inputs.
• Accommodating with the challenge of noise in brain network analysis, we provide a theoretical discussion on the failures of previous methods in effectively identifying the desired OOD generalizable and interpretable subgraph. Moreover, we also propose to incorporate a reconstruction-based loss to better tackle with influence of the noises.
• As previous OOD generalization and interpretability methods fail to identify the desired subgraph, we also empirically demonstrate the superior performance of our proposed approach.

While individual components may have roots in existing techniques, their integration into a unified GIB framework specifically for brain network analysis is novel. The emphasis on tackling OOD generalization—a critical but underexplored issue in this field—represents a meaningful technical advancement. Moreover, the interpretability aspect, enabled by feature selection and visualizable subgraphs, offers practical value for clinical applications.

[W2. Clarification of Feature Selection Mechanism.] Thank you for pointing this out. We acknowledge that the explanation could be clearer. In Eq. (8), the use of $Tanh()$ is not intended as a typical activation function. Instead, it serves to scale the range of the reconstructed features to align with the input connectivity matrix, which typically has values in the range [-1, 1]. Additionally, the self-multiplication operation $\hat{\boldsymbol{H}} \hat{\boldsymbol{H}}^\mathsf{T}$ is designed to ensure the output exhibits the symmetry property inherent in the connectivity matrix. This operation mimics the structure of the input data, making it easier for the model to capture meaningful patterns during reconstruction. For clarity, we have added this description in Section 4.2 of our revision.

[W3. Definition of OOD Problem and Comparison with Brain-Specific Models.] Thank you for your suggestion. To further evaluate how non-OOD methods perform in an OOD scenario, we report the ID and OOD accuracies for GCN and BrainNetCNN. As shown in the following table, while they achieve good performance on the ID set, a significant gap between ID and OOD results is observed, which indicates these methods fail to generalize to the OOD set. It demonstrates the challenge of OOD generalization for brain network datasets. We have added these results and discussions in Section 5.2 of our revision.

[W4. Grammar should be double checked.] Thank you for your feedback. We have thoroughly reviewed and revised the manuscript to address grammatical issues and enhance clarity throughout. We appreciate your attention to detail and believe the revised version now reads more smoothly.

[W1. Contribution of the Benchmark.] We appreciate the reviewer’s comment and would like to clarify our contribution regarding the benchmark setup. While it is true that the ABIDE and ADNI datasets are widely used in brain network analysis for evaluating brain disorder diagnosis models, existing studies typically utilize these datasets without explicitly addressing the challenge of distribution shifts across different acquisition sites. In most prior work, the focus has been on in-distribution (ID) evaluation, which assumes a uniform data distribution without taking into account the inherent domain shifts present across different data collection sites. In our work, we go beyond the conventional usage of these datasets by creating a specific OOD benchmark scenario that simulates real-world conditions where models encounter data from unseen sites during testing. This setup is motivated by the well-known issue of site-specific biases in neuroimaging data, which often lead to significant distribution shifts.

• To address this, we carefully partition the ABIDE and ADNI datasets across multiple folds, ensuring that each fold contains test sets with both ID and OOD subjects. By selecting the smallest/largest sites as OOD test sets based on the dataset characteristics (ABIDE: small sites; ADNI: large sites), we ensure a diverse and realistic OOD evaluation protocol.
• Our proposed benchmark is the first to systematically evaluate OOD generalization on brain network datasets with a focus on addressing site-specific variability, which is a critical challenge in clinical applications.
• The experimental results show that existing OOD methods fail to generalize well in this scenario, highlighting the necessity of specifically designed OOD algorithms for brain network data. We have revised the manuscript to further clarify this point in the introduction. We hope this explanation better conveys the significance of our benchmark in advancing OOD research for brain network analysis.

[W2. Alignment of Motivation, Method, and Analysis.]

The rationale of our proposed method to resolve the OOD generalization relies on the concept of graph information bottleneck, which aims to find a minimal sufficient subgraph to predict the label. Previous works show that predictions based merely on the minimal sufficient subgraph are able to generalize to OOD graphs [1,2]. Finding the minimal sufficient subgraph is essentially to remove the noisy and spuriously correlated substructures from the original input data.

However, the noisy nature of fMRI data and the limited encoding capacity of GNNs raise unique challenges in implementing the graph information bottleneck, as shown in our Theorem 4.1. Therefore, we propose several strategies to resolve the challenges and better identify the desired subgraphs.

For the analysis of the method’s generalizability, current visualization can provide some evidence. For the analysis of the method’s generalizability, current visualization can provide some evidence. As shown in Figure 3, the OOD model (top left) and the ID model (bottom right) highlight some common connections, particularly within regions such as the VIS, SMN, VAN, and DMN networks. While the specific scores differ, the overall patterns in these regions remain consistent. This observation suggests that BrainOOD captures critical features that are robust across both ID and OOD scenarios, indicating its strong generalization ability.

[W3. Paper Organization.] We included the background sections on GNNs and brain networks to cater to a broader audience, as our work intersects multiple domains (machine learning and neuroscience). However, we understand that this level of detail may not be necessary for readers already familiar with these topics. In response to your suggestion, we have streamlined these sections to focus on the most relevant aspects and reduce redundant explanations in our revision. Besides, we have moved the related work section after the experiment section in the revision to enhance the clarity and readability of the paper.

[Q1. ​​Balancing the Loss Functions and Risk of Overfitting.] Balancing multiple loss functions is indeed a critical aspect of our proposed method. We employ a series of trade-off hyperparameters to combine the four auxiliary losses. The weights are treated as hyperparameters and tuned using grid search on the validation set.

Regarding the risk of overfitting, we acknowledge that the complexity of the model and the inclusion of multiple modules may increase the likelihood of overfitting, particularly on small datasets. To mitigate this risk, we adopt several strategies:

• (1) Early Stopping: We monitor the validation loss and stop training when there is no improvement over a predefined number of epochs.
• (2) Regularization Techniques: We apply dropout and L2 regularization to reduce the model’s capacity and encourage generalization.
• (3) Cross-Validation: We use 10-fold cross-validation, ensuring that our results are robust and not dependent on a specific split.

These implementation details are included in Appendix D.2.

[Q2. Adaptability and Transferability to Other Models.] We appreciate the reviewer’s interest in the adaptability of our method. The BrainOOD framework is designed to be modular, making it straightforward to incorporate with various GNN backbones. Apart from the GIN backbone used in the main results, we also integrated BrainOOD with the GCN backbone, which shows consistent performance gains. This discussion is included in Appendix E.2. For broader applicability beyond the specific functional brain networks evaluated, we believe that the transferability of BrainOOD can further benefit models that encounter domain shifts in other types of neuroimaging data, such as DTI or EEG.

[Q3. Appropriateness of Using fMRI-Derived Brain Networks for OOD Generalization.] Our research is based on the brain network datasets proposed by [3], which focus on fMRI. We agree that evaluating our proposed method on brain networks that are constructed from various types of neuroimaging would enhance its robustness. However, our main focus is still developing an algorithm that could improve the OOD generalization ability of GNNs, rather than selecting the most appropriate modality for brian network construction. Moreover, fMRI-derived brain networks often exhibit greater variability across different sites and acquisition conditions, making them an ideal candidate for assessing the robustness of OOD methods.

[1] Interpretable and generalizable graph learning via stochastic attention mechanism. ICML 2022.

[2] Learning causally invariant representations for out-of-distribution generalization on graphs. NeurIPS 2022.

[3] Data-driven network neuroscience: On data collection and benchmark. NeurIPS 2023