Source-linked AI summary
Personalized Federated Learning with First Order Model Optimization
Michael Zhang, Karan Sapra, Sanja Fidler, Serena Yeung, Jose M. Alvarez
TL;DR
Federated learning’s single global model may not suit clients with different target distributions. FedFomo instead computes client-specific weighted combinations of uploaded models without assuming client similarities, and it outperforms alternatives while supporting targets outside local training distributions.
Problem
A single global model may not provide strong performance for every client, especially when decentralized data are non-IID and client target distributions differ.
Method
FedFomo evaluates uploaded models against each client’s objective and efficiently computes personalized weighted combinations without assuming knowledge of client data distributions.
Results
FedFomo outperforms other methods across non-IID settings, with up to 70% improvement in some settings and support for target distributions different from local training data.
Takeaways & Limitations
Personalized federated updates can use client-specified target objectives and enable optimization beyond each client’s local data distribution.
Takeaways & Limitations
Computing the optimal model weights directly can be cumbersome and prohibitively expensive in communication rounds when averages are computed only server-side.
Abstract
from arXiv · showhide
While federated learning traditionally aims to train a single global model across decentralized local datasets, one model may not always be ideal for all participating clients. Here we propose an alternative, where each client only federates with other relevant clients to obtain a stronger model per client-specific objectives. To achieve this personalization, rather than computing a single model average with constant weights for the entire federation as in traditional FL, we efficiently calculate optimal weighted model combinations for each client, based on figuring out how much a client can benefit from another's model. We do not assume knowledge of any underlying data distributions or client similarities, and allow each client to optimize for arbitrary target distributions of interest, enabling greater flexibility for personalization. We evaluate and characterize our method on a variety of federated settings, datasets, and degrees of local data heterogeneity. Our method outperforms existing alternatives, while also enabling new features for personalized FL such as transfer outside of local data distributions.
1 INTRODUCTION
The paper reframes federated learning around client-specific target distributions rather than one global model, using personalized weighted combinations of uploaded models. FedFomo efficiently computes these updates without assuming client similarities and outperforms alternatives in non-IID settings.
- Personalized federated learning targets stronger performance on each client’s target distribution instead of a single aggregate test distribution.
- FedFomo weights uploaded models for each client according to how well they align with that client’s interests and target task.
- Up to 70% improvement is demonstrated in some settings, with larger gains as the number of clients or non-IIDness increases.
- Each client can personalize to a specific target distribution irrespective of its available local training data.
- The method computes personalized federated updates by efficiently calculating optimal weighted combinations of uploaded models.
- The method strongly outperforms other methods in non-IID federated learning settings.
2 RELATED WORK
Prior work addresses non-IID federated learning through robust global models, multiple client models, or personalized components. These approaches generally do not directly optimize performance on client-level target distributions, especially when those differ from local training data.
- Federated Learning with Non-IID Data: Non-IID decentralized data can cause parameter divergence, distribution biases, and unguaranteed convergence.
- Federated Learning with Non-IID Data: FedProx limits local parameter drift with a proximal objective, while FedAvgM uses momentum to reduce oscillations during global aggregation.
- Personalized Federated Learning: Existing approaches do not directly address client-level performance when target distributions differ substantially from the global aggregate or local training data.
- Personalized Federated Learning: MOCHA frames clients as tasks and learns one model per client through distributed multi-task learning.
- Personalized Federated Learning: Mixture methods personalize by combining local and global model parameters, while other methods mix components across network layers or preserve local privacy.
3 FEDERATED FIRST ORDER MODEL OPTIMIZATION
FedFomo replaces a single global model with client-specific weighted combinations of available models, optimized against each client’s target objective. It uses first-order updates and model-selection mechanisms while limiting communication overhead through restricted downloads.
- Personalized optimization: FedFomo lets each client optimize a weighted combination of available server models for its own target objective.Clients evaluate received models on target-task validation data and use those performances to weight model parameters.
- Personalized optimization: Each federated update evaluates received models on a client’s target task, then uses their relative performance to construct a personalized parameter update.Uploaded parameters remain individually available so clients can independently download and combine model copies.
- Problem setting: The framework supports heterogeneous clients without assuming knowledge of client distributions or that local and target data share a distribution.Clients may have distinct local objectives and target test distributions, including targets outside their local training distributions.
- First-order update: The optimal model combination is approximated with first-order weights that favor models improving the client’s loss and models closer to its current parameters.Negative weights are clipped to zero, and positive weights are normalized for the final update.
- Communication and selection: The method keeps one download and one upload round between E local epochs, matching FedAvg’s total communication-round count while adding multiple-model bandwidth.The authors note that quantization and distillation can alleviate this additional bandwidth, while restricting M avoids worst-case N^2 overhead.
- Communication and selection: Model selection uses an affinity matrix whose entries estimate how likely each available model is to benefit each client, sending only the top M models.Affinity estimates are updated from prior target-objective performance, restricting downloads under limited bandwidth.
4 EXPERIMENTS
FedFomo is evaluated across pathological and latent non-IID settings, out-of-distribution personalization, model-selection ablations, and differentially private training. Across these experiments, it generally outperforms competing methods while learning client-to-client weighting patterns aligned with target distributions.
- Pathological Non-IID: FedFomo strongly outperforms alternatives in pathological non-IID settings with larger client populations and remains competitive with 15 clients.The pathological setup assigns each client two of ten classes; single-model robust methods show mixed success over FedAvg and do not exceed local training alone.
- Latent Distribution Non-IID: FedFomo strongly outperforms competing methods in highly non-IID latent-distribution settings, with the performance gap widening as heterogeneity increases.It is competitive with local training in the most extreme case and outperforms both local training and FedAvg at moderate heterogeneity, measured by EMD ∈[1, 2].
- Exploration with ε and number of models downloaded M: With fixed ε, downloading more models increases performance by enabling earlier identification of beneficial model-client assignments.The ablation varies ε-greedy exploration and the download limit M on CIFAR-10 over 100 training epochs.
- Out-of-local-distribution personalization: FedFomo outperforms alternatives when clients optimize for target distributions different from their local training distributions.The out-of-client-distribution evaluation uses CIFAR-10 and CIFAR-100 with 15 clients and five latent distributions; other methods do not convincingly optimize this shifted target.
- Personalized model weighting: FedFomo learns client-to-client weights that upweight clients sharing the relevant training or target distribution.The visualizations show weights increasing for clients belonging to the original client’s target distribution while other clients are downweighted.
- Locally Private FedFomo: Under differential privacy, FedFomo retains a sizeable improvement over FedAvg with minimal ε, while stronger privacy from greater σ reduces performance.Noisy updates still preserve the ability to discover larger latent distributions, although the initial weighting is noisier.
5 CONCLUSION
FedFomo is a flexible personalized federated learning framework for non-IID settings, including objectives on target distributions distinct from clients’ local training data. It uses personalized model combinations, discovers aligned client distributions, and preserves functionality under added privacy guarantees.
- 5 CONCLUSION: FedFomo personalizes federated updates by downloading weighted combinations of individual client models according to each client’s target objective.The method replaces previously used model averages with individually maintained client models and calculates combinations efficiently.
- 5 CONCLUSION: FedFomo achieves strong performance across various non-IID settings and enables optimization for target distributions different from local training data.This out-of-local capability extends personalization beyond objectives tied to each client’s available training data.
- 5 CONCLUSION: FedFomo empirically discovers local client data distributions and upweights models trained on data aligned with each client’s target objective.The framework uses client-specific model suitability to determine which other models should contribute to each update.
- 5 CONCLUSION: FedFomo preserves its core functionality and strong personalization when additional privacy guarantees are introduced.The conclusion reports this behavior as part of the framework’s evaluation in federated settings.
A.1 DERIVING THE FOMO UPDATE
The Fomo update treats personalized federated learning as a constrained model-combination problem and chooses weights that approximate a gradient-based update for a client’s objective. A first-order approximation yields practical normalized weights for combining available client models.
- A.1 DERIVING THE FOMO UPDATE: A federated model download can be viewed as an iterative update whose parameters are determined by the personalized weights and model deltas.This perspective connects the Fomo update to iterative optimization while retaining the fixed available models at each round.
- A.1 DERIVING THE FOMO UPDATE: Each client chooses weights over available model parameters to move its local model toward an objective defined by a target dataset or distribution.The available model parameters restrict the update paths, while the objective need not be global test performance.
- A.1 DERIVING THE FOMO UPDATE: The weights are constrained to sum to 1, forming a normalized budget allocated across models in the federated update.The same reparameterization also supports directly calculating a client’s weight for its own model.
- A.1 DERIVING THE FOMO UPDATE: The derivation seeks weights whose model-combination update closely approximates an optimal gradient update with respect to the client’s objective.The correspondence is established by matching the model-combination update with the gradient-based update in vector or matrix form.
- A.1 DERIVING THE FOMO UPDATE: For each client model, the method derives a first-order approximation of the objective derivative with respect to its weight and solves for the corresponding weight.The resulting value specifies how much client ci weights model θn in its federated update.
- A.1 DERIVING THE FOMO UPDATE: The practical weights normalize the solved values after distributing the negative α to select the update direction, with the constant canceling under normalization.The normalized weights are the values actually used in practice.
A.2 ADDITIONAL LATENT DISTRIBUTION NON-IID EXPERIMENTS
On CIFAR-100, the latent non-IID in-distribution personalization experiments compare FedFomo with recent alternatives as heterogeneity increases. FedFomo consistently outperforms the alternatives across increasingly heterogeneous client distributions.
- A.2 ADDITIONAL LATENT DISTRIBUTION NON-IID EXPERIMENTS: FedFomo consistently outperforms all alternatives on CIFAR-100 as the number of non-IID data distributions increases across 15 clients.The experiment reports accuracy averaged over all client runs and gives increasing EMD values in parentheses.
- A.2 ADDITIONAL LATENT DISTRIBUTION NON-IID EXPERIMENTS: FedFomo is more competitive when client data are more similar and handily outperforms other methods as statistical label heterogeneity increases.Similarity is characterized using lower EMD and fewer data distributions.
A.3 CLIENT WEIGHTING WITH PERSONALIZATION
The client-weighting experiments examine whether FedFomo identifies helpful clients for local and non-local target distributions. Across the visualized settings, it downweights unhelpful distributions and increases or maintains weights for aligned clients.
- A.3 CLIENT WEIGHTING WITH PERSONALIZATION: The experiments use client-weight trajectories colored by local training distribution to assess whether FedFomo identifies which clients should federate together.Helpful clients should gain or retain weight, while other client weights should decline over federation rounds.
- A.3 CLIENT WEIGHTING WITH PERSONALIZATION: FedFomo quickly downweights clients with unhelpful data distributions in both in-distribution and out-of-local personalization.For in-distribution targets, it increases or maintains higher weights for clients from the same distribution.
- A.3 CLIENT WEIGHTING WITH PERSONALIZATION: Under differentially private federated learning, FedFomo is evaluated with a model-average baseline in addition to the standard comparison methods.The privacy-focused comparison is reported in Table 6.
A.4 ADDITIONAL PRIVACY EXPERIMENTS
The multiple-model privacy variant averages uploaded models to make individual data membership harder to infer, but its personalization performance is lower than the single-model-download variant.
- The variant downloads averages excluding and including a client’s own model, then uses their comparison to calculate personalized weights.This changes the weighting calculation relative to the baseline model because the client evaluates model averages rather than individual models.
- The privacy evaluation uses 15-client CIFAR scenarios with δ = 1 × 10^-5 and Gaussian noise σ ∈{0, 1, 2}.The experiment compares accuracy and privacy loss ε at fixed δ.
- The model-average variant still outperforms non-private FedAvg on CIFAR-10 and CIFAR-100, but performs considerably worse than single-model-download FedFomo.The authors attribute the drop to noisier estimates of another model’s potential contribution when evaluating many models together.
- Balancing personalization and privacy across the two weighting schemes remains future work.
A.5 LATENT DISTRIBUTION NON-IID MOTIVATION AND SETUP
The latent-distribution setting models realistic client heterogeneity through feature and semantic variation, using learned image embeddings to construct client distributions.
- A.5 LATENT DISTRIBUTION NON-IID MOTIVATION AND SETUP: The authors caution that pathological non-IID setups may be useful experimentally but may not represent realistic or frequent data-partitioning patterns.They illustrate alternatives where location explains road-image variation or music combinations define meaningful groups.
- A.5 LATENT DISTRIBUTION NON-IID MOTIVATION AND SETUP: The latent-distribution method partitions datasets according to feature and semantic similarity before sampling client data distributions.
- A.5 LATENT DISTRIBUTION NON-IID MOTIVATION AND SETUP: Image representations come from a VGG-11 network trained to at least 85% classification accuracy, using 4096-dimensional second-layer embeddings.
- A.5 LATENT DISTRIBUTION NON-IID MOTIVATION AND SETUP: For CIFAR-10/100 and MNIST, training uses SGD with specified learning-rate, decay, momentum, and weight-decay settings; FedFomo uses 5 or 10 downloads per client.
A.7 ADDITIONAL DESIGN ABLATIONS
The ablation examines how clients should divide local data between training and validation for personalized updates. Validation fractions below one-half converge faster and reach higher accuracy, while larger fractions reduce performance.
- Size of the validation split: FedFomo requires a local validation split that reflects the client’s objective or target test distribution.The split determines how much data supports local training versus evaluating received models for personalized weighting.
- Size of the validation split: Clients average 3333 local data points, and final accuracy with standard deviation is reported over five runs.
- Size of the validation split: Figure 9 plots in-distribution accuracy over the validation split ratio.
- Size of the validation split: Using under half of local data for validation produces faster convergence and higher accuracy, while larger validation allocations cause a notable drop-off.The authors attribute this decline to less data remaining for local model training; performance also decreases slightly from fractions 0.05 to 0.1.