SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Sara Babakniya
Computer Science  
University of Southern California  
Los Angeles, CA  
babakniy@usc.edu
Zalan Fabian
Electrical and Computer Engineering University of Southern California Los Angeles, CA zfabian@usc.edu
Chaoyang He
FedML Sunnyvale, CA ch@fedml.ai
Mahdi Soltanolkotabi
Electrical and Computer Engineering University of Southern California Los Angeles, CA soltanol@usc.edu
Salman Avestimehr
Electrical and Computer Engineering University of Southern California Los Angeles, CA avestime@usc.edu
Deep learning models often suffer from forgetting previously learned information when trained on new data. This problem is exacerbated in federated learning (FL), where the data is distributed and can change independently for each user. Many solutions are proposed to resolve this catastrophic forgetting in a centralized setting. However, they do not apply directly to FL because of its unique complexities, such as privacy concerns and resource limitations. To overcome these challenges, this paper presents a framework for federated class incremental learning that utilizes a generative model to synthesize samples from past distributions. This data can be later exploited alongside the training data to mitigate catastrophic forgetting. To preserve privacy, the generative model is trained on the server using data-free methods at the end of each task without requesting data from clients. Moreover, our solution does not demand the users to store old data or models, which gives them the freedom to join/leave the training at any time. Additionally, we introduce SuperImageNet, a new regrouping of the ImageNet dataset specifically tailored for federated continual learning. We demonstrate significant improvements compared to existing baselines through extensive experiments on multiple datasets.
Federated learning (FL) [40, 29] is a decentralized machine learning technique that enables privacy-preserving collaborative learning. In FL, multiple users (clients) train a common (global) model in coordination with a server without sharing personal data. In recent years, FL has attracted tremendous attention in both research and industry and has been successfully employed in various fields, such as autonomous driving [17], next-word prediction [21], health care [13], and many more.
Despite its popularity, deploying FL in practice requires addressing critical challenges, such as resource limitation and statistical and system heterogeneity [27, 33]. While tackling these challenges is an essential step towards practical and efficient FL, there are still common assumptions in most FL frameworks that are too restrictive in realistic scenarios.
Figure 1: In the real world, users constantly change their interests, observe new data, or lose some of the old ones. As a result, the training dataset is divided into different tasks. For example, here, at  $\text{Task} = 1$ , the clients' datasets dominantly include pictures of animals, and by the end of the training  $\text{Task} = T$ , the trend shifts towards landscapes.
In particular, one of the most common assumptions is that clients' local data distribution is fixed and does not change over time. However, in real-world applications [49], clients' data constantly evolve due to changes in the environment, trends, or new interests. For example, [6] presents the real-world data of an online shop, suggesting interest in items shifts through seasons. Another example arises in healthcare, where a model trained on old diseases should be able to generalize to new diseases [58]. In such scenarios (Figure 1), the model must rapidly adapt to the incoming data while preserving performance on past data distributions to avoid catastrophic forgetting [28, 39].
In the centralized setting, such problems have been explored in continual learning [48, 34] (also called lifelong learning [3] or incremental learning [9, 7] based on the initial settings and assumptions). In recent years, various algorithms have been proposed in Continual Learning (CL) to tackle catastrophic forgetting from different angles and can achieve promising performance in different scenarios.
Despite all the significant progress, most CL methods are not directly applicable to the federated setting due to inherent differences (Table 1) between the two settings. For instance, experience replay [47] is a popular approach, where a portion of past data points is saved to maintain some representation of previous distributions throughout the training. However, deploying experience replay in FL has resource and privacy limitations. It requires clients to store and keep their data, which may increase the memory usage of already resource-limited clients. Furthermore, users may not be able to store data for more than a specific time due to privacy concerns. Finally, depending solely on the clients to preserve the past is not reliable, as clients leaving means losing their data.
Challenge	Limitation
Low memory	Clients cannot store many examples
Clients drop out	Causes loss of information stored in memory
New clients join	New clients only have access to new classes
Privacy	Limits data saving and sharing of the clients
Table 1: Challenges that limit the direct use of continual learning methods in federated settings.
To address the aforementioned problems, we propose MFCL, Mimicking Federated Continual Learning: a privacy-preserving federated continual learning approach without episodic memory. In particular, MFCL is based on training a generative model in the server and sharing it with clients to sample synthetic examples of past data instead of storing the actual data on the client side. The generative model training is data-free in the sense that no form of training data is required from the clients, and only the global model is used in this step. It is specifically crucial because this step does not require powerful clients and does not cause any extra data leakage. Finally, this algorithm has competitive performance; our numerical experiments demonstrate improvement by  $10\% - 20\%$  in average accuracy while reducing the training overhead of the clients.
Moreover, benchmarking federated continual learning in practical scenarios requires a large dataset to split among tasks and clients. However, existing datasets are not sufficiently large, causing most of the existing works in federated continual learning evaluating on a few clients (5 to 20) [45, 24, 52]. To enable more practical evaluations, we release a new regrouping of the ImageNet dataset, SuperImageNet. SuperImageNet enables evaluation with many clients and ensures all clients are assigned sufficient training samples regardless of the total number of tasks and active clients.
We summarize our contributions below:
- We propose a novel framework to tackle the federated class incremental learning problem more efficiently for many users. Our framework specifically targets applications where past data samples on clients are unavailable.
- We point out potential issues with relying on client-side memory for FCL. Furthermore, we propose using a generative model trained by the server in a data-free manner to help overcome catastrophic forgetting while preserving privacy.
- We modify the client-side training of traditional FL techniques in order to mitigate catastrophic forgetting using a generative model.
- We propose a new regrouping of the ImageNet dataset, SuperImageNet, tailored to federated continual learning settings that can be scaled to a large number of clients and tasks.
- We demonstrate the efficacy of our method in more realistic scenarios with a larger number of clients and more challenging datasets such as CIFAR-100 and TinyImageNet.
Continual Learning. Catastrophic forgetting [39] is a fundamental problem in machine learning: when we train a model on new examples, its performance degrades when evaluated on past data. This problem is investigated in continual learning (CL) [59], and the goal is for the model to learn new information while preserving its knowledge of old data. A large body of research has attempted to tackle this problem from different angles, such as adding regularization terms [31, 1, 41], experience replay by storing data in memory [2, 10, 4, 35], training a generative model [56, 53, 32], or architecture parameter isolation [16, 38, 19, 51].
In CL settings, the training data is presented to the learner as a sequence of datasets - commonly known as tasks. In each timestamp, only one dataset (task) is available, and the learner's goal is to perform well on all the current and previous tasks.
Recent work focuses on three main scenarios, namely task-, domain- and class-incremental learning (IL) [54]. In Task-IL, tasks are disjoint, and the output spaces are separated by task IDs provided during training and test time. For Domain-IL, the output space does not change for different tasks, but the task IDs are no longer provided. Finally, in Class-IL, new tasks introduce new classes to the output space, and the number of classes increases incrementally. Here, we work on Class-IL, which is the more challenging and realistic, especially in FL. In most of the FL applications, there is no task ID available, and it is preferred to learn a single model for all the observed data.
Class Incremental Learning. In standard centralized Class-IL, the model is trained on a sequence of non-overlapping  $T$  tasks  $\{\mathcal{T}^{(1)},\mathcal{T}^{(2)},\dots,\mathcal{T}^{(T)}\}$  where the data distribution of task  $t$ ,  $D^t$ , is fixed but unknown in advance, while all the tasks share the same output space  $(\mathcal{V})$ . For task  $t$ ,  $D^t$  consists of  $N^t$  pairs of samples and their labels  $\{(x_i^t,y_i^t)\}_{i = 1}^{N^t}\}$ , where all the newly introduced classes  $(y_{i}^{t})$  belong to  $\mathcal{V}^t (y_i^t\in \{\mathcal{V}^t\})$  and  $\bigcup_{j = 1}^{t - 1}\{\mathcal{V}^j\} \cap \{\mathcal{V}^t\} = \emptyset$ ). Moreover, a shared output space among all tasks means that at the end of task  $t$ , the total number of available classes equals  $q = \sum_{i = 1}^{t}|\mathcal{V}^i|$ .
Federated Continual Learning. In real-life scenarios, users' local data is not static and may evolve. For instance, users' interests may change over time due to seasonal variations, resulting in more examples for a given class. On the other hand, reliability issues or privacy concerns may lead to users losing part of their old data as well. In Federated Continual Learning (FCL), the main focus is to adapt the global model to new data while maintaining the knowledge of the past.
Even though FCL is an important problem, it has only gained attention very recently, and [58] is the first paper on this topic. It focuses on Task-IL, which requires a unique task ID per task during inference. Furthermore, it adapts separate masks per task to improve personalized performance without preserving a common global model. This setting is considerably different than ours as we target Class-IL with a single global model to classify all the classes seen so far. [37] employs server and client-side knowledge distillation using a surrogate dataset. [15] relaxes the problem as clients have access to large memory to save the old examples and share their data, which is different from the standard FL setting. Some works, such as [26, 44, 52], explore the FCL problem in domains other than image classification. [42] has proposed using variational embedding to send data to the server securely and then server-side training to rehearse the previous task for Domain-IL.
This work focuses on Class-IL for supervised image classification without memory replay, similar to [45, 24]. However, [24] allows overlapping classes between tasks and focuses on few-shot learning, which is different from the standard Class-IL. The most related work to ours is [45], where authors propose FedCIL. This work also benefits from generative replay to compensate for the absence of old data and overcome forgetting. In FedCIL, clients train the discriminator and generator locally. Then, the server takes a consolidation step after aggregating the updates. In this step, the server generates synthetic data using all the generative models trained by the clients to consolidate the global model and improve the performance. The main difference between this work and ours is that in our work, the generative model is trained by the server in a data-free manner, which can reduce clients' training time and computation and does not require their private data (detailed comparison in Appendix H).
Data-Free Knowledge Distillation. Knowledge distillation (KD) [25] is a popular method to transfer knowledge from a well-trained teacher model to a (usually) smaller student model. Common KD methods are data-driven, and at least a small portion of training data is required. However, in some cases, training data may not be available during knowledge distillation due to privacy concerns.
To tackle this problem, a new line of work [12, 22] proposes data-free knowledge distillation. In such methods, a generative model is used as a training data substitute. This generative model is trained to generate synthetic images such that the teacher model predicts them as their assigned label (Figure 2). This method has recently become popular in CL [57, 50] as well, mainly due to the fact that it can eliminate the need for memory in preserving knowledge. Data-free KD has been previously used in FL [60] to reduce the effect of data heterogeneity. However, to the best of our knowledge, this is the first work that adapted such a technique in the context of federated continual learning.
Figure 2: Data-Free Knowledge Distillation. The generator receives random noise as input labels and synthesizes images that are labeled correctly by the trained teacher model.
3 Federated Class Incremental Learning with MFCL
In federated Class-IL, a shared model is trained on  $T$  different tasks. However, the distributed and private nature of FL makes it distinct from the centralized version. In FL, users may join, drop out, or change their data independently. Besides, required data or computation power for some centralized algorithms may not be available in FL due to privacy and resource constraints.
To address the aforementioned problems, we propose MFCL, which is less reliant on the client-side memory and computational power. This algorithm includes two essential parts: first, at the end of each task, the server trains a generative model with data-free knowledge distillation methods to learn the representation of the seen classes. Second, clients can reduce catastrophic forgetting by generating synthetic images from the trained generative model obtained from the server side. This way, clients are not required to use their memory for storing old data. Moreover, this technique can address the problem of newly connected clients without past data. Furthermore, since the server trains the generative model training without additional information, this step does not introduce new privacy issues. Finally, MFCL can help mitigate the data heterogeneity problem, as clients can synthesize samples from classes they do not own [60] in memory. Next, we explain the two key parts of MFCL: server-side generative model (Figure 3 Left) and client-side continual learning (Figure 3 Right).
3.1 Server-Side: Generative Model
The motivation for deploying a generative model is to synthesize images that mimic the old tasks and to avoid storing past data. However, training these generative models on the client's side, where the training data exists, is computationally expensive, requires a large amount of training data and can be potentially privacy concerning. On the other hand, the server has only access to the global model and aggregated weights and no data. We propose training a generative model on the server, but in a data-free manner, i.e., utilizing model-inversion image synthesis [57, 50]. In such approaches, the goal is to synthesize images optimized with respect to the discriminator (global model). Then, the
Figure 3: Overview of MFCL. Left. The server aggregates the updates every round and trains a generator using data-free methods at the end of each task. Right. Clients train their models locally using their local data and synthetic images of past tasks from the generator.
generative model is shared with the clients to generate images during local training. To this aim, we utilize a generative model with ConvNet architecture,  $\mathcal{G}$ , that takes noise  $z\sim \mathcal{N}(0,1)$  as input and produces a synthetic sample  $\tilde{x}$ , resembling the original training input with the same dimensions. In order to train this model, we must balance the various training objectives we detail next.
Cross Entropy Loss. First, the synthetic data should be labeled correctly by the current discriminator model (global model or  $\mathcal{F}$ ). To this end, we employ cross entropy classification loss between its assigned label  $z$  and the prediction of  $\mathcal{F}$  on synthetic data  $\tilde{x}$ . Note that noise dimension can be arbitrary and greater than the current discovered classes of task  $t$ ; therefore, we only consider the first  $q$  dimension here, where  $q = \sum_{i=1}^{t} |\mathcal{V}^i|$  (which is equal to the total number of classes seen in the previous tasks). Then, we can define the cross-entropy loss as
$$
\mathcal {L} _ {C E} = C E (\operatorname {a r g m a x} (z [: q ]), \mathcal {F} (\tilde {x})). \tag {1}
$$
Diversity Loss. Synthetic images can suffer from a lack of class diversity. To solve this problem, we utilize the information entropy (IE) loss [12]. For a probability vector  $\mathfrak{p} = (p_1, p_2, \dots, p_q)$ , information entropy is evaluated as  $\mathcal{H}_{info}(\mathfrak{p}) = -\frac{1}{q}\sum_{i}p_{i}\log(p_{i})$ . Based on the definition, inputs with uniform data distributions have the maximum IE. Hence, to encourage  $\mathcal{G}$  to produce diverse samples, we deploy the diversity loss defined as
$$
\mathcal {L} _ {d i v} = - \mathcal {H} _ {i n f o} \left(\frac {1}{b s} \sum_ {i = 1} ^ {b s} \mathcal {F} \left(\tilde {x} _ {i}\right)\right). \tag {2}
$$
This loss measures the IE for samples of a batch (bs: batch size). Maximizing this term encourages the output distribution of the generator to be more uniform and balanced for all the available classes.
Batch Statistics Loss. Prior works [22, 57, 50] in the centralized setting have recognized that the distribution of synthetic images generated by model inversion methods can drift from real data. Therefore, in order to avoid such problems, we add batch statistics loss  $\mathcal{L}_{BN}$  to our generator training objective. Specifically, the server has access to the statistics (mean and standard deviation) of the global model's BatchNorm layers obtained from training on real data. We want to enforce the same statistics in all BatchNorm layers on the generated synthetic images as well. To this end, we minimize the layer-wise distances between the two statistics written as
$$
\mathcal {L} _ {B N} = \frac {1}{L} \sum_ {i = 1} ^ {L} K L (\mathcal {N} (\mu_ {i}, \sigma_ {i} ^ {2}), \mathcal {N} (\tilde {\mu} _ {i}, \tilde {\sigma} _ {i} ^ {2})) = \log \frac {\hat {\sigma}}{\sigma} - \frac {1}{2} (1 - \frac {\sigma^ {2} + (\mu - \hat {\mu}) ^ {2}}{\hat {\sigma} ^ {2}}). \tag {3}
$$
Here,  $L$  denotes the total number of BatchNorm layers,  $\mu_{i}$  and  $\sigma_{i}$  are the mean and standard deviation stored in BatchNorm layer  $i$  of the global model,  $\tilde{\mu}_i$ ,  $\tilde{\sigma}_i$  are measured statistics of BatchNorm layer  $i$  for the synthetic images. Finally,  $KL$  stands for the Kullback-Leibler (KL) divergence.
We want to note that this loss does not rely on the BatchNorm layer itself but rather on their stored statistics  $(\tilde{\mu}_i,\tilde{\sigma}_i)$ .  $\mathcal{G}$  aims to generate synthetic images similar to the real ones such that the global model would not be able to classify them purely based on these statistics. One way to achieve this is to ensure that synthetic and real images have similar statistics in the intermediate layers, and this is
the role of  $\mathcal{L}_{BN}$ . In our experiments, we employed the most common baseline model in CL, which already contains BatchNorm layers and measures those statistics. However, these layers are not a necessity and can be substituted by similar ones, such as GroupNorm. In general, if no normalization layer is used in the model, clients can still compute the running statistics of specific layers and share them with the server, and later, the server can use them in the training of the  $\mathcal{G}$ .
Image Prior Loss. In natural images, adjacent pixels usually have values close to each other. Adding prior loss is a common technique to encourage a similar trend in the synthetic images [22]. In particular, we can create the smoothed (blurred) version of an image by applying a Gaussian kernel and minimizing the distance of the original and Smooth  $(\tilde{x})$  using the image prior loss
$$
\mathcal {L} _ {p r} = \left| \left| \tilde {x} - S m o o t h (\tilde {x}) \right| \right| _ {2} ^ {2}. \tag {4}
$$
In summary, we can write the training objective of  $\mathcal{G}$  as Equation 5 where  $w_{div}, w_{BN}$  and  $w_{pr}$  control weight of each term.
$$
\min  _ {\mathcal {G}} \mathcal {L} _ {C E} + w _ {d i v} \mathcal {L} _ {d i v} + w _ {B N} \mathcal {L} _ {B N} + w _ {p r} \mathcal {L} _ {p r}, \tag {5}
$$
3.2 Client-side: Continual Learning
For client-side training, our solution is inspired by the algorithm proposed in [50]. In particular, the authors distill the stability-plasticity dilemma into three critical requirements of continual learning and aim to address them one by one.
Current Task. To have plasticity, the model needs to learn the new features in a way that is least biased towards the old tasks. Therefore, instead of including all the output space in the loss, the CE loss can be computed for the new classes only by splitting the linear heads and excluding the old ones, which we can write as
$$
\mathcal {L} _ {C E} ^ {t} = \left\{ \begin{array}{l l} C E \left(\mathcal {F} _ {t} (x), y\right), & i f y \in \mathcal {Y} ^ {t} \\ 0, & O. W. \end{array} \right. \tag {6}
$$
Previous Tasks. To overcome forgetting, after the first task, we train the model using synthetic and real data simultaneously. However, the distribution of the synthetic data might differ from the real one, and it becomes important to prevent the model from distinguishing old and new data only based on the distribution difference. To address this problem, we only use the extracted features of the data. To this aim, clients freeze the feature extraction part and only update the classification head (represented by  $\mathcal{F}_t^*$ ) for both real  $(x)$  and synthetic  $(\tilde{x})$  images. This fine-tuning loss is formulated as
$$
\mathcal {L} _ {F T} ^ {t} = C E \left(\mathcal {F} _ {t} ^ {*} ([ x, \tilde {x} ]), y\right). \tag {7}
$$
Finally, to minimize feature drift and forgetting of the previous tasks, the common method is knowledge distillation over the prediction layer. However, [50] proposed importance-weighted feature distillation: instead of using the knowledge in the decision layer, they use the output of the feature extraction part of the model (penultimate layer). This way, only the more significant features of the old model are transferred, enabling the model to learn the new features from the new tasks. This loss can be written as
$$
\mathcal {L} _ {K D} ^ {t} = \left| \left| \mathcal {W} \left(\mathcal {F} _ {t} ^ {1: L - 1} ([ x, \tilde {x} ])\right) - \mathcal {W} \left(\mathcal {F} _ {t - 1} ^ {1: L - 1} ([ x, \tilde {x} ])\right) \right| \right| _ {2} ^ {2}, \tag {8}
$$
where  $\mathcal{W}$  is the frozen linear head of the model trained on the last task  $(\mathcal{W} = \mathcal{F}_{t - 1}^{L})$
In summary, the final objective on the client side as
$$
\min  _ {\mathcal {F} _ {t}} \mathcal {L} _ {C E} ^ {t} + w _ {F T} \mathcal {L} _ {F T} ^ {t} + w _ {K D} \mathcal {L} _ {K D} ^ {t}, \tag {9}
$$
where  $w_{FT}$  and  $w_{KD}$  are hyper-parameters determining the importance of each loss term.
3.3 Summary of MFCL Algorithm