MAMBA-BASED RESIDUAL GENERATIVE ADVERSARIAL NETWORK FOR FUNCTIONAL CONNECTIVITY HARMONIZATION DURING INFANCY

Abstract

How to harmonize site effects is a fundamental challenge in modern multi-site neuroimaging studies. Although many statistical models and deep learning methods have been proposed to mitigate site effects while preserving biological characteristics, harmonization schemes for multi-site resting-state functional magnetic resonance imaging (rs-fMRI), particularly for functional connectivity (FC), remain undeveloped. Moreover, statistical models, though effective for region-level data, are inherently unsuitable for capturing complex, nonlinear mappings required for FC harmonization. To address these issues, we develop a novel, flexible deep learning method, Mamba-based Residual Generative adversarial network (MR-GAN), to harmonize multi-site functional connectivities. Our method leverages the Mamba Block, which has been proven effective in traditional visual tasks, to define FC-specified sequential patterns and integrate them with a multi-task residual GAN to harmonize multi-site FC data. Experiments on 939 infant rs-fMRI scans from four sites demonstrate the superior performance of the proposed method in harmonization compared to other approaches.

1. INTRODUCTION

Recent large-scale multi-site neuroimaging studies have demonstrated greater power to detect biologically relevant variability and has offered invaluable insights into the mechanisms underlying neurodevelopmental and neurodegenerative disorders [1]. However, the aggregation of neuroimaging data across different sites and scanners typically introduces non-biological variability, also known as site effects. These effects, stemming from site-related variations, can bias neuroimaging measures and distort the interpretation of clinical signals [2]. Consequently, many harmonization methods have been proposed to mitigate these unwanted site effects while preserving biological variability [3].

Drawing inspiration from image-to-image translation techniques in computer vision [4], many approaches have harmonized neuroimaging data in the image domain by synthesizing brain images across different sites [5]. However, due to the inherent complexity of neuroimaging data processing pipelines, these image-based harmonization techniques cannot guarantee that the final derived structural or functional features are free from site effects. Although some feature-based harmonization techniques have been proposed to directly mitigate site effects in derived volumetric [6], structural and functional [7, 8] magnetic resonance imaging (MRI) features, most of these methods rely on linear statistical models, such as ComBat [8] and its variants [6]. As data dimensionality increases and connections between brain regions become more complex, these methods struggle to capture subtle topological changes and effectively remove correlated noise. Moreover, when handling matrix data such as functional connectivity (FC), linear models and their parameter distribution assumptions may overlook the intricate topological structure of FC and the underlying connections between brain regions, leading to inadequate nonlinear mappings for FC. Although some deep learning methods have been applied to harmonize high-dimensional features such as cortical thickness maps [9], no learning-based FC-specified harmonization tools currently exist.

Recently, State Space Models (SSMs) [10], particularly structured state space sequence models (S4), have emerged as efficient and effective building blocks for deep networks, achieving state-of-the-art performance in continuous long-sequence data analysis. Mamba [11] enhances S4 by introducing a selective mechanism that enables the model to extract relevant information based on input data. With a hardware-aware implementation, Mamba outperforms Transformers in dense modalities such as language and genomics. State space models have also demonstrated promising results in vision tasks, such as image classification [12]. Given that image patches and features can be represented as sequences [12], these appealing properties of SSMs motivated us to explore Mamba blocks for feature extraction in FCs.

To this end, we propose a novel Mamba-based Residual Generative adversarial network (MR-GAN) for harmonizing early developing FC across sites. This model effectively removes site-related features and uniformly maps multi-site FC to a target site space, facilitating harmonization and enabling comprehensive FC analysis in subsequent studies. Specifically, 1) drawing inspiration from State Space Sequence Models (SSMs) [11, 10] and its variants [12], a family of deep sequence models renowned for their ability to handle long sequences, we first develop an FC-specified Mamba block capable of capturing both localized fine-grained features and long-range brain region dependencies in FC. 2) We leverage an adversarial strategy [13] to enhance the disentanglement of site-related and site-unrelated features. 3) Inspired by Cycle-GAN [4], we develop a cycle-consistent scheme to enforce consistency and employ a multi-task discriminator to ensure the authenticity and precision of the generated FC. Experiments on 939 infant rs-fMRI scans from four sites demonstrate the superior performance of the proposed method in harmonization compared to other approaches.

