Gradient Coreset for Federated Learning

Federated learning (FL) is an approach to machine learning in which clients collaborate to optimize a common objective without centralizing data [32]. The training dataset is distributed across a group of clients, and they contribute to the training process by sharing privacy-preserving updates with the central server across communication rounds until the model converges.

FL proves particularly valuable in situations where a central server lacks a sufficient amount of data for standalone model training but can leverage the collective data from multiple clients, including edge devices, sensors, or hospitals. For instance, a hospital aiming to develop a cancer prediction model can benefit from training their model using data from other hospitals, while maintaining the privacy of sensitive patient information. However, in scenarios where clients have limited computational resources or their data is noisy, it becomes imperative to design robust algorithms that enable their participation while minimizing computation and energy requirements. Our proposed solution addresses this challenge by identifying a subset of each client’s data, referred to as the ”coreset,” which reduces noise and facilitates effective training of the central server’s model.

Conventional coreset selection methods such as Facility Location, CRUST, CRAIG, Glister, and Gradmatch [34, 38, 33, 19] have been developed to enhance the efficiency and robustness of machine learning model training. However, adapting these strategies to Federated Learning (FL) settings is challenging due to the non-i.i.d. nature of clients’ datasets. Traditional coreset algorithms aim to select a representative subset that ensures a model trained on it performs similarly to the entire dataset. In FL, each client’s data originates from diverse distributions and is influenced by noise. For example, in a hospital context, data distributions differ due to varying demographics across locations and may be subject to varying amounts of different types of noise. As a result, biased updates from clients hinder the learning progress of FL algorithms. We demonstrate this obstacle through a motivating experiment in Section 2.

We introduce Gcfl (as shown in Figure  1), an algorithm specifically designed to address the aforementioned challenges. Our approach involves selecting a coreset every $K$ communication rounds to derive local updates. Similar to [49], we assume the server has access to a small validation dataset for guiding coreset selection. However, in contrast to [49], our validation dataset is not public; we exclusively use (last layer) gradients derived from it. Gcfl uses these gradients to identify a coreset at each client, effectively training the FL model. In our experiments, we demonstrate that our approach is robust to different types of noise, efficient, while preserving privacy and minimizing communication overhead¹¹1The code can be found at https://github.com/nlokeshiisc/GCFL_Release/tree/master.

In summary, our work makes the following contributions:

1. 1.

   We introduce Gcfl, a framework for efficiently selecting coresets for federated learning while preserving privacy.

2. 2.

   Through experiments, we show that Gcfl achieves the best tradeoff between accuracy and speed in a non-noisy setting.

3. 3.

   Furthermore, we demonstrate that Gcfl effectively filters out various types of noise, including closed-set label noise, open-set label noise, and attribute noise, resulting in improved performance compared to well-established baselines for FL and coreset selection.

To emphasize the need for a coreset algorithm like Gcfl in federated learning and to showcase its impact on performance, we conducted an experiment using a small toy dataset. We generated ten isotropic Gaussian blobs in $\mathbb{R}^{10}$ with varying standard deviations (ranging from $1$ to $8$) using scikit-learn’s make_blob() utility [39]. A test set was reserved, containing $15\%$ of the samples, while the remaining training data was divided among the server and ten clients. The training subset allocated to the server serves as a validation dataset, as will be explained in our algorithm, and is solely employed to guide coreset selection at the clients.

To simulate noise, we randomly flipped $40\%$ of the labels in each client’s samples. We trained logistic regression models under three settings: (i) FedAvg, (ii) Gcfl, and (iii) Skyline. In the Skyline approach, clients only computed updates from the clean ($60\%$) samples to establish an upper-bound performance benchmark using clean data. Figure 2 displays the results, revealing that FedAvg’s performance is adversely affected by the presence of noisy training samples. In contrast, Gcfl outperforms FedAvg and falls between Skyline, demonstrating its effectiveness in mitigating the impact of noisy FL data, likely due to its use of a small, server-guided training subset. Gcfl holds significant promise for FL, especially in noisy data settings, where it significantly enhances accuracy and model generalization.

