From Optimization to Generalization: Fair Federated Learning against Quality Shift via Inter-Client Sharpness Matching

1 Introduction

In light of escalating concerns regarding data privacy, federated learning (FL) [McMahan et al., 2017] has emerged as a promising approach for training deep neural networks in the realm of medical image analysis [Dou et al., 2021]. A significant challenge within FL is the inherent data heterogeneity [Ye et al., 2023; Huang et al., 2022; Huang et al., 2023a] observed across various medical institutions, primarily attributed to their independent data collection processes.

∗Corresponding author

This heterogeneity has been examined from several perspectives in existing research, e.g., domain shift [Li et al., 2021; Liu et al., 2021; Jiang et al., 2023b], label skew [Zhang et al., 2022; Wu et al., 2023b], and label quality variation [Wu et al., 2023a; Chen et al., 2023; Wu et al., 2024]. Nevertheless, the prevalent issue of quality heterogeneity [Fang et al., 2023] in medical imaging, a factor that could potentially raise new challenges for FL, remains under-explored.

Despite relatively rigorous protocols in medical imaging environments, uniformity in image quality across various institutions cannot be strictly assured. As illustrated in Fig. 1, images from some clients may exhibit noise due to equipment malfunctions or the necessity for low-dose imaging. Moreover, random movements from either patients or cameras in certain cases can result in motion blur. Typically, the proportion of low-quality (i.e., corrupted) images is smaller compared to high-quality (i.e., clean) images [Huang et al., 2023d]. In such scenarios, characterized by quality shifts across clients, FL tends to exhibit a bias towards the moreprevalent clean images, thereby compromising the performance on the less-frequent corrupted images. Such a bias can be a critical concern when applying federated models to clients with corrupted images. Therefore, it is imperative to address the issue of biased performance under quality shifts.

In this paper, we pioneer the identification and formulation of this significant challenge in FL within real-world medical contexts. In FL, especially under the mantle of privacy protection, there is an absence of prior knowledge about the imaging quality of each client, suggesting the necessity for a data-agnostic solution to this issue. We theoretically demonstrate improving performance on the poorest-quality images is equivalent to achieving client-level fairness in FL. Hence, we remodel this challenge as a problem of client-level fairness, i.e., how to ensure equitable FL performance among clients with clean and corrupted image distributions? This investigation marks the first effort to promote fair FL across clients with image quality shifts, a departure from previous research focused on fairness under domain shifts [Jiang et al., 2023a] or class distribution shifts [Li et al., 2020].

In this work, we observe that existing methods to foster client-level performance fairness in FL typically modify the importance weights in global aggregation to harmonize certain training metrics (e.g., loss) across clients [Mohri et al., 2019; Li et al., 2020; Jiang et al., 2023a]. However, as illustrated in Fig. 2, while these methods facilitate fairness regarding optimization during training, they do not necessarily ensure fairness regarding generalization during testing. This gap stems from the issue that equalizing training losses for all clients may result in a simplistic convergence towards a sharp minimum, particularly for the minority distributions (i.e., corrupted image distributions), thereby impairing generalization on testing sets. To counter this, our focus extends beyond fair optimization to fair generalization. Recent advancements in sharpness-aware minimization reveal an inverse relationship between the generalization capability and the sharpness of the loss surface [Foret et al., 2021; Zhuang et al., 2022]. Inspired by this finding, we introduce a novel Federated learning framework with Inter-client Sharpness Matching (FedISM), which aims to equalize the sharpness levels across clients for fair generalization. In this framework, local optimization in FL is made sharpnessaware, where each client's local update aims to minimize sharpness on its local data. Subsequently, weights for global aggregation are determined based on each client's sharpness level, with higher weights assigned to clients exhibiting greater sharpness. In this way, the global model update is more effective in minimizing sharpness in higher-level clients, leading to more uniform sharpness and fair generalization across both clean and corrupted image distributions as shown in Fig. 2. It is worth mentioning that FedISM involves only simple modifications to local optimization and global aggregation compared to FedAvg [McMahan et al., 2017], thus eschewing complex loss functions and additional regularization and offering ease of implementation.

The contributions are summarized into three folds:

• New Fairness Challenge in FL - Imaging Quality Shift: We identify and articulate a new challenge in FL, specifically concentrating on achieving performance fairness across clients with diverse imaging qualities.
• Innovative Solution - FedISM: To address this challenge, our focus extends beyond previous fair optimization to fair generalization. Our proposed solution, FedISM, integrates sharpness-aware local updates with sharpness-dependent global aggregation, promoting uniform sharpness across clients and achieving more equitable generalization.
• Extensive Validation: We validate the effectiveness of FedISM through a series of experiments on two realworld medical image classification datasets, i.e., RSNA ICH and ISIC 2019. FedISM demonstrates superior performance, outperforming several state-of-the-art methods in promoting fairness in FL.

