PaperlayerView exact revision on arXiv

Source-linked AI summary

Federated Reconstruction: Partially Local Federated Learning

Karan Singhal, Hakim Sidahmed, Zachary Garrett, Shanshan Wu, Keith Rush, Sushant Prakash

arXiv:2102.03448v6cs.LGcs.DC

TL;DR

Federated learning personalization is constrained by privacy, communication, client heterogeneity, and the impracticality of stateful clients at cross-device scale. Federated Reconstruction addresses these constraints by training global parameters while reconstructing local parameters on demand, outperforming standard baselines on unseen clients and enabling fast personalization with less communication.

  • Problem

    Fully global personalization can expose sensitive user-specific parameters, while prior partially local approaches may require stateful or always-available clients that are impractical at cross-device scale.

  • Method

    Federated Reconstruction partitions models into global and local parameters, trains global parameters federatively, and reconstructs local parameters on demand from client data.

  • Results

    FEDRECON outperforms standard centralized and federated training on unseen clients, enables fast personalization with less communication, and supports inference for unseen clients.

  • Takeaways & Limitations

    Federated Reconstruction provides model-agnostic partially local federated learning suitable for training and inference at scale, including a successful mobile-keyboard deployment.

  • Takeaways & Limitations

    The deployment discussion identifies fairness implications as an area for future work.

Abstract

from arXiv · show

Personalization methods in federated learning aim to balance the benefits of federated and local training for data availability, communication cost, and robustness to client heterogeneity. Approaches that require clients to communicate all model parameters can be undesirable due to privacy and communication constraints. Other approaches require always-available or stateful clients, impractical in large-scale cross-device settings. We introduce Federated Reconstruction, the first model-agnostic framework for partially local federated learning suitable for training and inference at scale. We motivate the framework via a connection to model-agnostic meta learning, empirically demonstrate its performance over existing approaches for collaborative filtering and next word prediction, and release an open-source library for evaluating approaches in this setting. We also describe the successful deployment of this approach at scale for federated collaborative filtering in a mobile keyboard application.

1 Introduction

Federated Reconstruction addresses privacy, communication, heterogeneity, and scalability limits of fully global or stateful personalized federated learning by partitioning models into global and local components. It reconstructs local parameters on demand and reports strong performance, fast personalization, reduced communication, and production deployment.

  • Motivation: Fully global training can underperform local models for heterogeneous clients and expose sensitive user-specific parameters.These concerns are especially acute for collaborative filtering models with user embeddings.
  • Motivation: Stateful or always-available-client approaches are impractical for large-scale cross-device settings with unreliable, partial participation.Clients may be sampled infrequently, making stored local state stale or unavailable.
  • Approach: Federated Reconstruction combines federated training of global parameters with local reconstruction to support model-agnostic, scalable, inferential, and fast personalization.The stated criteria include compatibility with partial participation, inference for new clients, and rapid adaptation to personal data.
  • Results: FEDRECON outperforms standard centralized and federated training on unseen clients while enabling fast adaptation and matching other personalization methods with less communication.The evaluations cover collaborative filtering and next word prediction.
  • Deployment and resources: The approach was successfully deployed at scale for federated collaborative filtering in a mobile keyboard application.The paper also releases an open-source evaluation library for partially local federated learning.

2 Related Work

Prior personalization and partial-model approaches reduce some federated-learning limitations, but they generally communicate all client parameters or require persistent client state. Federated Reconstruction is motivated as a way to aggregate only part of a model without those deployment constraints.

  • Existing approaches: Existing personalization approaches include finetuning, meta learning, bi-level optimization, model interpolation, heterogeneous-gradient methods, and resource-heterogeneity methods.These approaches target personalization, convergence, or client resource variation through different mechanisms.
  • Existing approaches: Many existing approaches require communicating all client parameters, creating privacy and communication concerns for some models.This motivates methods that aggregate only a subset of model parameters.
  • Partial-model methods: Earlier partial-model methods require clients to maintain local models across rounds and may assume continuous client availability.Those assumptions are problematic when clients are sampled from large populations.
  • Collaborative filtering: Federated collaborative-filtering methods can keep user factors local or protect aggregation cryptographically, but they require client state unlike Federated Reconstruction.The cited approaches focus on federated matrix factorization and related formulations.
  • Meta learning: Federated Reconstruction is connected to meta learning, while prior work related federated personalization and FEDAVG to MAML and REPTILE.The connection provides conceptual context for training models that personalize efficiently.

