Federated Learning with Instance-Dependent Noisy Label

We start with the standard $C$-class classification problem in FL, involving $K$ clients and a central server. For each client $k$, it contains the local dataset $\mathcal{D}_{k}=\left\{\left(\mathbf{x}_{i},y_{i}\right)\right\}_{i=1}^{N_{k}}$ where $N_{k}$ is denoted the size of dataset, i.e. $|\mathcal{D}_{k}|=N_{k}$, and $y_{i}\in\{1,2,\dots,C\}$. Each client $k$ has a local model $w^{k}$ and a model output $f(\mathbf{x}_{i},w^{k})\in\mathbb{R}^{C\times 1}$, which is the class probability distribution (a.k.a., confidence score vector). The global objective is defined as follows: $\min_{x}\sum_{k=1}^{K}{N_{k}\mathcal{L}_{k}}/{N}$, where $N=\sum_{k=1}^{K}N_{k}$ is the total dataset size. Denote $\boldsymbol{y_{i}}$ as the one-hot form of $y_{i}$, thus $\mathcal{L}_{k}={\sum_{i=1}^{N_{k}}}\mathcal{L}\left(f\left(\mathbf{x}_{i},w^{k}\right),\boldsymbol{y_{i}}\right)/N_{k}$ is the local objective function of client $k$, in which $\mathcal{L}(\cdot,\cdot)$ represents the loss function, and $(\mathbf{x}_{i},y_{i})\in\mathcal{D}_{k}$ is the $i$-th sample of client $k$.

In the classic FL setting, it is assumed that training samples in each client are labeled with ground-truth. However, this assumption is overly idealistic in real world. In this paper, to capture label noise in real-world data, we consider the IDN model, which assumes that the noise rate is not constant and varies based on each sample’s features, rather than varying on each class. Let $D$ be the overall distribution of pairwise random variables $(\mathbf{X},\mathbf{Y})$, where $\mathbf{X}$ denotes the variables of all training samples stored at clients, and $\mathbf{Y}$ represents the variable of corresponding labels. It can be treated as if each data sample in each client $k$ is independently drawn according to $D$.

Correspondingly, define $\hat{D}$ as the overall distribution of pairwise random variables $(\mathbf{X},\mathbf{\hat{Y}})$, where $\mathbf{\hat{Y}}$ represents the variable of corresponding noisy labels. The actual dataset held by each client $k$ is denoted as $\mathcal{\hat{D}}_{k}=\left\{\left(\mathbf{x}_{i},\hat{y}_{i}\right)\right\}_{i=1}^{N_{k}}$, which depends on the noisy distribution $\hat{D}$. To establish the bridge between the clean distribution $D$ and the noisy distribution $\hat{D}$ and build a statistically consistent classifier, the IDNTM (i.e. $T(\mathbf{x})=\left(T_{i,j}\left(\mathbf{x}\right)\right)_{i,j=1}^{C}\in[0,1]^{C\times C}$) is introduced. Here, we define $T_{i,j}(\mathbf{x})$ as the $i,j$-th element of IDNTM, representing the transition probability of a sample $\mathbf{x}$ with a clean label $i$ transitioning to a noisy label $j$. Thus the noisy class-posterior probability $P(\hat{Y}|\mathbf{X})$ can be inferred by the IDNTM $T(\mathbf{x})$ and the clean class-posterior probability $P(Y|\mathbf{X})$ as follows:

$\displaystyle P(\hat{Y}=j|X=\mathbf{x})=\sum_{i=1}^{C}T_{i,j}(\mathbf{x})P(Y=i|X=\mathbf{x}).$

(1)

As mentioned earlier, we need an IDNTM for each instance transiting from a noisy label distribution to a clean one. Typically, given a subset of instances that have both noisy and clean labels, training an IDNTM estimation network is a relatively straightforward task. However, in an FL context, clients face challenges in accurately labeling each instance, rendering it infeasible to directly build an estimation network for the IDNTM.

Therefore, our first step primarily focuses on training a weak global model from the noisy distribution $\hat{D}$ by incorporating the noisy labels during training to help isolate high-confidence clean data. This initial training lasts for a total of $T_{1}$ rounds and each client performs $P_{1}$ local steps during each round. In each round, FedBeat requires the server to disseminate the global model to all clients, which will then fully participate in training and uploading their local models. The server then aggregates the received local models to produce an updated global model by using approaches such as FedAvg [1] or FedProx [4]. We denote the local model uploaded by client $k$ during the final communication round as $w^{k}_{T_{1},P_{1}}$, and the global model as $\bar{w}=\sum_{k}N_{k}w^{k}_{T_{1},P_{1}}/N$.

