PaperlayerView exact revision on arXiv

Source-linked AI summary

Federated Learning with Partial Model Personalization

Krishna Pillutla, Kshitiz Malik, Abdelrahman Mohamed, Michael Rabbat, Maziar Sanjabi, Lin Xiao

arXiv:2204.03809v2cs.LGcs.DCmath.OC

TL;DR

The paper addresses limited convergence understanding for simultaneous and alternating federated updates in partially personalized models, especially under nonconvex objectives and partial participation. It analyzes both algorithms and evaluates them across image, text, and speech tasks, finding that small personal components capture most full-personalization benefits and that alternating updates usually have a small but consistent edge.

  • Problem

    Convergence properties of simultaneous and especially alternating partial-personalization algorithms remain incompletely understood under general nonconvex objectives with partial participation.

  • Method

    The paper analyzes FedSim and FedAlt under partial participation, using virtual full participation to decouple dependent random variables in the FedAlt analysis.

  • Results

    Across realistic image, text, and speech tasks, partial personalization obtains most full-personalization benefits with few personal parameters, and FedAlt usually outperforms FedSim by a small consistent margin.

  • Takeaways & Limitations

    Partial personalization can approach full personalization while using a small personalized parameter subset, and alternating updates are often preferable to simultaneous updates by a modest margin.

  • Takeaways & Limitations

    Personalization can improve average performance while reducing test accuracy for some devices, a fairness and performance issue that remains unresolved.

Abstract

from arXiv · show

We consider two federated learning algorithms for training partially personalized models, where the shared and personal parameters are updated either simultaneously or alternately on the devices. Both algorithms have been proposed in the literature, but their convergence properties are not fully understood, especially for the alternating variant. We provide convergence analyses of both algorithms in the general nonconvex setting with partial participation and delineate the regime where one dominates the other. Our experiments on real-world image, text, and speech datasets demonstrate that (a) partial personalization can obtain most of the benefits of full model personalization with a small fraction of personal parameters, and, (b) the alternating update algorithm often outperforms the simultaneous update algorithm by a small but consistent margin.

1 Introduction

The paper studies partial model personalization in federated learning as an alternative to one-model-fits-all and memory-intensive full personalization. It analyzes simultaneous and alternating updates theoretically and evaluates their empirical benefits and limitations.

  • Motivation: Non-identical device data can make one-model-fits-all federated learning ineffective or undesirable.The motivation is that devices may have diverse user characteristics and data distributions.
  • Motivation: Full model personalization can require twice the full model’s memory footprint, limiting the size of trainable models.Each device stores both personalized parameters and a reference model.
  • Problem formulation: Partial personalization partitions parameters into shared u and device-specific v_i components, whose dimensions may differ across devices.This permits personalized components to use different parameter counts or architectures.
  • Algorithms: FedSim updates shared and personal parameters simultaneously, whereas FedAlt updates them alternately under partial participation.Selected devices retain updated personal parameters locally while sending only updated shared parameters for server aggregation.
  • Contributions: The paper proves convergence guarantees for both algorithms in smooth nonconvex settings with partial participation.For FedAlt, the analysis decouples dependent random variables using virtual full participation.
  • Empirical findings: Experiments find that partial personalization captures most benefits of full personalization with few personal parameters, while FedAlt consistently outperforms FedSim by a small margin.Personalization can nevertheless worsen performance for some devices even when average performance improves.

2 Partially Personalized Models

Partial model personalization assigns selected model components to each device while keeping the remaining parameters shared. The paper describes input/output partitioning, adapters, and generalized additive models as concrete architectures.

  • Layer partitioning: Modern multilayer models can personalize either input layers or output layers while sharing the remaining layers.Lower layers generally support feature extraction, whereas upper layers focus on complex pattern recognition.
  • Layer partitioning: Input-layer partitioning can split the input horizontally into shared and personal branches whose outputs are concatenated before upper layers.Personal embeddings for user-specific features remain local to the device.
  • Adapters: Adapter personalization adds device-specific adapter modules to a shared model while treating the rest of the parameters as shared.The described adapter uses a skip connection and a middle bottleneck.
  • Generalized additive models: A generalized additive model fuses predictions from separate shared and personalized models to generate the final prediction.For classification, the outputs can be combined through a learnable convex combination.
  • General formulation: Full model personalization is a special case of the paper’s shared-personal formulation.The reduction identifies the shared component with a reference model and the personal component with each device’s model.

