FedGTEA: Federated Class-Incremental Learning with Gaussian Task Embedding and Alignment

TL;DR

This work tackles Federated Class Incremental Learning (FCIL) by introducing FedGTEA, which jointly models task-level knowledge and uncertainty in a privacy-preserving, scalable way. On the client, a Cardinality-Agnostic Task Encoder (CATE) produces Gaussian task embeddings that summarize task context without growing parameters with the number of tasks. On the server, a 2-Wasserstein distance-based regularization enforces inter-task separation and supports robust consolidation via a three-term server loss that includes knowledge distillation and an anchor term. Empirical results on CIFAR-10/100 sequences show that FedGTEA achieves higher accuracy and substantially less forgetting than strong baselines, illustrating the benefit of incorporating task-level context into FCIL. The approach offers a practical, privacy-conscious pathway to scalable continual learning in distributed settings with non-IID data distributions.

Abstract

We introduce a novel framework for Federated Class Incremental Learning, called Federated Gaussian Task Embedding and Alignment (FedGTEA). FedGTEA is designed to capture task-specific knowledge and model uncertainty in a scalable and communication-efficient manner. At the client side, the Cardinality-Agnostic Task Encoder (CATE) produces Gaussian-distributed task embeddings that encode task knowledge, address statistical heterogeneity, and quantify data uncertainty. Importantly, CATE maintains a fixed parameter size regardless of the number of tasks, which ensures scalability across long task sequences. On the server side, FedGTEA utilizes the 2-Wasserstein distance to measure inter-task gaps between Gaussian embeddings. We formulate the Wasserstein loss to enforce inter-task separation. This probabilistic formulation not only enhances representation learning but also preserves task-level privacy by avoiding the direct transmission of latent embeddings, aligning with the privacy constraints in federated learning. Extensive empirical evaluations on popular datasets demonstrate that FedGTEA achieves superior classification performance and significantly mitigates forgetting, consistently outperforming strong existing baselines.

Figures (3)

• Figure 1: Given different task contexts, identical input can yield contradicting answers. Multi-task models need task-level context to process different tasks accurately.
• Figure 2: This illustrates how client $\mathcal{C}_k$ is trained over task $\mathcal{T}^t$. Given locally collected dataset $\mathcal{D}_k^t$, CATE and CNN in the discriminator extract task embedding $\mathcal{E}$ and data feature $\mathcal{F}$ respectively. The class head outputs the prediction probability using both $\mathcal{F}$ and $\mathcal{E}$. At the same time, the Real/Fake head in the discriminator classifies whether the input image is genuine or synthesized. After each local iteration, $G$ generates fake images for replay purposes.
• Figure 3: This figure illustrates the aggregation and model regularization step on the server. We first obtain an initial aggregated model $\hat{\theta}_g^t$ by typical federate aggregation. Then, the regularization and consolidation use the server loss that comprises three components. Anchor loss prevents the global model from drifting too far away; Wasserstein loss promotes inter-task separation; knowledge distillation loss mitigates catastrophic forgetting.