Once the weak global model $\bar{w}$ is obtained, it is returned to the clients to generate peudo-labels on the training data samples. However, rather than returning only $\bar{w}$, we employ a Bayesian model ensemble [18] approach to boost the classification accuracy, thereby extracting a greater amount of high-confidence data. Particularly, after denoting the average model $\mu=\bar{w}$, the server then calculates the standard deviation $\Sigma$ of the local models:

$\displaystyle\Sigma=\mathrm{diag}\left(\sum\limits_{k}\frac{N_{k}}{N}\left(w^{k}_{T_{1},P_{1}}-\mu\right)^{2}\right).$

(2)

The standard deviation $\Sigma$ is also returned to each client. Here the square operation is element-wise. With both $\bar{w}$ and $\Sigma$, each client $k$ samples $M$ ensemble models ${w^{k,(m)}}\sim\mathcal{N}(\mu,\Sigma),m=1,\dots,M$. It is important to note that the $M$ ensemble models can vary among clients. Such an operation is to derive smoother predictive distributions and minimize the risk of overconfidence. Specifically, each client $k$ computes the average outputs $\tilde{f}(\mathbf{x}_{i})$ (i.e., the average confidence score vector) for each instance $i\in\mathcal{\hat{D}}_{k}$ using its $M$ ensemble models $w^{k,(m)}$ as follows:

$\displaystyle\tilde{f}(x_{i})=\frac{1}{M}\sum_{m=1}^{M}f\left(\mathbf{x}_{i};w^{k,\left(m\right)}\right),$

(3)

where $\tilde{f}(\mathbf{x}_{i})\in\mathbb{R}^{C\times 1}$. Let $\tilde{f}_{c}(\mathbf{x}_{i})$ be the $c$-th elements of $\tilde{f}(\mathbf{x}_{i})$, then the pseudo-label $\tilde{y}_{i}$ of $\mathbf{x}_{i}$ is generated according to $\tilde{y}_{i}=\arg\max_{c}\tilde{f}_{c}(\mathbf{x}_{i})$, with its associated confidence score being $\tilde{f}_{\tilde{y}_{i}}(\mathbf{x}_{i})$. Using the predetermined confidence threshold value $\tau$ (a hyper-parameter), client $k$ filters its local dataset to construct a extracted dataset containing feature inputs (e.g., images), noisy labels, and pseudo-labels:

$\displaystyle\mathcal{\tilde{D}}_{k}=\{(\mathbf{x}_{i},\hat{y}_{i},\tilde{y}_{i});i\in\mathcal{\hat{D}}_{k}~{}\text{and}~{}\tilde{f}_{\tilde{y}_{i}}(\mathbf{x}_{i})\geq\tau\},$

(4)

where we denote the size of the extracted dataset $\tilde{D}_{k}$ as $\tilde{N}_{k}$.

With a proper choice of $\tau$, the pseudo-label accuracy is high (see Table 5). However, the number of extracted instances can be small compared to the overall dataset size. This limited number presents two challenges: First, training an accurate classification model solely with these instances is impractical. Second, the limited number of filtered instances per client hinders the individual training of an effective IDNTM.

After extracting data with high-confidence pseudo-labels, an IDNTM estimation network $\theta$ is trained using the data extracted on each client. When provided with an input example $\mathbf{x}$, the estimation network $\theta$ generates the IDNTM $\tilde{T}(\mathbf{x};\theta)$, which contains the probability of transitioning from the clean distribution to the noisy distribution. The training is conducted for $T_{2}$ rounds, and at each round $t$, a subset of clients is selected to participate in the training. For each selected client $k$, the current version of the estimation network, $\theta_{t}$, is first received, and it is then set as the local network $\theta_{t,0}^{k}=\theta_{t}$. The client then performs $P_{2}$ steps of local training based on a loss function $\tilde{\mathcal{L}}(\theta_{t,p}^{k})$.

$\displaystyle\tilde{\mathcal{L}}(\theta_{t,p}^{k})=-\frac{1}{\tilde{N}_{k}}\sum_{i=1}^{\tilde{N}_{k}}\sum_{c=1}^{C}\boldsymbol{\hat{y}_{i}}_{,c}\log\left(\left[{\boldsymbol{\tilde{y}_{i}}^{\mathrm{T}}}\cdot\tilde{T}\left(\mathbf{x}_{i};\theta_{t,p}^{k}\right)\right]^{\mathrm{T}}_{c}\right),$

(5)