2. METHODS

2.1. Preliminaries and FC-specified Mamba Block

State Space Sequence Models (SSMs) [10] are a type of system designed to process 1-D function or sequence $u\left(t\right)\mapsto y\left(t\right)\in R$ and can be represented as the following linear Ordinary Differential Equation (ODE):

$$\begin{gathered} h^{\prime }(t)=Ah(t)+Bu(t), \\ y(t)=Ch(t) \end{gathered}$$ (1)

where state matrix $\mathbf{A}\in R^{N\times N}$ and $\mathbf{B},\mathbf{C}\in R^N$ are its parameters and $h\left(t\right)\in R^N$ denotes the implicit latent state.

Recently, Mamba [11] has further advances SSMs in discrete data modeling (e.g., text and genome). By merging SSM blocks with linear layers, the Mamba architecture demonstrates state-of-the-art performance on a variety of long sequence domains, with significant computational efficiency during both training and inference.

To effectively facilitate feature extraction and representation learning of FCs, we introduce the FC-specified Mamba Block, inspired by [11, 12]. As shown in Fig. 1, we developed an FC-specified sorting strategy to capture both localized fine-grained features and long-range dependencies between brain regions. The $n\times n$ FC matrix x is partitioned into n row-wise patches to preserve the FC map of each brain region. We defined connection strengths above 0.5 as strong functional connections and computed the number of such connections between each brain region and others at the group level. These numbers are then arranged in descending order to create a new brain region sequence s. The sorted brain region sequence features by Sort Block are fed into the Bidirectional Vision Mamba block [12], which processes the sequence in both forward and backward directions to capture richer features from the 2-dimensional data. During the decoder stage, the features are resorted back into their original sequence by Re-Sort Block. The encoder and decoder in this model utilize the Mamba block as their backbone.

Fig. 1.

The framework of the proposed Mamba-based Residual Generative Adversarial Network (MR-GAN). The FC input encoder extracts latent features and employs the Adversarial Disentangled Strategy to distinguish site-specific from non-site features. These features are then merged with the target site features to generate FC residuals, which are added to the original FC to achieve harmonization. The discriminator and cycle-consistent strategies enhance the precision and robustness of the generated outputs.

2.2. Adversarial Disentangled Strategy

To detect and remove site effects from multi-site FC, we introduce an adversarial disentangled strategy [13]. As shown in Fig. 1, for the latent feature z obtained from the input s across M sites, let $\mathit{z}=\left({\mathit{z}}_t,{\mathit{z}}_i\right)$, where ${\mathit{z}}_t$ is a 1 × M vector representing classification probabilities, which is site-related, and ${\mathit{z}}_i$ contains the remaining subject-specified and site-invariant features. We use an adversarial excitation and inhibition method with term:

$$\begin{gathered} {ℒ}_{Excitation}= \\ -\left[t_i\cdot \mathit{log}\left({\mathit{W}}_t\left({\hat{\mathit{z}}}_t\right)\right)+\left(1-t_i\right)\cdot \mathit{log}\left(1-{\mathit{W}}_t\left({\hat{\mathit{z}}}_t\right)\right)\right] \end{gathered}$$ (2)

$$\begin{gathered} {ℒ}_{Inhibition}= \\ {t}_i\cdot \mathit{log}\left({\mathit{W}}_i\left({\hat{\mathit{z}}}_i\right)\right)+\left(1-t_i\right)\cdot \mathit{log}\left(1-{\mathit{W}}_i\left({\hat{\mathit{z}}}_i\right)\right) \end{gathered}$$ (3)