3 Algorithms and Convergence Analysis

The paper analyzes FedAlt and FedSim under smooth nonconvex objectives with partial device participation, addressing the dependent random variables created by FedAlt’s alternating updates. It establishes convergence results for both methods and compares the regimes in which their rates differ.

  • Algorithms: FedAlt updates personal parameters first and shared parameters second, whereas FedSim updates both parameter groups simultaneously during local training.Both methods use the same outer federated protocol, including random device sampling, broadcast of shared parameters, local updates, and aggregation of updated shared parameters.
  • Assumptions: The convergence analysis assumes continuously differentiable losses with Lipschitz gradients, unbiased stochastic gradients with bounded variance, and bounded partial-gradient diversity.These assumptions cover smoothness within each parameter block, cross-sensitivity between blocks, stochastic-gradient noise, and participation-related variation.
  • Convergence Guarantees: FedAlt and FedSim both converge in expectation to a stationary point under the stated smooth nonconvex assumptions, with rates involving effective noise variances and a standard T^-1/2 term.The FedAlt rate depends on problem parameters, stochasticity from local gradients and participation, and coupling between shared and personal parameters; FedSim has an analogous bound with different noise terms.
  • Challenges of FedAlt: FedAlt’s u-update is difficult to analyze because the updated personal variables V (t+1) and shared variables u(t+1) both depend on the same device sampling S(t).Consequently, the expectation over S(t) cannot pass through the relevant inner product as in standard stochastic-gradient analyses.
  • Virtual Full Participation: Virtual full participation replaces the sampled personal update with an analysis-only iterate computed as if every device participated, decoupling it from S(t).Conditioned on the current iterates, the virtual personal iterate is deterministic, allowing expectation over the sampling to be taken over the shared update alone; the replacement introduces an error term.
  • Rate Comparison: FedAlt dominates FedSim in a regime where the shared-gradient noise is approximately zero, improving the rate by a factor of (1+χ^2)^1/2.This regime can arise from using a large or full batch on a small number of samples per device, and experiments support its practical relevance.

4 Experiments

Experiments across image, text, and speech tasks show that carefully chosen partial personalization can approach or exceed full personalization with fewer personal parameters, while FedAlt generally edges out FedSim. Personalization gains vary by task and device, and can reduce test accuracy for some low-data devices despite improving average performance.

  • Partial versus full personalization: Partial personalization nearly matches full personalization and can sometimes outperform it across the evaluated tasks.On StackOverflow, output-layer personalization reaches 25.05% versus 25.21% for full personalization; adapter personalization matches full personalization on EMNIST and exceeds it on GLDv2 by 3.5pp.
  • Parameter efficiency: 1.2% of parameters personalized with adapters captures 72% of StackOverflow’s full-personalization accuracy boost.Personalizing 14% of parameters through the output layer raises this to nearly 90%; adapters personalize 11.5–12.5% on EMNIST and GLDv2 while matching or exceeding full personalization.
  • Task-dependent partitioning: The best personalized architecture is model and task dependent: output layers lead StackOverflow, adapters lead GLDv2 and EMNIST, and input layers lead speech recognition.The authors relate these choices to the primary source of cross-device heterogeneity, such as user-specific next-word preferences or speaker-specific audio characteristics.
  • FedAlt versus FedSim: FedAlt almost always outperforms FedSim by a small margin, including 0.08pp for StackOverflow adapters and 0.3pp for GLDv2 input-layer personalization.The FedAlt–FedSim difference exceeds run-to-run standard deviations in the reported comparisons; the authors recommend FedAlt as the default, while noting FedSim remains competitive.
  • Per-device generalization: Personalization improves training accuracy across devices but reduces test accuracy for some devices, especially among devices with relatively few training samples.Increasing regularization can reduce per-device accuracy spread but degrades average accuracy; weight decay and dropout do not mitigate the issue.

