Personalized Federated Learning with Contextual Modulation and Meta-Learning

Machine learning and deep learning have revolutionized various domains by enabling models to learn from vast amounts of centralized data. However, the traditional approach of collecting data into a central server raises concerns regarding data privacy and security. To address these issues, federated learning (FL) [27] has emerged as a promising framework that enables collaborative model training while keeping the data decentralized on individual devices. This decentralized approach has gained significant attention in domains such as healthcare, Internet of Things (IoT), and mobile applications, where data cannot be easily transferred to a central server due to privacy regulations, connectivity limitations, or security risks. However, in the context of FL, data distributions across clients are often not “independent and identically distributed” (non-i.i.d.) since different devices generate and collect data separately. This poses a challenge for FL methods, such as FedAvg [27], that were primarily designed to handle i.i.d. data. Personalized federated learning (PFL) [9] takes a step further by customizing models to the needs and preferences of individual clients, an important aspect in domains such as recommendation systems and user behavior analysis. One way to accomplish this is by integrating meta-learning techniques into the FL framework. For example, Per-FedAvg [9] finds an initialization of a global model that performs well after each client’s update with respect to its own data and loss function, potentially by performing a few steps of gradient descent. This approach accounts for the heterogeneity of data distributions while benefiting from collaborative training. However, when there are substantial differences between client distributions, relying on a single global model may be insufficient. Some approaches have attempted to tackle this issue by learning multiple global models that simultaneously fit different data distributions [8, 11, 26, 33, 43], but this can impose a high computational cost at the server level, as storing multiple models typically requires significant memory resources. Moreover, selecting the most suitable model for each client based on the limited locally available data can be challenging.

In this paper, we propose CAFeMe, a context-aware federated learning solution that leverages meta-learning to facilitate personalization to each client. Our approach is based on a global model comprising a federated modulator and a base model. The federated modulator network captures contextual information from the data of each client and generates modulation parameters that dynamically adjust the activations of the base model. In each communication round, the global model is distributed to the clients, which personalize it with their locally available data. The modulation parameters generated by the federated modulator serve as additional parameters within the base network, allowing for the personalization of the model to the unique characteristics of each client’s data distribution. This allows our proposed approach to deal with high heterogeneity in the clients’ data while maintaining a single model at the server level.

We demonstrate the effectiveness of CAFeMe in improving model convergence, generalization, and personalized prediction on various datasets within the PFL framework. The results highlight the potential of leveraging only relevant information for personalization to each client. This also results in a fast learning process that is particularly advantageous in time-sensitive applications. Furthermore, we showcase the robustness of our framework under the concept shift scenario, where the same class-label may be related to different concepts across clients. The code is available online at https://github.com/annaVettoruzzo/CAFeMe.

The pioneering federated learning (FL) algorithm, FedAvg [27], performs well in scenarios where the local data across clients is i.i.d. However, in scenarios with heterogeneous (non-i.i.d.) data distributions, several techniques have been proposed to improve local learning. These include regularization methods [1, 21, 24], local bias correction techniques [6, 17, 28], data sharing mechanisms [47], data augmentation methods [44], and selective clients aggregation [36, 41]. These methods primarily focus on improving the performance of the global model without considering better customization for each client. An alternative approach is personalized federated learning (PFL), which aims to develop customized models for individual clients while leveraging the collaborative nature of FL. Popular PFL methods include L2GD [13], which combines local and global models, FedRep [5] that decouples the parameters of the feature extractor and the classifier to learn unique local heads for each client, methods that investigate client-specific model aggregations [15, 46], as well as multi-task learning methods such as pFedMe [35], Ditto [22], and FedPAC [42]. Recently, some frameworks propose to incorporate meta-learning into the PFL framework [38]. One notable approach is Per-FedAvg [9], which proposes a decentralized version of model-agnostic meta-learning (MAML) [10] for the FL setting. MAML is a meta-learning algorithm that learns models that can efficiently be adapted, through an optimization procedure, to new tasks. In the context of PFL, personalization to a client aligns with adaptation to a task [16]. Other papers [3, 16, 23] explore the combination of MAML-type methods with FL architectures to improve personalization for each client’s local dataset. Additionally, the authors of [18] propose ARUBA, a meta-learning algorithm inspired by online convex optimization, which enhances the performance of FedAvg.

