💻 computer science

ProgFormer: Hierarchical Voxel Diffusion Transformer for Longitudinal Brain MRI Prediction

ProgFormer is a hierarchical voxel-space Diffusion Transformer that addresses the challenge of longitudinal brain MRI prediction by employing a dual-pathway architecture to jointly preserve global structural consistency and capture fine-grained disease progression through conditional flow matching, outperforming existing state-of-the-art methods across multiple benchmarks.

Original authors: Dexuan Ding, Yuankai Qi, Luping Zhou, Jian Yang, Quan Z. Sheng, Ming-Hsuan Yang

Published 2026-07-31

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Dexuan Ding, Yuankai Qi, Luping Zhou, Jian Yang, Quan Z. Sheng, Ming-Hsuan 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 are a time traveler with a very special camera. You can take a picture of a person's brain today, but you want to know what it will look like five or ten years from now. This isn't just about guessing; it's about predicting how the brain changes as we age or as diseases like Alzheimer's slowly take hold. Scientists call this "longitudinal MRI prediction." The tricky part is that most of the brain stays exactly the same from year to year—it's like a sturdy house that doesn't change much. But the disease is like a tiny, slow leak in the roof or a crack in a specific wall. It happens in very small, specific spots. If you try to predict the future by looking at the whole house at once, you might miss those tiny cracks because the rest of the house looks so perfect and unchanged.

To solve this, researchers have tried two main tricks. One is to shrink the brain picture down into a tiny, blurry summary (a "latent space"), make the prediction, and then blow it back up. The problem is, when you shrink and blow up a picture, you often lose the fine details, like the texture of the wall or the exact shape of the crack. The other trick is to look at every single pixel (or "voxel" in 3D) directly. But if you try to predict the whole brain and the tiny cracks at the same time using one simple tool, the tool gets confused by the huge, stable parts of the brain and ignores the tiny, important changes.

Enter ProgFormer, a new AI model designed to be the ultimate time-traveling brain photographer. Instead of shrinking the image or using a one-size-fits-all tool, ProgFormer uses a clever "two-step" strategy to predict the future of a brain scan. Think of it like an artist painting a portrait of a person aging. First, the artist sketches the whole face and head to get the big picture right—this is the Coarse Pathway. It looks at chunks of the brain (called "patches") to understand the overall shape and how the whole brain has changed over time. It's like drawing the outline of the house.

Once the big picture is sketched, the artist switches to a magnifying glass for the Fine Pathway. This part zooms in on the specific areas where the "leaks" or cracks might be happening, like the hippocampus (a memory center) or the ventricles (fluid-filled spaces). It uses the sketch from the first step as a guide, asking, "Okay, I know the general shape of this room; now, exactly how does the wall here need to shift?" By combining these two steps, ProgFormer can keep the brain looking stable and realistic while still catching those tiny, subtle signs of disease progression.

The researchers tested this model on three large groups of real patient data (ADNI, AIBL, and OASIS). They found that ProgFormer is better at predicting future brain scans than previous methods. It creates clearer images and, more importantly, gets the tiny details right. When they looked at how well the model predicted the size of specific brain regions, it made fewer mistakes than other top models. The study also showed that looking at a patient's history of scans (not just the most recent one) helps the model make even better predictions.

In short, ProgFormer doesn't just guess the future; it builds it layer by layer. It respects the fact that the brain is mostly stable but is hyper-aware of the small, critical changes that signal disease. By separating the "big picture" work from the "fine detail" work, it manages to predict the future of a brain with a level of accuracy that suggests this two-step approach is a powerful new way to understand how our brains age.

Technical Summary: ProgFormer

Problem Statement

Predicting future structural MRI (sMRI) of the brain is a critical task for analyzing and forecasting neurodegenerative diseases like Alzheimer's. However, this presents a unique challenge: longitudinal changes are often subtle and confined to specific anatomical regions (e.g., hippocampus, lateral ventricles), while the majority of the subject's brain structure remains stable over time.

Existing methods generally fall into two categories, both with limitations:

Latent-space-based methods: These encode scans into low-dimensional representations before prediction. While computationally efficient, they rely on a separately learned autoencoder, which can introduce information loss during compression and reconstruction, potentially degrading the fidelity of fine-grained structural details.
Direct voxel-space methods: These predict outputs directly in voxel space to avoid latent reconstruction. However, they typically employ a unified prediction pathway to model both the stable global brain structure and the subtle, progression-related local changes. Consequently, the dominant stable content often overshadows the subtle local variations, making it difficult to resolve fine-grained disease progression.

Methodology: ProgFormer

The authors propose ProgFormer, a hierarchical voxel-space Diffusion Transformer designed to address these challenges without relying on a separately trained image autoencoder. The model operates within the Conditional Flow Matching framework, learning a velocity field to transport Gaussian noise to the target future MRI distribution.

