🤖 AI

CLIMB: Controllable Longitudinal Brain Image Generation using Mamba-based Latent Diffusion Model and Gaussian-aligned Autoencoder

The paper introduces CLIMB, a controllable longitudinal brain image generation framework that combines a Mamba-based latent diffusion model with a Gaussian-aligned autoencoder to efficiently synthesize high-quality, temporally evolving brain MRI scans conditioned on various clinical and demographic factors.

Original authors: Duy-Phuong Dao, Muhammad Taqiyuddin, Jahae Kim, Sang-Heon Lee, Hye-Won Jung, Jaehoo Choi, Hyung-Jeong Yang

Published 2026-04-20

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Duy-Phuong Dao, Muhammad Taqiyuddin, Jahae Kim, Sang-Heon Lee, Hye-Won Jung, Jaehoo Choi, Hyung-Jeong Yang

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 have a time machine, but instead of traveling through space, it travels through time inside a human brain. That is essentially what this paper, titled CLIMB, is trying to build.

Here is the story of how they did it, explained without the heavy math jargon.

The Big Problem: Predicting the Future of a Brain

Doctors often need to know how a patient's brain will change over the next few years, especially for diseases like Alzheimer's. They have a "snapshot" (an MRI scan) of the brain today, but they want to see what it will look like in 5 or 10 years.

Currently, they have to wait for the patient to come back for a new scan. But what if they could generate that future scan right now using a computer? This would help doctors plan treatments earlier and understand the disease better.

The Solution: CLIMB

The researchers created a new AI system called CLIMB (Controllable Longitudinal brain Image generation using MamBa-based latent diffusion model). Think of it as a high-tech "Future-Brain Simulator."

Here is how it works, broken down into three simple parts:

1. The "Magic Sketchbook" (The Autoencoder)

Before the AI can predict the future, it needs to understand the brain in a simplified way.

The Old Way: Previous AI models tried to memorize every single pixel of a brain scan, like trying to memorize a library by reading every single word on every page. This is slow and messy. Also, they used a method that added "static" or "noise" to the memory, making the final picture blurry (like a photo taken in fog).
The CLIMB Way: They built a special tool called GATE (Gaussian-aligned Autoencoder). Imagine this as a magic sketchbook. Instead of memorizing every pixel, it compresses the brain scan into a clean, organized "sketch" (a latent representation).
The Trick: Unlike other tools that add random noise to their sketches, GATE draws the sketch so perfectly that it fits a specific, clean pattern (a "Gaussian distribution"). This means when the AI tries to turn the sketch back into a photo later, the image is crisp and sharp, not blurry.

2. The "Efficient Brain" (The Mamba Model)

Once the AI has the "sketch," it needs to figure out how the brain changes over time.

The Old Way: Most AI models use a mechanism called "Transformers" (specifically Self-Attention). Think of this like a student trying to read a book by looking at every single word and comparing it to every other word in the book to understand the story. It's very accurate, but it takes forever and requires a massive amount of energy (like a supercomputer running a marathon).
The CLIMB Way: They used a new architecture called Mamba. Think of Mamba as a smart reader who can understand the story by reading linearly, remembering the context without needing to re-read the whole book every time. It is much faster, uses less energy, and works just as well (or better) for this specific task. It's like upgrading from a horse-drawn carriage to a high-speed train.

3. The "Time Traveler's Controls" (Conditioning)

To make the prediction accurate, the AI needs more than just the current brain scan. It needs to know the "rules" of the patient's life.

CLIMB takes the current MRI scan and asks: "How old is the patient? Are they male or female? Do they have Alzheimer's? How big are their brain parts?"
It then uses a separate AI (called IRLSTM) to guess what those numbers will be in the future (e.g., "The patient will be 5 years older, and their memory center will shrink slightly").
The main AI then uses these "future stats" to draw the new brain scan. It's like giving an artist a reference photo and a set of instructions: "Draw this person 10 years older, with a slightly smaller memory center, but keep their unique face shape."

The Results: Why It Matters

The team tested this on thousands of real brain scans from the Alzheimer's Disease Neuroimaging Initiative (ADNI).

Accuracy: When they compared the AI's "fake" future scans to real scans taken years later, CLIMB was incredibly close. It got a score of 0.94 (on a scale where 1.0 is perfect), beating all previous methods.
Speed & Cost: Because it uses the "Mamba" engine instead of the heavy "Transformer" engine, it runs faster and uses less computer memory. This means hospitals could potentially run this on standard equipment, not just massive supercomputers.

The Bottom Line

CLIMB is a new tool that lets doctors see a patient's brain aging in fast-forward. By using a "clean sketchbook" (GATE) to store brain data and a "high-speed train" (Mamba) to process it, the AI can generate realistic, sharp images of what a brain will look like in the future based on age, genetics, and disease status.

This doesn't just save time; it could help doctors catch diseases like Alzheimer's earlier and plan better treatments before the damage becomes irreversible. It's like having a crystal ball for brain health, but one built on math and data rather than magic.

1. Problem Statement

