MAGIC-Flow: Multiscale Adaptive Conditional Flows for Generation and Interpretable Classification

TL;DR

MAGIC-Flow tackles the need for trustworthy medical-image generation and classification by unifying these tasks in a conditional multiscale normalizing flow with an invertible backbone. It derives exact likelihoods via a conditional change-of-variables formulation and leverages a hierarchical architecture with flow steps, squeeze/split operations, and mask-based transformations to enable both high-fidelity conditional synthesis and likelihood-based classification. The model introduces task-specific affine couplings (generation vs classification) and likelihood attribution maps for faithful interpretability, achieving superior fidelity, diversity, and scanner-robust performance on MRI and PET data while providing interpretable explanations. This framework supports privacy-preserving augmentation and robust generalization in data-limited, acquisition-variant clinical settings, with planned extensions to 3D processing, uncertainty quantification, pathology-conditioned generation, and cross-institutional deployment.

Abstract

Generative modeling has emerged as a powerful paradigm for representation learning, but its direct applicability to challenging fields like medical imaging remains limited: mere generation, without task alignment, fails to provide a robust foundation for clinical use. We propose MAGIC-Flow, a conditional multiscale normalizing flow architecture that performs generation and classification within a single modular framework. The model is built as a hierarchy of invertible and differentiable bijections, where the Jacobian determinant factorizes across sub-transformations. We show how this ensures exact likelihood computation and stable optimization, while invertibility enables explicit visualization of sample likelihoods, providing an interpretable lens into the model's reasoning. By conditioning on class labels, MAGIC-Flow supports controllable sample synthesis and principled class-probability estimation, effectively aiding both generative and discriminative objectives. We evaluate MAGIC-Flow against top baselines using metrics for similarity, fidelity, and diversity. Across multiple datasets, it addresses generation and classification under scanner noise, and modality-specific synthesis and identification. Results show MAGIC-Flow creates realistic, diverse samples and improves classification. MAGIC-Flow is an effective strategy for generation and classification in data-limited domains, with direct benefits for privacy-preserving augmentation, robust generalization, and trustworthy medical AI.

Figures (6)

• Figure 1: Illustration of components of the model: (a) single flow step and (b) full multiscale flow.
• Figure 2: Qualitative comparison of generated neuroimaging slices across (a) scanner types and (b) imaging modalities.
• Figure 3: Likelihood attribution maps from MAGIC-Flow. Warm colors indicate regions with the strongest positive contributions to the predicted likelihood.
• Figure 4: Architectures of the task-specific affine coupling layers. The generation layer emphasizes multi-scale and global context features, while the classification layer focuses on integrating label information into the feature transformations.
• Figure 5: Examples of the masks used in MAGIC-Flow: from left to right, (1) application-specific coronal mask, (2) checkerboard mask, and (3) channel-wise mask. Together, these masks provide complementary spatial, feature-level, and application-adaptive transformations.