Harmonizing Flows: Leveraging normalizing flows for unsupervised and source-free MRI harmonization

TL;DR

This work tackles MRI domain shift across sites by proposing Harmonizing Flows, an unsupervised, source-free, and task-agnostic framework that uses normalizing flows to model the source distribution and a harmonizer network to map target-domain images onto it. The method trains a normalizing flow to capture $p_{\mathbf{x}}(\mathbf{x})$ and pre-trains a UNet-based harmonizer to reconstruct source-like images from augmented sources, then performs test-time adaptation by updating the harmonizer to maximize the NF-based likelihood of harmonized outputs, with stopping criteria based on entropy or bits-per-dimension. Across adult and neonatal segmentation and neonatal brain-age estimation, the approach yields state-of-the-art performance and demonstrates robust generalization to unseen domains and tasks, surpassing histograms-based methods, GAN/autoencoder baselines, and test-time adaptation strategies. The work further shows that focusing on distributional alignment via NF guidance, rather than merely histogram similarity, leads to tangible improvements in downstream clinical tasks, supporting broader adoption in multi-center MRI studies.

Abstract

Lack of standardization and various intrinsic parameters for magnetic resonance (MR) image acquisition results in heterogeneous images across different sites and devices, which adversely affects the generalization of deep neural networks. To alleviate this issue, this work proposes a novel unsupervised harmonization framework that leverages normalizing flows to align MR images, thereby emulating the distribution of a source domain. The proposed strategy comprises three key steps. Initially, a normalizing flow network is trained to capture the distribution characteristics of the source domain. Then, we train a shallow harmonizer network to reconstruct images from the source domain via their augmented counterparts. Finally, during inference, the harmonizer network is updated to ensure that the output images conform to the learned source domain distribution, as modeled by the normalizing flow network. Our approach, which is unsupervised, source-free, and task-agnostic is assessed in the context of both adults and neonatal cross-domain brain MRI segmentation, as well as neonatal brain age estimation, demonstrating its generalizability across tasks and population demographics. The results underscore its superior performance compared to existing methodologies. The code is available at https://github.com/farzad-bz/Harmonizing-Flows

Figures (12)

• Figure 1: Harmonizing Flows Pipeline. Our method comprises three primary steps. Initially, a harmonizer network undergoes pre-training to reconstruct the original images from augmented counterparts, facilitating initial harmonization. Subsequently, normalizing flow (NF) is utilized to capture the distribution of the source domain. In the third stage (test time), the trained NF is leveraged to update the parameters of the harmonizer network, ensuring maximal alignment between the harmonized outputs and the learned NF distribution. Notably, steps 1 and 2 are independent of each other and can be executed interchangeably.
• Figure 2: Cross-site brain MRI segmentation matrix across the compared methods. Each cell indicates the segmentation result (DSC %) when the source dataset (in the rows) is used to harmonize each target dataset (in columns).
• Figure 3: Effect of each component of the harmonizing flow. Particularly, we depict the improvement gained using the proposed pre-trained harmonizer network ($\sim$44 DSC%), and the adaptation using normalizing flows ($\sim$2.2 DSC%).
• Figure 4: The effect of different stopping criteria to stop the harmonizer network adaptation. Oracle represents selecting the best iteration based on the performance of the target task (i.e., segmentation in this example), which serves as the upper bound. Both criteria, source BPD and minimum entropy, provide good stopping points, with a slight advantage of the minimum entropy criterion.
• Figure 5: Which is the best metric as stopping criteria? This plot depicts different metrics during the adaptation of the harmonizer network (from HBNSI to NYU). Step zero corresponds to using the initial harmonizer network without adaptation. The vertical lines show the stopping time-points based on two proposed stopping criteria: minimum entropy of the predictions (red)and reaching source BPD purple).