2 RelatedWork

2.1 Fair Federated Learning

Fairness [Huang et al., 2023c] has become a crucial topic in FL, primarily concentrating on collaborative fairness [Lyu et al., 2020; Xu et al., 2021] and performance fairness [Li et al., 2020]. Collaborative fairness aims for the reward of each participant to be proportional to its contribution to the federation. In contrast, performance fairness advocates for unbiased performance across different devices/attributes/distributions. Existing solutions mainly achieve this goal by balancing weights of different objects to optimize fairly. For instance, qFedAvg [Li et al., 2020] addresses this by prioritizing clients that are harder to optimize; FedCE [Jiang et al., 2023a] tackles this issue by considering task-specific performance. Nevertheless, fairness respecting generalization has not been considered thoroughly till now, leading to sub-optimal outcomes.

2.2 Sharpness of Loss Surface

Bridging the gap between training optimization and testing generalization remains a pivotal challenge in machine learning. Recent developments in sharpness-aware minimization suggest that models tend to generalize better on flat minima than on sharp minima [Foret et al., 2021; Zhuang et al., 2022]. Motivated by this insight, various studies have incorporated the concept of loss surface sharpness to address poor generalization. SharpDRO [Huang et al., 2023d] merges sharpness with GroupDRO [Sagawa et al., 2020] to achieve robust generalization. ImbSAM [Zhou et al., 2023] focuses on enhancing the performance for tail classes in long-tailed recognition by minimizing sharpness in these classes. In addition to typical machine learning, it has recently been applied in FL [Qu et al., 2022; Caldarola et al., 2022; Sun et al., 2023]. However, the relationship between sharpness and fairness within FL has not yet been investigated.

3 Methodology

3.1 Preliminaries

For a typical image classification task with C classes, we consider a cross-silo FL scenario with K participants. Each k-th participant possesses a private dataset Dk = {(xi ∈ X, yi ∈ Y)}Nk i=1 , k ∈ [K], where X and Y = [C] represent the input image and label spaces, respectively. An imagelabel pair (xi, yi) is drawn from the client-specific distribution Pk(x, y | ak), with ak ∈ [A] indicating an attribute only influencing image quality in client k. This paper highlights the uneven distribution of clients across various quality attributes. Define f(·; θ) : X → ∆C-1 as a deep learning model parameterized by θ and ℓ : ∆C-1 × Y → R+ as the loss function, where ∆ denotes the probability simplex. In this paper, the goal of FL is to optimize θ for performance enhancement on images with the worst-performing quality:

When treating each quality as a group, Eq. 1 achieves group fairness which is similar to the aim in distributionally robust optimization [Sagawa et al., 2020]. However, in FL, there is no prior knowledge of group information (i.e., the imaging quality in each client), making group-wise design impractical. Therefore, we achieve Eq. 1 via client fairness, as stated in the following theorem following [Papadaki et al., 2022]:

Theorem 1 (Equivalence). Assuming class distributions of the testing set and all clients' training sets are identical, we have:

and

The proof is detailed in the Appendix. Notably, even when label distribution shifts occur, they can be theoretically addressed by logit adjustment [Menon et al., 2021; Zhang et al., 2022], thereby not considered to be addressed in this paper. Theorem 1 demonstrates that achieving client fairness (Eq. 2) inherently ensures group fairness (Eq. 3). Therefore, this paper subsequently concentrates on enhancing fairness across clients with various imaging qualities.

3.2 Previous Solution: Fair Optimization

Client-level unfairness in FL arises from a tendency to overlook certain clients, particularly those with limited data or outlier distributions. These clients are often overlooked because optimizing them does not significantly contribute to the overall optimization objective. To address this challenge, various approaches have been proposed. AFL [Mohri et al., 2019] concentrates efforts on the worst-performing client; qFedAvg [Li et al., 2020] and FairFed [Ezzeldin et al., 2023] give greater weights to those clients with higher training losses; and FedCE [Jiang et al., 2023a] focuses on clients with lower task-specific metrics. These methods collectively aim to balance optimization across the empirical distribution of each client, striving for more uniform risks among clients, as represented by the following optimization problem:

