Source-linked AI summary

Brain Network Transformer

Xuan Kan, Wei Dai, Hejie Cui, Zilong Zhang, Ying Guo, Carl Yang

arXiv:2210.06681v2cs.LGcs.CVcs.NE

TL;DR

Brain-network analysis needs Transformer designs that fit dense, fixed-order graphs and exploit their connection profiles without redundant positional encodings. The paper introduces BRAINNETTF with pairwise attention, OCREAD, and standardized evaluation, and reports superior performance over SOTA baselines on two large-scale datasets.

  • Problem

    Existing graph Transformer encodings are poorly suited to dense brain networks, while inaccurate functional-module labels and limited open-access datasets constrain analysis and comparison.

  • Method

    BRAINNETTF uses connection profiles as node features, learns pairwise attention weights, and applies OCREAD to softly cluster node embeddings with orthonormal projection.

  • Results

    BRAINNETTF achieves superior performance over SOTA baselines across experiments on the public ABIDE and restricted ABCD brain-network datasets.

  • Takeaways & Limitations

    The re-standardized ABIDE split provides a fair evaluation for new brain-network methods, while OCREAD supports cluster-aware graph representations without relying on fixed functional-module labels.

  • Takeaways & Limitations

    OCREAD is presented as potentially generalizable beyond graph-level brain-network prediction, but one theoretical integral has no known elementary antiderivative.

Abstract

from arXiv · show

Human brains are commonly modeled as networks of Regions of Interest (ROIs) and their connections for the understanding of brain functions and mental disorders. Recently, Transformer-based models have been studied over different types of data, including graphs, shown to bring performance gains widely. In this work, we study Transformer-based models for brain network analysis. Driven by the unique properties of data, we model brain networks as graphs with nodes of fixed size and order, which allows us to (1) use connection profiles as node features to provide natural and low-cost positional information and (2) learn pair-wise connection strengths among ROIs with efficient attention weights across individuals that are predictive towards downstream analysis tasks. Moreover, we propose an Orthonormal Clustering Readout operation based on self-supervised soft clustering and orthonormal projection. This design accounts for the underlying functional modules that determine similar behaviors among groups of ROIs, leading to distinguishable cluster-aware node embeddings and informative graph embeddings. Finally, we re-standardize the evaluation pipeline on the only one publicly available large-scale brain network dataset of ABIDE, to enable meaningful comparison of different models. Experiment results show clear improvements of our proposed Brain Network Transformer on both the public ABIDE and our restricted ABCD datasets. The implementation is available at https://github.com/Wayfear/BrainNetworkTransformer.

1 Introduction

Brain networks differ from the settings targeted by existing graph Transformers because they are dense, fixed-order graphs whose connection profiles already encode node identity and position. BRAINNETTF exploits these properties with pairwise attention, orthonormal clustering readout, and a standardized ABIDE evaluation.

  • Brain networks are built from ROIs connected by pairwise correlations between their BOLD signal series.
  • Existing graph Transformer designs are impractical because brain networks are typically fully connected, equal-degree graphs with redundant eigendecomposition-based positional embeddings.
  • BRAINNETTF uses connection profiles as node features and learns fully pairwise attention weights to model predictive brain-network structures.
  • OCREAD pools graph embeddings from functionally similar ROI clusters using self-supervised soft clustering and orthonormal projection.
  • ABIDE is the only large-scale publicly accessible dataset with extracted brain networks, but its 17-site acquisition introduces scanner and parameter heterogeneity.

2 Background and Related Work

Related work adapts attention and positional encoding to general graphs, while brain-network studies develop grouping, ROI-aware, interpretable, and network-generation approaches. These methods motivate a specialized Transformer design for brain-network structure.

  • Brain-network models use grouping, ROI-aware pooling, interpretable disorder-specific representations, or learnable network generation for analysis and explanation.
  • Graph Transformers incorporate edge information, positional embeddings, global structure, or molecule-specific encodings into graph representation learning.

3 BRAIN NETWORK TRANSFORMER

