Flashback: Understanding and Mitigating Forgetting in Federated Learning

## Abstract

Abstract In Federated Learning (FL), forgetting, or the loss of knowledge across rounds, hampers algorithm convergence, particularly in the presence of severe data heterogeneity among clients.
This study explores the nuances of this issue, emphasizing the critical role of forgetting in FL’s inefficient learning within heterogeneous data contexts. Knowledge loss occurs in both client-local updates and server-side aggregation steps; addressing one without the other fails to mitigate forgetting. We introduce a metric to measure forgetting granularly, ensuring distinct recognition amid new knowledge acquisition.
Leveraging these insights, we propose Flashback, an FL algorithm with a dynamic distillation approach that is used to regularize the local models, and effectively aggregate their knowledge.
Across different benchmarks, Flashback outperforms other methods, mitigates forgetting, and achieves faster round-to-target-accuracy, by converging in 6 to 16 rounds.

## 1 Introduction

Federated Learning (FL) is a distributed learning paradigm that allows training over decentralized private data. These datasets belong to different clients that participate in training a global model.
Federated Averaging (FedAvg) (McMahan et al., ) is a prominent training algorithm that uses a centralized server to orchestrate the process. At every round, the server samples a proportion of the available clients. Starting form the current version of the global model, each sampled client performs $E$ epochs of local training using their private data and sends its updated model to the server. Then, the server aggregates the models by averaging them to obtain the new global model. This process is typically repeated for many rounds until a desired model performance is obtained.

A main challenge in FL is the heterogeneity in distribution between the private datasets, which are unbalanced and non-IID (Kairouz et al., ).
Data heterogeneity causes local model updates to drift – the local optima might not be consistent with the global optima – and can lead to slow convergence of the global model – where more rounds of communication and local computation are needed – or worse, when the desired performance may not be reached.
Addressing data heterogeneity in FL has been the focus of several prior studies. For instance, FedProx (Li et al., ) proposes a proximal term to limit the distance between the global model and the local model updates, mitigating the drift in the local updates. MOON (Li et al., ) mitigates the local drift using a contrastive loss to minimize the distance between the feature representation of the global model and the local model updates while maximizing the distance between the current model updates and the previous model updates. FedDF (Lin et al., ) addresses heterogeneity in local models by using ensemble distillation during the aggregation step (instead of averaging the model updates).

