PaperlayerView exact revision on arXiv

Source-linked AI summary

Personalized Federated Learning through Local Memorization

Othmane Marfoq, Giovanni Neglia, Laetitia Kameni, Richard Vidal

arXiv:2111.09360v3cs.LGstat.ML

TL;DR

Heterogeneous client data can make a single global model suboptimal for personalized federated learning. kNN-Per interpolates a collectively trained global model with a client-local kNN model based on the global representation; the authors report higher accuracy and fairness than state-of-the-art methods.

  • Problem

    Clients’ differing data distributions can make a single global model suboptimal and may leave some clients with arbitrarily bad predictions.

    Source: §1 · p. 1

  • Method

    kNN-Per interpolates a collectively trained global model with a local kNN model using the representation provided by the global model.

    Source: §1 · p. 1, §3 · p. 4

  • Results

    The authors report significantly higher accuracy and fairness than state-of-the-art methods on federated benchmarks.

    Source: §1 · p. 1

  • Takeaways & Limitations

    Local memorization can use a global model’s representation to improve client-specific predictions, and local datastores can be updated as distributions change without retraining the global model.

    Source: §1 · p. 1, §7 · p. 9

    The reported advantage over FedRep and pFedGP may reflect conflicting goals from jointly learning shared representations and local models, but the interaction requires further study.

    Source: §7 · p. 9

Abstract

from arXiv · show

Federated learning allows clients to collaboratively learn statistical models while keeping their data local. Federated learning was originally used to train a unique global model to be served to all clients, but this approach might be sub-optimal when clients' local data distributions are heterogeneous. In order to tackle this limitation, recent personalized federated learning methods train a separate model for each client while still leveraging the knowledge available at other clients. In this work, we exploit the ability of deep neural networks to extract high quality vectorial representations (embeddings) from non-tabular data, e.g., images and text, to propose a personalization mechanism based on local memorization. Personalization is obtained by interpolating a collectively trained global model with a local $k$-nearest neighbors (kNN) model based on the shared representation provided by the global model. We provide generalization bounds for the proposed approach in the case of binary classification, and we show on a suite of federated datasets that this approach achieves significantly higher accuracy and fairness than state-of-the-art methods.

1. Introduction

Federated clients differ in their data distributions and resources, challenging the use of one common global model. The paper proposes kNN-Per, which combines a collectively trained global model with local memorization, and reports higher accuracy and fairness than state-of-the-art methods.

1. Introduction

  • Clients differ in local dataset size and distribution as well as storage and computational capabilities, challenging the assumption that all should train a common global model.Under statistical heterogeneity, a global model may be arbitrarily bad for some clients, raising fairness concerns.Source: §1 · p. 1
  • kNN-Per interpolates a collectively trained global model with a client-local kNN model using the global model’s shared representation.Local memorization can capture client-specific distribution shift, and the experiments report greater benefit when that shift is larger.Source: §1 · p. 1
  • Clients can update their local datastore as data distributions change without retraining the global model, and tune their local kNN to their resource capabilities.Personalization occurs after communication exchanges, and the method can be combined with differential privacy techniques.Source: §1 · p. 1
  • kNN-Per achieves significantly higher accuracy and fairness than state-of-the-art methods on federated benchmarks.The authors also provide generalization bounds for binary classification.Source: §1 · p. 1

2. Related Work

The related work spans personalization for statistical and system heterogeneity, as well as neural networks augmented with nearest-neighbor retrieval. kNN-Per distinguishes itself by using the global model both for interpolation and to provide the local kNN representation.

2.1. Statistical Heterogeneity

  • Earlier statistical-heterogeneity approaches fine-tune a global model, cluster clients, or learn related client models through multi-task optimization or interpolation.These approaches include FedAvg+ [Jiang et al., 2019; Yu et al., 2020], clustered methods, and multi-task learning; some restrict models or penalization terms.Source: §2.1 · p. 2
  • kNN-Per’s global model also supplies the representation for its local kNN, and its local models are learned separately from the globally trained representation.The authors note that pFedGP also learns a semi-parametric model in federated learning, but has higher training and inference computational cost than kNN-Per.Source: §2.1 · p. 2
  • Unlike FedRep, FedPer, and pFedGP, kNN-Per trains the global representation separately from local models, which the authors suggest may avoid conflicting learning goals.The authors report better performance than FedRep and pFedGP but say further study is needed to understand the interaction.Source: §2.1 · p. 2

