🤖 machine learning

Unsupervised learning of acquisition variability in structural connectomes via hybrid latent space modeling

This paper introduces an unsupervised hybrid latent space framework with architectural annealing that effectively separates acquisition-related variability from biological signals in large-scale dMRI structural connectomes without requiring manual capacity tuning.

Original authors: Gaurav Rudravaram, Lianrui Zuo, Karthik Ramadass, Elyssa McMaster, Jongyeon Yoon, Aravind R. Krishnan, Adam M. Saunders, Chenyu Gao, Nancy R. Newlin, Praitayini Kanakaraj, Lori L. Beason Held, Murat B

Published 2026-05-15

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Gaurav Rudravaram, Lianrui Zuo, Karthik Ramadass, Elyssa McMaster, Jongyeon Yoon, Aravind R. Krishnan, Adam M. Saunders, Chenyu Gao, Nancy R. Newlin, Praitayini Kanakaraj, Lori L. Beason Held, Murat Bilgel, Laura A. Barquero, Micah DArchangel, Tin Q. Nguyen, Laurie B. Cutting, Derek Archer, Timothy J. Hohman, Daniel C. Moyer, Bennett A. Landman

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to organize a massive library of books about the human brain. These books aren't made of paper; they are digital maps showing how different parts of the brain are connected (called "structural connectomes").

The problem is that these maps were drawn by 13 different libraries (research studies) using 25 different types of cameras and scanners. Even if two people have identical brains, their maps might look slightly different just because one was scanned on a Tuesday with a specific setting, and the other on a Friday with a different setting. This is like trying to sort books by genre, but the cover colors are changing based on which printer made them, not the story inside.

This paper introduces a new, smart way to sort these brain maps so that the "printer differences" don't get mixed up with the actual "story differences."

The Problem: The "Noisy" Library

In the past, scientists tried to clean up these maps using standard tools. Think of these tools like a basic photo filter. They try to smooth out the differences, but they often treat everything as a continuous blur. They can't easily tell the difference between a "continuous" change (like a brain getting older) and a "discrete" change (like switching from Scanner A to Scanner B). As a result, the scanner differences get hidden inside the brain data, making it hard to see the real biological patterns.

The Solution: A Hybrid Sorting Machine

The authors built a new kind of AI, which they call a Joint-VAE. Imagine this AI as a two-lane highway for data:

The Continuous Lane: This lane handles smooth, gradual changes, like the natural aging process or subtle biological differences between people.
The Discrete Lane: This lane handles "on/off" or "category" changes, like which scanner was used or which study the data came from.

The goal is to force the AI to put the "scanner noise" into the Discrete Lane and the "brain biology" into the Continuous Lane.

The Innovation: The "Training Coach" (Architectural Annealing)

Here is the tricky part. If you just let the AI drive, it tends to ignore the Discrete Lane entirely. It's easier for the AI to dump all the information (both the scanner noise and the brain biology) into the big, wide Continuous Lane because that lane has unlimited space.

Previous methods tried to fix this by manually telling the AI, "Hey, you can only put 50% of the data in the Continuous Lane." But this is like a coach shouting specific numbers at a runner; it's hard to get right, and if you get the number wrong, the runner fails.

This paper's breakthrough is a new "Architectural Annealing" strategy. Instead of just shouting instructions, the authors built a training coach that physically blocks the Continuous Lane at the start of training.

Early Training: The coach puts a "Do Not Enter" sign on the Continuous Lane. The AI is forced to put all the information into the Discrete Lane. It learns to recognize the different scanners and protocols because it has no other choice.
Late Training: As the AI gets smarter, the coach slowly lifts the barrier on the Continuous Lane. Now, the AI can start using the Continuous Lane for the smooth biological details, but it has already learned to keep the scanner noise separate.

This happens automatically. The AI learns to balance itself without the researchers having to manually tune complex knobs.

The Results: Sorting the Library

The team tested this on a huge dataset of 7,416 brain maps from people aged 2 to 102, including healthy individuals and those with cognitive issues.

The Test: They asked the AI to group the maps based on which scanner they came from (without telling the AI the scanner names).
The Winner: The new "Architectural Annealing" method was the best at sorting the maps. It correctly identified the 25 different scanner groups much better than the old methods.
The Discovery: The AI didn't just learn to separate the scanners; it learned how the scanners differed. It realized that some scanners used different "shell" settings, different angles, or different timing. It organized the data in a way that reflected the complex mix of these technical settings.

The Bottom Line

This paper shows that by using a special "two-lane" AI with a smart training schedule, we can automatically separate the "noise" of different medical scanners from the "signal" of the human brain. This allows scientists to mix and match data from different hospitals and studies without the results getting confused by the equipment used to take the pictures.

