ActPerFL: Active Personalized Federated Learning
Abstract
In the context of personalized federated learning (FL), the critical challenge is to balance local model improvement and global model tuning when the personal and global objectives may not be exactly aligned.
Inspired by Bayesian hierarchical models, we develop ActPerFL, a self-aware personalized FL method where each client can automatically balance the training of its local personal model and the global model that implicitly contributes to other clients’ training.
Such a balance is derived from the inter-client and intra-client uncertainty quantification .
Consequently, ActPerFL can adapt to the underlying clients’ heterogeneity with uncertainty-driven local training and model aggregation .
With experimental studies on Sent140 and Amazon Alexa audio data, we show that ActPerFL can achieve superior personalization performance compared with the existing counterparts.
Introduction
Federated learning (FL) (Konevcny et al., 2016;
McMahan et al., 2017) is transforming machine learning (ML) ecosystems from “centralized in-thecloud” to “distributed across-clients,” to potentially leverage the computation and data resources of billions of edge devices (Lim et al., 2020), without raw data leaving the devices.
As a distributed ML framework, FL aims to train a global model that aggregates gradients or model updates from the participating edge devices.
Recent research in FL has significantly extended its original scope to address the emerging concern of personalization, a broad term that often refers to an FL system that accommodates client-specific data distributions of interest (Dinh et al., 2020a; Fallah et al., 2020a).
In particular, each client in a personalized FL system holds data that can be potentially nonIID.
For example, smart edge devices at different houses may collect audio data of heterogeneous nature (Purington et al., 2017; Diao et al., 2020, 2021) due to, e.g., accents, background noises, and house structures.
Each device hopes to improve its on-device model through personalized FL without transmitting sensitive data.
While the practical benefits of personalization have been widely acknowledged, its theoretical understanding remains unclear.
Existing works on personalized FL often derive algorithms based on a pre-specified optimization formulation or model aggregation rule.
In this work, we start with a toy example and develop insights into the nature of personalization from a statistical uncertainty perspective.
In particular, we aim to answer the following critical questions regarding personalized FL.
(Q1)
The lower-bound baselines of personalized FL can be obtained in two cases,
i
e., each client performs local training without FL, or all clients participate in conventional FL training.
However, the upper-bound for the client is unclear.
(Q2)
Suppose that the goal of each client is to improve its local model performance.
How to design an FL training that interpret the global model, suitably aggregate local models and fine-tune each client’s local training automatically?
Both questions are challenging.
The question (Q1) demands a systematic way to characterize the client-specific and globally-shared information.
To this end, we draw insights from a simplified and analytically tractable setting: two-level Bayesian hierarchical models , where the two levels respectively describe inter-client and intra-client uncertainty.
We make the following technical contributions:
• Interpreting personalization from a hierarchical model-based perspective and providing theoretical analyses for FL training.
• Proposing ActPerFL, an active personalized FL solution that guides local training and global aggregation via inter- and intra-client uncertainty quantification .
• Presenting a novel implementation of ActPerFL for deep learning, consisting of auto-
mated hyper-parameter tuning for clients and an adaptive aggregation rule.
• Evaluating ActPerFL on Sent140 and Amazon Alexa audio data.
Empirical results show promising personalization performance compared with existing methods.
To our best knowledge, ActPerFL is the first work that utilizes uncertainty quantification to drive
FL personalization.
Bayesian View of Personalized FL
We discuss how ActPerFL approaches personalized FL with theoretical insights from the Bayesian perspective in this section.
To develop insights, we study a two-level Gaussian model.
Similar arguments can be derived for generic parametric models.
The notations are defined as follows.
Let N ( µ, σ 2 ) denote Gaussian distribution with mean µ and variance σ 2 .
For a positive integer M , let [ M ] denote the set { 1 , . . .
, M } .
Let (cid:80) m ̸
= i denote the summation over all m ∈
[ M ] except for m = i .
Suppose that there are M clients.
From the server’s perspective, it is postulated that data z 1 , . . .
, z M are generated from the following two-layer Bayesian hierarchical model:
[Equation]
for all clients with m = 1 , . . .
, M .
Here, σ 20 is a constant, and θ 0 ∼ π 0 ( · ) is a hyperparameter with a non-informative flat prior.
The above model represents both the connections and heterogeneity across clients.
In particular, each client’s data are distributed according to a client-specific parameter ( θ m ), which follows a distribution decided by a parent parameter ( θ 0 ) .
The parent parameter is interpreted as the root of shared information.
With- out loss of generality, we study client 1’s local model as parameterized by θ 1 .
Under the above model assumption, the parent parameter θ 0 that represents the global model has a posterior distribution p ( θ 0 | z 1: M ) ∼ N ( θ ( G ) , v ( G ) ) , where:
[Equation]
From the perspective of client m , we suppose that the postulated model is the same as above for m = 2 , . . .
, M , and θ
1 = θ 0 .
It can be verified that the posterior distributions of θ 1 without and with global Bayesian learning are p ( θ 1 | z 1 ) ∼
N ( θ ( L ) 1 , v ( L ) 1 ) and p ( θ 1 | z 1: M ) ∼ N ( θ ( FL ), v ( FL )) , respectively, which can be computed as:
[Equation]
The first distribution above describes the learned result of client 1 from its local data, while the second one represents the knowledge from all the clients’ data in hindsight.
Using the mean square error as risk, the Bayes estimate of θ 1 or θ 0 is the mean of the posterior distribution, namely θ ( L )and θ ( FL ).
The flat prior on θ 0 can be replaced with any other distribution to bake prior knowledge into the calculation.
We consider the flat prior because the knowledge of the shared model is often vague in practice.
The above posterior mean θ ( FL )can be regarded as the optimal point estimation of θ 1 given all the clients’ data, thus is referred to as “FL-optimal”.
θ ( G ) can be regarded as the “globaloptimal.”
The posterior variance quantifies the re- duced uncertainty conditional on other clients’ data.
Specifically, we define the following Personalized FL gain for client 1 as:
[Equation]
Remark 1 (Posterior quantity interpretations)
Each client, say client 1, aims to learn θ 1 in the personalized FL context.
Its learned information regarding θ 1 is represented by the Bayesian posterior of θ 1 conditional on either its local data z 1 (without communications with others), or the data z 1: M in hindsight (with communications).
For the former case, the posterior uncertainty described by v ( L )depends only on the local data quality σ 21 .
For the latter case, the posterior mean θ ( FL )is a weighted sum of clients’ local posterior means, and the uncertainty will be reduced by a factor of GAIN 1 .
Since a point estimation of θ 1 is of particular interest in practical implementations, we treat θ ( FL )as the theoretical limit in the FL context (recall question Q1).
Remark 2 (Local training steps to achieve θ ( FL ))
Suppose that client 1 performs ℓ training steps using its local data and negative log-likelihood loss.
We show that with a suitable number of steps and initial value, client 1 can obtain the intended θ ( FL ).
The local objective is:
[Equation]
which coincides with the quadratic loss.
Let η ∈ (0 , 1) denote the learning rate.
By running the gradient descent:
[Equation]
for ℓ steps with initial value θ INIT, client 1 obtains:
[Equation]
It can be verified that Eqn.
(5) becomes θ ( FL )in
Eqn.
(2) if and only if:
[Equation]
In other words, with a suitably chosen initial value θ INIT, learning rate η , and the number of (earlystop) steps ℓ , client 1 can obtain the desired θ ( FL ).
Proposed Solution for Personalized FL
Our proposed ActPerFL framework has three key components as detailed in this section: (i) proper initialization for local clients at each round, (ii) automatic determination of the local training steps, (iii) discrepancy-aware aggregation rule for the global model.
These components are interconnected and contribute together to ActPerFL’s effectiveness.
Note that points (i) and (iii) direct
ActPerFL to the regions that benefit personalization in the optimization space during local training, which is not considered in prior works such as
DITTO (Li et al., 2021) and pFedMe
(Dinh et al., 2020b).
Therefore, ActPerFL is more than imposing implicit regularization via early stopping.
In this section, we show how the posterior quan- tities of interest in Section 2 can be connected with
FL.
Recall that each client m can obtain the FLoptimal solution θ ( FL ) m with the initial value θ INIT m in Eqn.
(6) and tuning parameters η, ℓ in Eqn.
Also, it can be shown that θ INIT m is connected with the global-optimal θ ( G ) in Eqn.
(1) through θ
[Equation]
The initial value θ INIT m in Eqn.
(8) is unknown during training since θ ( L ) m , θ ( G ) are both unknown.
A natural solution is to update θ INIT m , θ ( L ) m , and θ ( G ) iteratively, leading to the following personalized
FL rule of our ActPerFL framework.
Generic ActPerFL.
At the t -th ( t ≥ 1 ) round:
• Client m receives the latest global model θ t −from the server (initialized as θ 0 ), and calculates:
[Equation]
where θ t − 1 m is client m ’s latest personal parameter at round t − 1 , initialized to be θ 0 .
Starting from the above θ t, INIT m , client m performs gradient descent-based local updates with optimization parameters following Eqn.
(7) or its approximations, and obtains a personal parameter θ tm .
• Server collects θ tm and calculates:
[Equation]
In general, the above σ 20 , σ 2 m represent “inter-client uncertainty” and “intra-client uncertainty,” respectively.
When σ 20 and σ 2 m ’s are unknown, they can be approximated asymptotically or using practical finite-sample approximations.
SGD-based practical algorithm for DL.
For the above training method, the quantities σ 20 and σ 2 m are crucial as they affect the choice of learning rate η m and the early-stop rule.
However, these two values are unknown in complex learning models.
To approximate the uncertainty quantities, we generally treat σ 2 m as “uncertainty of the local optimal solution θ ( L ) m of client m ”, and σ 20 as “uncertainty of clients’ underlying parameters.”
Assume that for each client m , we had u independent samples of its data and the corresponding local optimal parameter θ m, 1 , . .
, θ m,u .
We could then estimate σ 2 m by their sample variance .
In particular, at round t , we approximate σ 2 m with:
[Equation]
Likewise, at round t , we estimate σ 20 by:
[Equation]
For multi-dimensional parameters, we introduce the following counterpart uncertainty measures.
For vectors x 1 , . . .
, x M , their empirical variance is defined as the trace of (cid:80) m ∈
[ M ] ( x m − ¯ x )( x m − ¯ x ) T , which is the sum of entry-wise empirical variances.
(cid:99) σ 2 m and (cid:99) σ 20 are defined from such empirical variances similar to Eqn.
(11) and (12).
The above quantities can be calculated recursively online with constant memory (Han et al., 2017).
Alg.
1 outlines the workflow of ActPerFL.
Experimental Studies
Experimental setup.
We evaluate ActPerFL’s performance on two NLP datasets: {+Sentiment
+}
(Go et al., 2009) and private Amazon Alexa audio data.
Sent140 is a text sentiment analysis dataset with two output classes and 772 clients.
We generate non-i.i.d.
data following FedProx (Li, 2020).
The audio dataset is collected for wake-word detection task (i.e., binary classification).
This dataset contains 39 thousand hours of training data andthousand hours of test data.
We use a two-layer
Algorithm
Input: A server and M clients.
Communication rounds T , client activity rate C , client m ’s local data D m and learning rate η m .
for each communication round t = 1 , . . .
T do
Sample clients: M t ← max( ⌊ C · M ⌋ , 1) for each client m ∈ M t in parallel do
Distribute server model θ t − 1 to client m
Estimate (cid:99) σ 2 m using Eqn.
Compute local step l m from Eqn.
(7) and local initialization θ INIT m via Eqn.
(6) θ tm ← LocalT rain ( θ INIT m , η m , l m ; D m )
Server estimates (cid:99) σ 20 using Eqn.
Server updates global model θ t via Eqn.
LSTM model and an 11-layer CNN model for these two datasets, respectively.
For comparison, we also evaluate the personalization performance of FedAvg (McMahan et al., 2017), DITTO (Li et al., 2021), PerFedAvg (Fallah et al., 2020b), and pFedMe (Dinh et al., 2020b).
4. Results on Alexa Audio Data
For Alexa audio data, we use a CNN that is pretrained on the training data of different device types (i.e., heterogeneous data) as the initial global model to warm-start FL training.
The personalization task aims to improve the wake-word detection performance at the device type level .
We assume there are five clients in the FL system and all of them participate in each round.
Each client has the training data for a specific device type.
Evaluation metric.
We evaluate the performance using the pre-trained model (for warm-start) as the baseline .
To compare different FL algorithms, we use the relative false accept (FA) value of the resulting model when the associated relative false reject (FR) is close to one as the metric.
So a smaller relative FA is preferred.
Here, the relative
FA and FR are computed using the baseline.
For comparison, we implement FedAvg and
DITTO with both equal-weighted and sample sizebased model averaging (denoted by the suffix ‘-e’ and ‘-w’, respectively) during aggregation.
For
PerFedAvg (Fallah et al., 2020b), we use its firstorder approximation and the equal-weighted aggregation.
We did not report pFedMe
(Dinh et al., 2020b) due to its divergence with various hyperparameters.
Table 1 summarizes the performance of the updated global model.
The results show that
ActPerFL achieves the lowest relative FA, thus obtaining the best global model.
We further compare
[Table] [Table]
the personalization performance of local models obtained by different FL algorithms in Table 2.
4. Results on Sent140 Text Data
In this experiment, we also use warm-start by training a global model from scratch with FedAvg for 200 rounds for initializing other FL algorithms.
Then, we continue FL training with various FL methods for another 400 rounds.
Figure 1 compares the training and test accuracy of the personalized models obtained by different FL algorithms where the accuracy is aggregated across clients.
We can see that both ActPerFL and FedAvg demonstrate better convergence performance compared to DITTO (Li et al., 2021), pFedMe
(Dinh et al., 2020b), and PerFedAvg (Fallah et al., 2020b).
[Figure]
Concluding Remarks
We proposed ActPerFL to address the challenge of balancing local model training and global model aggregation in personalized FL.
Our solution adaptively adjusts local training with automated hyper-parameter selection and performs uncertainty-weighted global aggregation.
Empirical studies show that ActPerFL can achieve promising performance on NLP applications.