2.2. System Heterogeneity

  • Prior system-heterogeneity approaches use distillation, client-specific submodels, or shared prototypes to accommodate different device capabilities or model architectures.The paper assumes a shared global model but says kNN-Per can be combined with heterogeneous models to address both system and statistical heterogeneity.Source: §2.2 · p. 3

2.3. Nearest Neighbours Augmented Neural Networks

  • Unlike nearest-neighbor augmentation for memorizing rare patterns or adversarial robustness, this paper uses local memorization to capture client-specific distribution shift.The authors relate this use to few-shot learning, while noting that those methods do not interpolate models or have kNN-Per’s generalization guarantees.Source: §2.3 · p. 3

3. kNN-Per Algorithm

kNN-Per personalizes a global discriminative model by adding a client-local nearest-neighbor estimate based on the global model’s fixed-length representation. The client interpolates the local estimate with the global model’s prediction.

3. kNN-Per Algorithm

  • Each client’s data are drawn i.i.d. from its own distribution, and personalized FL fits a separate hypothesis to each client’s distribution.A hypothesis maps inputs to probability distributions over labels and estimates a conditional distribution.Source: §3 · p. 3
  • Clients collectively train the global model, then build local datastores from their data without additional training, requiring a forward pass to compute representations.The datastore stores representation-label pairs and is queried during inference.Source: §3 · p. 4
  • The global model provides a fixed-length representation that kNN-Per uses to retrieve neighbors from each client’s local datastore.The representation may come from a CNN’s last convolutional layer, an RNN’s last hidden state, or a transformer self-attention layer.Source: §3 · p. 4
  • At inference, the client embeds an input, retrieves its k nearest local examples, and estimates the label distribution with a local kNN method.The local hypothesis may use a Gaussian kernel.Source: §3 · p. 4
  • The final client prediction interpolates the local kNN distribution with the global model’s distribution using a client-tuned parameter λm.Clients tune λm with local validation data or cross-validation; although k and distance metrics could vary by client, the paper treats them as common across clients.Source: §3 · p. 4

4. Generalization Bounds

The paper bounds kNN-Per’s expected risk for binary classification with one-neighbor Euclidean kNN under stated assumptions. The bound incorporates distribution discrepancy, sample sizes, representation dimension, and hypothesis-class complexity.

4. Generalization Bounds

  • The generalization bound applies to binary classification with k=1 and Euclidean distance, under the paper’s stated assumptions.The analysis assumes bounded loss and convexity, along with a representation-based condition relating proximity to label similarity.Source: §4 · pp. 4–5
  • The bound depends on hypothesis-class VC dimension, local and global dataset sizes, representation dimension, and discrepancy between a client distribution and the average distribution.The paper quantifies discrepancy through label discrepancy associated with the hypothesis class.Source: §4 · p. 5
  • When clients use only the global model, the bound is analogous to a prior probabilistic bound; when they use only kNN, it recovers a standard kNN generalization bound.The global-only comparison is described for the non-agnostic case.Source: §4 · p. 5
  • The bound predicts that a client should weight kNN more than the global model once its local sample count exceeds a threshold, even when client distributions are identical.The threshold is difficult to quantify analytically; experiments report that clients with a few tens of samples give greater weight to kNN.Source: §4 · p. 5

5. Experimental Setup

The experiments span four federated datasets and compare kNN-Per with local, global, and personalized federated-learning baselines. Dataset partitioning, model representations, and training procedures vary by dataset.

