Isolating Nonlinear Independent Sources in fMRI with $\beta$-TCVAE Models

Technical Summary: Isolating Nonlinear Independent Sources in fMRI with β-TCVAE Models

Problem Statement
Learning meaningful latent representations from nonlinear functional MRI (fMRI) data remains a fundamental challenge in neuroimaging. Traditional Independent Component Analysis (ICA), while widely used for estimating interpretable functional brain networks, relies on a linear mixing assumption for latent sources. This limitation restricts its ability to capture the inherently nonlinear and complex organization of brain dynamics. Although deep representation learning methods have emerged as alternatives for modeling nonlinear latent structures, they have been primarily evaluated on simulated datasets or natural image benchmarks, with comparatively limited validation on real-world neuroimaging data. Furthermore, existing nonlinear ICA extensions often lack systematic validation in neuroimaging contexts where complex spatial-temporal dependencies and biological variability present additional challenges.

Methodology
To address these gaps, the authors adapt and modify the β-TCVAE (Total Correlation Variational Autoencoder) framework, originally proposed as a refinement of the β-VAE to learn disentangled representations by penalizing total correlation (TC) among latent variables.

• Model Architecture: The authors employ a subject-conditioned variational autoencoder. The model takes multi-subject fMRI data $\{x^{(s)}_t\}$ as input, where $s$ indexes subjects and $t$ indexes time. To incorporate inter-subject variability, each subject index is mapped to a learned embedding vector $e_s$, which is concatenated with the input fMRI feature vector.
• Encoder and Decoder: The encoder $q_\phi(z | x, s)$ infers a latent distribution $z$ from the input and subject identity, outputting a diagonal Gaussian distribution. The decoder $p_\theta(x | z)$ reconstructs the fMRI signal from the latent representation. The architecture utilizes multilayer perceptrons (MLPs) with specific hidden layer configurations (512/256 units for the encoder; 256/512 for the decoder) and ReLU activations.
• Objective Function: The model optimizes a decomposition of the Evidence Lower Bound (ELBO) into four terms:
  1. Reconstruction Loss ($L_{rec}$): A Gaussian likelihood measuring the difference between the input and reconstructed signal.
  2. Mutual Information (MI): Encourages the latent variables to retain information about the observed data.
  3. Total Correlation (TC): Penalizes the dependence between latent dimensions to enforce disentanglement.
  4. Dimension-wise KL Divergence ($KL_{dim}$): Regularizes the marginal distribution of each latent dimension against an isotropic standard normal prior.
• Training Details: The model was trained on the Human Connectome Project (HCP) dataset using the Adam optimizer. Input data underwent standard preprocessing (motion correction, normalization to MNI space, smoothing). The latent dimensionality was set to 80, and a linear warm-up schedule was applied to the $\beta$ coefficient over the first 10 epochs.

Key Results

• Disentanglement Dynamics: Training curves demonstrated that as the total loss decreased, both Mutual Information and Total Correlation were progressively minimized, suggesting successful separation of latent components over time.
• Biological Relevance: The framework successfully identified well-defined resting-state networks. Specifically, it recovered the Default Mode Network (DMN) and Orbitofrontal Cortex (OFC) with clear spatial coverage across canonical regions (e.g., posterior cingulate cortex, medial prefrontal cortex).
• Comparison with InfoMax ICA: When compared to the traditional InfoMax ICA approach:
  • The β-TCVAE-derived DMN exhibited more spatial coherence and complete spatial coverage compared to the InfoMax-derived DMN, which showed reduced spatial continuity.
  • Spatial correlation analyses between corresponding components from both methods yielded moderate correlations, indicating comparable functional network structures.
  • Crucially, β-TCVAE identified unique components that captured multi-network coupling patterns (e.g., simultaneous visual and frontal networks), which were not observed in the InfoMax results. This suggests β-TCVAE better preserves interactions between distributed brain networks, whereas InfoMax tends to enforce stronger independence assumptions that may separate these patterns.
• Functional Network Connectivity (FNC): The learned representations were used to compute an average FNC matrix, which displayed distinct clustered structures along the diagonal, indicating organized interactions among subsets of networks.

Significance and Claims
The paper positions this work as a pilot investigation bridging nonlinear representation learning and fMRI analysis. The authors claim that:

1. The adapted β-TCVAE framework can recover meaningful nonlinear spatial components with biological relevance, including established intrinsic connectivity networks.
2. The model effectively separates mixed spatial and temporal brain signals into interpretable components without requiring additional hyperparameters during training beyond the standard β-VAE framework.
3. The approach offers a potential alternative to linear ICA by capturing coherent and interpretable brain organization patterns, including multi-network couplings that traditional ICA might overlook due to strict independence assumptions.
4. The study provides a foundation for developing constrained and interpretable representation learning methods for fMRI, extending traditional ICA toward approaches that better capture hierarchical and nonlinear brain dynamics.

The authors maintain a modest tone, noting that while the initial results are promising, further work is needed to improve stability and generalizability across larger, more diverse cohorts and to explore inference on unseen subjects.