0653

Peer-to-Peer Generative Learning for Architecture-Agnostic Federated MRI Reconstruction

Valiyeh Ansarian Nezhad^1,2, Gökberk Elmas^1,2, and Tolga Çukur^1,2,3
¹Department of Electrical and Electronics Engineering, Bilkent University, Ankara, Turkey, ²National Magnetic Resonance Research Center (UMRAM), Bilkent University, Ankara, Turkey, ³Neuroscience Program, Bilkent University, Ankara, Turkey

Synopsis

Keywords: AI/ML Image Reconstruction, Machine Learning/Artificial Intelligence, Federated learning, multi-institutional, collaborative learning, image reconstruction

Motivation: Federated learning (FL) enables privacy-preserving training of deep reconstruction models across multiple sites to improve generalization at the expense of lower within-site performance. Yet, existing methods require a common model architecture across sites, limiting flexibility.

Goal(s): Our goal was to devise an architecture-agnostic method for collaborative training of heterogeneous models across sites.

Approach: We introduced a novel peer-to-peer generative learning method (PGL-FedMR), where individual sites share a generative prior for their MRI data with remaining sites, and prior-driven synthetic data are used to train reconstruction models at each site.

Results: PGL-FedMR improves across-site generalization over local models, and within-site performance over conventional FL.

Impact: Improvements in within-site and across-site performance for MRI reconstruction through PGL-FedMR, coupled with the ability to handle heterogeneous architectures, may facilitate privacy-preserving multi-institutional collaborations to build reliable reconstruction models for many applications where data are scarce including rare diseases.

Introduction

Deep network models have gained prominence in accelerated MRI reconstruction given their high image quality^1-16. Yet, deep models have limited representation for rare features in their training sets; so models trained on single-site data characteristically yield poor across-site generalization¹⁷. Thus, collaborative approaches for training reconstruction models on multi-site data are direly needed.

Federated learning (FL) is an emerging privacy-preserving framework for collaborative model training, based on sharing of model weights across sites as opposed to imaging data^17-21. In conventional FL, each site trains a copy of a global model on its local data, and a server then averages locally-trained copies to update the global model^18-19. While FL-based MRI reconstruction has recently been reported to enhance generalization^22-24, it can suffer from losses in within-site performance compared to single-site models²². Furthermore, conventional FL requires all sites to prescribe a common model architecture, preventing collaboration among sites with heterogeneous architectural preferences^25-26 (e.g., due to varying complexity of reconstruction tasks, compute resources).

To address these limitations, here we introduce a novel method, PGL-FedMR, for federated MRI reconstruction based on peer-to-peer generative learning. In PGL-FedMR, generative priors for local MRI data are trained at each site and communicated among sites. Site-specific reconstruction models are then trained on both local data and prior-generated synthetic data via a mixing approach to prevent catastrophic forgetting. PGL-FedMR outperforms single-site models in across-site, and conventional FL in within-site reconstruction. The proposed method further enables collaborations based on heterogeneous models across sites, a setting where conventional FL fails.

Methods

