Network-Aware Bilinear Tokenization for Brain Functional Connectivity Representation Learning

TL;DR

This paper introduces NERVE, a novel self-supervised learning framework that tokenizes brain functional connectivity matrices based on network modules, improving representation stability and transferability.

Contribution

NERVE employs a structured bilinear factorization for network-aware FC tokenization, addressing heterogeneity and preserving network identity, advancing brain connectomics analysis.

Findings

NERVE outperforms existing MAE and graph-based methods in behavior and psychopathology prediction.

Network-aware tokenization improves cross-cohort transferability of representations.

Ablation studies confirm the effectiveness of bilinear embedding and anatomically grounded parcellation.

Abstract

Masked autoencoders (MAEs) have recently shown promise for self-supervised representation learning of resting-state brain functional connectivity (FC). However, a fundamental question remains unresolved: how should FC matrices be tokenized to align with the intrinsic modular organization of large-scale brain networks? Existing approaches typically adopt region-centric or graph-based schemes that treat FC as structurally homogeneous elements and overlook the large-scale network brain organization. We introduce NERVE (Network-Aware Representations of Brain Functional Connectivity via Bilinear Tokenization), a self-supervised learning framework that redefines FC tokenization by partitioning FC matrices into patches of intra- and inter-network connectivity blocks. Unlike image-based MAE, where fixed-size patches share a common tokenizer, FC patches defined by network pairs are heterogeneous in size and correspond to distinct functional roles. To resolve this problem, NERVE embeds FC patches through a novel structured bilinear factorization. This formulation preserves network identity and reduces parameter complexity from quadratic to linear scaling in the number of networks. We evaluate NERVE across three large-scale developmental cohorts (ABCD, PNC, and CCNP) for behavior and psychopathology prediction. Compared to structurally agnostic MAE variants and graph-based self-supervised baselines, the proposed network-aware formulation yields more stable and transferable representations, particularly in cross-cohort evaluation. Ablation studies confirm that the proposed bilinear network embedding and anatomically grounded parcellation are critical for performance. These findings highlight the importance of incorporating domain-specific structural priors into self-supervised learning for functional connectomics. Code is available at: https://github.com/leomlck/NERVE.