3 Partially Local Federated Learning

Partially local federated learning partitions model parameters so sensitive or heterogeneous components remain on clients while other parameters are federated. This design addresses privacy, communication, and state-management problems in large, changing cross-device populations.

  • Federated objective: Standard federated learning optimizes an expected client objective over a distribution of heterogeneous clients.The objective uses client-specific losses f_i(x) and samples from client distribution P.
  • Motivating model: Matrix factorization separates user-specific parameters from item parameters, making it a natural partially local collaborative-filtering model.The ratings matrix is factorized into user and item matrices.
  • Privacy and communication: Aggregating all parameters can expose personal user embeddings or require unreasonable communication, even when raw data remain decentralized.Sharing individual user parameters can reveal sensitive user-item affinities.
  • Cross-device constraints: Stateful partial-model methods are unsuitable when clients are rarely resampled because local parameters become unavailable or stale.Large-scale cross-device systems have partial availability and a constantly changing client population.
  • Broader scope: Partially local learning can also improve robustness to heterogeneous objectives and reduce communication for models without user-specific parameters.The local-global partition is determined by the use case and communication limitations.
  • Broader scope: A key contribution is practical partially local learning for large, changing client distributions with stateless clients.This targets cross-device training under partial participation.

4 Federated Reconstruction

Federated Reconstruction trains global parameters while reconstructing local parameters whenever a client participates, allowing clients to discard local state afterward. Its meta-learning formulation targets fast reconstruction and supports inference for unseen clients.

  • Framework: The framework relaxes client statefulness because reconstructed local parameters can be discarded after participation.This avoids requiring clients to preserve local parameters across rounds.
  • Training procedure: Clients reconstruct local parameters from current global parameters and support data, then update only the global parameters using query data.The server aggregates the resulting global updates across selected clients.
  • Training procedure: Support and query sets are typically disjoint so reconstruction and global updating promote meta-generalization.The paper notes that this assumption can be relaxed when clients lack enough data to partition.
  • Algorithm: FEDRECON uses local gradient descent for reconstruction and global updates, with weighted server aggregation of client updates.The server weighting uses each client’s query-set size n_i.
  • Inference: RECONEVAL reconstructs local models offline for inference, including on clients unseen during training.Clients can store reconstructed parameters for repeated use and optionally refresh them with new data.
  • Meta-learning connection: The meta-learning view treats global parameters as meta-parameters and local parameters as task-specific parameters adapted from support data.Reconstruction corresponds to the inner loop and global updating to the outer loop.
  • Meta-learning connection: FEDRECON trains global parameters for fast reconstruction without requiring persistent client state, and one reconstruction step can recover most remaining performance.The paper reports this fast-reconstruction behavior empirically on unseen clients.

5 Experimental Evaluation

