FedET: A Communication-Efficient Federated Class-Incremental Learning Framework Based on Enhanced Transformer

1 Introduction

Federated learning (FL) enables each participating local client to benefit from other clients' data while ensuring client's data does not leave the local [Yang et al., 2019; Dong et al., 2023]. On the premise of ensuring the data privacy of all clients, the problem of data silos has been successfully solved [Hong et al., 2021; Qu et al., 2020]. However, most existing FL methods are modelled in static scenarios, meaning the models' classes are preset and fixed, which undoubtedly reduces the model's generality. Therefore, Federated Class-Incremental Learning (FCIL) is proposed. FCIL solves the problem that FL needs to retrain the entire model when meeting the new classes, saving time and computing costs. For FCIL, how to deal with catastrophic forgetting, seek the plasticity-stability balance of the model and ensure the cooperation of multiple parties are the keys to the problem.

∗Corresponding author: JianzongWang, jzwang@188.com

To date, less work has been done on FCIL studies. The research conducted by [Hendryx et al., 2021] focuses on global IL by facilitating knowledge sharing among diverse clients. However, the author overlooks the non-i.i.d distribution of classes across these distinct clients. The paper [Dong et al., 2022] draw on the regularization methods used in Incremental Learning (IL) and proposes two loss functions. One for addressing the issue of forgetting old classes after IL, and the other is concentrate on the global forgetting caused by the non-i.i.d (non-independent and identically distributed) distribution of classes among different clients. However, this method needs a proxy server to achieve its best performance, leading to high communication costs and some privacy issues. To raise the accuracy of the model in FCIL settings, a natural idea is to choose a more powerful backbone model. We note that there is still no work to apply transformers to FCIL, and the biggest obstacle is that the communication cost is extremely high and cannot be reduced, which makes this application unrealistic. From another perspective, the accuracy and application scope will be significantly improved if we solve the communication and non-i.i.d. problem of class distribution between different clients.

Driven by these ideas, we propose a new Federated Enhanced Transformer (FedET) framework. Compared with other existing FCIL methods, FedET has better prediction performance, lower communication volume, and more universality. It has achieved excellent performance in both Computer Vision (CV) and Natural Language Process (NLP) fields, also it is more efficient when dealing with catastrophic forgetting. FedET consists of four main components: Pretrained Transformer Blocks, Enhancer Select Module, Enhancer Pool and Sample Memory Module (only the local clients have the Sample Memory Module). FedET first divides the entire label space into multiple domains, each with its corresponding Enhancer Group. When new classes need to learn, Enhancer Select Module will determine which domain the new classes belong to and train a temporary Enhancer Group. The new Enhancer Group is obtained by performing distillation between the temporary Group and the corresponding old one. In this way, not only can the local clients have the capability of IL, but large-scale models (such as MAE [He et al., 2022]) can also be used. At the same time, because only the parameters of the chosen Enhancer Group need to be updated, the communication cost is significantly reduced.

We make the following contributions:

• We introduce FedET in order to address the FCIL problem, which mitigates the issue of catastrophic forgetting in both local and global models and effectively reduces communication overhead. According to our knowledge, it is the first effort to explore the FCIL problem in a large-scale model.
• We propose the first FCIL framework used in both CV and NLP fields. Using different transformers as backbones, FedET can handle problems in multiple fields. Compared with baseline models, FedET improves the average accuracy of image classification by 3% and text classification by 1.6%.
• We develop a new loss to handle global catastrophic forgetting named entropy-aware multiple distillation. This is the first time an FCIL model incorporating entropy as a factor when setting the loss function.
• We combine the IL problem of text classification with FL for the first time. By discussing the experimental design method and baseline selection, we think it is a new challenge for both NLP and FCIL fields.

2 Preliminary

In standard IL [Rebuffi et al., 2017; Simon et al., 2021; Shmelkov et al., 2017], the streaming task sequence is defined by T = {T t }Tt=1 , in which T represents the task order, the first t tasks T t = {xti , yt i }N t i=1 contains N t pairs the sample xti and the corresponding one-hot encoded label yt i ∈ Y t . Y t represents the label space of the t-th task, which includes the new classes M t = ⋃Bb=1 mtb that have not appeared in the previous t - 1 tasks, and B represents the number of new classes. At this time, the set of all classes that the model can judge is M A = ⋃ti=1 M i . Inspired by [Ermis et al., 2022; Liu et al., 2020], based on the unique architecture of FL, we construct a Sample Memory Module S located on every local client to store |S| MA exemplars of each class at local, and it satisfies |S| MA ≪ N t Mt .