One can consider the idea of mitigating the impact of noise by fine-tuning the model refined from each communication round using the server’s validation dataset. However, such an approach may prove ineffective when the sample size in the validation dataset is insufficient. To explore this approach, we conducted experiments using CIFAR-10 and CIFAR-100 datasets, both affected by $40\%$ closed-set label noise. We compared the performance of Gcfl and FedAvg, as detailed in Table 1. Our observations indicate that while fine-tuning does offer some improvements, its effectiveness is limited by the small size of the validation dataset. In our upcoming experiments, we will demonstrate how Gcfl, in contrast, efficiently harnesses the validation dataset to guide coreset selection at the client side, ultimately optimizing its performance.

Federated Learning (FL) is distinguished by data heterogeneity, where various clients possess data from different sources with diverse characteristics. As a result, aggregating updates from such a heterogeneous data source can impact the convergence rate of models. The challenge of addressing this heterogeneity within client data in the context of Federated Learning (FL) has garnered significant attention in the literature [12, 23, 30, 26, 15, 22].

For instance, FedProx, introduced by [26], incorporates a proximal term into the objective function. This term penalizes updates that deviate significantly from the server’s parameters, aiming to accommodate client heterogeneity. Similarly, Scaffold, proposed by [15], focuses on reducing the variance in the server’s aggregated updates by managing the drift in update computation across clients. However, FL presents a unique challenge when dealing with noisy data owned by clients, as neither the server nor the clients have a complete view of the entire training dataset. Traditional data cleaning approaches [40, 28, 9, 41] may not be directly applicable to FL.

Coreset selection is a well-established technique in machine learning that involves selecting a weighted subset of data to approximate a desired quantity across the entire dataset. Traditional coreset selection methods typically depend on the model and use submodular proxy functions for coreset selection [47, 20, 16, 10, 35]. Recent developments in the literature have explored coreset selection in conjunction with deep learning models [33, 34, 18, 5, 38]. Nevertheless, most existing coreset selection approaches are tailored for conventional settings where all data is readily accessible, requiring thoughtful adaptation for FL.

Coreset selection in Federated Learning (FL) is an underexplored domain, primarily due to the complexities tied to privacy and non-i.i.d. data distribution among clients [37, 48, 4, 36]. Notably, [1] picks a coreset of clients to collectively represent the global update across all clients using the facility location algorithm. In contrast, [37] uses Shapley values in a game-theoretic framework for again to perform client selection, while [36] explores reinforcement learning techniques for data selection. However, Nonetheless, training [36] is a challenging task, imposing an additional workload on local clients by necessitating the training of an extra private model. In comparison to the prior work, Gcfl is easy to implement and blends well with the FL framework.

FL has different paradigms: Personalized FL strives to train specialized models for individual clients, and substantial research has been conducted in this direction [6, 8, 31, 24, 13]. In contrast, our work is focused on building models that exclusively account for the server’s distribution.

We finally note that various techniques have been introduced to ensure privacy in FL, including differential privacy [7, 14], homomorphic encryption [2, 25], and more. As Gcfl is model-agnostic, these methods can be seamlessly integrated with our approach.

In our Federated Learning setup, a group of $N$ clients is represented by the set $\mathcal{C}=\{c_{1},c_{2},\cdots,c_{N}\}$. The training dataset $D_{T}=\bigcup_{i=1}^{N}D_{i}$ is divided among the clients, where each client $c_{i}$ has a data chunk $D_{i}$ consisting of $n_{i}$ samples $\{(x_{ij},y_{ij})\}_{j=1}^{n_{i}}$. Here, $x_{ij}\in\mathcal{X}$ denotes the input features of the $j^{th}$ data point at the $i^{th}$ client, and $y_{ij}\in\mathcal{Y}$ represents its corresponding target. It is important to note that the data chunks are disjoint, and the set of samples in each data chunk is not obtained independently and identically distributed ($i.i.d.\hbox{}$) from the ground truth target distribution $\Pr_{S}$.

Our objective is to train a machine learning model $f_{\theta}:\mathcal{X}\rightarrow{\mathcal{Y}}$ where $\theta$ represents the learnable parameters. The server $S$ defines the objective for the downstream task and has access to a small dataset $D_{S}$ consisting of samples obtained independently and identically from the ground truth target distribution $\Pr_{S}$. Our aim is to minimize the expected value of the loss function $\ell(f_{\theta}(x),y)$, over instances $(x,y)$ sampled from the distribution $\Pr_{S}$.

