Generalizable Learning Reconstruction for Accelerating MR Imaging via Federated Neural Architecture Search

Reconstructing high-quality MR images from undersampling $k$-space data is a typical ill-posed inverse problem. Deep learning-based methods learn the mapping relationship from undersampled to fully-sampled data offline and then enable fast online reconstruction[3]. These methods can avoid manual fine-tuning of optimization parameters. Deep learning-based MR image reconstruction methods can be broadly categorized into data-driven[5, 7, 9, 11, 12, 13, 14, 15, 16] and model-driven approaches[4, 6, 8, 10]. Data-driven methods typically rely on large amounts of data and directly learn the mapping relationship from undersampled data to fully-sampled data. However, the networks could be totally black-box. Model-driven approaches formulate an optimization problem and then unfold the iterative optimization algorithm into the network. The network architecture is constructed based on iterative steps, leading to a slightly longer reconstruction time. While these methods have made certain contributions to the development of MR image reconstruction, they all require the collection of large amounts of data for centralized training. This undoubtedly increases labor costs and resource consumption, and it also brings a serious issue—data privacy leakage. We utilize a distributed collaborative training scheme to effectively avoid the problems of data collection and patient privacy leakage.

Federated learning is a distributed collaborative training paradigm that protects data privacy by training shared model without sharing data[20, 21, 22, 23, 24, 25, 26, 27, 28, 29]. One of the most classical federated learning methods is FedAvg[21], which obtains the global model through an average aggregation method, and it obtains a relatively better performance when the data distributions of different centers are similar. However, in real-world scenarios, there is often statistical heterogeneity among different centers, and the traditional federated learning methods are difficult to adapt to such requirements. Some works start to reconcile the differences between global and local models [22, 23], aiming to obtain a global model with better performance. For example, Li et al.,[22] facilitate the learning of the global model by imposing an additional constraint on the client’s objective function. In addition, there are also some proposed personalized federated learning methods to preserve the specificity of local data. For example, Zhang et al.,[25] proposed a local adaptive aggregation approach to improve the performance of local clients by adaptively aggregating global and local models. As for magnetic resonance imaging, the problem of data heterogeneity may also exist in different centers due to scanning devices, imaging protocols, and other issues[33]. Some methods start to solve the problems of feature drift and domain generalization. For example, Guo et al.,[26] reduces the data heterogeneity of clients by repeatedly aligning the hidden feature distributions between the source and target domains. It selects a target site in each communication round and then aligns it with the representation of the remaining sites, which adds more communication consumption. In addition, Feng et al.,[27] effectively preserve local client-specific features by dividing the reconstruction model into a globally shared encoder and a locally personalized decoder. These approaches effectively mitigate the privacy leakage problem, but their model architectures are hand-designed by experts, which are often complex and require more computational resources. We utilize an automatically differentiable architecture search method to obtain an optimal model architecture, effectively reducing the number of model parameters and the need for computational resources.

Neural architecture search can obtain optimal network architectures by employing an automated network architecture search approach while keeping fewer model parameters[30]. Based on different search strategies, neural architecture search algorithms can be broadly categorized into three types: reinforcement learning (RL)-based[34, 35], evolutionary algorithm (EA)-based[36, 37], and gradient-based (GD) methods[38, 39]. Reinforcement learning-based and evolutionary algorithm-based methods are gradient-free methods, in contrast to gradient-based methods that tend to be faster in the searching process and consume slightly fewer computational resources[39]. DARTS is the most classical gradient-based method, which scales the selection of candidate operations into continuous variables and then utilizes efficient gradient back-propagation for learning, effectively reducing the search time[38]. In addition, there is also some NAS-based work in MR image reconstruction[31, 32]. For example, Yan et al., [31] used NAS for compressed sensing MR image reconstruction, achieving enhanced performance. These NAS-based MR image reconstruction methods are used for centralized training scenarios but have not yet been explored in federated learning scenarios. In this paper, we utilize federated learning and neural architecture search for MR image reconstruction and improve the generalization performance for both in-distribution and out-of-distribution data through a fairness adjustment approach.

The goal of MR image reconstruction is to reconstruct a high-quality image $\mathbf{x}\in\mathbb{C}^{N}(M<N)$ from undersampled $k$-space data $\mathbf{b}\in\mathbb{C}^{M}$, formulated as:

where $\mathbf{A}\in\mathbb{C}^{M\times N}$ denotes the encoding matrix and $\mathbf{\epsilon}\in\mathbb{C}^{M}$ denotes the measurement noise. Solving $\mathbf{x}$ from Eq. (1) is an ill-posed inverse problem that can be expressed as the following optimization problem:

where the first term is the data fidelity, the second term is the regularization term, and $\lambda$ is the regularization parameter used to balance the tradeoff between the data fidelity term and the regularization term. Based on the method of model unfolding[8], Eq. (2) can be transformed into the following alternating iteration process:

where $f_{\theta}$ denotes the deep neural network used for denoising, $\mathbf{A}^{H}$ denotes the conjugate operator of $\mathbf{A}$, and $j$ denotes the iterative control variable, $\lambda$ is learnable.