While the effectiveness of these approaches in PFL has been demonstrated, concerns remain regarding their ability to handle highly heterogeneous data distributions among clients. A special type of PFL approach that tries to address this concern is clustered or group-based FL [8, 11, 26, 33, 43], where multiple (i.e., $K$) group-based global models are learned after partitioning clients into $K$ groups. These methods can either perform clustering at the server level, [33, 12, 43], or let the clients compute the cluster identity based on their local empirical loss [11]. However, partitioning clients into groups presents challenges due to the ambiguity in the clustering structure, and the memory costs of storing $K$ models at the server level.

Our contribution consists of improving the performance and scalability of the PFL framework, even in scenarios where clients exhibit substantial differences in the data distribution. Our proposed method incorporates context information extracted from each client’s local data to drive the personalized model towards a better fit of the client-specific data distribution, while ensuring that the personalized base model does not deviate too far from the global version at the server level. This approach ensures a balance between local adaptation and global consistency, facilitating effective model personalization for each client.

This section provides a formal description of the PFL problem and its connection with the MAML algorithm [10], specifically designed for meta-learning problems. In PFL, we consider a distributed learning setting where a central server communicates with $C$ clients. Let $\mathbb{C}=\{1,\cdots,C\}$ denote the set of clients participating in the training process. Each client $i$ has its own data distribution $P_{\mathcal{X}\mathcal{Y}}^{(i)}$ on $\mathcal{X}\times\mathcal{Y}$, where $\mathcal{X}$ represents the input space and $\mathcal{Y}$ denotes the label space. Each client $i$ has access to $n_{i}$ data points, represented as $\mathcal{D}_{i}=\{x_{k},y_{k}\}$, which are samples drawn from $P_{\mathcal{X}\mathcal{Y}}^{(i)}$. In this context, it is assumed that the probability distributions $P_{\mathcal{X}\mathcal{Y}}^{(i)}$ and $P_{\mathcal{X}\mathcal{Y}}^{(j)}$ differ for any pair of clients $i$ and $j$, which is typically the case in PFL. This is referred to as the non-i.i.d setting. The differences in probability distributions can arise from variations in $P_{\mathcal{X}}$, in $P_{\mathcal{Y}}$ or in the conditional $P_{\mathcal{Y}|\mathcal{X}}$, resulting in covariate shift, label shift, or concept shift, respectively. While previous works have primarily focused on the first two, our proposed solution is designed to effectively handle the concept shift scenario as well.

The main goal of PFL is to learn a global model $f_{\omega}$ that performs well after personalization to each client’s local dataset. This is achieved through an update scheme that proceeds in rounds of communication. In each communication round of the training process, the global model is initialized with $\omega$ and sent to the clients. Each client $i\in\mathbb{C}$ then personalizes the global model using its local data $\mathcal{D}_{i}$ to derive personalized model parameters $\omega^{\prime}_{i}$. Typically, this personalization process involves performing a few steps of gradient descent on the client’s local dataset. Formally, let $\mathcal{L}_{\mathcal{D}_{i}}(f_{\omega^{\prime}_{i}})$ be the empirical risk of the local model $f_{\omega^{\prime}_{i}}$ on the local dataset $\mathcal{D}_{i}$. The objective of PFL can be described as follows:

where $\bm{W}$ denotes the collection of all local models, i.e., $\bm{W}=\{\omega^{\prime}_{i}\}_{i=1}^{C}$. However, in the non-i.i.d setting, it is challenging to find a global model that performs well after personalization to heterogeneous clients. One way to address this challenge is to incorporate meta-learning into the PFL setting.