Note: The paper explicitly states that in this unsupervised setup, the AI learned mostly about the scanners, not the biological differences between patients (like Alzheimer's). The authors suggest that future work could use this method to help separate those biological differences, but this specific study focused on proving the method works for cleaning up the scanner data.

Technical Summary: Unsupervised Learning of Acquisition Variability in Structural Connectomes via Hybrid Latent Space Modeling

Problem Statement
Diffusion MRI (dMRI) is essential for analyzing white-matter connectivity and microstructural integrity in neurodegenerative diseases, cognitive decline, and aging. However, dMRI data is heavily impacted by heterogeneity in acquisition protocols (e.g., scanner type, echo time, repetition time, b-values, and shell configurations). These acquisition differences introduce variability that often obscures underlying biological signals, making it difficult to draw generalizable conclusions across multi-site studies. While harmonization techniques exist, many rely on predefined site labels or metadata, which may be incomplete or inaccurate. Furthermore, conventional dimensionality reduction methods (e.g., standard Variational Autoencoders or PCA) typically model all variance as continuous, making it difficult to disentangle discrete acquisition effects from continuous biological variation. Existing hybrid latent space models that attempt to separate discrete and continuous factors often require manual capacity tuning to ensure the discrete component captures the desired variability, leading to instability and sensitivity to hyperparameters.

Methodology
The authors propose an Architectural Annealing Joint-VAE, an unsupervised framework designed to learn a hybrid latent space comprising both continuous ( $z_c$ ) and discrete ( $z_d$ ) variables to model structural connectomes.

Data Curation: The study utilizes a large, heterogeneous dataset ( $N=7,416$ ) derived from 13 neuroimaging studies, spanning ages 2 to 102. The data includes 25 unique acquisition parameter combinations (sites) and covers neurotypical, mild cognitive impairment (MCI), and Alzheimer's disease (AD) cohorts. Structural connectomes were constructed using standard preprocessing pipelines (PreQual, SLANT, MRtrix3) involving 121 ROIs and 10 million streamlines per subject.
Model Architecture: The model extends the traditional Variational Autoencoder (VAE). The encoder maps flattened connectome data to a joint posterior $q_\phi(z_c, z_d | x)$ $q_{ϕ} (z_{c}, z_{d} ∣ x)$ , while the decoder reconstructs the input from the concatenated latent vector $[z_c, z_d]$ $[z_{c}, z_{d}]$ .
- Continuous Latent ( $z_c$ ): Modeled as a Gaussian distribution.
- Discrete Latent ( $z_d$ ): Modeled as a categorical distribution (25 classes corresponding to the number of sites), sampled via Gumbel-Softmax.
Architectural Annealing: The core innovation is a staged annealing procedure applied directly to the encoder outputs rather than just the loss function.
- An annealing factor $\lambda$ ramps linearly from 0 to 1 over a fixed number of iterations.
- The encoder's mean ( $\mu$ ) and log-variance ( $\log \sigma^2$ ) for the continuous latent are scaled by $\lambda$ (i.e., $\mu' = \lambda \mu$ , $\log \sigma'^2 = \lambda \log \sigma^2$ ) before sampling and KL divergence calculation.
- Mechanism: At the start of training ( $\lambda \approx 0$ ), the continuous pathway is suppressed, forcing the model to encode structural information into the discrete latent space. As $\lambda$ increases, the continuous pathway is gradually unmasked, allowing for a balanced joint optimization. This architectural constraint prevents the continuous space from dominating early training dynamics, a failure mode observed in loss-based annealing approaches.

Key Contributions

Principled Unsupervised Framework: The paper introduces a method to separate discrete acquisition effects from continuous biological variation without requiring site labels or manual capacity tuning.
Architectural Annealing: Unlike previous approaches that anneal only the loss term (KL divergence), this method anneals the encoder outputs themselves. This ensures the decoder is conditioned on representations with controlled capacity throughout training, leading to more stable dynamics.
Large-Scale Validation: The framework is validated on a significantly larger and more diverse dataset (7,416 subjects, 25 sites) compared to prior hybrid latent space studies.

Results

Sensitivity to Manual Tuning: Experiments on the standard Joint-VAE (Eq. 5) demonstrated that the model is highly sensitive to manual continuous capacity ( $C_c$ ) settings. Increasing $C_c$ caused the discrete space to collapse, merging distinct site clusters and failing to recover site structure.
Superior Site Discovery: The proposed Architectural Annealing Joint-VAE significantly outperformed baselines, including PCA+k-means, standard VAE+k-means, and a Joint-VAE using only loss-based annealing.
- At 25 discrete classes (matching the true number of sites), the model-annealed approach achieved an Adjusted Rand Index (ARI) of 0.53 ( $p < 0.05$ ), significantly higher than all other methods.
- The model-annealed approach maintained stable performance as the number of classes approached the ground truth, whereas loss-annealed and PCA baselines showed marked performance drops.
Latent Space Interpretability: Visualization of the latent space revealed that the model successfully learned to separate multi-shell from single-shell acquisitions. Within single-shell groups, the continuous latent space smoothly organized subjects according to acquisition parameters (TE, TR, angular resolution), demonstrating a nonlinear factorization of acquisition variability.

Significance and Claims
The paper claims that the proposed hybrid continuous–discrete latent space provides a robust, unsupervised mechanism for capturing acquisition-related variability in diffusion MRI. By jointly modeling smooth and categorical structures, the model recovers meaningful clusters aligned with scanner and protocol differences without supervision. The authors emphasize that the learned latent structure reveals that no single acquisition parameter dominates variability; rather, the model captures a multi-factor dependence involving TE, TR, b-values, and shell structure.

The authors modestly note a limitation: in this fully unsupervised setting, the latent space is dominated by acquisition-related variation rather than biological differences between subjects. They suggest that future work could extend this to a semi-supervised framework where the discrete variable captures acquisition structure while dedicated continuous dimensions are encouraged to encode biological variability, thereby supporting large-scale, cross-study neuroimaging analyses that preserve meaningful biological signals.