For FCIL, we give the initial setting under the FL framework [Yoon et al., 2021]: we set K local clients C = {Ck}Kk=1 and a global server CG, the model structures on all clients and server are the same, from the perspective of parameters, including the frozen parameter Φ (that is, the parameters of the selected pre-trained backbone model) and the variable parameter θ. When a clients (a < K) send applications to server for Class-Incremental Learning (CIL), they will access the tth task, updated θ, and select some samples {xti , yt i } put into Sample Memory Module.

3 Methodology

While CIL and FCIL share similarities, the key distinction between them is that FCIL involves tackling two types of forgetting: local and global. FedET addresses local forgetting through a dual distillation loss and mitigates global forgetting through auxiliary data construction and an entropy-aware multiple distillation loss. Figure 1 shows the general outline of the FedET approach.

3.1 Solution of Local Forgetting

In FedET, a local model mainly includes four parts: Pretrained Transformer Blocks, Enhancer Select Module, Enhancer Pool and Sample Memory Module. We show the local model's specific structure and predicting process in Figure 2.

Enhancer Pool and Enhancer Group

Enhancer is the core of FedET, so it is introduced here first. An Enhancer Group contains some Enhancers and a prediction head. And the number of Enhancers is decided by the frozen Pre-trained Transformer Blocks. The mainstream methods of IL fall into three categories [Lange et al., 2022]: playback, regularization, and parameter isolation. In FedET, we use a combination of three approaches: for each client, we set up an Enhancer Pool, which contains multiple Enhancer Groups H = {Hj }Jj=1 , each group is dedicated to being proficient in part of all existing classes. That is, for an Enhancer Group, the class it is responsible for is M Hj , and ⋃Jj=1 M Hj = M A. Setting the parameters of an Enhancer Group and a frozen Pre-trained transformer model as θ and Φ respectively, An Enhancer Group includes many Enhancers

θE and a prediction head θH. During training, only θ is modified, Φ is still frozen. Thereby greatly reducing the number of parameters that need to be adjusted without dropping accuracy. An Enhancer is comprised of three components: a down-projection with Wdown ∈ Rn×m, an activation function f(·), and an up-projection with Wup ∈ Rm×n . Since the encoder structures of transformers are almost the same, after completing FedET's experiments in the NLP and CV fields, we believe this framework can be used for most of the currently known transformers.

Why choose "Enhancer + Freeze the backbone model" instead of Freeze the underlying encoder to adjust the upper encoder? We draw two perspectives from experiments and the literature [R¨uckl´e et al., 2021]. First, Enhancer is freely pluggable, and its internal structure can keep the input of the original encoder, so it can retain the maximum amount of the knowledge that the backbone model has learned during the pre-training stage. Meanwhile, through the design of the Enhancer bottleneck structure and the freezing of the pretraining model, the entire model can learn the downstream tasks better, while the number of parameters that need to be adjusted is significantly reduced. Second, we note that direct fine-tuning can easily lead to overfitting during training on downstream tasks, whereas inserting the Enhancer module performs much better. Although it can be compensated by carefully tuning hyperparameters such as learning rate and batch size, it is undoubtedly time-consuming and labourintensive.

Enhancer Select Module and Sample Memory Module

After a sample is preprocessed, it will be input to the Enhancer Select Module Gs(x). The Enhancer Select Module is a pre-trained frozen classifier. The function of this module is to select a suitable Enhancer Group to handle the input sample. The output of this classifier tells FedET which group is the right group to call up. In t-th task T t , the Enhancer Select Module will first select an Enhancer Group (e.g. j-th) Hj according to the judgement that the new class M t is the most similar to the class M Hj , then Hj will participate in distillation as Hjold . There will be a randomly initialized temporary Enhancer Group Ht aiming to study M t . After the study is completed, Ht will perform distillation with Hjold to obtain a new Enhancer Group Hjnew which covers Hjold , and the specialized class of the group change to M Hjnew = M Hj ∪M t . The judgment methods of the Enhancer Select Module corresponding to different task fields are also different. For CV fields, we design the Enhancer Select Module as an EfficientNet [Tan and Le, 2019], and for NLP fields, we use textRCNN [Lai et al., 2015]. During the distillation process, the required old class samples SMHjold are provided by the Sample Memory Module. After the new class is learned, this module will also store one typical sample of each new type SMt , get new SMHjnew = SMHj ∪ SMt , to ensure that the subsequent distillation can proceed smoothly.