The overall architecture of GAutoMRI is shown in Fig. 1, which consists of two phases: the searching phase and the training phase. In the searching phase, we adjusted the search space according to the reconstruction task, details of which can be found in III.B.2 and Fig. 2, and the whole search process draws on the method in [40] to obtain the optimal network architecture. In the training phase, the architecture obtained in the searching phase is employed as the base model. In addition, we introduce a fairness adjustment approach to improve the generalization of the global model over out-of-distribution data. The specific search process and fairness adjustment approach will be described later.

After the searching phase, we can obtain an optimal model architecture, which is subsequently utilized for the training phase. However, some existing federated MR image reconstruction methods seldom consider inter-client fairness issues, leading to performance degradation of the trained global model on out-of-distribution data, especially when faced with statistical heterogeneity among clients. Intuitively, if the global model does not discriminate in performance on known clients (i.e., the performance gap is small on each client), then it logically should exhibit respectable performance even when faced with out-of-distribution data. Based on this notion, we introduce a fairness adjustment approach that utilizes the difference between the empirical risks of the global and local models in each communication round. This approach dynamically adjusts the aggregation weights of the clients, aiming to progressively diminish the performance gap of the global model across the known clients. As a result, this approach gradually improves the generalization performance of the global model on out-of-distribution data. Formally, the empirical risk gap between global and local models can be expressed as:

where $\ell_{D_{c}}(\theta^{r})$ and $\ell_{D_{c}}(\hat{\theta}_{c}^{r-1})$ denote the empirical risk of the global and local model, respectively. After completing local model training, the client sends the empirical risk gap $\mathcal{G}_{D_{c}}(\theta^{r})$ and local model parameters $\hat{\theta}_{c}^{r}$ to the server. After the server receives $\mathcal{G}_{D_{c}}(\theta^{r})$ and $\hat{\theta}_{c}^{r}$ from each client, it first updates the aggregation weight of each client according to the fairness adjustment approach, which is formulated as:

where $\acute{c}=\{1,2,\cdots,C\}$, $a_{c}^{r-1}$ denotes the aggregation weight of the $c^{th}$ client at the $(r-1)^{th}$ communication round, $a_{c}^{r}$ is the weight updated by the fairness adjustment approach, $\gamma$ is utilized to control the magnitude of the weight increment.

In each communication round, for each client $c$, if the empirical risk gap $\mathcal{G}_{D_{c}}(\theta^{r})\leq 0$, we do not change its aggregation weight; if the empirical risk gap $\mathcal{G}_{D_{c}}(\theta^{r})>0$, which indicates that the global model does not generalize very well for local client, we increase its aggregation weight to reduce the empirical risk gap. In this way, we can continuously adjust the empirical risk gap on each client, and finally make the empirical risk gap smaller on each client, so that the global model can adapt to out-of-distribution data. The details of the training phase are shown in Algorithm. 2.

To fully compare the performance of our method, we set up two scenarios: In-Distribution and Out-of-Distribution (different sampling pattern, acceleration rate, contrast, and unseen center) experiments. In this section, only the results of the In-Distribution scenario are presented, and the results of the Out-of-Distribution scenario will be presented in IV. C. The In-Distribution experiments are conducted using a 1D random 4$\times$ mask.

Table I provides the quantitative results of the different methods under In-Distribution scenario. From Table I, we can see that our method is superior to other FL methods in terms of PSNR and SSIM metrics. Compared with the classical FedAvg method, the average value of the PSNR metric of our GAutoMRI improves from 32.3288 dB to 35.5253 dB, and the average value of the SSIM metric increases from 0.8958 to 0.9384. This is primarily attributed to the fact that FedAvg uses an average aggregation approach, which does not adapt well to the statistical heterogeneity of the client. Comparatively, personalized federated learning methods (FedMRI and FedALA) perform better than FedAvg. They exhibit improved adaptability to local data distributions by maintaining client specificity. The In-house dataset comprises data with three different contrasts, resulting in greater heterogeneity, as illustrated in Fig. 3. Therefore, the performance of all methods on this dataset is slightly lower compared to their performance on the other two datasets. However, relatively speaking, our method still maintains a better performance, proving the superiority of our methods. Furthermore, our fairness adjustment approach contributes to achieving an obviously improved performance balance for each client.

In addition to the quantitative results, the qualitative results of the different methods are also presented in Fig. 4. The reference images and reconstruction results from the comparative methods are displayed from left to right, with the corresponding error maps shown in the second and fourth rows. The corresponding quantitative metrics are also provided in the upper-right corner of the reconstructed images. From Fig. 4, our method obtains the best reconstruction quality and the smallest error, which also indicates that our method is superior to other methods.

To demonstrate the performance of our proposed GAutoMRI on out-of-distribution data, we compare the model performance at different sampling pattern (1D random 4$\times$ for training; 1D equispaced 4$\times$ for testing), acceleration rate (1D random 4$\times$ for training; 1D random 6$\times$ for testing), contrast (fastMRI T2) and unseen center (MoDL-Brain). We divided it into two parts: 1) different sampling pattern and acceleration rate, and 2) different contrast and unseen center.

In this section, we conduct ablation studies and utilize a radar chart to analyze the performance differences among different methods within the source domain.