5 Discussion

Partial model personalization offers memory, communication, and potential privacy advantages, but FedAlt requires special analysis because partial participation creates dependent random variables. Virtual full participation restores an analyzable expectation structure while adding lower-order error terms.

  • Advantages: Partial personalization communicates only shared parameters, reducing transmitted information and potentially protecting sensitive features kept in personal components.The passage presents privacy implications as intuitive and speculative, including the possibility that less noise may be needed for differential privacy.
  • FedAlt analysis: FedAlt is difficult to analyze under partial participation because the updated personal parameters and shared parameters depend on the same sampled devices.This prevents directly moving the sampling expectation through the relevant first-order inner product.
  • FedAlt analysis: Virtual full participation replaces the sampled personal update with an independent virtual iterate, allowing the sampling expectation to be taken over the shared update.The replacement introduces an error term that is bounded using Young’s inequality and smoothness.
  • FedAlt analysis: The resulting virtual-iterate error is lower order, with the analysis reporting O(γ^2) in the single-local-update setting and O(γuγv^2) for a related translation.The bounds rely on smoothness, bounded stochastic-gradient variance, and partial gradient diversity.

A.3 Convergence Analysis of FedAlt

The FedAlt analysis handles alternating personal and shared updates under partial participation by decoupling their sampling dependence with virtual personal iterates. Under smoothness, variance, and gradient-diversity assumptions, this yields convergence bounds with controlled lower-order errors.

  • Convergence result: Theorem 3 establishes FedAlt convergence under smoothness, bounded variance, and partial gradient diversity assumptions with learning rates γu = η/(Luτu) and γv = η/(Lvτv).The theorem requires learning-rate conditions in addition to the stated assumptions.
  • Algorithm: FedAlt alternates local updates of personal parameters and shared parameters within each communication round.The algorithm uses separate local-update counts and step sizes for the two parameter blocks.
  • Challenge: Partial participation makes FedAlt’s shared-update analysis difficult because personal and shared iterates depend on the same device sample.The resulting first-order term is not directly unbiased with respect to client sampling.
  • Proof strategy: Virtual personal iterates remove this dependence from the shared-update first-order term, after which the remaining replacement error is bounded using smoothness and Young’s inequality.The virtual iterates are defined as if every device had participated and are independent of the sampling set.
  • Convergence result: The final FedAlt rate follows by combining descent bounds for the alternating blocks, virtual-iterate error control, and telescoping over communication rounds.The analysis also gives an optimized-learning-rate corollary and an asymptotic 1/T regime.

A.4 Convergence Analysis of FedSim

The FedSim analysis treats shared and personal parameters as simultaneous updates and derives non-asymptotic convergence under partial participation. Its proof uses block smoothness, unbiased sampling arguments, bounded variance, and gradient diversity.

  • Algorithm: FedSim jointly updates shared and personal parameters during each local-update sequence.The analysis uses separate shared and personal learning rates while applying the updates simultaneously.
  • Convergence result: Theorem 11 establishes FedSim convergence under smoothness, bounded variance, and partial gradient diversity assumptions with γu = η/(Luτ) and γv = η/(Lvτ).The theorem includes additional learning-rate conditions.
  • Convergence result: The final FedSim rate is stated in Corollary 12 using optimized learning rates, with an asymptotic 1/T regime under the corresponding conditions.The analysis also notes that T ≥ Ω(n/m) is necessary for every device to be sampled at least once on average.
  • Proof strategy: FedSim avoids FedAlt’s alternating-update dependence challenge by analyzing simultaneous block updates with virtual sequences independent of device selection.The proof bounds the resulting terms through block smoothness and sampling arguments.
  • Convergence result: The FedSim proof controls O(γu^2 + γv^2) noise terms and telescopes the resulting descent condition to obtain the final bound.The analysis separately bounds terms associated with personal and shared updates.