After selection, the parameters of the chosen Enhancer Group θold are linked with the frozen pre-trained model parameters Φ. The θold contains two parts: θold E , the parameters of the Enhancers, and θold H , the parameters of the prediction head.

Distillation of Enhancers

For a new task T t , the new classes need to learn is M t . After the temporary Enhancer Group Ht finish studying, it will be linked with the frozen transformer model as the temporary model ft, which contains Φ and θt . The frozen transformer model with Hold is called the old model fold. We use the following objective to distill fold and ft:

we set

where m and n is the number of M old and M t respectively. To ensure that the consolidated model's output f(x; θnew,Φ)

approximates the combination of outputs from ft and fold, we utilize the output of ft and fold as supervisory signals during the joint training of the consolidated parameters θnew .

To achieve this goal, we employ the double distillation loss proposed by [Zhang et al., 2020] to train fnew. The distillation process is as follows: ft and fold are frozen, and run a feedforward pass with every sample in training set to collect the logits of ft and fold:

respectively, where the superscript is the class label. Then main requirements is to reduce the gap between the logits generated by fnew and the logits generated by ft and fold. Based on prior work[Zhang et al., 2020], We choose L2 loss [Ba and Caruana, 2014] as the distance metric. Specifically, the training objective for consolidation is:

where U denotes the training samples from Sample Memory Module used for distillation. And Ldd is the double distillation loss:

in which yi are the logits produced by fnew for the t-th task, and

where ˆ y is the concatenation of ˆ yold and ˆ ynew. After the consolidation, the Enhancers parameters θnew are used for fold in the next round. The pseudo code for local forgetting solution is shown in Algorithm 1.

3.2 Solution of Global Forgetting

Global catastrophic forgetting primarily arises from the heterogeneity forgetting among local clients participating in incremental learning. Which means the non-i.i.d. classimbalanced distributions across local clients lead to catastrophic forgetting of old classes on a global scale, further exacerbating local catastrophic forgetting. Therefore, it is necessary to solve the heterogeneity forgetting problem across clients in global perspective. To ensure precision and speed, FedET handles this problem with double distillation loss and the difference of average entropy across different clients.

Distillation of Enhancers of Different Clients

FedET changed the stereotype of having to queue up for updates and proposed a new way to update the model. The newmethod is more scientific, reasonable, and time-effective. When a single client uploads the new Enhancer Group obtained after distillation to the server, all the server needs to do is update the parameters of the corresponding group.

Algorithm 1 Local ICL

Input: Enhancer Select Module Gs(x)

• Input: Enhancer Pool H = {Hj }Jj=1
• Input: Sample Memory Module S

Input: Parameters of pre-trained transformer model Φ

Input: t-th task data T t = {xti , yt i }N t i=1

• 1: for i = 1 →N t do
• 2: Group number j ←Gs(xti )
• 3: Put {xti , yt i } with the same j into a list Lj
• 4: end for
• 5: for every selected j do
• 6: while Non-convergence do
• 7: Randomly initialize temporary group Ht
• 8: train ft(x; θt , Φ) with Lj
• 9: Sample from Lj and add to S
• 10: fnew(x; θnew, Φ) = DISTILLATION(ft, fold, S)
• 11: θj ← θnew
• 12: Communicate Hj with Server to get Global best jth group in this turn
• 13: end while
• 14: end for

15:

• 15: function DISTILLATION(ft, fold, S)
• 16: Get ˆ yold from fold and S
• 17: Get ˆ yt from ft and S
• 18: Compute loss function as in Eq.4 and train fnew
• 19: return fnew

When many clients upload the same Enhancer simultaneously, queuing is unscientific because only the last client's update is critical, and this is how global catastrophic forgetting happens. FedET sets a server waiting time limitation. Within a specific time, multiple schemes for an Enhancer Group will be aggregated by the server to perform global model distillation.

For global distillation, the server will distill some Enhancer Group at same time, which means there will be f 1 t , f 2 t , · · · , f q t (q < the number of clinets) and a fold distill together. Note that the new classes are learned by all uploaded groups. Suppose the class which the distilled Enhancer Group major in is M t . For every group-uploaded client, they also upload the information entropy H(M t ) of M t to the server. The server uses H(M t ) to judge the importance of each ft, in detail, the consolidated model of global distillation is:

where Hsum is the sum of information entropy H(M t ) of all uploaded clients. Noted that all output of f here are logits, not hard-label. To get θnew , the entropy-aware multiple distillation loss Lemd is:

in which ¨ y is:

Because of the nature of FL, we cannot rely solely on the sampled data to consolidate the updated Enhancers. Therefore, auxiliary data must be used. During local Enhancer distillation, we generate U using the Sample Memory Module, which stores one representative sample per class and utilizes data augmentation to create the dataset. For global distillation, we construct an equivalent dataset to approximate the training samples. After the local clients send the label distribution Py to the server, the server can construct the auxiliary datasets using available data of a similar domain. Notably, these auxiliary datasets are dynamically fetched and inputted in mini-batches, reducing the storage burden, and discarded after distillation is complete.

Communication Cost Analysis

The parameter quantity of a single Enhancer is 2mn+n+m. For a single local model, if there are D Encoder modules in one Enhancer Group, after adding a group of Enhancers, the increased parameter quantity is:

Other parameters are frozen except for the Enhancer Group and prediction head in the model. As shown in Table 1, the network parameters that need to be transmitted are reduced by more than 70% compared with the various FCIL models previously proposed.

Computation Cost Analysis

The computation FLOPs for each Enhancer in the forward pass are 2 × m × n × sequence length (normalized to a single data sample). The overhead incurred in this way is negligible compared to the original model complexity, e.g., less than 1% on BERT. In the meantime, since all other parameters are fixed during the training period, the computation during backpropagation is reduced by skipping the gradient that computes most of the weights. As shown in Table 1, the use of Enhancers reduces the training time by about 40%.

4 Experiments

As discussed in Section 1, since transformers are widely used in both NLP and CV fields, we test the performance of FedET

on image and text classification tasks. The complete setup will be described in the following subsections.

4.1 Natural Language Processing (NLP)

Datasets and Baselines