BRAINNETTF transforms fixed-size brain-network matrices into node and graph representations through multi-head attention and OCREAD. OCREAD softly clusters learned node embeddings with orthonormal centers, while the paper provides theoretical support and discusses broader applicability and limitations.

  • 3.1 Problem Definition: Given X ∈ R^V×V, MHSA maps the brain-network matrix to attention-enhanced node embeddings Z^L ∈ R^V×V.
  • 3.2 Brain Network Transformer: The model uses connection profiles as initial node features, removes extra positional embeddings, and applies vanilla pairwise attention on the complete graph.
  • 3.3 ORTHONORMAL CLUSTERING READOUT: OCREAD assigns nodes softly to K clusters using a Softmax projection, then computes graph embeddings as Z_G = P^⊤Z^L.
  • 3.3 ORTHONORMAL CLUSTERING READOUT: OCREAD uses orthonormal cluster centers initialized through Xavier uniform initialization followed by Gram–Schmidt orthonormalization.
  • 3.3.1 Theoretical Justifications: Theoretical analysis states that the relevant integral is maximized for orthonormal centers and that orthonormal sampling yields a lower significance level α_Ek· for well-estimated pooling.
  • 3.4 Generalizing OCREAD to Other Graph Tasks and Domains: The paper notes that OCREAD may extend beyond graph-level brain-network prediction to node or edge tasks, subgraph embeddings, and biomedical graphs, while reporting an unclear elementary antiderivative for one integral.

4 Experiments

Experiments evaluate BRAINNETTF across datasets, baselines, readout choices, hyperparameters, and attention patterns. Results report strong baseline performance, OCREAD benefits, and attention consistency with functional modules.

  • Experimental Settings: ABIDE evaluation uses site-aware stratified splitting because multi-site heterogeneity can destabilize training and create validation-testing gaps.The authors share the re-standardized data split as a benchmark pipeline for future methods.
  • Experimental Settings: BRAINNETTF is evaluated on ABIDE and ABCD using binary classification tasks, with AUROC, accuracy, Sensitivity, and Specificity as performance metrics.ABIDE targets ASD diagnosis, whereas ABCD evaluates biological sex prediction; both datasets are class-balanced.
  • Performance Analysis (RQ1): BRAINNETTF outperforms SAN and Graphormer by up to 6% absolute on both datasets and consistently surpasses BrainGNN, BrainGB, and BrainnetCNN.The reported BRAINNETTF-over-baseline gains pass a t-test with p-value<0.03.
  • OCREAD with varying readout functions: OCREAD is the most effective readout across SAN, Graphormer, and VanillaTF, improving prediction power across the Transformer architectures.The comparison varies the readout function while holding the evaluated Transformer architectures in view.
  • OCREAD with varying cluster initializations: With orthonormal cluster centers, performance rises from 2 to 10 clusters and then falls through 100, indicating a relatively small optimal cluster number.The observed range is consistent with typical functional-module counts being smaller than 25 and also reduces computation.
  • In-depth Analysis of Attention Scores and Cluster Assignments (RQ3): Average first-layer self-attention scores in the ABCD test set match available functional-module divisions, supporting model explainability.The corresponding comparison cannot be visualized for ABIDE because functional-module labels are unavailable for its atlas.

5 Discussion and Conclusion

Brain imaging supports brain-behavior research and diagnosis prediction, while the paper presents BRAIN NETWORK TRANSFORMER with an orthonormal clustering readout for brain network analysis.

  • Neuroimaging studies investigate relationships between brain biology, demographic variables, and clinical outcomes such as mental disorders.
  • Using brain imaging data to predict diagnosis may support mechanisms for more effective preemption and treatment.
  • BRAIN NETWORK TRANSFORMER combines a specialized graph Transformer with ORTHONORMAL CLUSTERING READOUT for brain network analysis.
  • The paper reports superior performance over various state-of-the-art baselines on two large-scale brain network datasets and provides a re-standardized ABIDE split for fair evaluation.

Checklist