where $p\in\{1,\dots,P_{2}\}$ denotes the local step, the subscript $c$ denotes the $c$-th element of the vector which is usually corresponding to class $c$, $\tilde{T}(\mathbf{x}_{i};\theta_{t,p}^{k})$ is the estimated IDNTM for the example $\mathbf{x}_{i}$ generated by the local network $\theta_{t,p}^{k}$. Here $(\cdot)^{T}_{c}$ means we pick the $c$-th element from the $C$-dimensional column vector transposed by the row vector. It is noteworthy that both the noisy label $\boldsymbol{\hat{y}_{i}}$ and the pseudo-label $\boldsymbol{\tilde{y}_{i}}$ are hard labels, and in the one-hot form of $\hat{y}_{i}$ and $\tilde{y}_{i}$ respectively. $\tilde{\mathcal{L}}$ can be viewed as a traditional cross-entropy loss function which treats the noisy label as the ground-truth and $\boldsymbol{\tilde{y}_{i}}^{\mathrm{T}}\tilde{T}(\mathbf{x}_{i};\theta_{t,p}^{k})$ as the predicted output. Such a loss function ensures that the estimated IDNTM, which is generated based on the features of each instance, can successfully map the pseudo-label to the noisy label. Once the selected client completes their local updates, both the updated parameter $\theta_{t,P}^{k}$ and the size of their extracted dataset $\tilde{N}_{k}$ are sent back to the server for further processing. The server then updates the global $\theta_{t}$ as follows:

$\displaystyle\theta_{t+1}=\frac{1}{\tilde{N}}\sum\limits_{k}\tilde{N}_{k}\theta_{t,P_{2}}^{k},$

(6)

where $\tilde{N}$ is the sum of the number of all extracted examples. Note that it is crucial to upload $\tilde{N}_{k}$ for aggregation on the server instead of using $N_{k}$ directly as in the federated data extraction step. Though the ratio of aggregation weights among the clients for this process is similar to the previous step in most of IID settings, IID with a high noise rate or non-IID case either causes a huge imbalance in the amount of data extracted. Upon completing $T_{2}$ rounds, we obtain the final IDNTM estimation network $\theta_{T_{2}}$, denoted as $\tilde{\theta}$, which generates the IDNTM that approximates the ground-truth transition matrix.

Once the estimation network $\tilde{\theta}$ has been trained, it can be employed to correct the weak global classification model $\bar{w}$ derived from the data extraction step. Since $\bar{w}$ is trained on the noisy dataset $\hat{\mathcal{D}}$, simply continuing to train it in the same way as before would result in overfitting for low noise rates or difficulty in converging for high noise rates. Furthermore, the output distribution would still be based on the noisy dataset.

To obtain a strong classification model, we perform federated classifier correction by conducting $T_{3}$ communication rounds for global model aggregation, with each round consisting of $P_{3}$ local steps for fine-tuning the local models similar to the data extraction step. Thus, to minimize the impact of noisy labels, we use the IDNTM estimation network $\tilde{\theta}$ to map the original predictions to noisy outputs. We define the classifier correction loss function as follows:

$\displaystyle\bar{\mathcal{L}}(\bar{w}_{t,p}^{k})=-\frac{1}{N_{k}}\sum_{i=1}^{N_{k}}\sum_{c=1}^{C}\boldsymbol{\hat{y}_{i}}_{,c}\log\left(\left[{f\left(\mathbf{x}_{i};\bar{w}^{k}_{t,p}\right)^{\mathrm{T}}}\tilde{T}\left(\mathbf{x}_{i};\tilde{\theta}\right)\right]^{\mathrm{T}}_{c}\right),$

(7)

where $\bar{w}^{k}_{t,p}$ represents the local model on client $k$ during the $p$-th epoch of round $t$, $f(\mathbf{x}_{i};\bar{w}^{k}_{t,p})^{\mathrm{T}}\tilde{T}(\mathbf{x}_{i};\tilde{\theta})$ is the adjusted prediction obtained by mapping the original prediction $f(\mathbf{x}_{i};\bar{w}^{k}_{t,p})$ through the IDNTM estimated by the network $\tilde{\theta}$.

After calculating the loss function for all clients, we aggregate the local models using a global aggregation procedure, similar to the one employed in the federated data extraction step:

$\displaystyle\bar{w}_{t+1}=\frac{1}{N}\sum\limits_{k}N_{k}\bar{w}_{t,P_{3}}^{k}.$

(8)

We argue that the federated classifier correction step works by utilizing the IDNTM in reverse. Specifically, the IDNTM produced by $\tilde{\theta}$ adjusts the clean label distribution to approximate the noisy label distribution. Therefore, we can adjust the output distribution in reverse to approximate a clean label distribution by ensuring that the adjusted output closely matches the noisy label distribution.