The evaluation tests FEDRECON on MovieLens matrix factorization and Stack Overflow next word prediction, including unseen-client reconstruction, communication, personalization, and step-count ablations. Across these experiments, FEDRECON matches or improves relevant baselines while supporting local-parameter reconstruction without communicating those parameters.

  • MovieLens matrix factorization: FEDRECON uses local user embeddings and global item parameters for MovieLens, evaluating both seen users with STANDARDEVAL and held-out users with RECONEVAL.The model reports RMSE and rating prediction accuracy against centralized training and FEDAVG.
  • MovieLens matrix factorization: FEDRECON matches CENTRALIZED + STANDARD EVAL on MovieLens despite reconstructing embeddings for held-out users in the more difficult RECONEVAL setting.Reconstructing embeddings for unseen users performs significantly worse for centralized training than FEDRECON and server evaluation on seen users.
  • Stack Overflow next word prediction: 500 local OOV embeddings significantly boost FEDRECON accuracy over FEDYOGI for every tested Stack Overflow vocabulary size.FEDRECON accuracy improves with smaller vocabularies in this setting, unlike FEDYOGI and FEDRECON with one OOV embedding; the passage attributes this to more training data for the local component.
  • Stack Overflow next word prediction: FEDRECON reconstructs local embeddings with performance similar to OOV FINETUNING while avoiding communication of the local parameters.FULL FINETUNING performs better because all parameters are finetuned; FEDRECON+FINETUNE obtains comparable results to the corresponding personalization comparison.
  • Ablations: One reconstruction step recovers the majority of remaining performance across MovieLens and Stack Overflow, whereas zero steps leave MovieLens accuracy at 0.0.Stack Overflow relative accuracy remains above 90% without reconstruction, while MovieLens user embeddings are randomly initialized.
  • Ablations: One client update step recovers almost all accuracy, while additional steps gradually improve accuracy and can accelerate training compared with the fixed-round ku = 1 setting.The ku = 1 procedure corresponds to a first-order meta-learning objective, but its performance was still increasing after the fixed number of rounds.

6 Open-Source Library

The paper releases an Apache-2.0 code framework built on TensorFlow Federated for expressing and evaluating practical partially local federated models. It supports reproducing the experiments and applying the setting to user-selected models and tasks.

  • Open-source framework: The open-source framework uses TensorFlow Federated and provides a flexible, documented interface for simulations with models and tasks of users’ choice.Existing Keras models can be integrated into the framework, which is released under Apache License 2.0.

7 Deployment in a Mobile Keyboard Application

FEDRECON was deployed for federated collaborative filtering in a mobile keyboard application with hundreds of millions of clients. Production challenges included sparse data and client dropout, while expression recommendations increased click-through rate by 29.3%.

  • FEDRECON was deployed for federated collaborative filtering in a mobile keyboard application with hundreds of millions of clients.
  • Users’ heterogeneous expression-sharing styles made collaborative filtering suitable for predicting expressions they might share.
  • Sparse client data prompted oversampling and filtering contributions from clients with too few examples.
  • Client resource and availability heterogeneity caused some participating clients to drop out before sending updates.
  • The authors identify fairness implications of ignoring dropped-client updates as a valuable area for future work.
  • 29.3% increase in click-through rate followed deployment of the resulting model for expression recommendations.

8 Conclusion

The paper presents Federated Reconstruction as a model-agnostic framework for fast partially local federated learning at training and inference scale. It connects FEDRECON to meta learning, validates it across tasks, releases an open-source library, and reports a production deployment.

  • Federated Reconstruction is a model-agnostic framework for fast partially local federated learning suitable for training and inference at scale.
  • FEDRECON improves performance on unseen clients and enables fast personalization with less communication across collaborative filtering and next message prediction.
  • The work releases an open-source library and describes a successful production deployment.
  • Future work includes studying the balance of local and global parameters and applying differential privacy to global parameters.

A Frequently Asked Questions

The FAQ explains why FEDRECON uses stateless reconstruction rather than relying on remembered client parameters, and how its global-local split supports privacy and communication goals. It distinguishes the method from prior partially local approaches and reports empirical rather than theoretical convergence evidence.

  • Cross-device training samples only a small fraction of millions of clients, so most clients lack usable state from previous rounds.
  • Stateful algorithms can suffer stale state, while non-sampled clients would lack trained local parameters in partially local federated learning.
  • Reconstruction is practical because clients are unlikely to be reached repeatedly and lets unseen clients create partially local models offline.
  • FEDRECON differs from prior work by addressing partially local learning with a stateless algorithm for large-scale cross-device settings.
  • The authors report empirical convergence and a 29.3% recommendation click-through-rate improvement, but no additional theoretical convergence results.
  • User-specific or privacy-sensitive parameters can remain local, while making more variables local can reduce communication cost.