Core Architecture

ProgFormer structures the prediction into two distinct but coupled pathways:

Coarse Pathway (Patch-Level Reasoning):
- Input: Processes observed scans and the current flow state as non-overlapping 3D patch tokens ( $p \times p \times p$ ).
- Mechanism: Utilizes Spatial Multi-Head Self-Attention (MHSA) within each scan and Causal Temporal MHSA across observed scans. This allows the model to capture global brain structure and longitudinal context.
- Output: Generates a primary volumetric velocity ( $\nu_{coarse}$ ) and coarse representations ( $C$ ) that encode whole-brain structure and temporal context.
Fine Pathway (Voxel-Level Refinement):
- Input: Operates on individual voxel tokens within each 3D patch.
- Mechanism: Applies Voxel MHSA within patches and, crucially, a Coarse-to-Fine Multi-Head Cross-Attention (MHCA). In this cross-attention, voxel tokens act as queries, while the coarse representations ( $C$ ) from the corresponding patch serve as keys and values.
- Function: This mechanism uses the coarse pathway's output as "spatio-temporal grounding" to guide the refinement of local structural details and intensity variations.
- Output: Generates a refined velocity ( $\nu_{fine}$ ).

Integration and Inference

Velocity Estimation: The final predicted velocity field is the sum of the coarse and fine velocities: $\nu_\theta = \nu_{coarse} + \nu_{fine}$ .
Generation: During inference, the model starts from Gaussian noise and integrates the estimated velocity field over a sequence of Euler steps to generate the predicted future MRI.
Variants: The authors evaluate two compositions:
- Sequential (ProgFormer-S): All coarse blocks are completed before any fine blocks begin.
- Interleaved (ProgFormer-I): Coarse and fine blocks are interleaved in stages, allowing fine blocks to receive progressively refined coarse features.

Conditioning and Training

Conditioning: The model is conditioned on flow time, target age, and covariates (diagnosis, sex) via AdaLN-Zero (Adaptive Layer Normalization with zero initialization), which modulates scale, shift, and gate parameters for attention and MLP modules.
Loss Function: A weighted flow-matching objective is used. An anatomical label map (derived from SynthSeg) assigns higher weights to progression-relevant regions (e.g., hippocampus, ventricles) and lower weights to stable tissue/background, ensuring the model prioritizes disease-related changes.

Key Contributions

End-to-End Hierarchical Architecture: ProgFormer is the first unified architecture to jointly perform global patch-level spatio-temporal reasoning and fine-grained intra-patch voxel modeling for longitudinal MRI prediction without a separate autoencoder.
Coarse-to-Fine Grounding: The design explicitly decouples the modeling of global stability (coarse pathway) from local progression (fine pathway), using the former to ground the latter, thereby preventing subtle changes from being overshadowed by stable structures.
State-of-the-Art Performance: Extensive experiments on three benchmarks (ADNI, AIBL, OASIS) demonstrate superior performance in both pairwise and trajectory prediction settings compared to existing latent-space and direct volumetric methods.

Experimental Results

The model was evaluated on ADNI, AIBL, and OASIS datasets using PSNR, SSIM, and Region-level Mean Absolute Error (R-MAE).

Pairwise Prediction: ProgFormer-S achieved the highest PSNR across all three cohorts (e.g., 26.39 on ADNI, outperforming the next best by 0.20). ProgFormer-I achieved the best regional accuracy (lowest R-MAE) and the highest SSIM on ADNI.
Trajectory Prediction: Using observed history (multiple scans) significantly improved performance over single-input baselines. ProgFormer-S using history achieved the highest PSNR across all cohorts.
Ablation Studies:
- Removing the coarse pathway caused prediction to collapse (PSNR dropped to ~7), confirming its necessity for establishing global structure.
- Removing the fine pathway or cross-attention significantly degraded both image quality and regional accuracy, highlighting the importance of voxel-level refinement guided by coarse context.
- Regional weighting was shown to be critical; removing it led to significant performance drops, proving that emphasizing progression-relevant regions improves overall fidelity.
Visual Analysis: Qualitative results show that ProgFormer preserves ventricular configurations and cortical outlines more consistently than baselines, with errors concentrated in progression-relevant regions rather than accumulating globally.

Significance and Claims

The paper claims that ProgFormer demonstrates the effectiveness of hierarchical voxel-space modeling for longitudinal volumetric prediction. By separating the modeling of stable global structure from subtle local changes, the method achieves anatomically faithful and progression-sensitive forecasting. The authors position this as a significant step forward from latent-space methods (which suffer from reconstruction loss) and unified direct methods (which struggle to resolve subtle changes against stable backgrounds). The work establishes that direct Transformer-based modeling, when structured hierarchically with spatio-temporal grounding, is a viable and superior approach for generating future brain MRIs.