[Paper Review] Generative Autoregressive Transformers for Model-Agnostic Federated MRI Reconstruction

Valiyeh Ansarian Nezhad, Gokberk Elmas+4 more|ArXiv.org|Feb 6, 2025

Advanced MRI Techniques and Applications5 citations

TL;DR

This paper introduces FedGAT, a model-agnostic federated learning approach for MRI reconstruction that uses a global generative autoregressive transformer prior to enable heterogeneous-site models and improves both within-site and across-site reconstruction performance.

ABSTRACT

While learning-based models hold great promise for MRI reconstruction, single-site models trained on limited local datasets often show poor generalization. This has motivated collaborative training across institutions via federated learning (FL)-a privacy-preserving framework that aggregates model updates instead of sharing raw data. Conventional FL requires architectural homogeneity, restricting sites from using models tailored to their resources or needs. To address this limitation, we propose FedGAT, a model-agnostic FL technique that first collaboratively trains a global generative prior for MR images, adapted from a natural image foundation model composed of a variational autoencoder (VAE) and a transformer that generates images via spatial-scale autoregression. We fine-tune the transformer module after injecting it with a lightweight site-specific prompting mechanism, keeping the VAE frozen, to efficiently adapt the model to multi-site MRI data. In a second tier, each site independently trains its preferred reconstruction model by augmenting local data with synthetic MRI data from other sites, generated by site-prompting the tuned prior. This decentralized augmentation improves generalization while preserving privacy. Experiments on multi-institutional datasets show that FedGAT outperforms state-of-the-art FL baselines in both within- and cross-site reconstruction performance under model-heterogeneous settings.

Motivation & Objective

• Motivate the need for privacy-preserving, multi-site MRI reconstruction with diverse site-specific architectures.
• Develop a model-agnostic FL framework that decouples knowledge transfer from local model design.
• Propose a site-prompted generative autoregressive transformer prior to synthesize multi-site MR images.
• Enable site-specific reconstruction models to be trained on hybrid datasets consisting of local data and synthetic data from other sites.
• Evaluate FedGAT on multi-institutional MRI datasets and compare against standard FL baselines.

Proposed method

• Introduce FedGAT, a two-tier FL framework where a global GAT prior learns the multi-site image distribution (Tier 1).
• Use a site-prompted GAT prior that combines a VAE encoder/decoder with an autoregressive transformer to generate site-conditioned MR images (site prompts steer synthesis).
• Train site-specific reconstruction models on a hybrid dataset comprising local data plus synthetic data generated by the GAT prior (Tier 2).
• Decouple cross-site knowledge transfer from local model training to support model heterogeneity across sites.
• Employ a two-part loss for the GAT prior: a VAE reconstruction/latent loss and a cross-entropy loss for the autoregressive transformer, with federated averaging for the global GAT (Alg. 1).
• Synthesize MR images via a multi-scale token-based autoregressive process with nucleus sampling for inference, and use AdaLN to condition transformer outputs on site tokens.

Figure 1: FedGAT devises a two-tier strategy to collaboratively train heterogeneous models for MRI reconstruction. (a) The first tier conducts decentralized training of a global prior $\text{GAT}_{\theta_{\text{GAT}}}$ that captures the distribution of multi-site MR images. The global prior is built

Experimental results

Research questions

• RQ1 Can model-heterogeneous sites participate in federated MRI reconstruction without constraining architectures to a single global model?
• RQ2 Does a globally learned generative autoregressive prior enable effective cross-site knowledge transfer while preserving site-specific performance?
• RQ3 Do site-specific reconstruction models trained on hybrid local+synthetic data outperform conventional FL baselines in MRI reconstruction?
• RQ4 Is the proposed site-prompted GAT prior capable of producing high-fidelity, site-consistent MR images across multiple scales?
• RQ5 How does FedGAT perform in within-site and across-site reconstruction tasks compared to state-of-the-art FL baselines?

Key findings

• FedGAT enables model-heterogeneous FL for MRI reconstruction by decoupling global knowledge transfer from local model training.
• A site-prompted GAT prior (VAE + autoregressive transformer) synthesizes multi-site MR images conditioned on site prompts.
• Local reconstruction models are pre-trained on local data and then fine-tuned on a hybrid dataset that includes synthetic data from other sites.
• Comprehensive experiments show FedGAT achieves superior within-site and across-site reconstruction performance versus state-of-the-art FL baselines.
• Code for FedGAT is available at the authors’ repository.

Figure 2: Architecture of the proposed site-prompted GAT prior. (a) The GAT prior embodies a variational autoencoder (VAE), whose encoder module maps an input MR image onto a set of discrete token maps $\mathbf{f}_{1},\mathbf{f}_{2},\dots,\mathbf{f}_{S}$ across $S$ spatial scales, and whose decoder