Essentially, these solutions seek to achieve a balance in optimization by prioritizing clients that are performing worse. However, this strategy may result in clients with poorer performance rapidly converging to sharp minima, primarily because the rate of loss minimization is most pronounced when moving towards them. Due to the discrepancy between empirical and expected risks, this strategy to optimize fairly does not necessarily lead to client fairness in a strict sense, as denoted by Eq. 2. A conceptual illustration of this issue is presented in Fig. 2. Thus, it becomes crucial to extend the exploration of client fairness to include considerations of generalization, beyond mere optimization.

3.3 Measurement of Generalization: Sharpness

Since it is sub-optimal to measure generalization capacity by a single value of empirical risk (i.e., training loss), promoting fairness by merely seeking uniformity in such a metric is inappropriate. To tackle this problem, we should identify a new indicator more closely correlated with generalization ability.

A key limitation of using single training loss as a metric is its insensitivity to the geometric properties of the loss landscape, treating sharp and flat minima indiscriminately. To overcome this, we propose focusing on a range of loss values, specifically the sharpness of the loss surface [Foret et al., 2021; Zhuang et al., 2022]. It is defined as the largest loss change in the vicinity of the initial model parameters:

where ρ is a positive step parameter controlling the search radius. Calculating sharpness as per Eq. 5 poses a challenge due to the continuous and infinite nature of the perturbation. To simplify, the difference can be approximated linearly through the Taylor series when ρ is sufficiently small:

This approximation enables us to identify the optimal perturbation by maximizing the right-hand side of the equation:

Therefore, sharpness can be computed more feasibly as:

Unlike a single training loss, sharpness reflects the rate of change in training loss across the loss surface. Prior research has shown that this rate of change is closely related to generalization capacity. To be specific, models generally perform better on a flat minimum (i.e., with smaller sharpness) than on a sharp minimum (i.e., with larger sharpness) [Foret et al., 2021; Zhuang et al., 2022]. Therefore, we select this indicator to measure generalization capacity.

3.4 Sharpness Matching for Fair Generalization

As the sharpness of the loss surface is indicative of a model's generalization ability, we propose an FL method with interclient sharpness matching (FedISM) to establish a uniform sharpness distribution across clients, thereby achieving fairer generalization. The overview of FedISM is depicted in Fig. 3, and details are as follows.

Our proposed method is structured around the following objective:

where S(f (x; θ), y) represents the sharpness for a given sample and model, as defined in Eq. 8. Unlike previous solutions which focus on a uniformly-low loss (Eq. 4), our key objective is to achieve a uniformly-low sharpness across clients, thereby avoiding convergence to sharp minima while aligning empirical risks. Considering that lower sharpness typically reflects better generalization, this strategy emphasizes fairness regarding generalization more.

In conventional local training in FL, e.g., FedAvg [McMahan et al., 2017], the objective is to minimize the loss on local data via gradient descent:

with η as the learning rate. However, it tends to find an isolated point of low loss, without aiming for a flat minimum as desired in Eq. 9. To address this, local training should incorporate sharpness-awareness. Inspired by sharpness-aware minimization [Foret et al., 2021], the update rule is adapted to:

where ϵ∗ represents the optimal perturbation that maximizes the change in the training loss as Eq. 7. This modification directs the optimization towards a point where the surrounding area exhibits the lowest loss, effectively minimizing the sharpness of the loss surface.

After each round of local training, gradients from all clients are transmitted to the server for global aggregation. Notably, Algorithm 1 Pseudocode of FedAvg and our FedISM

Input: number of clients K, local datasets {D1, . . . ,DK}, local dataset size, total communication rounds T, learning rate of local training η.

Output: final global model θT+1

1: Initialize the global model θ1 2: for t = 1, 2, . . . , T do 3: for Client k = 1, 2, . . . , K in parallel do 4: θ(t,k) ← θt ▷ download the global model 5: for (xi, yi) ∈ Dk do 6: Update θ(t,k) with (xi, yi) by Eq. 10 7: Update θ(t,k) with (xi, yi) by Eq. 11 8: end for 9: end for 10: θt+1 ←Aggregate {θ(t,k)}Kk=1 with wAvg (Eq. 12) 11: θt+1 ←Aggregate {θ(t,k)}Kk=1 with w (Eq. 14) 12: end for 13: return θT+1

under sharpness-aware minimization (Eq. 11), each client's gradient is representative of the direction to minimize sharpness with respect to local data. The crux of achieving uniform sharpness across clients lies in how these gradients are fairly aggregated. Traditional FedAvg [McMahan et al., 2017] assigns aggregation weights based on the quantity of data each client contributes, as defined by:

However, such a weighting policy may not facilitate uniform sharpness. Given that clients with corrupted images typically have less data, the global update tends to be dominated by clients with clean images. Consequently, this might only reduce the sharpness for clean clients, leading to a disparity in sharpness similar to the loss disparity addressed by fair optimization [Mohri et al., 2019; Li et al., 2020; Ezzeldin et al., 2023; Jiang et al., 2023a]. To solve this, we propose sharpness-aware aggregation:

where Sk,t denotes the sharpness calculated on the entire local dataset Dk at any round t by Eq. 8, and q is a predetermined positive parameter. This strategy emphasizes clients with higher sharpness levels during aggregation. As q increases, the aggregation process focuses more on these highsharpness clients. In extreme cases, aggregation will exclusively focus on the client exhibiting the highest sharpness given q → +∞. To conclude, this method ensures that the sharpness levels of clients with initially higher sharpness are primarily reduced, thus promoting a uniform sharpness distribution across all clients. In this way, it accomplishes the objective defined in Eq. 9. To maintain stability in federated training, a moving average is further employed for rounds t > 1, formulated as:

In particular, we set w1 = ˜w1 for the first round.

It is important to note that FedISM only requires clients to share their sharpness values instead of any information regarding data distributions (e.g., imaging quality of each client), which is privacy-preserving.

To facilitate a clearer understanding, the processes of both FedISM and the conventional FedAvg [McMahan et al., 2017] are summarized in Algorithm 1. Compared to FedAvg, FedISM modifies only the local optimizer and global aggregation weights, without the need to introduce complex loss functions or additional regularization techniques. Such a streamlined approach enhances the ease of implementation.

4 Experiments

We conduct a series of experiments to assess the effectiveness of FedISM. More results and discussion can be found in Appendix.

4.1 Experimental Setup

Datasets

Two medical datasets are used for evaluation, in line with prior FL research [Jiang et al., 2022; Wu et al., 2023a]:

• RSNA ICH [Flanders et al., 2020]: The task is to classify each CT slice into five intracranial hemorrhage (ICH) subtypes. Following [Jiang et al., 2022], we randomly select 25000 images for experiments.
• ISIC 2019 [Tschandl et al., 2018; Codella and others, 2018; Combalia et al., 2019]: This dataset contains 25331 images for developing models to classify eight skin diseases.

Both datasets are split into training and test sets in an 8:2 ratio and resized to 224 × 224 pixels following the standard operation [Jiang et al., 2022]. For data partitioning, training sets are partitioned into 20 clients using a Dirichlet distribution (i.e., Dir(1.0)), simulating the prevalent label distribution shifts though this paper does not focus mainly on it. Imaging quality shifts are created via Gaussian noise added on a subset of clients following [Hendrycks and Dietterich, 2019].

Model

For standard evaluation, pretrained ResNet-18 [He et al., 2016] is used as the base model for all experiments.

Implement Details

To mitigate label distribution shifts between local training and testing sets, we incorporate logit adjustment [Menon et al., 2021] in the training of local models. We train batches of 32 images using the Adam optimizer, with a constant learning rate of 0.0003, beta values of (0.9, 0.999), and a weight decay of 0.0005. For FL specifics, we set a maximum of 300 communication rounds and the local epoch as 1. These parameters are consistent across all experiments to facilitate fair comparison. In FedISM, we implement GSAM [Zhuang et al., 2022] for sharpness-aware minimization, with default settings of q = 2.0 and β = 0.5.

Evaluation Strategy

The federated model's performance is evaluated on both an unaltered clean testing set and a generated corrupted testing set (with the identically distributed Gaussian noise in corrupted clients). We use class-balanced accuracy (ACC) and the area under the receiver operating characteristic curve (AUC) as our evaluation metrics. To ensure the robustness of results, three independent experiments are conducted, and performance is averaged over the last five communication rounds, in line with [Huang et al., 2023b].

4.2 Comparison to State-of-the-Arts

To validate the superiority, we compare FedISM with several leading methods, including the basic FL approach FedAvg [McMahan et al., 2017], and five advanced fair FL methods: Agnostic-FL [Mohri et al., 2019], q-FedAvg [Li et al., 2020], FairFed [Ezzeldin et al., 2023], FedCE [Jiang et al., 2023a], and FedGA [Zhang et al., 2023]. Comprehensive details on these methods and implementation strategies are available in the Appendix. In our experiments, 4 out of 20 clients are equipped with corrupted images, constituting a corrupted client ratio of 20%.