The checklist records affirmative responses for reporting contributions, limitations, assumptions, proofs, reproducibility details, and asset documentation, with several items marked not applicable.

  • The authors report that they described their work’s limitations in Section 5.
  • The checklist states that theoretical assumptions and complete proofs were provided in Appendix C.
  • The authors report providing code, data, instructions, training details, error bars, and compute information for reproducing experiments.
  • The implementation repository is shared under an MIT license, while ABIDE and ABCD are reported as providing no license.
  • The checklist marks participant instructions, participant risks, compensation, and potential negative societal impacts as not applicable.

A Training Curves of Different Models with or without StratifiedSampling

Figure 5 compares training curves for models with and without stratified sampling based on ABIDE site information.

  • Stratified sampling produces a smaller validation–test performance gap on ABIDE than training without stratified sampling.
  • Stratified sampling stabilizes training on ABIDE, especially for VanillaTF and other evaluated model variants.
  • The curves of different variants show similar patterns across three model architectures in a single run.
  • Graphormer is omitted from the figure because its performance is much worse than that of the other models.

B Transformer Performance with Different Node Features

The section compares Transformer node features and reports that augmenting connection profiles with identity or eigen features does not improve performance.

  • Connection profiles use each node’s corresponding adjacency-matrix row as its feature representation.
  • Identity features assign each node a unique one-hot vector, while eigen features derive k-dimensional vectors from k adjacency-matrix eigenvectors.
  • Adding identity or eigen node features to connection profiles cannot improve the Transformer’s performance according to empirical observations.
  • The section motivates statistical methods for assessing the orthonormal initialization because direct performance comparison with non-orthonormal initialization is impractical.

C.1 Proof of Theorem 3.1

The proof shows that the variance of the Softmax readout is maximized when cluster centers are orthonormal, using symmetry and directional-derivative arguments.

  • Theorem 3.1 establishes that the variance of Softmax projection averaged over a round ball attains its maximum with orthonormal cluster centers.
  • In two dimensions, spherical symmetry reduces the comparison to the angle between two cluster centers.
  • Periodic and symmetry properties of the integrand enable comparison of readout integrals without requiring an elementary antiderivative.
  • A nonnegative directional derivative along v = (1, 2) shows that the integral increases with the angle until the orthonormal case.
  • Higher-dimensional cases follow by analogous spherical and hyperspherical-coordinate arguments.

C.2 Proof of Theorem 3.2

The proof compares orthonormal and non-orthonormal cluster centers statistically, showing that orthonormality improves regression accuracy and reduces estimation error in the readout analysis.

  • Theorem 3.2 treats readout probabilities as regression outputs and compares estimation using orthonormal versus non-orthonormal cluster centers.
  • Non-orthonormal cluster centers create correlated categorical variables, producing overlap that complicates decision-boundary identification.
  • The analysis uses nonlinear regression, Gauss–Newton linearization, goodness-of-fit through R2, and hypothesis testing.
  • The proof argues that orthonormal centers yield higher regression accuracy, lower MSE, and more accurate appraisal in hypothesis testing.
  • VIF ≥ 1, with equality if and only if the variables are orthogonal; orthonormal centers therefore avoid variance inflation from variable correlation.

D Running Time

BRAINNETTF and VanillaTF run faster than Graphormer and SAN because their computations are optimized for the small, densely connected graphs typical of brain networks.

  • Graphormer and SAN are much slower than BRAINNETTF and VanillaTF in the running-time comparison.
  • Brain networks typically have fewer than 500 nodes and dense connectivity, with e ≃ v2 rather than e ≪ v2.
  • Sparse graph operations optimized for e ≪ v2 do not work properly for these dense brain-network graphs.
  • BRAINNETTF and VanillaTF accelerate computation through direct matrix multiplication, while Graphormer also incurs edge-feature-generation costs.
  • The comparison includes graph Transformers and other neural-network baselines, with implementation and tuning details reported separately.

H The Difference between Various Initialization Methods

The initialization experiment compares random and orthonormal cluster-center initialization using a difference score, finding significant margins on both ABIDE and ABCD.

  • The difference score d quantifies the contrast between initialization methods using the model outputs P.
  • The experiment compares random and orthonormal initialization methods on the ABIDE and ABCD datasets.
  • The margins between random and orthonormal initialization are significant on both datasets according to a t-test.
Loading 2210.06681v2…