Local generative priors: Unlike conventional FL that uses a server topology, PGL-FedMR leverages peer-to-peer communications to transfer local generative priors among $$$K$$$ sites (Fig.1). Priors are operationalized as K generators $$$\theta_G^{k}$$$ and discriminators $$$\theta_D^{k}$$$, $$$k\in[1,2,...,K]$$$, trained to capture the distribution of local MR images derived as $$$x_{ref}^{k}=\mathcal{C}^*\mathcal{F}^{-1}\{y_{fs}\}$$$ $$$(\mathcal{C}^*$$$: adjoint of coil sensitivities, $$$\mathcal{F}^{-1}$$$: inverse Fourier transform) from fully-sampled acquisitions $$$y_{fs}$$$. The training losses over $$$n_g$$$ epochs are:$$L^k_G(\theta^k_G)=-\mathbb{E}_{p(z)}\left\{\mathrm{log}(f(D_{\theta_D^k}(G_{\theta_G^k}(z))))\right\},\\L^k_D(\theta^k_D)=-\mathbb{E}_{p(z)}\left\{\mathrm{log}(1-f(D_{\theta_D^k}(G_{\theta_G^k}(z))))\right\}-\mathbb{E}_{p(x^k_{ref})}\left\{\mathrm{log}(f(D_{\theta_D^k}(x^k_{ref}))\right\},$$where $$$\mathbb{E}$$$ denotes expectation, $$$p$$$ denotes probability distribution, $$$f$$$ denotes sigmoid activation, $$$z$$$ is a random normal variable. Each site transfers its generator parameters $$$\theta_G^{k}$$$ to remaining sites:$${Site}_k\xrightarrow{\theta_G^{k}}Site_j,\mbox{ }k,j\in[1,..,K];\mbox{ }j\neq k.$$

Site-specific reconstruction models: Following generative learning, PGL-FedMR trains site-specific reconstruction models $$$H^k_{\phi^k}$$$, $$$k\in[1,2,...,K]$$$, which can employ distinct architectures across sites, and reconstruct images as $$$\bar{x}=H^k_{\phi^k}(x^{k}_{us},y^{k}_{us})$$$ where the input image $$$x_{us}^{k}=\mathcal{C}^*\mathcal{F}^{-1}\{y^{k}_{us}\}$$$ denotes the zero-filled reconstruction of the acquisition $$$y^{k}_{us}$$$ undersampled with mask $$$\mathcal{M}.$$$

To ensure high within-site performance, the reconstruction model for each site is first trained for $$$n_{r}$$$ epochs on local MRI data:$$\hat{\phi}^k:=\min_{\phi^k}\mathbb{E}_{p(x^k_{ref},x^{k}_{us},y^{k}_{us})}\left\{\left\|x_{ref}^{k}-H^k_{\phi^k}(x^{k}_{us},y^{k}_{us})\right\|^2\right\}.$$ To improve generalization, the model is then trained using prior-generated synthetic data from all sites: $$x^j_{syn}=G_{\theta_G^{j}}(z),\mbox{ }j\in \{1,..,K\}.$$ For adaptation to remaining sites without forgetting previously learned knowledge on the local site, local data from site $$$k$$$ are mixed in equal proportion with synthetic data from other sites over $$$n_{s}$$$ epochs:\[\mathrm{for}\mbox{ }e\in[1,..,n_{s}]:\\\quad\quad\quad j=\mathrm{modulus}(e,K)+1,\\\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad X_{ref}=\{X^k_{ref}\cup X^j_{syn}\},X_{us}=\{X_{us}^k\cup X^j_{syn,us}\},\\\quad\quad Y_{us}=\{Y_{us}^k\cup Y^j_{syn,us}\},\\\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad \hat{\phi}^k:=\min_{\phi^k}\mathbb{E}_{p(X_{ref},X_{us},Y_{us}[i])}\left\{\left\|X_{ref}[i]-H^k_{\phi^k}(X_{us}[i],Y_{us}[i])\right\|^2\right\}\]where $$$i$$$ is sample index, and $$$X_{ref},X_{us},Y_{us}$$$ are mixed reference images, zero-filled reconstructions, and undersampled acquisitions respectively.

Analyses: Single-coil data from IXI (https://brain-development.org/ixi-dataset/), BRATS²⁷, fastMRI²⁸, and multi-coil data from fastMRI-brain, fastMRI-knee, Calgary-brain²⁹ datasets were analyzed. A (70%,10%,20%) split of (training,validation,test) subjects was used. Variable-density undersampling was performed at rate R=3-4. Coil sensitivities were estimated via ESPIRiT³⁰. Adam optimizer was used with $$$n_g$$$=100, $$$n_r$$$=15, $$$n_s$$$=15, 0.01 learning rate.

Results

Fig.3a compares PGL-FedMR against a privacy-violating centralized benchmark with the same training strategy as the proposed method, albeit using actual instead of synthetic data. PGL-FedMR performs on par with the benchmark, indicating that synthetic data are representative of actual data. Fig.3b-c compare PGL-FedMR against single-site models trained using only local data (Single-site¹⁷) and conventional FL (FedAvg²²). Compared to Single-site, PGL-FedMR performs similarly for within-site, and yields a (PSNR, SSIM) improvement of (0.9dB, 3.0%) for across-site reconstruction. Compared to FedAvg, it yields improvements of (3.4dB, 2.4%) for within-site, and (0.4dB, 2.8%) for across-site reconstruction. Improved image quality with PGL-FedMR is also evident in representative reconstructions (Fig.4). Lastly, Fig.5 examines performance under heterogeneous model settings across sites, where conventional FL is inapplicable. Compared to Single-site, PGL-FedMR performs similarly for within-site, and yields a (1.2dB, 0.9%) improvement for across-site reconstruction.

Discussion

Here we introduced the first architecture-agnostic technique for federated MRI reconstruction to our knowledge. PGL-FedMR uses peer-to-peer learning of generative MRI priors, and trains reconstruction models on a mixture of local data and prior-generated synthetic data. Going beyond the reach of conventional FL, PGL-FedMR enables collaborative learning of heterogeneous models across multiple sites. Thus, it shows great promise for expanding the scope of multi-site collaborations in learning-based MRI reconstruction.

Acknowledgements

This work was supported in part by a TUBITAK 1001 Grant No. 123E142, and by TUBA GEBIP 2015 and BAGEP 2017 fellowships.

References

1. Lustig, M., Donoho, D., Pauly, J.M., “Sparse MRI: The application of compressed sensing for rapid MR imaging,” Magn. Reson. Med., vol. 58, no. 6, pp. 1182–1195, 2007.

2. Haldar, J.P., Hernando, D., Liang, Z.P., “Compressed-sensing MRI with random encoding,” IEEE Trans. Med. Imaging, vol. 30, no. 4, pp. 893–903, 2010.

3. Wang, S., Su, Z., Ying, L., Peng, X., Zhu, S., Liang, F., Feng, D., Liang, D., “Accelerating magnetic resonance imaging via deep learning,” in IEEE 13th Int. Symp. Biomed. Imaging (ISBI), 2016, pp. 514–517.

4. Hammernik H., Klatzer T., Kobler R., Recht M.P., Sodickson D.K., Pock T., Knoll F., “Learning a variational network for reconstruction of accelerated MRI data,” Magn. Reson. Med., vol. 79, no. 6, pp. 3055–3071, 2018.

5. Zhu, B., Liu, J.Z., Rosen, B.R., Rosen, M.S., “Image reconstruction by domain transform manifold learning,” Nature, vol. 555, no. 7697, pp. 487–492, 2018.

6. Qin, C., Schlemper, J., Caballero, J., Price, A.N., Hajnal, J.V., Rueckert, D., “Convolutional recurrent neural networks for dynamic MR image reconstruction,” IEEE Trans. Med. Imaging, vol. 38, no. 1, pp. 280–290, 2018.

7. Mardani, M., Gong, E., Cheng, J.Y., Vasanawala, S.S., Zaharchuk, G., Xing, L., Pauly. J.M., “Deep Generative Adversarial Neural Networks for Compressive Sensing MRI,” IEEE Trans. Med. Imaging, vol. 38, no. 1, pp. 167-179, 2019.

8. Akçakaya, M, Moeller, S, Weingärtner, S, Uğurbil, K., “Scan-specific robust artificial-neural-networks for k-space interpolation (RAKI) reconstruction Database-free deep learning for fast imaging,” Magn. Reson. Med., vol. 81, pp. 439–453, 2019.

9. Tamir, J.I., Yu S., Lustig. M., “Unsupervised deep basis pursuit: Learning reconstruction without ground-truth data,” in Proceedings of ISMRM, 2019, p. 0660.

10. Aggarwal H.K., Mani, M.P., Jacob, M., “MoDL: Model-Based Deep Learning Architecture for Inverse Problems,” IEEE Trans. Med. Imaging, vol. 38, no. 2, pp. 394–405, 2019.

11. Peng, X., Sutton, B.P., Lam, F., Liang, Z.P., “DeepSENSE: Learning coil sensitivity functions for SENSE reconstruction using deep learning,” Magn. Reson. Med., vol. 87, no. 4, pp. 1894–1902, 2020.

12. Kuestner, T., Fuin, N., Hammernik, K., Bustin, A., Qi, H., Hajhosseiny, R., Masci, P. G., Neji, R., Rueckert, D., Botnar, R. M., Prieto, C., “CINENet: deep learning-based 3D cardiac CINE MRI reconstruction with multi-coil complex-valued 4D spatio-temporal convolutions,” Scientific Reports, vol. 10, no. 1, 2020.

13. Polak, D., Cauley, S., Bilgic, B., Gong, E., Bachert, P., Adalsteinsson, E., Setsompop, K., “Joint multi-contrast variational network reconstruction (jVN) with application to rapid 2D and 3D imaging,” Magn. Reson. Med., vol. 84, no. 3, pp. 1456–1469, 2020.

14. Eo, T., Jun, Y., Kim, T., Jang, J., Lee, H. J., Hwang, D., “KIKI-net: cross-domain convolutional neural networks for reconstructing undersampled magnetic resonance images,” Magn. Reson. Med., vol. 80, no. 5, pp. 2188–2201, 2018.

15. Dar, S.U., Yurt, M., Shahdloo, M., Ildız, M.E., Tınaz, B., Cukur, T., “Prior-Guided Image Reconstruction for Accelerated Multi-Contrast MRI via Generative Adversarial Networks,” IEEE J. Sel. Top. Signal Process., vol. 14, no. 6, pp. 1072–1087, 2020.

16. Liu, F., Feng, L., Kijowski, R., “MANTIS: Model-Augmented Neural neTwork with Incoherent k-space Sampling for efficient MR parameter mapping,” Magn. Reson. Med., vol. 82, no. 1, pp. 174–188, 2019.

17. Guo, P., Wang, P., Zhou, J., Jiang, S., Patel, V.M., “Multi-institutional Collaborations for Improving Deep Learning-based Magnetic Resonance Image Reconstruction Using Federated Learning,” arXiv:2103.02148, 2021.

18. G. A. Kaissis, M. R. Makowski, D. Rueckert, R. F. Braren, “Secure, privacy-preserving and federated machine learning in medical imaging,” Nat. Mach. Intelli., vol. 2, no. 6, pp. 305–311, 2020.

19. N. Rieke, J. Hancox, W. Li, F. Milletar, H. R. Roth, S. Albarqouni, S. Bakas et al., “The future of digital health with federated learning,” NPJ Digit. Med., vol. 3, no. 1, p. 119, 2020.

20. X. Li, Y. Gu, N. Dvornek, L. H. Staib, P. Ventola, J. S. Duncan, “Multi-site fmri analysis using privacy-preserving federated learning and domain adaptation: Abide results,” Med. Image Anal., vol. 65, p. 101765, 2020.

21. Rajagopal A, Redekop E, Kemisetti A, Kulkarni R, Raman S, Sarma K, Magudia K, Arnold CW, Larson PEZ, “Federated Learning with Research Prototypes: Application to Multi-Center MRI-based Detection of Prostate Cancer with Diverse Histopathology,” Acad Radiol., vol. 30, no. 4, pp. 644-657, 2023.

22. Elmas, G, Dar, SUH, Korkmaz, Y, Ceyani, E, Susam, B, Ozbey, M, Avestimehr, S, Çukur, T. Federated Learning of Generative Image Priors for MRI Reconstruction. IEEE Trans. Med. Imaging. vol. 42, no. 7, pp. 1996-2009, 2023.

23. Feng C.M., Yan Y., Wang S., Xu, Y., Shao, L., Fu, H. “Specificity-Preserving Federated Learning for MR Image Reconstruction,” arXiv:2112.05752v3, 2022.

24. Levac, B.R., Arvinte, M., Tamir, J.I., "Federated End-to-End Unrolled Models for Magnetic Resonance Image Reconstruction," Bioengineering, vol. 10, no. 3, p. 364, 2023.

25. Dalmaz, O., Mirza, U., Elmas, G., Ozbey, M., Dar, SUH., Ceyani, E., Avestimehr, S., Çukur, T, “One model to unite them all: Personalized federated learning of multi-contrast MRI synthesis,” arXiv:2207.06509, 2022.

26. Wu, R., Li, C., Zou, J., Wang, S., “FedAutoMRI: Federated Neural Architecture Search for MR Image Reconstruction,” arXiv:2307.11538, 2023.

27. Menze, B.H., Jakab, A., Bauer, S., Kalpathy-Cramer, J., Farahani, K., Kirby, J., et al. "The Multimodal Brain Tumor Image Segmentation Benchmark (BRATS)," IEEE Trans. Med. Imaging, vol. 34, no. 10, pp. 1993-2024, 2015.

28. F. Knoll, J. Zbontar, A. Sriram, M. J. Muckley, M. Bruno, A. Defazio, M. Parente et al., “fastMRI: A publicly available raw k-space and DICOM dataset of knee images for accelerated MR image reconstruction using machine learning,” Rad. Artif. Intell., vol. 2, no. 1, p. e190007, 2020.

29. R. Souza, O. Lucena, J. Garrafa, D. Gobbi, M. Saluzzi, S. Appenzeller, L. Rittner, R. Frayne, R. Lotufo, “An open, multi-vendor, multi-field strength brain MR dataset and analysis of publicly available skull stripping methods agreement,” NeuroImage, vol. 170, pp. 482–494, 2018.

30. Uecker, M., Lai, P., Murphy, M.J., Virtue, P., Elad, M., Pauly, J.M., Vasanawala, S.S., Lustig, M., “ESPIRiT--an eigenvalue approach to autocalibrating parallel MRI: where SENSE meets GRAPPA,” Magn. Reson. Med., vol. 71, no. 3, pp. 990–1001, 2014.

Figures

Figure 1. PGL-FedMR trains local generative priors operationalized as generators $\theta_G^{k}$ and discriminators $\theta_D^{k}$, $k\in[1,2,...,K]$. Generator parameters are shared via peer-to-peer communication among sites: $Site_k\xrightarrow{\theta_G^{k}}Site_j,\mbox{ }k,j\in[1,..,K];\mbox{ }j\neq k$. A site-specific reconstruction model $H^k_{\phi}$ is then trained at each site using a mixture of local MRI data and synthetic MRI data produced by the received generators: $x^j_{syn}=G_{\theta_G^{j}}(z),\mbox{ }j\in \{1,..,K\}$.

Figure 2. Training algorithm for PGL-FedMR. Local generative priors are trained to capture the distribution of local MRI data. The trained generators are peer-to-peer communicated among sites. At a given site $k$, priors received from all sites are leveraged to generate synthetic MRI data. Site-specific reconstruction models are first trained on local data to achieve high within-site performance. They are then trained on a mixture of local data at site $k$ and synthetic data from other sites to ensure reliable generalization without catastrophic forgetting.

Figure 3. PSNR/SSIM in three-site FL setups with a common reconstruction model across sites. Within-site (e.g., IXI-specific model tested on IXI) and across-site (e.g., IXI-specific model tested on fastMRI and BraTS) performances are listed. (a,b) FL on (IXI, BraTS, fastMRI) datasets, R=3 and MoDL-5 cascade architecture¹⁰ for reconstruction. PGL-FedMR against (a) a privacy-violating centralized benchmark, (b) locally-trained single-site models and global FedAvg model. (c) FL on (fastMRI-brain, fastMRI-knee, Calgary-brain) datasets, R=4 and rGAN architecture¹⁵.

Figure 4. Reconstructions from zero-filled Fourier method (Zero-filled), locally-trained single-site models (Single-site), conventional FL models (FedAvg), and PGL-FedMR, along with the reference image derived from fully-sampled acquisitions. A three-site FL setup with (IXI, fastMRI, BraTS) was considered for R=3 and a common MoDL-5 architecture¹⁰. (a) Within-site reconstructions in IXI, (b) across-site reconstructions in fastMRI, using IXI-specific Single-site, PGL-FedMR models and global FedAvg model. Zoom-in windows are included to emphasize method differences.

Figure 5: PSNR/SSIM in a three-site FL setup with heterogeneous reconstruction models across sites. Within-site (e.g., IXI-specific model tested on IXI) and across-site (e.g., IXI-specific model tested on fastMRI and BraTS) performances are listed. The FL-setup is based on (IXI, BraTS, fastMRI) datasets, and (MoDL-7 cascade¹⁰, MoDL-3 cascade, rGAN¹⁵) architectures at R=3. Note that conventional FedAvg is not applicable in heterogeneous model settings.

Proc. Intl. Soc. Mag. Reson. Med. 32 (2024)

0653

DOI: https://doi.org/10.58530/2024/0653