Meta-learning approaches aim to learn models that can be efficiently adapted to new tasks by leveraging the knowledge acquired from previous tasks. Optimization-based meta-learning methods, such as MAML [10], use a bi-level optimization process to learn an initial set of parameters $\omega$ that can be effectively adapted with gradient descent on new tasks. During the meta-training phase, a set of training tasks $\{\mathcal{T}_{i}\}_{i=1}^{T}$ is used. Each task $\mathcal{T}_{i}$ corresponds to data generating distributions $\mathcal{T}_{i}\triangleq\{P_{\mathcal{X}}^{(i)},P_{\mathcal{Y}|\mathcal{X}}^{(i)}\}$. The data sampled from each task is divided into a support set $\mathcal{D}_{i}^{(sp)}$ containing $K$ training examples, and a query set $\mathcal{D}_{i}^{(qr)}$ used for evaluating the model’s performance. MAML meta-trains a neural network model $f_{\omega}$ parameterized by $\omega$ using a two-stage procedure consisting of an inner and an outer loop. In the inner loop, the initial parameters $\omega$ are adapted to each training task $\mathcal{T}_{i}$ by taking a few gradient descents steps on the support set $\mathcal{D}_{i}^{(sp)}\triangleq\{x_{k},y_{k}\}_{k=1}^{K}$. This process results in task-specific parameters $\omega^{\prime}_{i}$, as illustrated in Eq. 3.2 when a single gradient step is used:

The initial parameters $\omega$ are then optimized in the outer loop by minimizing the loss achieved by the task-specific parameters $\omega^{\prime}_{i}$ on the query set $\mathcal{D}_{i}^{(qr)}$ of the same task $\mathcal{T}_{i}$. This is shown in Eq. 3.3:

The result is a model initialization $\omega$ that performs well when updated with a few training examples $\mathcal{D}_{new}^{(sp)}\triangleq\{x_{k},y_{k}\}_{k=1}^{K}$ from a new task $\mathcal{T}_{new}$. The goal is to train the model on $\mathcal{D}_{new}^{(sp)}$ while leveraging the previous knowledge acquired during meta-training in order to achieve good generalization performance on new unlabeled test examples from $\mathcal{T}_{new}$.

As suggested in [16, 9], optimization-based meta-learning algorithms, such as MAML, bear relevance to the PFL setting due to the inherent heterogeneity present both in tasks (for meta-learning) and in clients (for federated learning). While MAML assumes access to only a few labeled data from a new task at test time, PFL does not have this constraint. Consequently, an estimation of the empirical loss (in Equations 3.2 and 3.3) can be computed using independent batches of data (e.g., $\mathcal{D}_{i}^{\prime}$, $\mathcal{D}_{i}^{\prime\prime}$, etc.) that belongs to $\mathcal{D}_{i}$, as suggested in [9]. By incorporating the MAML formulation into PFL, the objective of PFL shifts towards finding a model initialization shared among all clients that performs well once each client updates it based on the loss function computed with their own local data.

In this section, we present our proposed CAFeMe approach in detail. The aim is to improve the performance and scalability of the PFL framework by leveraging meta-learning techniques and extracting contextual information from each client’s local data. To achieve this, we introduce a federated modulator network $g_{\mu}$ followed by a base network $f_{\psi}$, as shown in Figure 1. We define the complete set of model parameters encompassing both the modulator and the base network as $\omega=\{\mu,\psi\}$. Additionally, we incorporate modulation layers in the base network to enable efficient personalization to each client’s local dataset. These layers aim to condition the base network $f_{\psi}$ on that client’s data by employing a feature-wise gating mechanism to the activations of the preceding layers, allowing only the relevant features (or activations) to propagate forward while disregarding the non-relevant ones. The parameters of the modulation layers are predicted by the federated modulator. The latter extracts some context from batches of data available at the client level and transforms it to generate the modulation parameters. The context is designed to be representative of the client’s local dataset and invariant to the permutation of data samples, ensuring the predicted modulation parameters are independent of the ordering of the data. We delve into the architecture of the federated modulator in Section 4.1 and provide the complete algorithm in Section 4.2.