B Experiments: Detailed Setup and Hyperparameters

The experiments cover four real-world federated tasks spanning image, text, and speech modalities, using naturally non-i.i.d. client data. The setup reports task metrics, baselines, hyperparameters, and client-level data distributions.

  • Experimental scope: The experiments use four datasets from image, text, and speech modalities with natural non-i.i.d. client splits.The authors state that the splits reflect data heterogeneity encountered in federated learning.
  • Experimental scope: The experimental section specifies task datasets and metrics, the comparison pipeline and baselines, and algorithm hyperparameters.Device weights αk are proportional to the number of datapoints available on each device.
  • Data distribution: Figure 6 shows the distribution of training samples per device for each experimental task.For GLDv2, the displayed distribution omits a long right tail whose maximum is 1000 data points per device.

B.1 Datasets, Tasks and Models

The experiments cover image, text, character-recognition, and speech-to-text tasks, using task-specific models and partial personalization architectures. These architectures personalize input layers, output layers, or lightweight adapters while sharing the remaining model components.

  • Partial Personalization Architectures: Input layer personalization keeps the first layer client-specific while sharing the rest of the model.For the next-word transformer, the first transformer layer replaces the word embedding layer because of its size.
  • Partial Personalization Architectures: Output layer personalization shares a global representation while personalizing the prediction layer.Large transformer components are substituted for the prediction layer in the next-word and speech-to-text models.
  • Partial Personalization Architectures: Adapter personalization shares the full model and adds lightweight client-specific modules between its layers.The experiments use transformer adapters and residual adapters.
  • Datasets and Tasks: The study evaluates four federated tasks: next-word prediction, visual landmark recognition, character recognition, and speech-to-text.The tasks use StackOverflow, GLDv2, EMNIST, and LibriSpeech datasets, respectively.
  • Models: The models include a 6-million-parameter transformer for StackOverflow and ResNet-18 models with 12 million and 11 million parameters for GLDv2 and EMNIST.The GLDv2 model uses ImageNet pretraining, while the EMNIST model is modified for 28 × 28 grayscale inputs.

B.2 Experimental Pipeline and Baselines

The experiments use a staged federated pipeline: non-personalized training, optional personalized federated training, and final local finetuning. Partial personalization is compared with non-personalized training and three full-model personalization baselines.

  • Experimental Pipeline: The training pipeline begins with non-personalized federated training using a one-model-fits-all global model.This stage uses FedAvg variants before any personalization is introduced.
  • Experimental Pipeline: Personalized federated training jointly trains shared and personal parameters, initialized from the non-personalized global model.This optional stage uses a personalized federated learning approach and warm-starts both parameter groups.
  • Experimental Pipeline: Final finetuning updates only the personalized parameters while leaving the shared parameters fixed.For input/output personalization, personal parameters start from the corresponding global-model components; adapters use shared random initialization across devices.
  • Baselines: Full-model baselines are Finetune, Ditto with ℓ2 regularization, and pFedMe, while partial personalization uses FedAlt or FedSim.Each partial architecture follows personalized federated training and final local finetuning.

B.3 Hyperparameters and Evaluation Details