Longitudinal medical image generation aims to synthesize follow-up medical scans (e.g., MRI) based on a baseline scan to simulate anatomical changes over time. This is critical for modeling disease progression (e.g., Alzheimer's), planning treatments, and handling missing data. However, existing methods face three major challenges:

Computational Inefficiency: Conventional Latent Diffusion Models (LDMs) often rely on Transformer-based backbones with self-attention mechanisms. These have quadratic complexity, making them computationally expensive and memory-intensive, especially for high-dimensional 3D medical volumes.
Latent Space Fidelity: Standard LDMs use Variational Autoencoders (VAEs) to compress images into a latent space. The stochastic sampling inherent in VAEs (reparameterization trick) introduces noise, often resulting in reconstructed images that lack sharpness or structural clarity.
Limited Controllability: Many existing models struggle to effectively incorporate multiple clinical conditional variables (e.g., projected age, gender, disease status, genetic info, and specific brain structure volumes) to drive realistic temporal evolution.

2. Methodology

The authors propose CLIMB (Controllable Longitudinal brain Image generation via MamBa-based latent diffusion model), a framework consisting of two primary stages:

A. Gaussian-Aligned Autoencoder (GATE)

To address the blurring issues of standard VAEs, the authors introduce GATE.

Deterministic Alignment: Unlike VAEs that inject stochastic noise, GATE uses a Sliced Wasserstein Distance (SWD) approach. It aligns the distribution of encoded latent features directly with a standard Gaussian prior ( $\mathcal{N}(0, I)$ ) using sliced cumulative density function (CDF) alignment. This is done deterministically without sampling noise.
Loss Function: The training objective combines four components:
1. Reconstruction Loss ( $\mathcal{L}_{rec}$ ): Pixel-level accuracy.
2. Sliced Distance Loss ( $\mathcal{L}_{SD}$ ): Enforces the latent distribution to match the Gaussian prior.
3. Perceptual Loss ( $\mathcal{L}_{perc}$ ): Ensures high-level feature similarity (using a pre-trained VGG/CLIP network).
4. Adversarial Loss ( $\mathcal{L}_{adv}$ ): Uses a discriminator to ensure photo-realistic outputs.
Architecture: Based on a ResNet encoder-decoder structure.

B. Mamba-based Latent Diffusion Model

Instead of using Transformers with self-attention, CLIMB employs Mamba, a State Space Model (SSM), as the backbone for the denoising network.

Efficiency: Mamba replaces self-attention with a recurrent hidden state mechanism that can be computed via convolution. This reduces complexity from quadratic to linear time, significantly lowering memory usage and inference time while maintaining the ability to capture long-range dependencies.
Two-Stage Training Strategy:
1. Stage 1: Train the diffusion model using only projected clinical variables (age, disease status, volumes) to learn the relationship between these conditions and the latent target.
2. Stage 2: Freeze the Stage 1 weights and initialize a parallel image-feature processing network. Train the full model with both image features and clinical variables, updating only the image network parameters. Zero-initialized convolutions are used to ensure smooth information flow.
Conditioning: The model takes a baseline MRI, its acquisition age, and projected future data (predicted by a separate IRLSTM model) as inputs. It uses Cross-Attention to inject these conditional cues into the denoising process.

C. Inference Process

Extract features from the baseline MRI using the GATE encoder.
Predict future clinical variables (disease status, brain volumes) using the IRLSTM model.
Feed the baseline features and projected variables into the Mamba-based denoising network.
Use Denoising Diffusion Implicit Models (DDIM) with 25 steps and latent average sampling (10 samples) to generate the future latent representation.
Decode the latent representation into the final MRI scan using the GATE decoder.

3. Key Contributions

CLIMB Framework: A novel architecture for controllable longitudinal brain image synthesis that integrates baseline scans with multi-modal conditional variables (age, gender, disease status, genetics, volumes).
Gaussian-Aligned Autoencoder (GATE): A deterministic autoencoder that aligns latent features to a Gaussian distribution using Sliced Wasserstein Distance, eliminating the stochastic noise and blurring associated with traditional VAEs.
Mamba Integration: The first application of the Mamba (SSM) architecture as the backbone for a latent diffusion model in medical imaging, offering superior computational efficiency (linear complexity) compared to Transformer-based self-attention mechanisms.
Two-Stage Training: A specialized training strategy that decouples the learning of clinical condition relationships from image feature processing to stabilize optimization.

4. Experimental Results

Dataset: The model was trained and evaluated on the ADNI (Alzheimer's Disease Neuroimaging Initiative) dataset, comprising 6,306 MRI scans from 1,390 participants (CN, MCI, and AD categories).
Baselines: Compared against SADM (State-space Attention Diffusion Model) and BrLP (Brain Longitudinal Prediction). Note: SADM was adapted to use latent diffusion to fit memory constraints.
Quantitative Performance: CLIMB achieved state-of-the-art results across all metrics:
- SSIM: 0.9433 (Highest, indicating superior structural similarity).
- PSNR: 27.82 (Highest).
- MSE: 2.01e-3 (Lowest error).
- LPIPS: 0.0587 (Lowest perceptual distance, indicating highest realism).
Efficiency:
- Inference Time: ~2.92 seconds per image (faster than baselines).
- Memory Usage: ~5,699 MB (significantly lower than Transformer-based counterparts).
Ablation Studies:
- Replacing GATE with a standard VAE resulted in a higher LPIPS (0.1080 vs. 0.0587), confirming the benefit of the Gaussian-aligned approach.
- Replacing Mamba with Self-Attention (Res+SA) increased memory usage and inference time while slightly reducing quality.
- Removing global context mechanisms (Res+CA) led to the worst performance.

5. Significance

Clinical Utility: CLIMB enables the generation of realistic future brain scans, aiding in early diagnosis, prognosis, and treatment planning for neurodegenerative diseases like Alzheimer's.
Scalability: By leveraging Mamba, the model makes 3D longitudinal diffusion modeling feasible on standard hardware, overcoming the memory bottlenecks of Transformer-based models.
Image Quality: The GATE module solves the "blurring" problem common in medical image generation, preserving critical anatomical details required for clinical interpretation.
Controllability: The ability to condition generation on specific clinical trajectories (e.g., "what will this patient's brain look like in 5 years if they progress to AD?") offers a powerful tool for personalized medicine research.

The authors note that while quantitative metrics are strong, future work requires validation by medical experts (radiologists/neurologists) to ensure the generated images are clinically accurate for diagnostic purposes.