The experimental results are summarized in Tables 1 and 2 for the synthetic and realistic datasets, respectively. Across all datasets, CAFeMe consistently outperforms the other baseline methods. The performance gap becomes even more evident when the heterogeneity in clients’ data distribution is large (as observed in RMNIST and Meta-Dataset), highlighting the effectiveness of the proposed approach in personalizing to each client’s local data distribution. This observation is supported by the results presented in Figure 3, which illustrates the performance of different PFL methods under varying levels of data heterogeneity in the RMNIST dataset.

When the data are (almost) balanced and i.i.d, as indicated by the “None” column in the histogram and when $\alpha=1000$, the performance of CAFeMe is comparable to other baselines, such as Ditto and FedRep. However, as the data heterogeneity across clients increases (lower $\alpha$ values), CAFeMe consistently maintains strong performance, while the performance of the other baselines significantly deteriorates. This is particularly evident when using the shards partition scheme, where most of the clients have data from only two classes. In this challenging scenario, personalizing the global model to each local dataset becomes difficult due to the narrow local data distribution compared to the one used for learning the global model. Nevertheless, CAFeMe overcomes this challenge by effectively modulating the global model through meta-learning, resulting in robust personalization even in the presence of concept shifts. This observation is further supported by the results on Meta-Dataset (Table 2). The variability in the conditional distribution $P_{\mathcal{Y}|\mathcal{X}}$ across different clients hurts the performance of other methods, which struggle to adapt the global model to such variability. Even using multiple global models, as in IFCA-FT, is insufficient for effective personalization on the client side, also due to the limitations on the estimation of the cluster identity. In contrast, CAFeMe leverages meta-learning and the federated modulator to achieve superior performance in personalizing the global model among clients with concept shifts.

Furthermore, Figure 4 shows that CAFeMe achieves efficient personalization with a small number of personalization steps at test time across all datasets. This demonstrates the efficiency of the proposed approach, enabling quick personalization for clients with limited computational resources. We also conducted additional experiments to investigate the impact of different dataset sizes at the client level. As depicted in Figure 5, when the local data size is too small (e.g., 100), all baselines exhibit low performance due to the inability of personalizing the global model effectively to the client’s local data distribution. On the other hand, as the size of data increases, CAFeMe shows a substantial improvement in performance, suggesting its applicability even in scenarios with small amounts of data at the client level.

In this work, we presented CAFeMe, a novel approach for PFL that leverages meta-learning and a federated modulator network to achieve efficient and effective personalization of the global model to each client’s local data distribution. Our proposed approach overcomes the limitations of traditional FL methods in scenarios with non-i.i.d data distributions by introducing modulation layers that condition the base model on each client’s data.

Through extensive experiments on various synthetic and realistic datasets, we demonstrated the superiority of CAFeMe over state-of-the-art baselines. The performance gap was particularly evident when the data heterogeneity among clients was significant, showing the robustness and adaptability of our approach in such challenging scenarios. The ability to personalize the global model effectively, even in the presence of concept shifts, was one of the key strengths of CAFeMe, setting it apart from traditional approaches that struggled to handle such scenarios. We believe that CAFeMe holds great promise for advancing personalized federated learning in various real-world applications, and we hope that our work inspires further research in this exciting and rapidly evolving field. Future research directions may involve investigating more sophisticated modulator architectures as well as diverse meta-learning techniques and algorithms to optimize the personalization process for clients with varying data distributions. Furthermore, it would be interesting to explore the robustness of CAFeMe under adversarial settings and develop mechanisms to defend against potential security and privacy threats.

This work was supported by the “Knowledge Foundation” (KK-stiftelsen).