B Proof of Connection to Meta Learning

The meta-learning connection treats global parameters as meta-parameters optimized for rapid reconstruction of client-specific parameters. FEDRECON uses first-order updates that omit costly second-order terms, and the evaluation applies reconstruction to unseen users in matrix factorization and next-word prediction.

  • Meta-learning formulation: FEDRECON treats global parameters g as meta-parameters and client-specific parameters l_i as task-specific parameters.
  • Meta-learning formulation: The framework minimizes an objective evaluated after reconstructing local parameters from client data.
  • First-order approximation: FEDRECON’s reconstruction and update steps correspond to first-order meta learning when reconstruction uses gradient descent on support data and the update uses query data.
  • First-order approximation: Treating reconstruction output as constant drops Jacobian-gradient terms, avoiding computation and memory quadratic in the number of local parameters.
  • Empirical connection: One reconstruction step is sufficient to recover useful local parameters, supporting fast personalization on unseen clients.
  • Evaluation: Matrix-factorization evaluation tests standard performance on seen users and reconstruction-based generalization to unseen users.
  • Evaluation: For matrix factorization, FEDRECON keeps user embeddings local and the item matrix global, using reconstruction on support data.
  • Evaluation: For next-word prediction, FEDRECON keeps out-of-vocabulary embeddings local while the remaining model is global.

D Additional Empirical Results

Additional experiments show that FEDRECON remains effective across training variants, server optimizers, vocabulary configurations, and unseen-user evaluation settings.

  • FEDRECON results are reported for matrix factorization and next word prediction, including ablations and additional comparisons.The tables cover both tasks and include previously discussed results for ease of comparison.
  • Sharing Support and Query Data: Using the full client dataset for both support and query causes only a minimal performance drop across MovieLens and Stack Overflow.This relaxation also helps combat data sparsity in the real-world deployment.
  • Joint Training after Reconstruction: Joint training after reconstruction produces similar but slightly degraded performance compared with alternating updates.Freezing local parameters during global training helps ensure that global parameters are updated significantly.
  • Adaptive Optimizers: Different server optimizers converge well for MovieLens, and ADAGRAD improves FEDRECON results; Stack Overflow uses Yogi, which outperforms SGD.FEDRECON treats aggregated global updates as an antigradient that can be supplied to server optimizers.
  • Stack Overflow Vocabulary and OOV Sizes: Lowering core vocabulary coverage reduces Stack Overflow loss, while adding out-of-vocabulary buckets improves accuracy most for smaller core vocabularies.Performance plateaus around 500 out-of-vocabulary buckets for all vocabulary sizes.
  • MovieLens Centralized Evaluation on Unseen Users: RECONEVAL produces reasonable unseen-user results, and FEDRECON outperforms centralized training because its global parameters support reconstruction.Randomly initialized user embeddings perform poorly under standard evaluation.

E Privacy Implications and Limitations

FEDRECON reduces the need to centralize privacy-sensitive parameters, but it still sends gradients to a central server and offers no formal leakage guarantee.

  • FEDRECON keeps privacy-sensitive parameters on clients while providing a partially local alternative to centralized and fully global federated training.The paper also releases the deployment framework to support partially local federated learning.
  • Like other federated algorithms, FEDRECON sends gradient updates to a central server, which can leak client information or reconstruct training data.The paper notes that leakage has been observed in realistic settings, not only simple one-example cases.
  • Communicating only global-parameter updates may provide some protection against gradient leakage, but vanilla FEDRECON offers no formal guarantee.The authors suggest differential privacy or secure aggregation for provable guarantees in highly privacy-sensitive applications.
  • Differentially private training typically clips and noises gradients, with performance depending on the dimensionality of the model representation.The discussion relates this dimensionality dependence to regret-based analyses of private optimization.
  • FEDRECON can reduce the dimensionality of aggregated parameters, lowering the variance of noise needed for user-level differential privacy.The paper presents this as a systems-focused application rather than a new differential privacy algorithm.
Loading 2102.03448v6…