As $D_{S}$ is small, it is insufficient for training $f_{\theta}$ and using it alone can lead to overfitting. To overcome this, the server seeks assistance from the clients to learn $f_{\theta}$ while respecting their privacy constraints. Federated Learning is a promising solution to this problem, where the learning progresses through $T$ communication rounds. In each round $t$, the server selects a subset of clients $\mathcal{C}_{sel}^{t}$ and shares the current FL model parameters $\theta^{t}$ with them. The selected clients initialize their local model with $\theta^{t}$ and train it for a few epochs with their respective private data chunks $D_{i}$ to arrive at the updated model parameters $\theta_{i}^{\prime}$. The difference in model parameters, computed as $\delta_{i}^{t}=\theta^{\prime}_{i}-\theta^{t}$, is then transmitted back to the server. The server then averages the parameter updates received from clients and updates the FL model as follows:

The global learning rate used by the server is denoted as $\eta_{g}$, while $\eta_{l}$ represents the local learning rate used by each client to train Gcfl. Although updates generated using equation (2) can minimize the objective (1) when the clients’ datasets are independently and identically distributed according to $\Pr_{S}$, computing updates from the entire dataset is not recommended when the data contains noise. Moreover, for resource-constrained clients such as edge devices, it is crucial to compute updates in an energy-efficient manner. To address these challenges, we propose using adaptive coreset selection, which involves selecting a weighted subset that approximates the characteristics of the entire dataset. The selected coreset should reduce computation costs without compromising the performance of the FL model, and also prevent the updates from only minimizing the client’s local loss, especially when the client’s data distribution significantly differs from the ground truth distribution $\Pr_{S}$. In this regard, we aim to answer the following question: How can clients in $\mathcal{C}$ select an effective coreset that facilitates the computation of $\delta^{t}_{i}$ and also helps minimize the objective (1)?

Let us denote the coreset selected by a client $c_{i}$ in communication round $t$ as $\mathcal{X}_{i}^{t}$ and its associated weight as ${\mathbf{w}}_{i}^{t}$. We begin the exposition by listing certain desiderata for coreset selection in FL and then proceed to explain how Gcflmeets them.

1. 1.

   The algorithm should align with the current data distribution of clients while also approximating the ground truth distribution $\Pr_{S}$, guaranteeing a coreset that mirrors the desired target.

2. 2.

   The coreset algorithm should adapt and update $\mathcal{X}_{i}^{t}$ as FL model $f_{\theta}$ evolves, maintaining relevance.

3. 3.

   Assumptions in the coreset approach should uphold FL privacy constraints, safeguarding client data and confidentiality.

To meet the first requirement, relying solely on signals from a client’s local dataset, denoted as $D_{i}$, is insufficient, as these datasets are not identically and independently distributed with respect to the global distribution $\Pr_{S}$. However, $D_{S}$ contains samples drawn from the target distribution and can potentially aid in the selection of a coreset. Previous research [49] has demonstrated that making $D_{S}$ publicly accessible enables clients to choose an effective coreset. Nevertheless, in privacy-sensitive domains like healthcare, even the inclusion of $D_{S}$ may raise privacy concerns, making it unsuitable for sharing.

Hence, the task of coreset selection within the input feature space becomes challenging. As an alternative, we shift our focus towards coreset selection in the gradient space, as Federated Learning (FL) allows the server to disseminate gradients computed from $D_{S}$. It is important to note that any cryptographic techniques applied to secure clients’ gradients, such as differential privacy, can also be employed to safeguard the server’s gradients. Additionally, to minimize communication overhead, we opt to transmit an aggregated gradient (average) derived from $D_{S}$. This choice is grounded in the Information Bottleneck theory [43], which suggests that parameters in the final layers contain crucial discriminatory class information $\mathcal{Y}$, while those in the initial layers primarily encapsulate feature-specific information $\mathcal{X}$.

