cs.CVSep 24, 2026

Atlases Are Already Inside: Recovering Population Templates from Pretrained Diffusion Models

Authors: Jian Shi, John Femiani, Peter Wonka

Organizations: KAUST · Miami University

Abstract

We present a new inference-time sampler for diffusion models that gives a pretrained model a capability it was never trained for: constructing the atlas of the population it synthesizes. The sampler converges from every random seed to the population's central anatomy, which we call the intrinsic atlas. The advantage is threefold. (1) It requires no retraining. A diffusion model that has already learned a coherent population, including the released ones, yields its atlas in a single inference pass without involving deformable registration. (2) It applies to multiple domains, such as brain MRI, chest X-ray, faces, and 3D shapes. (3) It extends to subpopulations. One age-conditioned model gives an atlas at any age in its training range, and the resulting family reproduces the CSF expansion of healthy aging. Evaluated as a registration target, the intrinsic atlas is best or second-best on every dataset against classical and learned templates, and the most central template on held-out brain MRI cohorts. Atlas construction can be reframed as a byproduct of generative modeling: a diffusion model is a learned representation of population structure, and the atlas is what it already contains.

Figures & tables

Appendix figures & tables9 assets

Supplementary material from the paper’s appendix.

Appendix

Explore similar work

Oct 29, 2025

MedForj: An open, large-scale foundational generative prior for high-resolution 3D brain MRI

This work introduces MedForj, a suite of 3D foundational generative priors based on diffusion models. The MedForj models were trained on 72,65972{,}659 1~mm isotropic 3D T1T_1-weighted MRI human brain image volumes from 38,17438{,}174 subjects, drawn from a curated corpus of 80,67580{,}675 volumes from 42,50642{,}506 subjects spanning 3838 publicly available datasets. These training images were manually inspected to exclude those with poor quality and excessive pathology, and otherwise were minimally processed. The models include six different diffusion training strategies: rectified flow, latent diffusion rectified flow, flow matching, velocity prediction, clean prediction, and noise prediction. Image samples produced by each of these models were compared to each other and against real, ground truth data under downstream segmentation distributions, FID, five inverse problems, and blind human inspection in an observer study. Flow matching was the strongest strategy overall, achieving the best inverse problem solving results at 28.8028.80~dB PSNR and 0.8740.874 SSIM averaged over the five forward problems, the highest rate of reconstructions judged real by blind human raters at 72.6%72.6\%, and the closest per-structure match to real segmented anatomy in a permutation test. It was not best everywhere: rectified flow produced the most convincing unconditional samples in the observer study and the best FID, and the latent rectified-flow model achieved the smallest joint distributional distance to real anatomy. No other strategy, however, performed consistently well across all four evaluations. We therefore recommend flow matching as the default MedForj prior, while releasing every strategy so that the choice can be revisited per application. All model weights and corresponding code are publicly available at https://github.com/piksl-research/medforj.
May 11, 2026

INFANiTE: Implicit Neural representation for high-resolution Fetal brain spatio-temporal Atlas learNing from clinical Thick-slicE MRI

Spatio-temporal fetal brain atlases are important for characterizing normative neurodevelopment and identifying congenital anomalies. However, existing atlas construction pipelines necessitate days for slice-to-volume reconstruction (SVR) to generate high-resolution 3D brain volumes and several additional days for iterative volume registration, thereby rendering atlas construction from large-scale cohorts prohibitively impractical. We address these limitations with INFANiTE, an Implicit Neural Representation (INR) framework for high-resolution Fetal brain spatio-temporal Atlas learNing from clinical Thick-slicE MRI scans, bypassing both the costly SVR and the iterative non-rigid registration steps entirely, thereby substantially accelerating atlas construction. Extensive experiments demonstrate that INFANiTE outperforms existing baselines in subject consistency, reference fidelity, intrinsic quality and biological plausibility, even under challenging sparse-data settings. Additionally, INFANiTE reduces the end-to-end processing time (i.e., from raw scans to the final atlas) from days to hours compared to the traditional 3D volume-based pipeline (e.g., SyGN), facilitating large-scale population-level fetal brain analysis. Code: https://github.com/hu2274898/INFANiTE
Jul 20, 2026

DiFA: Inference-Time Forward-Process Alignment for Diffusion Models

The prevailing inference framework for diffusion models formulates generation fundamentally as a problem of numerical integration. This perspective casts the model as an exact estimator, neglecting the inherent statistical uncertainty of the denoising process. In this work, we propose Forward-Process Aligned Diffusion prediction (DiFA), a training-free framework that reframes inference-time data prediction refinement as a sequential state estimation problem. Rather than reusing past outputs solely for numerical integration, DiFA treats iterative data predictions along the reverse trajectory as correlated observations to build a forward-aligned temporal consensus. Inspired by Kalman filtering, this consensus aggregates historical predictions according to structural consistency and noise-level compatibility. To counteract the over-smoothing tendency of temporal consensus, we introduce a deviation guidance mechanism to adaptively preserve residual details. Empirically, DiFA yields significant improvements on CIFAR-10 and ImageNet across the evaluated metrics, including FID, IS, and FD-DINOv2, demonstrating that aligning inference with the forward statistical structure substantially improves generative fidelity.