However, we experimentally observe that under severe data heterogeneity, these proposals provide little or even
no advantage over FedAvg. [Figure 1(a)](#S1.F1.sf1) illustrates the test accuracy of FedAvg and other baselines while training a DNN over the CIFAR10 dataset (Krizhevsky, ) (details in [§ 5](#S5)).

This motivates us to better understand how data heterogeneity poses a challenge for FL and devise a new approach of handling non-IID datasets.
We investigate the evolution of the global model accuracy broken down by its per-class accuracy. [Figure 1(b)](#S1.F1.sf2) shows as a heatmap the per-class accuracy for FedAvg; each rectangle represents the accuracy of the global model on a class at a round.
Other baseline methods show similar results.
Our key observation is that there is a significant presence of *forgetting*: i.e., cases where some knowledge obtained by the global model at round $t$ is dropped at round $t+1$, causing the accuracy to decline (as shown by the prominent number of light-shaded rectangles at the right side of darker ones in the figure).

Figure: (a) Global model accuracy of FedAvg and baselines.
Refer to caption: /html/2402.05558/assets/x1.png

A similar phenomenon is known as *catastrophic forgetting* in Continual Learning (CL) literature (Parisi et al., ).
CL addresses the challenge of sequentially training a model on a series of tasks, denoted as $\{T_{1},T_{2},\ldots,T_{n}\}$, without revisiting data from prior tasks. Formally, given a model with parameters $\theta$ and task-specific loss functions $L_{t}(\theta)$ for each task $T_{t}$, the objective in CL is to update $\theta$ such that performance on the current task is optimized without significantly degrading the model’s performance on previously learned tasks. This is non-trivial, as naïve sequential training often leads to catastrophic forgetting, where knowledge from prior tasks is overridden when learning a new task. An inherent assumption in this paradigm is that once the model transitions from task $T_{i}$ to task $T_{i+1}$, data from $T_{i}$ becomes inaccessible, amplifying the importance of knowledge retention strategies (De Lange et al., ).

While the premises and assumptions of FL differ from those of traditional machine learning and continual learning, forgetting remains an issue. This can be viewed as a side effect of data heterogeneity, a commonality FL shares with CL. In FL, the global model evolves based on a fluctuating data distribution. Specifically, in each communication round, a diverse set of sampled clients, each with distinct data distributions, contribute with a model update.
Furthermore, these model updates need to be aggregated to obtain a global model.
This situation presents dual levels of data heterogeneity. Firstly, at the intra-round level, heterogeneity arises from the participation of clients with varied data distributions within the same round. This diversity can inadvertently lead to “forgetting” specific data patterns or insights from certain clients. Secondly, at the inter-round level, the participating clients generally change from one round to the next. As a result, the global model may “forget” or dilute insights gained from clients in previous rounds.

To remedy this issue, we propose Flashback, a FL algorithm that employs a dynamic distillation approach to mitigate the effects of data heterogeneity.
Flashback’s dynamic distillation ensures that the local models learn new knowledge while retaining knowledge from the global model during the client updates by adaptively adjusting the distillation loss. Moreover, during the server update, Flashback uses a very small public dataset as a medium to integrate the knowledge from the local models to the global model using the same dynamic distillation.
Flashback performs these adaptations by estimating the knowledge in each model using label counts as a proxy of the model knowledge.
Overall, Flashback results in a more stable and faster convergence compared to existing methods.

Our contributions are the following:
$\bullet$ We investigate the forgetting problem in FL. We show that under severe data heterogeneity, FL sufferers from forgetting. We dissect how and where forgetting happens ([§ 3](#S3)).
$\bullet$ We propose a new metric for measuring forgetting over the communication rounds([§ 3](#S3)).
$\bullet$ We introduce *Flashback*, a FL algorithm that employs a dynamic distillation during the local updates and the server update ([§ 4](#S4)). By addressing the forgetting issue, Flashback not only mitigates its detrimental effects but also converges to the desired accuracy faster than existing methods ([§ 5](#S5)).

## 2 Background

We consider a standard cross-device FL setup in which there are $N$ clients. Each client $i$ has a unique dataset $D_{i}=\{(x_{j},y_{j})\}_{j=1}^{n_{i}}$ where $x_{j}$ represents the input features and $y_{j}$ is the ground-truth label for $j$-th data point and $n_{i}$ represent the size of the local dataset of client $i$. The goal is to train a single global model that minimizes the objective:

$$ $\min_{w\in\mathbb{R}^{d}}\sum^{N}_{i=1}\frac{|D_{i}|}{|\cup_{i\in[N]}D_{i}|}\left\{L_{i}(w)=\frac{1}{|D_{i}|}\sum_{j=1}^{|D_{i}|}l(w;(x_{j},y_{j}))\right\},$ $$

where $L_{i}(w)$ represents the local loss for client $i$, and $l(w;(x_{j},y_{j}))=\mathcal{L}_{\text{CE}}(F_{w}(x),y)$ is the cross-entropy loss for a single data point,
where $F_{w}$ denotes the model parameterized by learnable weights $w$.

FedAvg provides a structured approach to efficiently address this distributed problem. At each communication round $t$, the server randomly selects $K$ clients from the total available $N$ clients. These clients (denoted with ${\mathbb{S}}_{t}$) receive the previous global model, $w_{t-1}$. Then, they update this model based on their local data using their local loss function $L_{i}$. After updating, each client sends their modified model $w_{k,t}$ back to the server that updates the global model using a weighted average of local models, i.e.,
$w_{t}=\sum_{k\in{\mathbb{S}}_{t}}\frac{|D_{k}|w_{k,t}}{|\cup_{k\in[K]}D_{k}|}.$
To accommodate the intrinsic heterogeneity in client data, various FL algorithms introduce modifications either at the local update level or during the global aggregation. The nuances of these variations are further explored in [§ 6](#S6).

Figure: Figure 2: Local (client) & Global Forgetting in some of the baselines using CIFAR10. The first row represents the global model per-class test accuracy at round $t-1$; then, the rows in the middle are all the clients that participated in round $t$, and finally, in the last row, the global model at the end of round $t$. Local forgetting happens when clients at round $t$ lose the knowledge that the global model had at round $t-1$. The global forgetting happens when the global model at round $t$ loses the knowledge that in the clients’ models at round $t$.
Refer to caption: /html/2402.05558/assets/x3.png

Among these, Knowledge Distillation (KD) is a training method wherein a smaller model, referred to as the student, is trained to reproduce the behavior of a more complex model or ensemble called the teacher. Let $F_{w_{s}}$ denote the student model with weights $w_{s}$ and $F_{w_{t}}$ represent the teacher model with weights $w_{t}$. For a given input $x$, the student aims to minimize the following distillation loss:

$$ $\displaystyle\mathcal{L}_{\text{KD}}((x,y);w_{s},w_{t})=\mathcal{L}_{\text{CE}}(F_{w_{s}}(x),y)(1-\alpha)$ (1) $\displaystyle+\mathcal{L}_{\text{KL}}(F_{w_{t}}(x),F_{w_{s}}(x))\alpha$ $$

Here, $\mathcal{L}_{\text{CE}}$ is the standard cross-entropy loss with true label $y$, and $\mathcal{L}_{\text{KL}}$ represents the Kullback-Leibler (KL) divergence between the teacher’s and the student’s output probabilities. It is defined as $\mathcal{L}_{\text{KL}}({\bm{p}},{\bm{q}})=\sum_{c=1}^{C}{p}^{c}\log\left(\frac{{p}^{c}}{{q}^{c}}\right)$, where $C$ is the number of classes, ${\mathbf{p}}$ is the target output probability vector, and ${\mathbf{q}}$ is the predicted output probability vector.
The hyperparameter $\alpha\in[0,1]$ balances the importance between the learning from the true labels and the teacher’s outputs.

While distillation originally emerged as a method for model compression (Hinton et al., ; Buciluǎ et al., ; Schmidhuber, ), its utility extends to FL. In the federated context, distillation can combat challenges like data heterogeneity  (Lin et al., ; Lee et al., ) and communication efficiency  (Jeong et al., ).
Specifically, the global model can act as a guiding teacher during local updates, directing the training process for each client. Additionally, distillation techniques streamline the aggregation step, assisting in the incorporation of varied knowledge from diverse clients to update the global model. Furthermore, employing distillation for aggregation mitigates model heterogeneity, allowing for the use of different model architectures. Perhaps most interestingly, by using distillation, FL systems can potentially bypass the traditional method of transmitting weight updates. This is accomplished by sending soft labels that encapsulate the essence of local updates–communication becomes more efficient, reducing bandwidth usage.

## 3 Forgetting in FL

We now investigate where forgetting happens and devise a metric to quantify this phenomenon.
Recall that in FL, the models are updated in two distinct phases: 1) during local training – when each client $k$ starts from global model $w_{t-1}$ and locally trains $w_{k,t}$ – and 2) during the aggregation step – when the server combines the client models to update the new global model $w_{t}$.

Intuitively, forgetting in FL is when knowledge contained in the global model will be lost after the completion of communication round $w_{t-1}\rightarrow w_{t}$.
We observe that forgetting may occur in the two phases of FL.
We refer to the former case as local forgetting, where some knowledge in the global model will be lost during the local training $w_{t-1}\rightarrow w_{k,t}$. This is due to optimizing for the clients’ local objectives, which depend on their datasets. Local forgetting is akin to the form of forgetting seen in CL, where tasks change over time (as with clients in FL) and, consequently, the data distribution.
We refer to the latter case as aggregation forgetting, where some knowledge contained in the clients’ model updates will be lost during aggregation $\sum\{w_{k,t}\mid k\in{\mathbb{S}}_{t}\}\rightarrow w_{t}$. This might be due to the coordinate-wise aggregation of weights as opposed to matched averaging in the parameter space of DNNs (Wang et al., ).

We illustrate forgetting in [Fig. 2](#S2.F2) based on actual experiments with several baseline methods. The figure shows the per-class accuracy of $w_{t-1}$, all local models $w_{k,t}$, and the new global model $w_{t}$. The local forgetting is evident in the drop in accuracy (lighter shade of blue) of the local models $w_{k,t}$ compared to the global model $w_{t-1}$. The aggregation forgetting is evident in the drop in accuracy of the global model $w_{t}$ compared to the local models $w_{k,t}$. The figure also previews a result of our method, Flashback, which shows a significant mitigation of forgetting. Local and aggregation forgetting lead to the main forgetting problem in FL, which we term as round forgetting, affecting $w_{t-1}\rightarrow w_{t}$.

In CL, forgetting is often quantified using Backward Transfer (BwT) (Chaudhry et al., ).
Lee et al. () adapted this metric for FL as:

$$ $\textstyle\mathcal{F}=\frac{1}{C}\sum_{c=1}^{C}\operatorname*{arg\,max}_{t\in{1,T-1}}(A^{c}_{t}-A^{c}_{T}),$ (2) $$

where $C$ is the number of classes and $A^{c}_{t}$ is the global model accuracy on class $c$ at round $t$.

However, $\mathcal{F}$ is a coarse-grain score that evaluates forgetting in aggregate across all rounds.
We seek a finer-grain metric that measures forgetting round-by-round. Furthermore, we wish to account for knowledge replacement scenarios, such as when a decline in accuracy for one class might be accompanied by an increase in another, essentially masking the negative impact of forgetting in aggregate measures.
Thus, for our evaluation results ([§ 5](#S5)), we propose to measure round forgetting by focusing only on drops in accuracy using the following metric:

$$ $\textstyle\mathcal{F}_{t}=-\frac{1}{C}\sum_{c=1}^{C}\min(0,(A^{c}_{t}-A^{c}_{t-1}))$ $$

where $t>1$ is the round at which forgetting is measured.

Our metric accounts for the pitfalls of the previously proposed forgetting metric. It discounts knowledge replacement scenarios that can happen between rounds by only focusing on the negative changes in accuracy. Furthermore, it provides a granular view of forgetting, because it measures forgetting within the rounds (whereas [Eq. 2](#S3.E2) measures how much the global model has forgotten by the end of training).

### 3.1 Forgetting in Federated Learning and Continual Learning

Forgetting is a prominent problem in CL, where tasks change over time, and consequently, the data distribution, places models at risk of overriding previously learned knowledge. Looking at FL with the same perspective, we have the intra-round and inter-round data distribution heterogeneity.
Intra-round clients with different data distributions participate, by updating the current global model weights $w_{t-1}$. The server obtains a set of $\{w_{k,t}\mid k\in S_{t}\}$ where $S_{t}$ is the set of clients participating in round $t$. The goal is after the aggregation step, a new global model $w_{t}$ is obtained containing all the knowledge that was present in $\{w_{k,t}\mid k\in S_{t}\}$.
Inter-round the global model $w_{t}$ has learned knowledge that over the prior rounds $1\ldots t$. At each round, different sets of clients participate. The goal is to carry this knowledge to the next round $t+1$ even though the new set $S_{t+1}$ most likely is different with respect to the data distribution than the previous set $S_{t}$.
Presenting distinct forgetting challenges in FL compared to CL.

## 4 Forgetting-Robust Federated Learning

Our key idea to mitigate forgetting is to leverage a dynamic form of knowledge distillation, which is fine-tuned in response to the evolving knowledge captured by the different models in the training process.
During local training, distillation ensures that each local model learns from the client’s local dataset while retaining knowledge from the current global model.
On the server side, after the clients’ updates, Flashback begins by aggregating the locally updated models—much in the vein of FedAvg. Then, Flashback distills the knowledge of the freshly updated global model using our dynamic distillation, learning from both its immediate predecessor—the global model obtained at the previous round—and the ensemble of locally updated models, which all play the role of teachers.
The Flashback algorithm is detailed in [Algorithm 1](#alg1).
The remainder of this section discusses our distillation approach in detail.

Figure: Algorithm 1 Flashback algorithm.

### 4.1 Dynamic Distillation

As established in [§ 3](#S3), a client’s local model can forget and override model knowledge with what is present in its private data.
Moreover, even the global model can be imperfect for two reasons:
i) As we established before, the global model is susceptible to forgetting in the aggregation step.
ii) Assuming no forgetting in the aggregation step, the knowledge contained in the clients who participated so far might not represent all the available knowledge, especially in the early rounds.
Overall, both local models and the global model can be imperfect. Therefore, the logits of all the different classes cannot be treated equally.

We propose to use the label count as an approximation of the knowledge within a model.
Here, the label count refers to the occurrences of each class in the training data that the model saw during training.

In machine learning, a model’s knowledge is fundamentally tied to the data it has been exposed to. If certain classes have higher representation (or label counts) in the training data, it’s intuitive that the model would have more opportunities to learn the distinguishing features of such classes. Conversely, underrepresented classes might not offer the model sufficient exposure to learn their nuances effectively.
Our experimental results suggest that per-class performance on the test set correlates highly with the label counts in the training data. In scenarios where certain classes were more abundant, the model demonstrated higher proficiency in predicting those classes on the test set. As an example, [Fig. 5](#S4.F5) in the appendix, illustrates for a randomly-chosen client that the client’s model performance on the test set well reflects the label count distribution of its private data. From this and many similar observations, we conclude that the label count can be indicative of a model’s knowledge.

In standard knowledge distillation ([Eq. 1](#S2.E1)), all logits are treated equally since it is assumed that the teacher model has been trained on a balanced dataset. Owing to the heterogeneity of data distribution in local datasets, this assumption does not hold in FL.
As a result, we cannot directly treat the current global model nor the local model updates as equally reliable teachers across all classes.
Instead, we propose weighting the logits using the label count as an approximation of the per-class knowledge within a model.

We now revisit the distillation loss in [Eq. 1](#S2.E1) and transform the scalar $\alpha$ to a matrix form that is automatically tuned according to the label count of both the student and the teachers and used directly within the KL divergence loss.
Namely, the dynamic $\bm{\alpha}$ parameter (defined below) will change during the training as the label counts change. Flashback maintains the global model counts over the rounds; this mechanism is detailed in the next section.

We consider a single student model $F_{w_{s}}$ with weights $w_{s}$ and a set ${\mathbb{T}}$ of $K$ teacher models; the $i$-th teacher model is denoted as $F_{w_{i}}$ with weights $w_{i}$.
Let $\bm{\nu}\in\mathbb{R}^{C}$ be the relative label count vector of the student model, where $\nu^{c}$ is the relative occurrences of class $c$ in the dataset.
Similarly, let $\bm{\mu}_{i}\in\mathbb{R}^{C}$ be the relative label count vector of the $i$-th teacher model.

The dynamic $\bm{\alpha}\in[0,1]^{K\times C}$ is defined as $[\bm{\alpha}_{1}^{\intercal},\ldots,\bm{\alpha}_{K}^{\intercal}]$, with $\alpha^{c}_{i}=\frac{\mu^{c}_{i}}{\nu^{c}+\sum_{k}\mu^{c}_{k}}$.

Then, we embed $\bm{\alpha}$ directly in the KL divergence loss ($\mathcal{L}_{\text{KL}}$ in [Eq. 1](#S2.E1)) as follows:

$$ $\mathcal{L}_{\text{dKL}}({\bm{p}},{\bm{q}};\bm{\alpha}_{i})=\sum_{c=1}^{C}\alpha_{i}^{c}\cdot{p}^{c}\log\left(\frac{{p}^{c}}{{q}^{c}}\right)$ $$

Similar to standard distillation, to account for the student model knowledge with respect to the ground-truth class $y$, we define $\alpha^{c}_{s}=\frac{\nu^{c}}{\nu^{c}+\sum_{k}\mu^{c}_{k}}$. Thus, $\alpha^{c}_{s}+\sum_{k=1}^{K}\alpha^{c}_{k}=1$ for all classes $c\in[C]$.

Finally, the dynamic knowledge distillation loss ($\mathcal{L}_{\text{dKD}}$) is:

$$ $\displaystyle\mathcal{L}_{\text{dKD}}((x,y);w_{s},{\mathbb{T}},\bm{\alpha})=\alpha^{y}_{s}\mathcal{L}_{\text{CE}}(F_{w_{s}}(x),y)$ (3) $\displaystyle\;+\sum_{{w_{i}}\in{\mathbb{T}}}\mathcal{L}_{\text{dKL}}(F_{w_{i}}(x),F_{w_{s}}(x);\bm{\alpha}_{i})$ $$

The dynamic $\bm{\alpha}$ will weigh the divergence between the logits of different classes, making the student model focus more on learning from the teacher’s strengths while being cautious of its weaknesses.
This is of great importance in FL because of the data heterogeneity problem.
For instance, in the initial training rounds, the global model may not encounter certain classes. Our distillation approach assigns a zero weight to the divergence of these classes, shielding the client model from adopting unreliable knowledge from the global model. Similarly, if a client possesses significantly larger data for a specific class compared to what the global model has encountered, the weight assigned to that class’s divergence will be small. This implies that the client model’s remains more grounded in classes where it has more comprehensive data.

A property of our distillation is that it will ignore the global model as a teacher in the first communication round. Since the initial global model label counts are all zeros, [Eq. 3](#S4.E3) reduces to just the cross-entropy: $\mathcal{L}_{\text{dKD}}=\mathcal{L}_{\text{CE}}$.

Figure: (a) CIFAR10.
Refer to caption: /html/2402.05558/assets/x4.png

Figure: (a) CIFAR10.
Refer to caption: /html/2402.05558/assets/x8.png

### 4.2 Estimating the Global Model Knowledge

Note that to apply the dynamic distillation loss [Eq. 3](#S4.E3), we require to obtain the student’s and teachers’ label count vectors.
While the label count of local models can be easily obtained (from the class frequency of local datasets), the label count of the global model is not readily available.
We construct $\bm{\pi}$, the global model’s relative label count, as follows.
Let $r_{k}$ denote the number of rounds that client $k$ has participated in.
For every client $k$ that participates at round $t$, Flashback adds a fraction $\gamma\in(0,1]$ of client $k$ label count ($\bm{\mu}_{k}$) to $\bm{\pi}$, unless $\gamma r_{k}>1$, in which case $\bm{\pi}$ is not updated based on $k$’s label count. The latter case means that client $k$ has participated enough times that its label count is fully accounted for in $\bm{\pi}$.

Intuitively, the parameter $\gamma$ indicates the rate at which we rely on the global model. When $\gamma$ is set to 1, it implies complete trust in the global model’s ability to incorporate the clients’ knowledge after just one round of participation. However, expecting such immediate and full assimilation is unrealistic, so we typically set $\gamma<1$.
The gradual build-up of the global label count plays a vital role in maintaining a balanced distillation in [Eq. 3](#S4.E3) during local updates. This progressive approach mirrors our growing trust in the global model’s capabilities. It prevents the risk of assigning excessively high weights too soon, which could otherwise hurt the learning process.

### 4.3 Label Count Motivation

[Fig. 5](#S4.F5) shows, on the left, the per-class accuracy of a randomly-chosen example client from a FedAvg training experiment. On the right, the figure shows the corresponding label count at the client. This example suggests that label count can be representative of the model performance.

Figure: Figure 5: (left) Per-class accuracy of a client model on all the rounds where it participated. (right) Data distribution of that client.
Refer to caption: /html/2402.05558/assets/x11.png

## 5 Experiments & Results

We outline and analyze our experimental findings to investigate whether mitigating forgetting successfully addresses the issues of slow and unstable convergence observed in the initial problem laid out in [§ 1](#S1). The experimental results stem from three settings: CIFAR10 and CINIC10, where heterogeneous data partitions are created using Dirichlet distribution with $\beta=0.1$ and FEMNIST with 3,432 clients, following the natural heterogeneity of the dataset.
Furthermore, we do an ablation study on the different components of the algorithm.
We use the same neural network architecture that is used in  Lee et al. (); McMahan et al. (), which is a 2-layer Convolutional Neural Network (CNN).
Summaries of the datasets, partitions, and more details on the experimental setup, as well as additional results, are reported in [Appendix A](#A1).

We compare Flashback against several baseline methods, namely:
1) FedAvg (McMahan et al., ),
2) FedDF (Lin et al., ),
3) FedNTD (Lee et al., ),
4) FedProx (Li et al., ),
5) FedReg (Xu et al., ),
6) MOON (Li et al., ).
It is noteworthy that both FedNTD and FedReg target forgetting in FL (discussed further in [§ 6](#S6)).
Flashback server distillation is performed until early stopping is triggered on the public validation set (details in [Appendix A](#A1)). Moreover, Flashback only introduces one additional hyperparameters $\gamma$, which represents how fast trust is built in the global model. We analyze the effect of $\gamma$ later in the section.
We start by evaluating Flashback performance by showing round-to-accuracy, round forgetting, and the local-global loss over the rounds.

Improved round-to-accuracy.
We evaluate the learning efficiency of Flashback and other baselines by showing the accuracy over rounds in [Fig. 3](#S4.F3). Flashback consistently shows a faster convergence to a high accuracy. Furthermore, we show the number of rounds it takes to reach a target accuracy and fractions of that target accuracy in [Table 1](#S5.T1); we include the result of a central training on the public dataset. Flashback shows a much faster convergence than other baselines.
This indicates that addressing forgetting on the clients’ local update and at the aggregation step does provide training stability and indeed a faster convergence.

**Table 1: Number of rounds to reach accuracy $A_{x}=A\cdot x$ where $A$ is the target accuracy and $x$ is a fraction of it.**
|  | CIFAR10, $A=48.2$% | CINIC10, $A=43.5$% | FEMNIST, $A=69.5$% |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
|  | $A_{0.5}$ | $A_{0.75}$ | $A_{0.95}$ | $A_{0.5}$ | $A_{0.75}$ | $A_{0.95}$ | $A_{0.5}$ | $A_{0.75}$ | $A_{0.95}$ |
| FedAvg | 12 | 82 | - | 13 | - | - | 49 | 75 | 138 |
| FedDF | 7 | 40 | 112 | 2 | 30 | - | - | - | - |
| FedNTD | 12 | 41 | 112 | 13 | 46 | - | - | - | - |
| FedProx | 35 | 93 | - | 13 | - | - | 142 | - | - |
| FedReg | 35 | 108 | - | 16 | - | - | - | - | - |
| MOON | 82 | - | - | 124 | - | - | - | - | - |
| Flashback | 2 | 4 | 10 | 4 | 5 | 6 | 3 | 5 | 16 |

Less round forgetting.
We show the empirical cumulative distribution function (ECDF) of the round forgetting in [Fig. 6](#S5.F6). We see that Flashback successfully reduces round forgetting. Also, FedNTD has less round forgetting than the remaining baselines. In the appendix, we show the round forgetting over the rounds ([Fig. 12](#S5.F12)).

Figure: (a) CIFAR10.
Refer to caption: /html/2402.05558/assets/x12.png

Minimizing local models divergence.
To further understand the effect of Flashback on the training behavior, we show the transition of the mean loss of the local models to the loss of the global model over the rounds in [Fig. 4](#S4.F4). This gives us an insight into the effect of the regularization made by our dynamic distillation.
We see that the mean loss of the local models of the other baselines always spikes, signifying a divergence of these models from the global training objective, while Flashback has a much more stable loss. This shows that [Eq. 3](#S4.E3) regularizes the local models well so that they do not diverge too much from the global training objective.

In the rest of the section, we delve deeper into Flashback to understand its behavior and validate its performance gains.

Dissecting the distillation.
To show the importance of performing dynamic distillation during clients’ updates and at the server’s aggregation step, we conduct an experiment where we run Flashback with local distillation only and with server distillation only (c.f. [Fig. 7](#S5.F7)). We observe that doing dynamic distillation at either side of the algorithm – client update and aggregation step – doesn’t address forgetting or gets a similar performance to Flashback.

Figure: Figure 7: Performing distillation at only one side of the algorithm (client & server) on CIFAR10.
Refer to caption: /html/2402.05558/assets/x15.png

To validate the *importance of dynamic distillation at client- and server sides* towards Flashback’s performance gains, we create a baseline where we replace the local distillation loss with not-true distillation (NTD) loss Lee et al. ().
From [Fig. 8](#S5.F8), we observe that this baseline doesn’t perform as well as Flashback, and performs similarly to FedNTD.
This indicates that performing the distillation at the server needs well-regularized local models (teachers), which is further supported by our previous experiment contrasting Flashback to single-side distillation (c.f. [Fig. 7](#S5.F7)).

Figure: Figure 8: Using NTD loss instead of Flashback’s dynamic distillation at the local update on CIFAR10.
Refer to caption: /html/2402.05558/assets/x18.png

Effect of public dataset size.
Flashback requires the availability of a public labeled dataset. This may be a limiting assumption in some cases. To study this limitation, we explore a few scenarios for the size of the public dataset in [Fig. 9](#S5.F9):
1) 9000 samples ($15$% of CIFAR10),
2) 1125 samples ($1.88$% of CIFAR10),
3) 1283 samples that have unbalanced class distribution ($2.14$% of CIFAR10),
4) 450 samples ($0.75$% of CIFAR10).
For all of these scenarios, we train a model centrally on the public dataset.
We find that Flashback can benefit from a large balanced public dataset. Most importantly, Flashback can work well with a small public dataset (1125 samples is the default in all experiments). Furthermore, even if the public dataset has a class imbalance Flashback still performs relatively well. In all of the cases, Flashback always outperforms central training on the public dataset.
Overall, Flashback requires the availability of a public dataset, however, it does not require a huge amount of data or hard requirements for the class distribution to be very balanced.

Figure: Figure 9: Flashback using different public datasets, and the results of central training on public datasets of similar sizes.
Refer to caption: /html/2402.05558/assets/x21.png

Training on the public dataset.
We also experiment to answer the following question: *does the performance improvement of Flashback come from the fact that we train the global model on a public labeled dataset?* To answer this question we create a naive baseline, where we extend FedAvg to fine-tune the global model after the aggregation step at every communication round. From [Fig. 10](#S5.F10), we see that FedAvg with fine-tuning quickly reaches a stale model and eventually collapses. We believe this collapse happens due to the local models diverging too much after the local update such that aggregating those models fails. This is evident by the spike of the local model’s loss.

Figure: Figure 10: Comparing Flashback to FedAvg with fine-tuning: (left) test accuracy over 3 runs on CIFAR10; (right) the transition of the local models’ loss to the global model loss over the rounds.
Refer to caption: /html/2402.05558/assets/x23.png

The importance of $\gamma$.
As mentioned, Flashback has a single hyperparameter $\gamma$, which dictates how fast models will trust the global model as a competent teacher. We explore the effect of this hyperparameter in [Fig. 11](#S5.F11). We find that setting this parameter to a larger value leads the learning process to get to a stale solution quickly. This is intuitive since large $\gamma$ leads the global model label count to grow faster, therefore, this dominates the loss term in [Eq. 3](#S4.E3) during both the client- and the server update.
Smaller $\gamma$ gives the best performance because it gives the local models time to learn from their own private data by having a small weight to the distillation term in [Eq. 3](#S4.E3). However, too small value for $\gamma$ such as $0.001$ slows the training process, since during the local training the distillation term in [Eq. 3](#S4.E3) will be very small in the early rounds.

Figure: Figure 11: Varying the trust of the global model parameter $\gamma$.
Refer to caption: /html/2402.05558/assets/x25.png

### 5.1 Additional Results

In [Fig. 12](#S5.F12), we show the round forgetting score computed over the rounds. We see that the baselines have very flaunting round forgetting score.

In [Fig. 13](#S5.F13), we show the transition of the average loss of the local models to the global model loss on the test set. Interestingly, we can see that even though performing local distillation only doesn’t have the same performance as Flashback, it does mitigate the local forgetting. That is, we do not see a spike in the loss after the clients perform their local update.

In [Fig. 14](#S5.F14), we see the ECDF of the round forgetting. The Flashback variant with NTD even shows worse forgetting than Flashback and FedNTD. Further showing that just performing distillation at both sides doesn’t address the forgetting problem.

Figure: (a) CIFAR10.
Refer to caption: /html/2402.05558/assets/x27.png

Figure: Figure 13: Performing distillation at only one side of the algorithm (client & server) on CIFAR10.
Refer to caption: /html/2402.05558/assets/x30.png

Figure: Figure 14: Using Not-True Distillation instead of Flashback’s dynamic distillation at the local update on CIFAR10.
Refer to caption: /html/2402.05558/assets/x32.png

## 6 Related Work

Federated learning. FL is commonly viewed as a ML paradigm wherein a server distributes the training process on a set of decentralized participants that train a shared global model using local datasets that are never shared (Konečnỳ et al., ; Shokri & Shmatikov, ; Konečnỳ et al., ; Konečnỳ, ; Li et al., ; McMahan et al., ; Kairouz et al., ). FL has been used to enhance prediction quality for virtual keyboards among other applications (Bonawitz et al., ; Yang et al., ). A number of FL frameworks have facilitated research in this area (Caldas et al., ; OpenMined, ; tensorflow.org, ; Abdelmoniem et al., ).

Heterogeneity in FL. A key challenge in FL systems is uncertainties stemming from learner, system, and data heterogeneity. The non-IID distributions of Learners’ data can significantly slow down convergence (McMahan et al., ; Kairouz et al., ) and several algorithms are proposed as means of mitigation (Fourati et al., ; Wang et al., ; Karimireddy et al., ; Li et al., , ).

In this section, we review works that targeted the issue of forgetting in FL.
Forgetting in FL is an under-studied area that poses significant challenges, leading to slow model convergence and loss of crucial knowledge acquired during the learning process (Chaudhry et al., ; Dupuy et al., ). There have been some notable attempts to mitigate the impact of forgetting on the learning process (Lee et al., ; Xu et al., ).

FedReg (Xu et al., ) addresses the issue of slow convergence in FL, asserting it to be a result of forgetting at the local update phase. They demonstrate this by comparing the loss of the global model $w_{t-1}$ on specific client data points with the averaged loss of updated clients’ models $\{w_{t,k}\mid k\in{\mathbb{S}}_{t}\}$ on the same data points, highlighting a significant increase in the average loss, indicative of forgetting.
However, in our work, we propose a systematic way of measuring forgetting using a metric designed to capture it. Furthermore, we show that forgetting doesn’t only occur in the local update, but it also happens at the aggregation step ([§§ 3](#S3) and [2](#S2.F2)).
FedReg proposes to generate fake data that carries the previously attained knowledge. During the local update, Fast Gradient Sign Method (Goodfellow et al., ) is used to generate these data using the global model $w_{t-1}$ and the client data. Then, the loss of the generated data is used to regularize the local update.
While FedReg employs regularization using synthetic data during local updates, our work, Flashback, leverages dynamic distillation to ensure knowledge retention at both local updates and aggregation steps.

FedNTD (Lee et al., ) makes a connection between CL and FL, suggesting that forgetting happens in FL as well. Similarly to FedReg, their analysis shows that forgetting happens at the local update, where global knowledge that lies outside of the local distribution of the client is susceptible to forgetting.
To address this, they propose to use a new variant of distillation [Eq. 1](#S2.E1) named Not-True Distillation (NTD), that masks the ground-truth class logits in the KL divergence as $\mathcal{L}_{\text{KL}}({\bm{p}},{\bm{q}})=\sum_{i=c,c\neq y}^{C}{p}^{c}\log(\frac{{p}^{c}}{{q}^{c}})$,
where $y$ is the ground-truth class. NTD is used at the local update, while all the other steps in the algorithm remain the same as FedAvg. FedNTD aims to preserve global knowledge during the local update.

Both FedReg and FedNTD diagnose the issue of forgetting primarily within the realm of local updates, asserting that this stage risks losing valuable global knowledge. Consequently, both works present innovative solutions specifically tailored to counteract this local update forgetting. However, their perspective overlooks a pivotal aspect of the forgetting problem: the occurrence of forgetting during the aggregation step. As we delve into in [§ 3](#S3), this oversight in recognizing and addressing forgetting during aggregation has repercussions on the later local updates. In contrast, Flashback takes a holistic approach, targeting forgetting comprehensively across both the local updates and the aggregation phase, leading to faster convergence.

## 7 Conclusion

We explored the phenomenon of forgetting in FL. Our investigation revealed that forgetting occurs during both local and global update phases of FL algorithms.
We presented Flashback, a novel FL algorithm explicitly designed to counteract forgetting by employing dynamic knowledge distillation. Our approach leverages data label counts as a proxy for knowledge, ensuring a more targeted and effective forgetting mitigation.
Our empirical results showed Flashback’s efficacy in mitigating round forgetting, thereby supporting the hypothesis that the observed slow and unstable convergence in FL algorithms is closely linked to forgetting. This result underlines the importance of addressing forgetting, paving the way for the advancement of more robust and efficient FL algorithms.

## Appendix A Experiments Details

### A.1 Datasets

In this section, we provide an overview of the datasets used, the data split, and the specific experimental setups. For each dataset, we perform two sets of experiments to analyze the effects of data heterogeneity on the algorithms’ performance. The datasets used are CIFAR10, CINIC10, and FEMNIST.

CIFAR10  (Krizhevsky, ). A famous vision dataset that includes 50k training images and 10k testing images. We emulate a realistic, heterogeneous data distribution by using a Dirichlet distribution with $\beta=0.1$. A $\beta$ value of $0.1$ is chosen to simulate a more heterogeneous, and challenging data distribution. A $2.5$% random sample of the training set creates a public dataset, further divided into training and validation sets. This yields a very small public training dataset with size of $1.88$%. The remaining $97.5$% is distributed among 100 clients, with each client’s data being split into training (90%) and validation (10%) subsets.

CINIC10 (Darlow et al., ). A drop-in replacement of CIFAR10, this dataset comprises 90k training, 90k validation, and 90k test images. We merge the training and validation sets and adopt a similar approach as with CIFAR10, taking out $2.5$% of the data to be the public dataset, similar to the CIFAR10 case this $2.5$% is further divided into training set and validation set. Then employing Dirichlet distribution with $\beta$ value of 0.1 to split the $97.5$% remaining data into 200 clients, with each client’s data further divided into training (90%) and validation (10%) sets.

Figure: (a) CIFAR10, Dir$(\beta=0.1)$
Refer to caption: /html/2402.05558/assets/x34.png

FEMNIST (Caldas et al., ). This federated learning dataset is based on extended MNIST with natural heterogeneity, where each writer is considered a client. From the 3597 total writers, those with less than 50 samples are excluded. We randomly selected 150 writers to form a public dataset. The remaining 3432 writers’ data is divided into train (approx. 70%), validation (approx. 15%), and test (approx. 15%) sets. The collective test sets from all writers form the overall test set. At every round, 32 clients are randomly selected for participation.

For CIFAR10 and CINIC10, we chose $\beta$ values of 0.1 and client participation value of 10. While for FEMNIST we have 3432 clients (writers) with client participation vale of 32
In all cases, the training data distribution among clients is illustrated in [Fig. 15](#A1.F15).

### A.2 Baselines & Hyperparameters

We evaluate the following algorithms as baselines:
1. FedAvg (McMahan et al., );
2. FedDF (Lin et al., );
3. FedNTD (Lee et al., );
4. FedProx (Li et al., );
5. FedReg (Xu et al., ); and
6. MOON (Li et al., )
.
Both FedNTD and FedReg target forgetting in FL (discussed in [§ 6](#S6)). We use the same neural network architecture that is used in  Lee et al. (); McMahan et al. (), which is a 2-layer CNN. Note that for MOON  (Li et al., ) we add an additional layers to the model as described in their source code for the projection head.
Moreover, for the optimizer, learning rate, and model we follow Lee et al. (); McMahan et al. (), and when a baseline has different hyperparameters we use their proposed values. For example, in FedDF the number of local epochs is set to 40, while in the other baselines and Flashback, it is set to 5 epochs. As for Flashback hyperparameters, during the server distillation, we train until early stopping gets triggered using the validation set; we set the label count fraction $\gamma=0.025$ for CIFAR10, i.e., we add $2.5$% of the client label count each time it participates, while we set $\gamma=0.1$ for CINIC10 and FEMNIST. As for distillation-specific hyperparameters, we have one fewer hyperparameter since $\bm{\alpha}$ is computed automatically, and for temperature, we use the standard $T=3$.

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 3
presentation: 3
questions: 1.	Federated Continual Learning (FCL) also focuses on the catastrophic forgetting problem. What’s the difference between the problem of this paper and that of FCL?
soundness: 2
strengths: 1.	The NonIID problem in FL is important and using distillation to solve this issue achieves promising results.
2.	The writing is easy to follow.
summary: This paper formulates the NonIID problem in FL as the catastrophic forgetting from both local and global perspective. To solve this problem, they propose using knowledge distillation in both local and global sides to maintain knowledge to mitigate forgetting.
weaknesses: The major concern is that the experiments are too limited: 
1.	All baselines are too old. The newest baseline FedReg which is proposed in 2022. In fact, many recent baselines for soving the NonIID problem are proposed in 2023 and 2024.
2.	The evaluations should be conducted on different NonIID settings. This paper only adopt a fixed NonIID setting across all datasets, which is not sufficient to demonstrate the effectiveness of the proposed method.
3.	Larger models such as ResNet18 can be included. Merely 2-layer CNN may not be enough.
4.	Varying fraction of one-round selected clients should be evaluated. This factor is strongly related to the global knowledge forgetting. 
5.	Current datasets only cover 10-classes settings, which may be limited to make a comprehensive comparison because $\alpha$ proposed by the method is strongly related to the number of classes. Cifar100 is recommended to add to the experiments.

---

confidence: 5
contribution: 1
initial_rating: 3
presentation: 2
questions: What are the differences between forgetting in federated learning and forgetting in continual learning?

What are the differences between the federated continual learning scenario and the federated learning scenario in this paper?

Since the paper states that forgetting occurs in both local updating and global aggregation phases, shouldn't there be more detailed forgetting metrics (e.g., phased forgetting metrics)?
soundness: 2
strengths: This paper deals with data heterogeneity in federated learning from the perspective of forgetting which is very novel and gives a comprehensive framework from observation, derivation, solution design and experimental proof.

It is well written and easy to understand.
summary: This work focuses on addressing the problem of data heterogeneity in federated learning. This work analyzes the mechanism of data heterogeneity in multiple rounds of iterations of federated learning from the novel perspective of forgetting. In addition, this paper proposes a new metric for measuring the degree of forgetting between training rounds. The major contribution of this paper is the proposed Flashback algorithm, which utilizes knowledge distillation to mitigate forgetting during the local update and global update phases. In conclusion, this paper presents a novel perspective to study data heterogeneity in federated learning and proposes methodological solutions, and the article ideas are easy to understand.
weaknesses: 1 The researched work related to data heterogeneity in federated learning is insufficient, especially for personalized federated learning and clustered federated learning.

2 The elaboration of the concept of forgetting is not detailed enough, lacking comparisons between different categories of continual learning (class-CL, task-CL, domain-CL), and lacks a detailed elaboration of forgetting mechanisms (e.g., weight drift, activation drift, inter-task confusion, and task-recency bias).

3 The forgetting metric between training rounds proposed in this paper is also computed based on accuracy, which is not different enough from the forgetting rate metric in continual learning to be considered as a contribution point.

4 The idea of using knowledge distillation in Flashback has been widely used, and the approach in this paper lacks innovation.

5 Flashback's dependence on labeled datasets on the server side is the major limitation and is not feasible in real scenarios, perhaps try unlabeled datasets or data-free distillation.

6 The experiments are less persuasive:

The datasets used (CIFAR10, CINIC10, FEMNIST) are too simple, it would be more persuasive to use more complex datasets such as CIFAR100 or TinyImageNet.

Comparison methodology lacks the most recent work (including 2023 and 2024) and work related to personalized federated learning.

7 The font sizes in images is inconsistent, e.g., Figures 4 and 5.

---

confidence: 5
contribution: 2
initial_rating: 5
presentation: 2
questions: In addition to the weaknesses, one more question is as below.

1. In the experiments, do all baselines use the public dataset on the server?
soundness: 2
strengths: 1. The authors provide a detailed analysis and illustration of the key issues - where forgetting occurs in non.iid FL, which is an important problem in FL.

2. The proposed Flashback method is simple and easy to understand, and it outperforms a variety of FL baselines

3. This paper introduces fine-grained evaluation metrics for forgetting in FL.
summary: his paper analyzes the problem of forgetting in a non.iid federated learning environment, and concludes that the forgetting occurs in local updating and aggregation processes. For this, the authors propose metrics to measure the degrees of forgetting in the training process. In addition, the authors propose Flashback, a knowledge distillation (KD)-based model to mitigate forgetting. Flashback allows both clients and the server to perform dynamic KD according to the relative label count, and achieve fast convergence against baselines.
weaknesses: 1. This paper proposes a fine-grained metric for assessing forgetting, but this does not seem to be reflected in Flashback. The authors should emphasize the connection between this metric and Flashback.

2. I have concerns about the use of the public dataset on the server. If the distribution of that dataset is similar to the distribution of data on each client, does this mean that there is already a leakage problem in that environment? Also, I'm curious what happens if a different dataset is used than the training dataset (e.g., cifar10 as the public dataset and cinic10 as the training data).

3. The authors should reorganize Sections 3, 4 for clearer understanding, and placing the algorithm in 4.1 would make it easier to understand.

4. Flashback involves a variety of additional parameters (e.g., label counts for teacher and student models, dependency on the global model, etc.), do these parameters make the algorithm less robust to some extent? This paper would benefit from an analysis of this problem.

5. Missing some related works, e.g. test-time FL [1-3].

[1] L Jiang, T Lin, Test-Time Robust Personalization for Federated Learning, ICLR 2023
[2] Y Tan, et al., ​​Is heterogeneity notorious? taming heterogeneity to handle test-time shift in federated learning, NeurIPS 2023
[3] W Bao, et al., Adaptive Test-Time Personalization for Federated Learning. NeurIPS 2023

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 2
questions: See above.
soundness: 2
strengths: 1. The NonIID problem in FL is important and using distillation to solve this issue achieves promising results.

2. The writing is easy to follow.
summary: This paper proposes a new distillation method to solve the NonIID problem in FL. Specifically, they consider the NonIID problem will cause the forgetting of model knowledge in both local update and global aggregation processes. Then, they apply the distillation over both local and global processes to mitigate the forggeting.
weaknesses: 1. The novelty is limited. In fact, there have been massive methods using knowledge distillation to solve the NonIID problem, e.g., [1,2]. The experiments should also include them for the comprehensiveness of comparison.

[1] DaFKD: Domain-Aware Federated Knowledge Distillation. CVPR 2023.

[2]  Data-free knowledge distillation for heterogeneous federated learning. ICML 2021.

2. It is better to include a figure to illustrate the method framework for ease of understanding.

3. It would be better if the theoretical advantages were provided.

4. The ablation study about the proposed method of local and global distillation should be provided.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: (a) Why are so many values missing in Table 1, particularly for the FEMNIST dataset?

(b) Why is the baseline performance on FEMNIST so low, appearing to learn almost nothing, and why does performance fluctuate significantly across all datasets? This fluctuation suggests that hyperparameters may not be well-tuned for optimal convergence. Providing the training and validation loss trends for each baseline would help demonstrate convergence.
soundness: 3
strengths: (a) The method includes comparisons with a variety of baselines that incorporate regularization and distillation techniques. 

(b) The paper is well-written, with clear presentation and structure that is easy to follow.
summary: The motivation of this paper is to address the problem of forgetting in Federated Learning (FL), which slows down convergence, particularly in settings with high data heterogeneity among clients. The main contribution is the introduction of Flashback, a new FL algorithm that uses dynamic distillation to mitigate forgetting by regularizing both client-local updates and server aggregation.
weaknesses: (a) The investigation of forgetting is less systematic than claimed. While the paper frames forgetting as a key factor in FL underperformance, it lacks detailed analysis and comparisons with baselines from continual learning, where many regularization- [1] based methods effectively mitigate forgetting. Exploring whether these methods can similarly address forgetting in FL would be valuable.

(b) Fairness in performance comparison on public datasets is an issue, as the proposed method uses a portion of data for validation while other methods do not, making the comparisons unfair. Including simple baselines that utilize similar validation strategies, such as server-side distillation [2] or model selection, would improve fairness.

(c) The motivation and application of the new forgetting metric in FL are not clearly explained, making it difficult to fully understand this contribution’s impact.

[1] Kirkpatrick, J., et al. “Overcoming catastrophic forgetting in neural networks.” Proceedings of the National Academy of Sciences 114.13 (2017): 3521-3526.

[2] Huang, Chun-Yin, et al. "Overcoming Data and Model Heterogeneities in Decentralized Federated Learning via Synthetic Anchors." arXiv preprint arXiv:2405.11525 (2024).