where ${\mathit{W}}_t$ and ${\mathit{W}}_i$ are site classifiers utilizing latent variables ${\mathit{z}}_t$ and ${\mathit{z}}_i$, respectively. The one-hot vector $t_i$ encodes the ground-truth site label. We apply multi-class binary cross-entropy loss to align ${\mathit{z}}_t$ with the correct site label as in Eq. 2, while an inhibition process ensures ${\mathit{z}}_i$ remains site-unrelated with loss shown in Eq. 3. This approach allows us to extract non-site features, which are then combined with site-specific features in the next step and passed to the decoder to generate the target FC.

2.3. Cycle-Consistent Residual GAN

We develop a Cycle-Consistent Residual GAN to make sure the generated FC contains the target site features and closely matches the ground truth, as illustrated in Fig. 1. For a FC from site a, we generate its residual FC using the encoder E and decoder D, which is then added to the original $\mathrm{F}\mathrm{C}{\mathit{x}}_i^a$ to obtain the $\mathrm{F}\mathrm{C}{\hat{\mathit{x}}}_i^b$ in target site b. After that, we perform a reverse mapping to predict the original input $\mathrm{F}\mathrm{C}{\mathit{x}}_i^a$ as follows:

$$\begin{gathered} ({\hat{\mathit{z}}}_t^a,{\hat{\mathit{z}}}_i)=\mathit{E}({\mathit{x}}_i^a),{\hat{\mathit{x}}}_i^b=\mathit{D}({\mathit{z}}_t^b,{\hat{\mathit{z}}}_i)+{\mathit{x}}_i^a \\ ({\widehat{\stackrel{ˆ}{\mathit{z}}}}_t^b,{\widehat{\stackrel{ˆ}{\mathit{z}}}}_i)=\mathit{E}({\hat{\mathit{x}}}_i^b),{\widehat{\stackrel{ˆ}{\mathit{x}}}}_i^a=\mathit{D}({\mathit{z}}_t^a,{\widehat{\stackrel{ˆ}{\mathit{z}}}}_i)+{\hat{\mathit{x}}}_i^b \end{gathered}$$ (4)

We then impose the cycle-consistency loss to ensure that the reconstructed $\mathrm{F}\mathrm{C}{\widehat{\stackrel{ˆ}{\mathit{x}}}}_i^a$ and latent non-site feature ${\mathit{z}}_i$ are precise representations of the original $\mathrm{F}\mathrm{C}{\mathit{x}}_i^a$ and reverse latent non-site feature ${\hat{\mathit{z}}}_i$:

$${ℒ}_{\mathit{cycle}}={\Vert {\widehat{\stackrel{ˆ}{\mathit{x}}}}_i^a-{\mathit{x}}_i^a\Vert }_2^2+{\Vert {\widehat{\stackrel{ˆ}{\mathit{z}}}}_i-{\hat{\mathit{z}}}_i\Vert }_2^2$$ (5)

Besides, we use a multi-task discriminator to ensure the authenticity and precision of the generated FC. We label all ground-truth FC as real and the generated FC as fake to train the discriminator. Additionally, the discriminator also performs site classification to ensure that the generated FC preserves the target site’s information.

3. EXPERIMENTS

3.1. Datasets and Implementation Details

We evaluated our method on four infant datasets: UNC/UMN Baby Connectome Project (BCP) [14], Multi-visit Advanced Pediatric Brain Imaging Study (MAP) [15], Calgary study [16] and Pixar study [17]. Detailed information is provided in Table 1. For infant structural and functional MRI processing, we followed the methodologies detailed in [18] to extract the fMRI time-series for each vertex on middle cortical surfaces and all-time series within each cortical region were then averaged based on [19]. We selected 200 cortical regions per hemisphere to construct the functional connectivity matrix by computing the Pearson’s correlation coefficient between the time series of each pair of region of interests (ROIs), followed by applying Fisher’s r-to-z transformation.

Table 1.

Data information of four infant datasets.

	Subject Number (Males)	Total Scans	Age Range (Months)
BCP	305 (143)	661	0-73
MAP	79 (39)	328	0-75
Calgary	67 (37)	156	31-85
Pixar	61 (27)	61	42-72