Evaluation uses weighted client-level performance with task-specific metrics, fixed training schedules, and validation-based hyperparameter tuning. Training partial personalization requires less memory than full personalization in the reported settings.

  • Hyperparameters: Hyperparameters are tuned on validation sets formed by holding out 20% of each device’s training data.After tuning, experiments are rerun using the full training data.
  • Evaluation: The primary metric is client-test accuracy weighted by each client’s number of test examples; speech recognition uses weighted word error rate.The weighted accuracy corresponds to pooling all client data locally.
  • Training Configuration: Personalized training runs for 500 rounds on StackOverflow and EMNIST and 600 rounds on GLDv2 after non-personalized warm-start training.Non-personalized training uses 1000, 2500, and 2000 rounds for those datasets, respectively.
  • Training Configuration: Training uses 50 devices per round for StackOverflow and GLDv2 and 10 for EMNIST, with uniform random selection.Each selected device runs one local epoch with minibatches of 64 or 32, depending on the task.
  • Memory Requirement: 34% to 42% lower memory consumption is reported for partial personalization across the evaluated models and datasets.During deployment, partial and full personalization have the same memory footprint because one full model is deployed.
  • Reported Results: Input layer personalization with FedAlt achieves the smallest word error rate among the speech-recognition models considered.Speech recognition comparisons use WER, where smaller values indicate better predictive performance.

C.2 Ablation: Final Finetuning for FedAlt and FedSim

Final local finetuning has little effect on partial personalization overall, although it benefits FedSim more than FedAlt. FedAlt nevertheless remains ahead of FedSim after finetuning.

  • Per-device Effects: The StackOverflow personalization distributions compare per-device train and test accuracy changes, including regularized adapter and output-layer variants.The figure distinguishes adapter personalization with ℓ2 regularization from output-layer personalization with dropout.
  • Effect of Final Finetuning: Final finetuning contributes less than 0.5 percentage points on GLDv2, compared with 5 to 15 percentage points from personalization overall.For FedAlt on StackOverflow, the improvement from final finetuning is close to zero.
  • FedAlt versus FedSim: Final finetuning is more important for FedSim than for FedAlt.The ablation reports that finetuning helps FedSim more, but does not close its performance gap with FedAlt.

C.3 Effect of Personalization on Per-Device Generalization

Personalization consistently improves training accuracy but can reduce test accuracy for some devices. Neither regularization nor dropout eliminates this per-device generalization issue, and the outcome can depend on the random seed.

  • Regularization and dropout: Regularization tuned for best average accuracy leaves nearly the same test-accuracy distribution, while stronger regularization narrows spread at the cost of a lower median.Dropout has a similar effect, with best dropout improving the median without mitigating all per-device degradation.
  • Per-device generalization: Personalization improves training accuracy on all devices but reduces test accuracy on some devices.This pattern appears for both partial and full personalization and across StackOverflow and GLDv2.
  • Per-device generalization: Personalization sometimes helps or hurts the same device depending on the random seed.The recommended practical response is A/B testing personalized and non-personalized models after deployment.
  • Regularization: Best-tuned weight decay matches the unregularized case in per-device statistics and does not prevent personalization from hurting individual devices.Increasing the regularization weight reduces the spread of per-device accuracy but worsens mean accuracy.
  • Dropout: Best dropout increases average accuracy by 0.18 pp and reduces hurt devices from 256 out of 1000 to 193, but does not eliminate the issue.Increasing dropout further degrades per-device statistics.

C.4 Partial Personalization for Stateless Devices

Stateless-device variants reinitialize personalized parameters whenever a device participates, with FedAlt showing mild improvements over FedSim. Final local finetuning supplies most of the stateless-device gains, while personalized federated training still contributes 1 to 1.34pp.

  • Stateless-device setup: Stateless-device variants reinitialize personalized parameters at each participation round instead of retaining them across rounds.The methods warm-start these parameters from the non-personalized model; adapter initializations are fixed and reused.
  • Algorithm comparison: FedAlt is better than FedSim for stateless devices, although the improvement is smaller and all algorithms perform similarly overall.FedAlt obtains mild improvements over both FedSim and finetuning for GLDv2, including 0.24pp with adapters.
  • Contribution of finetuning: For GLDv2, final finetuning contributes 11.68 and 10.42pp of total gains of 12.67 and 11.76pp for FedAlt and FedSim, respectively.These values are reported for the stateless setting and are measured over five random seeds.
  • Contribution of federated training: Personalized federated training still increases accuracy by 1 to 1.34pp in the stateless setting.Thus, finetuning accounts for most of the improvement without accounting for all of it.
Loading 2204.03809v2…