Tab. 1 summarizes the mean and standard deviation of ACC and AUC for both clean and corrupted images, as well as their average. Although most fair FL methods improve the performance on corrupted images, excluding the less stable Agnostic-FL, they focus primarily on fair optimization rather than generalization, which can be suboptimal. Comparatively, our FedISM, with its emphasis on sharpness minimization, achieves better generalization on corrupted images. For instance, on the ICH dataset, FedISM outperforms FedAvg and the second-best method (FedGA) by 10.85% and 3.78%, respectively. Notably, previous methods often boost the performance on corrupted images at the expense of performance degradation on clean images in this setting. For instance, FedGA's ACC on clean ICH images falls by 4.84% compared to FedAvg. This trade-off can be problematic in medical scenarios, potentially discouraging high-quality institutions from participation in FL. In contrast, FedISM not only enhances the performance on corrupted images but also maintains top results on clean images. This balanced improvement is crucial in medical scenarios, ensuring accurate diagnostics and encouraging broader participation in FL.

4.3 Ablation Study

FedISM is comprised of two key components: SharpnessAware Local Training (SALT, Eq. 11) and Sharpness-Aware Global Aggregation (SAGA, Eqs. 13 and 14). To evaluate the effectiveness of each component, we conduct ablation studies by integrating them individually with FedAvg, with performance summaries provided in Table 2. SALT, by directing the model towards flat minima, enhances FedAvg's performance on both clean and corrupted images. Nonetheless, this approach is not fully optimized for fairness, as it fails to address the disparity in sharpness between clean and corrupted image distributions. This is where SAGA comes into play; it helps the model prioritize distributions with poorer generalization capacity, typically those of corrupted images. As a result, the combination of both SALT and SAGA yields the best performance on corrupted images, demonstrating the comprehensive effectiveness of FedISM's designs.

4.4 SALTHelps Fair Optimization

A key motivation for developing Sharpness-Aware Local Training (SALT) is to address the issue of fair optimization methods converging to sharp minima, often resulting in suboptimal performance (see Fig. 2). To validate this premise and demonstrate SALT's effectiveness, we integrate SALT with existing fair optimization methods and analyze the impact on performance enhancement. The results, presented in Tab. 3, reveal that, aside from Agnostic-FL [Mohri et al., 2019] which shows large standard deviations, other methods exhibit performance improvements when combined with SALT. It not only confirms that SALT is an effective and adaptable component for fair FL, but also validates our initial motivation for its design.

4.5 Robustness to Different Imbalanced Ratios

Our experiments also demonstrate the robustness of FedISM, showcasing stable performance enhancements across varying ratios of clean and corrupted clients. Quantitative results under various ratios of clients with corrupted images through {0.1, 0.2, 0.3} are illustrated in Fig. 4. Across all ratios and on both image distributions, FedISM consistently outperforms other methods in the two evaluation metrics. Such findings further highlight FedISM's strong adaptability and effectiveness irrespective of the corruption ratio.

4.6 Discussion on Parameters

The parameter q is closely related to how strongly FedISM concentrates on clients with higher sharpness levels. We examine this impact by varying q through {0.1, 0.5, 1.0, 2.0, 5.0, 10.0} and assess the performance on the ICH dataset with 20% corrupted clients, as shown in Fig. 5. For comparison, we also report the performance of FedAvg and the secondbest fair FL method. As q increases, FedISM increasingly focuses on clients with greater sharpness (typically those with corrupted images), which results in a decrease in performance on clean images and an improvement on corrupted images. It is important to note that using very high values of q, such as 5 and 10, would slightly degrade the performance on corrupted images, likely due to training instability. Overall, FedISM demonstrates improved performance across a wide range of q values, reducing the burden of parameter tuning in practice.

5 Conclusion

In this paper, we pioneer the identification and formulation of a new fairness challenge in FL, specifically concerning imaging quality shifts across clients. To address this problem, existing FL approaches have primarily concentrated on balancing the empirical risk among distinct client distributions. Despite their effectiveness for fair optimization, they often neglect the crucial aspect of fair generalization. To address this overlooked area, we introduce Federated Learning with Inter-client Sharpness Matching (FedISM). FedISM innovatively refines both local training and global aggregation by integrating sharpness-awareness, effectively harmonizing sharpness levels across clients to achieve fair generalization in FL. Extensive empirical evaluations, on the well-recognized RSNA ICH and ISIC 2019 datasets, clearly demonstrate FedISM's superiority over current state-of-theart FL methods in terms of promoting fairness. It highlights the effectiveness of FedISM in resolving fairness issues stemming from imaging quality shifts across medical datasets. We are confident that our proposed challenge and solution will pave the way for more equitable and effective FL system designs in medical applications and beyond.

Acknowledgments

This work was supported in part by the National Natural Science Foundation of China under Grants 62202179 and Grant 62271220, and in part by the Natural Science Foundation of Hubei Province of China under Grant 2022CFB585. The computation is supported by the HPC Platform of HUST.