5. Experimental Setup

  • Shakespeare and FEMNIST use natural client partitions, while CIFAR-10 and CIFAR-100 use constructed heterogeneous splits, including Dirichlet and Pachinko allocation variants.The CIFAR splits also include v2 variants in which clients receive samples from only a limited number of classes.Source: §5 · pp. 5–6
  • MobileNet-v2 supplies 1280-dimensional representations for image and character-recognition tasks, while Shakespeare uses a two-layer stacked LSTM with 1024-dimensional representations.The local kNN uses the representation provided by the global model.Source: §5 · pp. 5–6
  • The benchmarks include local training, FedAvg, and personalized methods FedAvg+, ClusteredFL, Ditto, FedRep, APFL, and pFedGP.The implemented first-order MAML approach performed worse than FedAvg+.Source: §5 · p. 6
  • The number of neighbors was fixed at k = 10, while each client tuned its interpolation parameter from a seven-value grid.The models were retrained on concatenated training and validation data after selecting hyperparameters.Source: §5 · p. 6
  • Training schedules differed by dataset, with full participation on CIFAR, 10% participation on Shakespeare, and 5% participation on FEMNIST.The reported training spans 200 rounds for CIFAR, 300 for Shakespeare, and 1000 for FEMNIST.Source: §5 · p. 6

6. Experiments

Across the tested datasets, kNN-Per achieves the highest average accuracy and benefits clients in the bottom accuracy decile. Experiments also examine unseen clients, datastore capacity, heterogeneity, global-model quality, and adaptation to distribution shift.

6. Experiments

  • kNN-Per achieves the highest average weighted accuracy across all datasets.Local training performs much worse than federated methods on CIFAR-10, while it can outperform some methods on CIFAR-10/100 v2 because each client’s task is easier.Source: §6 · p. 7
  • Even clients in the bottom decile have better personalized-model accuracy with kNN-Per.The minimum accuracy is especially noisy when local test datasets are very small.Source: §6 · p. 7
  • With 20% of clients joining after training, kNN-Per outperforms all other methods on average test accuracy for unseen clients.The comparison reports training accuracy in parentheses.Source: §6 · p. 7
  • New clients retain most kNN-Per benefits even when their local datastore is reduced by a factor of three.The global model was held fixed to isolate the effect of local datastore size.Source: §6 · p. 7, Figure 1 · p. 8
  • The memorization benefit is larger under greater data heterogeneity, and kNN-Per appears to address heterogeneity more effectively than other methods.The reported comparison associates greater heterogeneity with smaller α; identical local distributions would leave no advantage for personalization.Source: §6 · p. 8
  • Accuracy is not highly sensitive to neighbor count or distance scaling, while larger local datasets lead clients to rely more on kNN.For datasets with 100 samples, clients assign the kNN model more weight than the global model (λ > 1/2).Source: §6 · p. 8, Figure 2 · p. 8
  • With only the kNN predictor used, its test accuracy is almost linearly related to the global model’s test accuracy across heterogeneity levels.Additional experiments, including CIFAR-100, confirm the relationship.Source: §6 · pp. 8–9, Figure 3 · p. 9
  • After a distribution shift, FIFO datastore updates recover accuracy as new-distribution samples replace old ones, while fixed datastores show a significant drop.Concatenating new samples also improves accuracy after the shift, but more slowly than FIFO because old samples are never evicted.Source: Figure 4 · p. 9, §6 · p. 9

7. Conclusion

Local memorization can use a global model’s representation to address statistical heterogeneity even when its predictions are inaccurate for some clients. The comparison with FedRep and pFedGP also raises an unresolved question about jointly learning representations and local models.

7. Conclusion

  • A global model may predict poorly for some clients yet still provide representations that a local kNN can use effectively.The paper presents local memorization as a simple way to address statistical heterogeneity in federated learning.Source: §7 · p. 9
  • kNN-Per outperforms FedRep and pFedGP, whose joint learning of shared representations and local models may involve conflicting goals.The authors say further study is needed to understand this interaction.Source: §7 · p. 9
  • A semi-parametric framework could formalize the interaction, but its extension to federated learning remains unexplored.Source: §7 · p. 9

A. Proofs