Owing to the limited label space of a single dataset, we integrated five text classification datasets [Chen et al., 2020] to evaluate FedET. Table 2 displays the details of the dataset. Considering the domain similarity of Yelp and Amazon, we merge their label spaces for a total of 33 classes. We followed specific task sequences as outlined in Table 3 during training. To alleviate the impact of sequence length and task order on experiment results, we test task sequences of length-3 and length-5 in different orders. The first three tasks of length3 sequences are a cyclic shift of AGnews→Yelp→Yahoo, which belong to three distinct domains (news classification, sentiment analysis, Q&A classification). The remaining three task sequences of length-5 follow the experimental design proposed by [de Masson d'Autume et al., 2019]. During validation, the validation set comprise all classes.

Currently, there is no text classification in the FCIL field, so we choose the baseline of text classification in the ClassIncremental Learning (CIL) field and federate it to form the baseline of this experiment. We compare FedET with five baselines:

• Finetune [Yogatama et al., 2019] + FL: Only new tasks are used to fine-tune the BERT model in turn.
• Replay [de Masson d'Autume et al., 2019] + FL: Replay some old tasks examples during new-tasks-learning to Finetune the model.
• Regularization + Replay + FL: On the foundation of Replay, add an L2 regularization term to the hidden state of the classifier following BERT.
• IDBR [Huang et al., 2021] + FL: On the basis of Regularization + Replay + FL, replace the L2 regularization term with an information disentanglement-based regularization term.

• Multi-task Learning (MTL): Train the model with all class in one task. This approach represents an upper bound on the performance achievable through incremental learning.

Implementation Details

In this NLP experiment, we set J = 3 Enhancer Groups and K = 10 local clients. The prediction head for each group is a linear layer with a Softmax activation function. For sample collection (i.e. experience replay), we stored 3% (store ratio γ = 0.03) of observed examples in the Sample Memory Module, which is used for local Enhancer distillation. We choose the pre-trained Bert-Base-Uncased from HuggingFace Transformers [Wolf et al., 2020] as our backbone model. All experiments utilized a batch size of 16 and a maximum sequence length of 256. During training, ADAMW [Loshchilov and Hutter, 2019] is used as the optimizer, with a learning rate lr = 3e-5 and a weight decay 0.01 for all parameters. For each round of global training, three clients are randomly selected for ten epochs of local training. Selected clients are randomly given 60% of the classes from the label space of its seen tasks.

Results

As shown in Table 1 and Table 4, we can directly see the importance of experience replay for FCIL in NLP. Moreover, the simple regularization approach based on experience replay consistently improves results across all six orders. In most cases, FedET achieves higher performance in incremental learning compared to other baseline methods, while significantly reducing the communication cost. Specifically, compared to IDBR+FL, FedET's Enhancer structure adds a segmentation step to the regularisation and empirical replay, further improving the performance of the model.

4.2 Computer Version(CV)

Datasets and Baselines

We use ImageNet-Subset [Deng et al., 2009] and CIFAR-100 [Krizhevsky et al., 2009] to evaluate our method. We follow the same protocol as iCaRL [Rebuffi et al., 2017] to set incremental tasks and we use the same order generated from iCaRL for a fair comparison. In detail, we compare FedET with the following baselines in the FL scenario: iCaRL, BiC [Wu et al., 2019], PODNet [Douillard et al., 2020], SSIL [Ahn et al., 2021], DDE+ iCaRL [Hu et al., 2021] and GLFC [Dong et al., 2022].

Implementation Details

In this CV experiment, we set J = 10 for Enhancer Groups and K = 30 for local. The prediction head for each group is a linear layer with a Softmax activation function. We collected samples at a store ratio of γ = 0.01. In the CIL baselines, we choose ResNet18 [He et al., 2016] to be the backbone with cross-entropy as the classification loss. On the other hand, FedET uses frozen pre-trained ViT-Base [He et al., 2022] as the backbone. All experiments have a batch size of 64. The training of the Enhancer used an SGD optimizer with minimum learning rate lrmin = 1e-5 and a base learning rate lrb = 0.1. In each round of global training, ten clients are randomly selected for ten local-training epochs. Selected clients are randomly given 60% of the classes from the label space of its seen tasks.

Results

Table 6 and Table 5 show that FedET consistently outperform all the baselines by 3.3% ∼ 10.5% in terms of average accuracy and reduces the communication cost to 11.0% of baseline models. These results demonstrate that FedET can cooperatively train a global class-incremental model in conjunction with local clients more efficiently. Furthermore, for all incremental tasks, FedET has steady performance improvement over other methods, validating the effectiveness in addressing the forgetting problem in FCIL.

(b) Same incremental tasks with different number of clients 𝐾 on CIFAR-100

Figure 3: Ablation study

Ablation Studies

Table 6 and Figure 3 illustrate the results of our ablation experiments on the number of Enhancers and local clients.

The number of Enhancers (J) Figure 3 (a) shows the model's performance in four cases, J = 1, J = 5, J = 10, J = 15, respectively. Compared with J = 10, the performance of J = 1 and J = 5 is worse but J = 15 is better. Since the Enhancer Select Module is frozen, in FedET, we cannot set J = the number of classes. It has been verified that a suitable value of J will make FedET powerful and efficient. We observed that increasing the value of J makes FedET perform better at the beginning of incremental learning but degrades faster as learning progresses. We believe that the reasons may be: (1) With the increase of J, the requirements for the Enhancer Select Module are higher. When the Enhancer Select Module cannot precisely perform rough classification, the model accuracy is bound to decrease. (2) As J increases, the learning cost on a single Enhancer Group is lower, which means that as long as the Enhancer Select Module selects the correct group, the possibility of accurate judgment will be significantly improved.

The number of clients (K) We tested the performance of FedET with client numbers of 30, 50, and 100, as depicted in Figure 3(b). It is clear that FedET's capacity declines as K increases. The performance decrease is most noticeable when K = 100. We believe that with an increase in K, the central server needs to perform distillation on more enhancers simultaneously, which causes the model to not fully converge within the specified number of iterations, resulting in a decrease in model performance.

5 Conclusion

FedET is an FCIL framework that can be used in many fields. Based on previous work, it introduces transformers to improve the accuracy of FCIL and increase the application field of the framework. To reduce communication costs and streamline training, only the Enhancer and its related components are designated as trainable parameters. Our experiments on datasets in NLP and CV demonstrate that FedET outperforms existing methods in FCIL while decreasing communication information by up to 90%.

Acknowledgements

This paper is supported by the Key Research and Development Program of Guangdong Province under grant No.2021B0101400003. Corresponding author is Jianzong Wang from Ping An Technology (Shenzhen) Co., Ltd.