MAGNET.

Accepted at NeurIPS 2026

MAGNET: Manifold-Aware Graph Diffusion Network for Functional Brain Connectome Generation

Protyay Dey*1 Ayush Roy*1 Hyuna Cho2 Won Hwa Kim3 Vishnu Suresh Lokhane1

* Equal contribution

1 University at Buffalo, SUNY  ·  2 National Institutes of Health, United States  ·  3 POSTECH, South Korea

Valid correlation geometry is only the beginning. MAGNET generates class-conditional functional connectomes while preserving finer clinical structure.

Code & reproduction Explore the method

Abstract

Recovering class-conditional geometry in functional connectomes.

Functional brain connectome represents neural connectivity as a matrix of pairwise interactions between brain regions. Generation of functional connectomes is not only a question of validity; rather, having satisfied the constraints on the correlation matrix, the next step is to recover the class-conditional geometry buried under coarse labels. We propose MAGNET, a Manifold-Aware Graph Diffusion Network which uses a normalized-Cholesky representation of the manifold of correlation matrices that guarantees validity. MAGNET lifts noisy latent states into ROI-level region tokens and performs denoising with a relational inductive bias over brain atlas regions. To deal with structural problems induced by coarse labels, MAGNET employs class-anchored conditioning, amortized structural bridge, and relevance-preserving corruption. Across ABIDE, ADNI, and OASIS-3, MAGNET consistently achieves favorable results compared to previous manifold-aware approaches, demonstrating improvements of 7-21% in class-conditional fidelity (α,β-F1) across three cohorts and better sampling efficiency. Moreover, while training only with strict binary labels, MAGNET is capable of maintaining clinical heterogeneity through fine substructure of the connectomes in a zero-shot setting, improving subclass covariance alignment (λ-MSE) by over 30%. These results suggest that geometric validity is a necessary but insufficient condition for clinical utility in connectome synthesis. Moreover, efforts in making the diffusion denoising class-conditional manifold aware finds utility beyond the highly curved brain connectome generation as this is a critical problem in various general settings.

The approach

Built around connectome geometry.

MAGNET treats a connectome as a structured point on the correlation-matrix manifold. A normalized-Cholesky coordinate preserves this geometry before ROI-level graph denoising, class-anchored conditioning, an amortized structural bridge, and relevance-preserving corruption shape the diffusion process.

For a correlation matrix C, MAGNET writes C = LLT and removes the diagonal scale from L before diffusing only through its strict lower triangle. With 39 MSDL regions this gives 741 latent coordinates. Decoding reconstructs the unit-lower-triangular factor, forms its Gram matrix, and renormalizes the diagonal, yielding a symmetric positive-definite correlation matrix with unit diagonal by construction.

MAGNET architecture showing correlation-matrix encoding, class-anchored conditioning, structural bridge, graph transformer denoising, and valid correlation decoding View full size
6sampling steps
7-21%improvement in class-conditional α,β-F1 across three cohorts
30%+improvement in zero-shot subclass covariance alignment

The evidence

Structure beyond validity.

Across ABIDE, ADNI, and OASIS-3, MAGNET is evaluated for class-conditional fidelity and preservation of finer connectome structure.

Explore the code