The encoder, decoder, and discriminator all employ the Vision Mamba (Vim) structure [12]. We adopt the Vim-base model, consisting of 24 encoder blocks, 8 decoder blocks, a 384-dimensional encoder embedding, and a 256-dimensional decoder embedding. Additionally, a masked autoencoder [20] is employed as a pretraining strategy to enhance the feature-capturing ability of the Mamba block. For the discriminator’s identification task and the site classification task, we use fully connected layers with dimensions of (256, 2) and (256, 4), respectively. During the disentanglement process, a simple sigmoid layer outputs the probability of z_t, while a fully connected layer of dimension (380,4) is used fo z_i to adversarially train the encoder to extract site-unrelated features. We set BCP dataset as the target dataset, with all four datasets mapped to the BCP domain for harmonization.

3.2. Validation on Removing Site Effects

To validate whether site effects are successfully removed while preserving the essential non-site FC features, we present the population-level developmental trajectory of functional networks in Fig. 2. We use the Combat algorithm and use age as covariates to compare the generated results with MR-GAN. We calculate the intra- and inter-network FC strength between 17 functional networks of each subject by averaging the FC connection related to the corresponding functional networks. We selected intra-network connections of Visual Central (VisCent), SomatomotorA (SomMotA), and Dorsal Attention A (DorsAttnA) as the representative connections to demonstrate the population-level developmental trajectory. As shown in Fig. 2, the linear statistical model, ComBat, struggles to effectively handle extreme values and outliers, often diminishing the overall strength of functional connections. In contrast, our MR-GAN achieves better harmonization while preserving both general trajectory trends and the unique characteristics for individual trajectory.

Fig. 2.

The visualization of the functional network developmental trajectories. Intra-network connections of VisCent, SomMotA, and DorsAttnA were selected as the representatives. Solid lines: growth chart; shaded area: 95% confidence intervals.

3.3. Validation on Downstream Tasks

We further investigate the impact of different harmonization methods on downstream tasks. We use Brain Connectivity based Graph Convolutional Networks (BC-GCN) [21], a deep learning based model designed specifically for FC based age prediction, to predict the scan age. Given the significant differences in age distributions across the four datasets, we designed four age prediction tasks: (1) using BCP as the training set and MAP as the test set (0–800 days); (2) using MAP as the training set and BCP as the test set (0–800 days); (3) using MAP and Pixar as the training set and Calgary as the test set (800–1600 days); (4) using Calgary as the training set and MAP and Pixar as the test set (800–1600 days). We used mean absolute error (MAE) to evaluate performance. The results, presented in Table 2, indicate that ComBat does not improve prediction accuracy compared to non-harmonized data, likely due to the limitations of linear statistical models in preserving the detailed information and intricate connection structures in two-dimensional FC data. As a result, such harmonization may cause the loss of many local features. In contrast, our MR-GAN significantly improves the accuracy and reliability of analysis across multi-site data compared to both non-harmonized and ComBat-based harmonization method.

Table 2.

Experimental results of age prediction by BC-GCN in terms of MAE (days).

	Task 1	Task 2	Task 3	Task 4
Raw	129.77	103.22	141.85	186.71
ComBat	127.07	137.96	147.80	185.27
Mamba	123.07	98.20	135.63	183.28

4. CONCLUSION

In this study, we propose a novel Mamba-based Residual Generative Adversarial Network (MR-GAN) to harmonize multi-site infant brain functional connectivity (FC). Unlike the previous harmonization methods, MR-GAN includes an FC-specified Mamba block, designed to preserve detailed features of FC. The Mamba Block effectively captures the relational characteristics between brain regions in FC and supports subsequent FC generation. To the best of our knowledge, this is the first application of the Mamba structure in the FC field. Additionally, we employed an adversarial disentanglement strategy to separate site-related and non-site features and utilized the Cycle-GAN structure to ensure that the generated FCs contain the target site feature and closely resemble the real FC. The promising results across four infant datasets underscore the superiority of our model in harmonizing functional connectivity during infancy.