Baseline Methods. Given that no previous work has considered IDN in the FL setting, we incorporate several baseline methods to demonstrate the superior performance of our proposed method. We consider three types of baselines. First, we examine some traditional FL methods that learn under the assumption of clean labels (e.g. FedAvg [1] and FedProx [4]). Then, we look at a scenario where each client individually applies the state-of-the-art IDN method developed in centralized learning (BLTM [12]). Lastly, we consider FedCorr [15] which is a state-of-the-art FL approach solving CCN.

Implementation Details. Our experiments encompass a range of noise levels (IDN rates: 30%, and 50%) across the CIFAR-10 and SVHN datasets. Specific details vary based on the dataset: CIFAR-10:We consider 10 clients, each having 5000 balanced training data samples in the IID setting. In the non-IID scenario, the sample count is scaled down to 60% of the original. SVHN: In the IID setting, we conisder 10 clients, where each client has 4659 balanced training data samples. In the non-IID setting, the number of samples is reduced to 50% of the original. We use the Dirichlet distribution ($\alpha_{Dir}$) to simulate the non-IID case among clients.

For the initial step, we employ ResNet-34 [19] for CIFAR-10 and ResNet-18 [19] for SVHN as the backbone model with an initial learning rate of 0.01. We warm-up the classification network for a total of 15 rounds, with 10 clients participating in each round, and each client performing 2 local epochs. We then set the number of model ensembles, denoted by $M$, to 10 and the threshold $\tau$ to 0.65. For the second step, which involves training the estimation network, we utilize ResNet-34 for CIFAR-10 and ResNet-18 for SVHN with a 100-dimensional output as the backbone model with an initial learning rate of 0.01. We conduct $T_{2}$ rounds, set to 85, with 10 clients participating in each round. For the final step, we use the model obtained from the final round of the first step as the backbone model with an initial learning rate of 0.0005. We conduct $T_{3}$ rounds, set to 50, with 10 clients participating in each round.

We first evaluate FedBeat against the baseline methods at different noise levels under the IID setting. The results in Tables 1 and 2 from both datasets indicate that our method significantly surpasses the performance of the baselines. Specifically, on CIFAR-10, FedBeat outperforms the baseline methods by at least $8.48\%$ in the IDN-$30\%$ scenario and the performance gap widens to $12.52\%$ in the IDN-$50\%$ scenario. Similar observations can be made from the results on the SVHN dataset. Note that BLTM-local exhibits the worst test accuracy performance under various noise levels. This highlights that directly applying IDN methods developed for centralized learning without suitable cooperation may not yield the desired results in the FL setting. The lackluster performance of FedCorr and BLTM-local emphasizes the effectiveness of our proposed method in FL.

We also perform experiments under the non-IID setting. The results in Tables 1 and 2 show FedBeat, consistently outperforms the baseline methods. Except for BLTM-local, all methods demonstrate a drop in test accuracy on both datasets compared to their performance under IID settings. The stability of BLTM-local’s performance can be attributed to its design, as it does not involve any inter-client cooperation.

Further, we conducted additional non-IID experiments with varying Dirichlet parameters, while keeping the IDN rate fixed at 30%. These results, presented in Table 3, validate that our method consistently performs best under different degrees of non-IID settings. Specifically, FedBeat outperforms all baselines by at least 9.89% with $\alpha_{Dir}=0.5$, and by at least 7.49% with $\alpha_{Dir}=5$.

We compare our method with an ablated variant where the first step uses only the average of local models to extract instances with clean pseudo labels. We set the same extraction threshold $\tau$ and ablated variant to isolate the differences between them. As shown in Table 4, the results (the accuracy of pseudo labeling / the number of extracted data / the final classification accuracy) demonstrate the use of Bayesian model ensembles allows us to filter a similar number of instances with a higher accuracy. This improved accuracy in pseudo-labeling aids the second step in training a more reliable IDNTM estimation network, thereby enhancing the final model’s performance.

We further examine the impact of threshold $\tau$. We present the results (the accuracy of pseudo labeling / the number of extracted data) under different $\tau$ values in Table 5. If $\tau$ is set too high, the requirement for confidence scores to pass the threshold and be included in the dataset becomes stringent. This could result in too few samples being filtered into the dataset, despite a high correct rate of pseudo-labeling. Conversely, setting a low $\tau$ can result in an excessive number of instances with incorrect pseudo-labels being included in the extracted dataset. Both an excessively small dataset (even with a high accuracy) and a low correct rate (even with a large dataset) can impair the learning performance of the transition network. Therefore, we suggest using a moderate threshold $\tau=0.65$, which ensures a sufficient number of accurate extracted data.