We begin Gcfl by defining the server’s objective and then illustrate how we incorporate it within the Federated Learning framework. We use $\ell_{S}$ to signify the loss incurred by the server concerning the validation data $D_{S}$, and denote the loss of client $c_{i}$ w.r.t. its local dataset $D_{i}$ as $\ell_{i}$.

where $\ell:\mathcal{Y}\times\mathcal{Y}\rightarrow\mathbb{R}^{+}$ is a loss function that is pertinent to the problem.

We define $\nabla_{\theta}\ell_{S}(\theta)$ as the average gradient of $\ell_{S}$ at $\theta$ on the validation dataset $D_{S}$, and $\{\nabla_{\theta}\ell_{i}^{j}\}_{j=1}^{n_{i}}$ as individual data gradients of client $c_{i}$ for all $j\in[n_{i}],i\in[N]$. The objective for each client $c_{i}$ is to select a coreset $\mathcal{X}_{i}^{t},{\mathbf{w}}_{i}^{t}$ of size $b$ such that the gradients derived from it closely match $\nabla\theta\ell_{S}(\theta)$. Our coreset selection objective is:

Here, $\lambda$ is a hyper-parameter that regulates the weights of selected items in the coreset. Due to the combinatorial nature of the optimization objective, it is known to be NP-Hard [19]. Therefore, we employ a greedy approximation method, which we will explain in more detail later.

Assuming that the client has solved Eq. 5, we now describe how it computes an update to share with the server. Let $\ell_{i}^{gm}$ denote the loss on the selected coreset, defined as

To minimize the above loss, the client runs several epochs of stochastic gradient descent on the coreset. From the updated model, the client derives its update $\delta^{t}_{i}$ as follows:

Where $E$ is the number of local gradient update steps performed by the client on the coreset, and $\theta^{t}_{k-1}$ denotes the model parameters at the $k^{th}$ intermediate step, with $\theta^{t}_{0}=\theta^{t}$. In our experiments, we observed that the coreset weight ${\mathbf{w}}_{i}$ had a minimal impact on computing the update, so we omitted it in Eq  (7).

We present experiments on various real world datasets with a range of noise settings to demonstrate the efficacy of Gcfl over state of the art approaches.

In this work, we introduced a new approach called Gcfl to address the challenge of learning in a federated setting where the distribution of data across the client nodes is non-i.i.d. , and ingested with noise. Our proposed approach selects a coreset from each client that best approximates the server’s last layer gradient. Our experimental results illustrate that Gcfl outperforms state-of-the-art methods, and achieves the best accuracy vs. efficiency trade-off when the datasets are not noisy. In case of noise, Gcfl was able to achieve significant gains compared to other FL and coreset selection baselines.

Durga Sivasubramanian and Ganesh Ramakrishnan express their gratitude to Google Research award for their assistance and funding. Ganesh Ramakrishnan is also grateful to the IIT Bombay Institute Chair Professorship for their support and sponsorship.

We have released the code in github: https://github.com/nlokeshiisc/GCFL_Release/tree/master.

We list the notations and their corresponding description used in Gcfl in the Table 2.

Algorithm 1 presents the pseudocode of our proposed approach. The algorithm depicts the actions performed by both the server and clients at each round. The server’s actions are are highlighted in orange and the clients’ actions shown in blue color.

At a communication round $t$, first the server samples $m$ clients (line 3), and then broadcasts $\theta^{t}$ to them. Additionally, if the clients need to perform coreset selection, the server also broadcasts the validation gradients. The clients subsequently use these to construct the coreset $\mathcal{X}_{i}^{t}$ iteratively by solving our label-wise OMP problem (line $8$ of the algorithm). Then clients use the selected coreset to train the model for $E$ epochs (line 10–13). Finally, the parameters of the updated model are uploaded back to the server (line $14$). The server waits until it receives the updated model from all the sampled clients $C^{t}$, then it averages them and the updates the global model with the average(line $17$). This completes one round of Gcfl execution.

In table 3 we list the details about the datasets used in the experiments. Further, to showcase the challenging nature of datasets used in our experiments, we show the non-i.i.d.  nature of the partition across clients using heatmaps in Figure 10. We observe that each client is characterized by an allocation of data points $D_{i}$ such that it is skewed in favor of one or two classes.