The paper’s general description and experiments use each client’s full dataset both to train the shared base model and to populate its local datastore. For the analysis, the dataset is instead split into disjoint parts, with equal sizes in the simplified case.

A. Proofs

  • In the general description and experiments, each client uses its whole dataset to train the shared base model and populate the local datastore.Source: A. Proofs · p. 14
  • For the analysis, each local dataset is split into disjoint parts for training the base model and populating the datastore.Source: A. Proofs · p. 14
  • The simplified analysis assumes the two parts have equal size; the result also holds when their relative size is fixed across clients.Source: A. Proofs · p. 14

A.1. Proof of Theorem 4.1

Theorem 4.1 provides a bound for the personalized model under Assumptions 1–4 and mixing weights in (0, 1). Its proof combines bounds for the shared model and local kNN retrieval before applying convexity of the loss.

A.1. Proof of Theorem 4.1

  • Theorem 4.1 gives a bound under Assumptions 1–4 for each client and mixing weight λm in (0, 1).Source: §A.1 · pp. 14–15
  • The proof bounds the shared base model’s expected error and local kNN retrieval error, then uses loss convexity to bound the personalized model’s error.Source: §A.1 · p. 14
  • The proof assigns constants c1 = 2, c2 = max{12γ2, 6γ1}, c3 = δ1, c4 = δ2, and c5 = 6γ2δ1.Source: §A.1 · p. 15

A.2. Intermediate Lemmas

The intermediate lemmas bound the shared model’s error and the local 1-NN rule’s expected error. The appendix also examines test accuracy across datastore capacity and data heterogeneity.

A.2. Intermediate Lemmas

  • Lemma A.1 bounds the shared model’s deviation from its expected loss using the hypothesis class’s VC dimension and growth function.The proof invokes Sauer’s lemma to bound the growth function.Source: §A.2 · pp. 15–16
  • Lemma A.2 provides an upper bound on the expected error of the local 1-NN rule under Assumptions 1, 2, and 4.Source: §A.2 · p. 16
  • Figure 5 plots test accuracy against the number of neighbors k.Source: Figure 5 · p. 18
  • With greater data heterogeneity, the global model performs worse, while a sufficiently large datastore can compensate for that effect at least partially.Figure 6 plots test accuracy against local datastore capacity, normalized by the client’s initial dataset partition size.Source: Figure 6 · p. 18

B. Additional Experiments

Additional experiments examine compression, datastore management under distribution shift, and sensitivity to model and client conditions. They report tradeoffs in accuracy and memory, alongside relationships between accuracy, global-model quality, and local sample count.

Effect of kernel scale parameter σ

Effect of datastore’s size on the optimal λ

  • The optimal mixing parameter λopt is positively correlated with the local sample count nm, consistent with clients with more samples relying more on memorization.The reported relationship is evaluated on each client’s test dataset.Source: B. Additional Experiments · p. 19, Figure 8 · p. 20

Effect of hardware heterogeneity

  • As system heterogeneity increases, average test accuracy decreases because gains for strong clients do not offset losses for weak clients.The experiment varies how a fixed total memory is divided between weak and strong clients.Source: B. Additional Experiments · p. 19, Figure 9 · p. 21

Adding compression techniques

  • 82.3% versus 83.0% on CIFAR-10 with ProtoNN uses one-third as many prototypes and reduces memory fourfold; CIFAR-100 falls from 55.0% to 52.1% with 20% less memory.On CIFAR-10, the compressed approach still outperforms the other methods; on CIFAR-100, compression is less advantageous.Source: B. Additional Experiments · p. 19, Figure 11 · p. 21

Effect of global model’s quality

Robustness to distribution shift

  • After a shift at t0=50, fixed datastores lose accuracy, FIFO recovers as old samples are evicted, while concatenation improves more slowly after the shift.Concatenation also improves accuracy before the shift as the datastore grows; FIFO can fluctuate before the shift as its contents change.Source: B. Additional Experiments · pp. 20–21, Figure 14 · p. 23
Loading 2111.09360v3…