Source-linked AI summary
Adaptive Personalized Federated Learning
Yuyang Deng, Mohammad Mahdi Kamani, Mehrdad Mahdavi
TL;DR
Heterogeneous client data can make a single global federated model generalize poorly locally, motivating adaptive personalization. APFL mixes local and global models, derives generalization and convergence guarantees, and reports empirical support for adaptive mixing and the theory.
Problem
Statistical heterogeneity can cause the global model to generalize poorly on individual clients’ data, raising the question of how much each client should personalize.
Method
APFL trains local models while contributing to a global model, learns personalized local-global mixtures, and provides generalization and convergence analyses.
Results
Experiments empirically support the personalization schema and established generalization theories, including adaptive tuning that can outperform fixed α.
Takeaways & Limitations
The supported conclusion is that adaptive per-client mixing provides a way to balance shared and local knowledge when client data are non-IID.
Takeaways & Limitations
The nonconvex convergence analysis assumes K = n, and alleviating this assumption remains open.
Abstract
from arXiv · showhide
Investigation of the degree of personalization in federated learning algorithms has shown that only maximizing the performance of the global model will confine the capacity of the local models to personalize. In this paper, we advocate an adaptive personalized federated learning (APFL) algorithm, where each client will train their local models while contributing to the global model. We derive the generalization bound of mixture of local and global models, and find the optimal mixing parameter. We also propose a communication-efficient optimization method to collaboratively learn the personalized models and analyze its convergence in both smooth strongly convex and nonconvex settings. The extensive experiments demonstrate the effectiveness of our personalization schema, as well as the correctness of established generalization theories.
1 Introduction
Federated learning seeks a shared global model, but heterogeneous client data can make that model generalize poorly locally. APFL addresses this by adaptively mixing local and global models to balance collaboration with personalization.
- Motivation: Federated learning trains on decentralized client data while aiming to learn a shared global model with reduced privacy risks and communication costs.Clients retain data and training locally while contributing to server-side aggregation.
- Motivation: High statistical heterogeneity can cause FedAvg to diverge or generalize poorly on individual clients’ local data.FedAvg is sensitive to hyperparameters because local models may drift substantially from one another.
- Evidence: Increasing client-data diversity produces poor local generalization for FedAvg and SCAFFOLD global models, while the effect diminishes for the proposed personalized model.Figure 1 evaluates MNIST logistic regression as diversity increases.
- Contribution: APFL learns a client-specific mixture of optimal local and global models while theoretically analyzing mixing, distribution divergence, and local and global sample sizes.The approach is intended to balance collaboration benefits against statistical heterogeneity.
2 Related Work
Related personalization work includes fine-tuning, multi-task learning, contextualization, regularization, and model interpolation. APFL focuses on adaptively selecting the local-global mixture according to client data distributions.
- Existing approaches: Federated personalization has been approached through local fine-tuning, multi-task learning, contextualization, model regularization, and model interpolation.These approaches differ in whether they adapt models locally, group clients, use context, regularize parameters or predictions, or mix local and global models.
- Existing approaches: Some personalization methods require client meta-features, creating privacy concerns that make them less feasible for federated learning.This limitation is noted for client clustering and data interpolation approaches.
- Model interpolation: APFL’s model interpolation generalization analysis makes the bound depend on the mixture parameter, data diversity, and local and global optimal models.The paper contrasts this with prior interpolation theory that did not expose the advantage of mixing models.
- Open question: When client data are IID, personalization does not provide a benefit, whereas non-IID data can make personalization useful for balancing shared and local knowledge.The paper frames the remaining question as choosing the appropriate personalization degree for each client.
- Open question: APFL answers the per-client personalization question by adaptively tuning the degree of local-global mixing.The stated goal is to make the choice agnostic to local data distributions.
3 Personalized Federated Learning
The paper models personalization as a convex mixture of global and local empirical-risk minimizers and derives generalization guarantees for that mixture. The optimal mixing depends on distribution divergence and sample sizes, while the theory remains bounded by stated loss and scope assumptions.
- 3.1 Personalized model: The motivation is to use global information to compensate for limited local data while reducing harm from heterogeneity between local and global distributions.When heterogeneity is high, global and local optimal models may drift substantially.
- 3.1 Personalized model: Each client’s personalized model is a convex combination of the global and local empirical-risk minimizers.The global model minimizes empirical risk over the average client distribution, while each client incorporates part of it through a mixing weight α_i.
- Scope and limitations: The analysis is presented for squared hinge and mean-squared losses, with extension to convex smooth losses, while broader divergence theory for common losses remains open.The authors specifically identify hinge loss, cross-entropy loss, and optimal transport as targets for future theory.
- 3.2 Generalization guarantees: Theorem 1 bounds the mixed model’s risk on a client distribution under finite VC dimension, Lipschitz bounded loss, and probability 1 − δ.The bound includes local and total sample sizes and the divergence between the average and client distributions.
- 3.2 Generalization guarantees: The generalization risk mainly depends on total global data, average-to-local distribution divergence, and the client’s local sample count.The paper notes that global data are usually much larger than an individual client’s data, while high divergence weakens the global model’s local relevance.
- Optimal mixing: The optimal mixing parameter is bounded in [0, 1], increases toward local emphasis as distribution divergence grows, and increases global emphasis when local data are scarce.The parameter minimizes the quadratic generalization bound in α_i.
4 Optimization Method
Local Descent APFL jointly learns a shared global model and client-specific models by mixing global and local parameters. Its adaptive mixing coefficient is updated during optimization using local gradient information and model differences.
- Personalized formulation: The personalized model mixes the optimal global model with a modified local empirical risk minimizer, with client-specific α_i controlling the balance.When local and global distributions align, α_i can favor shared information; under larger distribution drift, α_i should approach one.
- Optimization procedure: Local Descent APFL uses a two-phase optimization that updates the shared model globally and client models locally.Clients maintain local, global, and mixed personalized models while optimizing on their own data.
- Optimization procedure: At each communication round, selected clients perform τ local updates, send their updated shared-model copies, and the server averages them before broadcasting the result.Unselected clients retain their previous local models.
- Adaptive α update: The optimal mixing coefficient is generally infeasible to compute directly because the distance between a client distribution and the average distribution is unknown.The method therefore estimates α_i empirically during optimization.
- Adaptive α update: α is updated from the correlation between the personalized-to-global model difference and the gradient evaluated at the personalized model.When personalized and global models are close, as with IID data, α changes less.
5 Convergence Analysis
The convergence analysis characterizes APFL under smooth strongly convex and nonconvex objectives using quantities that measure gradient and local-global heterogeneity. It explains how personalization, synchronization, and residual heterogeneity affect convergence guarantees.
- Analytical quantities: The analysis introduces gradient diversity and the local-global optimality gap to quantify heterogeneity across client objectives and between local and global optima.Gradient diversity is zero exactly when all local functions are identical.
- Assumptions: The convergence theory relies on smoothness, bounded stochastic-gradient variance, and strong convexity where required.The quantities used in the analysis depend on local-data distributions and loss-function geometry.
- Strongly convex convergence: The strongly convex analysis studies convergence of both the global and personalized local models under Local Descent APFL.The personalized model is analyzed relative to each client’s local optimum, while the global model is analyzed relative to the global optimum.
- Strongly convex convergence: O(1/KT) is recovered when τ = T/K, matching the convergence rate of local SGD with periodic averaging.The analysis also notes that setting α_i = 1 recovers the single-machine SGD convergence rate.
- Strongly convex convergence: The personalized-model residual error depends on gradient diversity and the local-global optimality gap, scaled by 1 − α_i.Thus, mixing in the global model can leave the empirical-risk-minimizing local model suboptimal.
- Strongly convex convergence: Removing the lower-bound constraint on α_i yields a looser residual error depending on the bounded-gradient parameter G.This relaxation requires gradients bounded by G and parameters restricted to a bounded domain, with projection during updates.
- Nonconvex convergence: In the nonconvex setting, the analysis assumes all clients participate, and only partial terms achieve speedup with respect to the number of clients.Using different learning rates for local and global updates and removing the full-participation assumption remain open problems.
6 Experiments
The experiments evaluate APFL across strongly convex and nonconvex models, heterogeneous datasets, client-sampling rates, adaptive personalization, and competing personalization methods. Across these settings, APFL generally improves local training or validation performance, with adaptive α benefiting heterogeneous data.
- Experimental setup: Experiments span logistic regression, CNN, and MLP models across MNIST, CIFAR10, EMNIST, and synthetic datasets.The setup uses 100 clients, except for EMNIST, which uses 1000 clients.
- Strongly convex loss: On highly non-IID MNIST, personalized APFL models outperform global FedAvg and SCAFFOLD models in local validation performance.The best generalization performance is obtained with α = 0.25 in both highly non-IID and IID cases, while increasing α lowers training loss.
- Effect of sampling: APFL outperforms localized FedAvg in both training and generalization across sampling rates K ∈ {0.3, 0.5, 0.7}.Larger α values can reduce the negative impact of client sampling on personalized models in highly non-IID data.
- Adaptive personalization: Adaptive α yields nearly equal generalization performance across datasets with different diversity levels and substantially exceeds FedAvg.The experiments motivate adaptive α when the degree of client-data diversity is unknown.
7 Discussion and Extensions
The discussion connects APFL’s adaptive mixing to distributional divergence and compares its personalized model with local ERM, global, and FedAvg alternatives. It also describes new-device personalization and identifies agnostic APFL and tighter drift bounds as open extensions.
- Theoretical connections: The generalization analysis links the mixed-model guarantee to divergence between local and global distributions and to gradient diversity in convergence.The discussion identifies these as corresponding information-theoretic and optimization quantities.
- Adaptive mixing: When local distributions drift from the average, APFL theory and convergence analysis require increasing the local mixing weight α.The adaptive update depends on local-gradient correlation and the deviation between local and global models.
- Comparison with local ERM model: The personalized model can have lower risk than the local ERM model under the paper’s stated high-probability bound.Corollary 3 establishes that the local ERM risk can be strictly worse than the personalized model’s risk.
- New participant nodes: A new device can personalize a trained global model through a few local stochastic-gradient updates using the global model as initialization.The paper also evaluates this setting by comparing APFL’s personalized model with FedAvg’s local model after previously unseen clients join.
- New participant nodes: APFL outperforms FedAvg for new-client personalization in both training loss and validation accuracy, including when increased diversity makes personalization harder.The comparison uses 100 communication rounds with 97 clients followed by fine-tuning on data from 3 joining clients.
- Extensions: Agnostic APFL remains future work because its existing risk bound does not apply to the different optimization domain and its minimax convergence analysis is more involved.The limitation concerns both theoretical guarantees and computational convergence analysis.
8 Conclusion and Future Work
The paper concludes that APFL learns personalized mixtures of local and global models, provides generalization and convergence analyses, and supports these results empirically. It identifies client-distribution diversity and per-device sample size as important personalization factors.
- Contributions: APFL learns personalized models as mixtures of local and global models while allowing clients to contribute to global learning.The algorithm is designed to combine personalization with federated collaboration.
- Theory: The paper derives generalization guarantees that depend on client-distribution diversity and the number of samples available per device.These factors are presented as key determinants in the personalization analysis.
- Optimization: The proposed communication-reduced optimization method is analyzed for smooth strongly convex and nonconvex objectives.The convergence analysis covers both settings stated in the conclusion.
- Empirical validation: Federated experiments empirically support the proposed personalization scheme and the established theoretical results.The conclusion describes the experiments as evidence for both algorithmic effectiveness and theoretical correctness.
A Proof of Generalization Bound
The proof decomposes the mixed model’s local risk, bounds its components using convexity and uniform generalization results, and quantifies the drift from the local empirical minimizer. The analysis assumes bounded loss and acknowledges that tighter drift bounds remain open.
- Proof strategy: The proof bounds the mixed model’s risk on a local distribution by decomposing it into two parts using the convexity of the risk function.This decomposition is the high-level starting point for the generalization analysis.
- Drift control: The mixed local model is not the sole local empirical risk minimizer, so the proof measures its drift from that minimizer through hypothesis capacity.The quantity λH(S) quantifies empirical loss discrepancy among pairs of hypotheses in H.
- Limitations: The paper leaves tighter bounds for the drift caused by incorporating the global model as future work.The stated limitation is that the current drift bound does not depend more precisely on how the global model is incorporated.
- Risk bounds: Uniform VC-dimension error bounds and a bounded-loss lemma are used to control the proof’s empirical and distributional risk terms.The proof applies these tools to bound T1 and T2 before combining them.
- Assumptions: The analysis assumes bounded loss functions, which can follow from bounded parameters and data in examples such as constrained linear models.The paper gives a linear-model setting with ∥w∥≤1 and bounded data as an example.
B.1 Proof without Sampling
The no-sampling analysis studies Local Descent APFL when all clients participate in each averaging stage, then states convergence results for local and global models under objective and mixing-weight conditions.
- Setting: Without client sampling, all clients participate in the averaging stage every τ iterations in the warm-up analysis.This is the setting used before extending the proof to uniformly sampled clients.
- Local model: Theorem 7 gives a local-model convergence result under client objective assumptions and a lower bound on each mixing weight α_i.The theorem applies the no-sampling Local Descent APFL algorithm.
- Local model: The local convergence analysis uses a decaying learning rate η_t = 16/[μ(t+a)] with a = max{128κ, τ}.The theorem also specifies an averaged local-model scheme.
- Global model: Theorem 8 states convergence for the global model under the same objective and mixing-weight framework.Its result uses an averaged global model and identifies w* as the optimal global solution.
B.1.1 Proof of Useful Lemmas
This section establishes technical bounds for local-global gradient differences and model deviations, then uses them in convergence proofs for sampled and unsampled updates.
- Lemma 2 bounds the gap between local and global gradients at each iteration for Algorithm 2.
- Lemma 3 bounds each local version’s deviation from the global model without sampling.
- The proof repeatedly combines smoothness, strong convexity, Jensen’s inequality, and telescoping sums.
- For sampled updates, Lemmas 5 and 6 address gradient gaps and local-model deviations under partial participation.
- The proofs use learning rate ηt = 16/[µ(t+a)] under the stated smooth strongly convex assumptions.
- The sampling analysis introduces an indicator for client selection and accounts for additional variance from updating only K of n devices.
C Proof of Convergence without Assumption on αi
This section provides the proof of Theorem 4 by applying the preceding bound and then following the established proof procedure.
- The section is devoted to proving Theorem 4.
- The proof invokes equation (43) as its starting relation.
- An expectation bound is applied before the argument follows the same procedure used earlier.
D.1 Proof of Technical Lemmas
This section develops technical lemmas for the nonconvex convergence analysis, including bounds across synchronization stages and their combination into the final result.
- Lemmas 8, 9, and 10 state intermediate bounds under Theorem 6’s assumptions.
- The analysis uses the updating rule, smoothness of local objectives, expectations, and Jensen’s inequality to derive these bounds.
- Synchronization-stage bounds are summed over intervals and divided by T to obtain aggregate statements.
- The final proof re-arranges and sums the intermediate inequalities, then plugs Lemmas 9 and 10 into the result.