Personalized Federated Learning through Local Memorization

Othmane MarfoqGiovanni NegliaRichard VidalLaetitia Kameni

article2022ICML132 citations

Proposes kNN-Per, a personalized federated learning method that couples a shared global model with client-specific k-nearest neighbors classifiers over local datastores to effectively handle statistical heterogeneity with theoretical generalization bounds and minimal computational overhead.

Listen

Federated learning enables multiple client devices to collaboratively train a shared machine learning model while keeping sensitive data stored locally. However, standard federated approaches produce a single global model that frequently performs poorly when client data distributions vary widely or when client hardware capabilities differ. This lack of customization creates accuracy bottlenecks and unfair performance disparities across participating devices. The article evaluates and demonstrates a personalized federated learning method, called kNN-Per, which couples a shared global neural network with a local nearest-neighbor memorization mechanism to tailor predictions to individual client data.

To address this challenge, the authors propose training a standard global neural network across clients and using its internal feature representations to build a private, local datastore on each device. During inference, the final prediction linearly blends the global model output with a local nearest-neighbor search over the stored representations, requiring no extra distributed retraining. The approach was evaluated through mathematical generalization bounds and extensive empirical testing across four standard federated benchmark datasets covering image classification, handwritten character recognition, and language modeling, across hundreds to thousands of simulated client devices.

The findings establish that this local memorization strategy outperforms existing personalization methods. Across all evaluated benchmarks, kNN-Per consistently achieved higher average test accuracy and delivered superior results for the bottom ten percent of clients, improving overall system fairness. Second, the method generalized seamlessly to newly arriving clients who did not participate in training, outperforming competing adaptation techniques. Third, performance gains scaled with data heterogeneity: clients benefited most when their local distributions diverged sharply from the global average. Fourth, the system showed strong operational adaptability, successfully navigating local data distribution shifts through simple first-in-first-out updates to the local datastore without triggering costly global retraining.

These results demonstrate that organizations can achieve highly effective personalization and improved fairness without increasing communication overhead, retraining costs, or data privacy risks. By decoupling global feature learning from local memorization, devices retain privacy since local datastores are never transmitted to the central server. Organizations deploying federated systems should consider implementing local representation caching rather than running complex, jointly trained personalization architectures. When client devices face memory constraints, teams can deploy approximate nearest-neighbor compression techniques or sliding-window update strategies to maintain real-time personalization.

The primary limitations include increased local storage demands for maintaining feature datastores and potential computational latency during nearest-neighbor retrieval on highly resource-constrained devices. In addition, the theoretical generalization guarantees are strictly proven for binary classification under specific bounded-loss assumptions, though multi-class empirical results remain strong. Overall, confidence in the findings is high given the consistent empirical performance across diverse tasks and alignment with formal theoretical bounds.

Accenture INRIA Université Côte d’Azur binary classification correlation analysis federated learning personalized federated learning local memorization agnostic meta-learning

arXiv: 2111.09360

Paper: Federated Learning with Personalization Layers, Manoj Ghuhan Arivazhagan et al. (2019). It introduces the foundational FedPer baseline of splitting shared base neural representations from local personalization layers, which kNN-Per directly builds upon and outperforms.
Paper: Ditto: Fair and Robust Federated Learning Through Personalization, Tian Li et al. (2020). It establishes key personalization objectives and benchmarks balancing local customization with global fairness across heterogeneous client distributions.
Paper: Personalized Federated Learning with Theoretical Guarantees: A Model-Agnostic Meta-Learning Approach, Alireza Fallah et al. (2020). It establishes the theoretical foundations and benchmark formulations for personalized federated learning under non-convex client heterogeneity.
Paper: Personalized Federated Learning with Moreau Envelopes, Canh T. Dinh et al. (2020). It formalizes personalized federated optimization and localized adaptation against which local memorization methods are comparatively benchmarked.
Paper: Communication-Efficient Learning of Deep Networks from Decentralized Data, H. B. McMahan et al. (2016). It introduces the standard FederatedAveraging algorithm that serves as the global feature-learning foundation for kNN-Per.
Paper: Towards Personalized Federated Learning, Alysa Ziying Tan et al. (2021). It provides a comprehensive taxonomy of personalized federated learning techniques and the trade-offs of decoupling global and local representations.
Paper: Federated Learning with Non-IID Data, Yue Zhao et al. (2018). It demonstrates the fundamental accuracy degradation caused by non-IID client distributions that local personalization techniques are designed to resolve.

Paper: FedALA: Adaptive Local Aggregation for Personalized Federated Learning, Jianqing Zhang et al. (2023). It advances beyond static feature memorization by adaptively aggregating element-wise local and global model parameters on edge devices.
Paper: CD2-pFed: Cyclic Distillation-guided Channel Decoupling for Model Personalization in Federated Learning, Yiqing Shen et al. (2022). It extends representation-level personalization by using cyclic distillation and channel decoupling across network layers to handle multi-faceted heterogeneity.
Paper: FedTGP: Trainable Global Prototypes with Adaptive-Margin-Enhanced Contrastive Learning for Data and Model Heterogeneity in Federated Learning, Jianqing Zhang et al. (2024). It builds on localized representation sharing by learning global class prototypes via margin-enhanced contrastive learning for heterogeneous clients.
Paper: DisPFL: Towards Communication-Efficient Personalized Federated Learning via Decentralized Sparse Training, Rong Dai et al. (2022). It extends personalized federated learning to fully decentralized peer-to-peer topologies using customized sparse training.
Paper: Federated Domain Generalization with Generalization Adjustment, Ruipeng Zhang et al. (2023). It investigates federated domain generalization to unseen target distributions, complementing client-side personalization strategies.
Paper: An Aggregation-Free Federated Learning for Tackling Data Heterogeneity, Yuan Wang et al. (2024). It develops an aggregation-free paradigm to mitigate client drift under severe data heterogeneity without standard centralized model aggregation.

Cover for Personalized Federated Learning through Local Memorization

Abstract

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.

Topics

binary classification correlation analysis federated learning k-nearest neighbors neural networks privacy preserving application servers personalized service distribution functions hierarchical model

Keywords

personalized federated learning local memorization agnostic meta-learning metric learning differential privacy multi-task learning local kNN few-shot learning

1. Introduction
2. Related Work
2.1. Statistical Heterogeneity
2.2. System Heterogeneity
2.3. Nearest Neighbours Augmented Neural Networks
3. kNN-Per Algorithm
4. Generalization Bounds
5. Experimental Setup
6. Experiments
7. Conclusion
Acknowledgments
References
A. Proofs
A.1. Proof of Theorem 4.1
A.2. Intermediate Lemmas
B. Additional Experiments

Knowls

Knowl 1 — kNN-Per Algorithm for Personalized Federated Learning
algorithm
The kNN-Per algorithm provides federated personalization without altering the shared global model representation during local adaptation. It first trains a standard global deep neural network across all participating clients using Federated Averaging (FedAvg). Each client then builds a local datastore of intermediate feature representations (embeddings) from its local dataset in a single forward pass. At inference time, queries are classified by interpolating the softmax output of the global network with a non-parametric $k$ -nearest neighbors ( $k$ NN) retrieval distribution computed over the client's local datastore.
Input: Number of clients $M$ , local client datasets $S_m = \{(x_m^{(i)}, y_m^{(i)})\}_{i=1}^{n_m}$ for $m \in [M]$ , feature extractor and global network architecture $h(x) = g(\phi_h(x))$ , number of neighbors $k$ , distance metric $d$ , interpolation parameters $\lambda_m \in [0, 1]$
Output: Personalized decision rules $h_{m,\lambda_m}$ for all clients $m \in [M]$
1. Server and clients collaboratively train global model $h_S$ via FedAvg
2. for each client $m \in [M]$ in parallel do
3. Initialize datastore $(K_m, V_m) \leftarrow \emptyset$
4. for each local sample $(x_m^{(i)}, y_m^{(i)}) \in S_m$ do
5. Compute representation $z_m^{(i)} \leftarrow \phi_{h_S}(x_m^{(i)})$
6. Append key-value pair $(z_m^{(i)}, y_m^{(i)})$ to $(K_m, V_m)$
7. end for
8. Tune interpolation weight $\lambda_m$ using local validation data
9. end for
10. Function Inference( $x$ , client $m$ ):
11. Compute global prediction $h_S(x)$ and embedding $\phi_{h_S}(x)$
12. Retrieve $k$ nearest neighbors $\mathcal{N}_m^{(k)}(x)$ from $K_m$ using distance $d(\phi_{h_S}(x), \cdot)$
13. Compute local $k$ NN distribution $h_{S_m}^{(k)}(x)$ over classes $y \in \mathcal{Y}$ via Gaussian kernel weighting:
14. $[h_{S_m}^{(k)}(x)]_y \propto \sum_{i=1}^k \mathbf{1}_{\{y = y_{\pi_m^{(i)}(x)}\}} \exp(-d(\phi_{h_S}(x), \phi_{h_S}(x_{\pi_m^{(i)}(x)})))$
15. return $h_{m,\lambda_m}(x) = \lambda_m h_{S_m}^{(k)}(x) + (1 - \lambda_m) h_S(x)$
Knowl 2 — Local kNN Estimation and Hypothesis Interpolation in kNN-Per
model/method
In kNN-Per, personalization is formulated as an improper learning problem where a client $m \in [M]$ combines a parametric global model $h_S: \mathcal{X} \to \Delta^{|\mathcal{Y}|}$ and a non-parametric local memory estimator $h_{S_m}^{(k)}: \mathcal{X} \to \Delta^{|\mathcal{Y}|}$ , where $\Delta^{|\mathcal{Y}|}$ is the probability simplex over discrete label set $\mathcal{Y}$ .

Given the intermediate representation map $\phi_{h_S}: \mathcal{X} \to \mathbb{R}^p$ extracted from the trained global network $h_S$ , client $m$ constructs a key-value datastore $(K_m, V_m)$ from its local dataset $S_m = \{(x_m^{(i)}, y_m^{(i)})\}_{i=1}^{n_m}$ :

$(K_m, V_m) = \left\{\left(\phi_{h_S}\left(x_m^{(i)}\right), y_m^{(i)}\right) \mid \left(x_m^{(i)}, y_m^{(i)}\right) \in S_m\right\}$

For an input query $x \in \mathcal{X}$ , let $\pi_m(x) = (\pi_m^{(1)}(x), \dots, \pi_m^{(n_m)}(x))$ denote the permutation of indices $[n_m]$ sorted by increasing distance in representation space:

$d\left(\phi_{h_S}(x), \phi_{h_S}\left(x_{\pi_m^{(i)}(x)}\right)\right) \le d\left(\phi_{h_S}(x), \phi_{h_S}\left(x_{\pi_m^{(i+1)}(x)}\right)\right) \quad \forall i \in [n_m - 1]$

The local $k$ -nearest neighbor distribution $h_{S_m}^{(k)}(x)$ computes the probability for each class $y \in \mathcal{Y}$ using a Gaussian kernel over the $k$ nearest representations:

$\left[h_{S_m}^{(k)}(x)\right]_y = \frac{\sum_{i=1}^k \mathbf{1}_{\left\{y = y_{\pi_m^{(i)}(x)}\right\}} \exp\left(-d\left(\phi_{h_S}(x), \phi_{h_S}\left(x_{\pi_m^{(i)}(x)}\right)\right)\right)}{\sum_{i=1}^k \exp\left(-d\left(\phi_{h_S}(x), \phi_{h_S}\left(x_{\pi_m^{(i)}(x)}\right)\right)\right)}$

The final prediction $h_{m, \lambda_m}(x)$ is the linear interpolation between the local $k$ NN distribution and the global model distribution parameterized by a client-specific hyperparameter $\lambda_m \in [0, 1]$ :

$h_{m, \lambda_m}(x) = \lambda_m \cdot h_{S_m}^{(k)}(x) + (1 - \lambda_m) \cdot h_S(x)$
Knowl 3 — Generalization Error Bound for kNN-Per in Binary Classification
theoretical result
For binary classification where $\mathcal{Y} = \{0, 1\}$ , $k = 1$ neighbor is used, and $d(\cdot, \cdot)$ is the Euclidean distance, let $\mathcal{H}$ be a hypothesis class of VC-dimension $d_{\mathcal{H}}$ , $M$ be the total number of clients, $n_m$ be the size of the local dataset $S_m \sim \mathcal{D}_m^{n_m}$ , $n = \sum_{m=1}^M n_m$ be the total number of samples across all clients, $\bar{\mathcal{D}} = \sum_{m=1}^M \frac{n_m}{n} \mathcal{D}_m$ be the mixture average data distribution, $p$ be the dimension of the embedding space, and $\text{disc}_{\mathcal{H}}(\bar{\mathcal{D}}, \mathcal{D}_m) = \max_{h \in \mathcal{H}} |L_{\bar{\mathcal{D}}}(h) - L_{\mathcal{D}_m}(h)|$ be the label discrepancy between distributions $\bar{\mathcal{D}}$ and $\mathcal{D}_m$ .

Under bounded representation $\phi_{h_S}: \mathcal{X} \to [0, 1]^p$ , bounded loss $l: \Delta^2 \times \{0, 1\} \to [0, 1]$ with $l(e_y, y') = \mathbf{1}_{\{y \neq y'\}}$ , convex loss $l$ on its first argument, and representation-label Lipschitz regularity with constants $\gamma_1, \gamma_2 > 0$ , for any client $m \in [M]$ and mixing weight $\lambda_m \in (0, 1)$ , there exist constants $c_1, c_2, c_3, c_4, c_5 \in \mathbb{R}$ such that:

$\mathbb{E}_{S \sim \bigotimes_{j=1}^M \mathcal{D}_j^{n_j}} \left[L_{\mathcal{D}_m}\left(h_{m, \lambda_m}\right)\right] \le (1 + \lambda_m) L_{\mathcal{D}_m}\left(h_m^*\right) + c_1(1 - \lambda_m) \cdot \text{disc}_{\mathcal{H}}\left(\bar{\mathcal{D}}, \mathcal{D}_m\right) + c_2 \lambda_m \frac{\sqrt{p}}{n_m^{1/(p+1)}} \left(\text{disc}_{\mathcal{H}}\left(\bar{\mathcal{D}}, \mathcal{D}_m\right) + 1\right) + c_3(1 - \lambda_m) \sqrt{\frac{d_{\mathcal{H}}}{n}} \sqrt{c_4 + \log\left(\frac{n}{d_{\mathcal{H}}}\right)} + c_5 \lambda_m \sqrt{\frac{d_{\mathcal{H}}}{n}} \sqrt{c_4 + \log\left(\frac{n}{d_{\mathcal{H}}}\right)} \frac{\sqrt{p}}{n_m^{1/(p+1)}}$

where $h_m^* = \arg\min_{h \in \mathcal{H}} L_{\mathcal{D}_m}(h)$ .

When $\lambda_m = 0$ , the bound matches standard agnostic federated bounds converging at rate $\tilde{O}\left(\sqrt{d_{\mathcal{H}}/n}\right)$ plus discrepancy. When $\lambda_m = 1$ , it recovers standard $1$ -NN bounds converging at rate $O(n_m^{-1/(p+1)})$ . For large local dataset sizes $n_m$ , the bound prescribes $\lambda_m > 0.5$ .
Knowl 4 — Representation-Label Regularity Condition
assumption
Let $\eta_m(x) = \mathcal{D}_m(y=1 \mid x)$ denote the true conditional label probability under the local data distribution $\mathcal{D}_m$ , $\phi_{h_S}: \mathcal{X} \to \mathbb{R}^p$ denote the representation extractor obtained from the global federated model $h_S$ , and $h_m^* = \arg\min_{h \in \mathcal{H}} L_{\mathcal{D}_m}(h)$ denote the optimal model in hypothesis class $\mathcal{H}$ on distribution $\mathcal{D}_m$ .

There exist strictly positive constants $\gamma_1, \gamma_2 > 0$ such that for any aggregated dataset $S$ sampled from $\mathcal{X} \times \mathcal{Y}$ and any input points $x, x' \in \mathcal{X}$ :

$|\eta_m(x) - \eta_m(x')| \le d\left(\phi_{h_S}(x), \phi_{h_S}(x')\right) \cdot \left(\gamma_1 + \gamma_2 \left(L_{\mathcal{D}_m}(h_S) - L_{\mathcal{D}_m}\left(h_m^*\right)\right)\right)$

This condition ensures that inputs mapped close to each other in the global representation space share similar conditional label probabilities, and that this representation fidelity improves linearly as the excess true risk of the global model $L_{\mathcal{D}_m}(h_S) - L_{\mathcal{D}_m}(h_m^*)$ decreases.

Knowl 5 — Test Accuracy and Fairness across Federated Benchmarks

data/table

kNN-Per was evaluated across standard personalized federated learning benchmarks: FEMNIST (character recognition, MobileNet-v2), CIFAR-10 and CIFAR-100 (image classification, MobileNet-v2, partitioned via Dirichlet distribution with $\alpha=0.3$ ), CIFAR-10 (v2) and CIFAR-100 (v2) (partitioned such that each client holds 2 and 10 classes respectively), and Shakespeare (next-character prediction, stacked 2-layer LSTM). In all cases, $k=10$ , distance $d$ is Euclidean, and $\lambda_m$ is tuned per client via grid search over $\{0.0, 0.1, 0.3, 0.5, 0.7, 0.9, 1.0\}$ .

Dataset	Local	FedAvg	FedAvg+	ClusteredFL	Ditto	FedRep	APFL	kNN-Per (Ours)
FEMNIST	71.0 / 57.5	83.4 / 68.9	84.3 / 69.4	83.7 / 69.4	84.3 / 71.3	85.3 / 72.7	84.1 / 69.4	88.2 / 78.8
CIFAR-10	57.6 / 41.1	72.8 / 59.6	75.2 / 62.3	73.3 / 61.5	80.0 / 66.5	77.7 / 65.2	78.9 / 68.1	83.0 / 71.4
CIFAR-10 (v2)	82.4 / 71.3	67.9 / 60.1	85.0 / 79.6	79.9 / 72.3	86.3 / 80.6	89.1 / 85.3	82.6 / 76.4	93.8 / 88.2
CIFAR-100	31.5 / 19.8	47.4 / 36.0	51.4 / 41.1	47.2 / 36.2	52.0 / 41.4	53.2 / 41.7	51.7 / 41.1	55.0 / 43.6
CIFAR-100 (v2)	45.7 / 38.2	42.3 / 34.8	48.1 / 41.9	43.5 / 37.2	48.7 / 40.3	70.1 / 65.2	48.3 / 42.1	74.6 / 67.3
Shakespeare	32.0 / 16.0	48.1 / 43.1	47.0 / 42.2	46.7 / 41.4	47.9 / 42.6	47.2 / 42.3	45.9 / 42.4	51.4 / 45.4

Values indicate test accuracy (weighted average across clients / bottom decile, i.e., the $(M/10)$ -th worst client accuracy to assess fairness). On CIFAR-10 (v2) and CIFAR-100 (v2), pFedGP achieved $88.9 / 84.1$ and $61.1 / 50.0$ , respectively. kNN-Per consistently achieves the highest average accuracy and the highest bottom decile accuracy across all benchmarks.

Knowl 6 — Generalization to Unseen Clients Joining Post-Training

data/table

In realistic federated settings, new clients arrive after collaborative training has finished. In kNN-Per, a new client simply downloads the frozen global model $h_S$ and builds its local datastore $(K_m, V_m)$ via a single forward pass over its local dataset without performing gradient descent or communicating with the server.

In experiments where $80\%$ of clients participate in training and $20\%$ join post-training as new clients, test accuracy across unseen clients is reported below (with training accuracy in parentheses):

Dataset	FedAvg	FedAvg+	ClusteredFL	Ditto	FedRep	APFL	pFedGP	kNN-Per (Ours)
FEMNIST	83.1 (83.3)	84.2 (88.5)	83.2 (86.0)	83.9 (86.9)	85.4 (88.9)	84.2 (85.5)	–	88.1 (90.5)
CIFAR-10	72.9 (72.8)	75.3 (78.2)	73.9 (76.2)	79.7 (84.3)	76.4 (79.5)	79.2 (80.6)	–	82.4 (87.1)
CIFAR-10 (v2)	67.5 (68.1)	85.1 (85.0)	79.6 (79.9)	85.9 (86.0)	89.0 (89.1)	82.3 (82.5)	89.0 (88.8)	93.0 (93.1)
CIFAR-100	47.1 (47.5)	50.8 (53.4)	47.1 (48.2)	52.1 (57.3)	53.5 (58.2)	49.1 (52.7)	–	56.1 (59.3)
CIFAR-100 (v2)	42.1 (42.2)	47.9 (48.1)	43.2 (43.4)	48.8 (48.5)	69.8 (70.0)	48.2 (48.4)	61.3 (61.0)	74.3 (74.5)
Shakespeare	49.0 (48.3)	49.3 (48.1)	49.4 (46.7)	48.1 (49.2)	48.7 (47.8)	46.1 (52.7)	–	50.7 (64.2)

kNN-Per outperforms all competing personalization mechanisms on unseen clients without requiring backpropagation updates or clustering assignments.

Knowl 7 — Datastore Adaptation Strategies under Local Distribution Shifts
empirical result
When client local distributions dynamically change at time $t_0$ from $\mathcal{D}_m$ to a shifted distribution $\mathcal{D}'_m \neq \mathcal{D}_m$ , kNN-Per can adapt without retraining the global model $h_S$ by updating the local datastore $(K_m, V_m)$ using one of three strategies:
1. First-In-First-Out (FIFO): Newly received samples at step $t$ replace the oldest samples in the datastore. Accuracy drops immediately at $t = t_0 = 50$ , but rapidly recovers as new distribution instances populate the datastore and old ones are evicted, stabilizing at the new distribution's equilibrium accuracy.
2. Concatenate: Newly received samples are appended to the datastore without eviction. For $t < t_0$ , accuracy increases steadily as datastore capacity grows. For $t > t_0$ , accuracy drops at $t_0$ and recovers, but at a significantly slower pace than FIFO because obsolete representations from $\mathcal{D}_m$ remain in memory.
3. Fixed Datastore: The datastore is never updated. Accuracy drops abruptly at $t = t_0$ (e.g., from $\approx 82\%$ to $\approx 73\%$ on CIFAR-10) and remains degraded permanently.
This demonstrates that dynamic local memorization enables rapid post-training domain adaptation solely through local buffer replacement.
Knowl 8 — Empirical Relationship between Local Sample Size and Mixing Parameter
empirical result
Evaluating the optimal mixing parameter $\lambda_{\text{opt}} = \arg\max_{\lambda} \text{Acc}(h_{m, \lambda})$ against local dataset size $n_m$ demonstrates a strong positive correlation between client data volume and reliance on the local $k$ NN component:
1. When clients have as few as $100$ local samples on CIFAR-10 and CIFAR-100, the optimal parameter satisfies $\lambda_{\text{opt}} > 0.5$ , indicating that the non-parametric local memory contributes more to the final prediction than the parametric global model.
2. As local dataset capacity increases from $n_m = 5$ to $n_m = 250$ , peak accuracy shifts monotonically toward higher $\lambda$ values.
3. The performance advantage of kNN-Per over standard FedAvg is inversely proportional to the Dirichlet concentration parameter $\alpha$ : smaller $\alpha$ (higher statistical heterogeneity across clients) yields substantially higher accuracy gains from local memorization.
Knowl 9 — Datastore Memory Compression with ProtoNN
empirical result
To reduce memory and compute overhead for hardware-constrained clients, kNN-Per can be combined with ProtoNN, which learns a reduced set of prototype vectors and a low-dimensional feature projection matrix for the local datastore.

On CIFAR-10 with a baseline MobileNet-v2 embedding dimension of 1280:
- Compressing the local datastore to $n_m / 3$ prototypes and a projection dimension of 1000 reduces the memory footprint by a factor of 4 while incurring a minor accuracy reduction from $83.0\%$ to $82.3\%$ , which still exceeds all baseline personalization algorithms.
- On CIFAR-100, ProtoNN compression is less favorable, experiencing a 2.9 percentage point drop (from $55.0\%$ to $52.1\%$ ) for only a $20\%$ reduction in memory footprint.
Knowl 10 — Robustness of kNN-Per to Neighbor Count and Distance Scaling
empirical result
kNN-Per classification accuracy is largely insensitive to the hyperparameters of the nearest-neighbor search:
1. Number of Neighbors $k$ : Varying $k$ between 7 and 14 for CIFAR-10, and between 5 and 12 for CIFAR-100, results in less than $0.2$ percentage points of test accuracy variation (peaking near $k = 10$ ).
2. Distance Scaling Parameter $\sigma$ : In the scaled Euclidean distance metric $d_\sigma(z, z') = \|z - z'\|_2 / \sigma$ , varying $\sigma$ across orders of magnitude between $0.1$ and $1000$ produces virtually indistinguishable test accuracy curves over $\lambda$ .
Consequently, client-side tuning in kNN-Per can be restricted to the scalar interpolation weight $\lambda_m$ , with $k = 10$ and standard Euclidean distance serving as robust defaults.

Coverage note — Omitted the intermediate proof lemmas (Lemma A.1 and Lemma A.2) from Appendix A that derive intermediate bounds for the global model risk and 1-NN risk respectively, as well as centralized training ablation curves for linear accuracy correlation, keeping the main standalone theorem, algorithm, and key empirical results.

References

1.Tian Li, Anit Kumar Sahu, Ameet Talwalkar, and Virginia Smith. Federated learning: Challenges, methods, and future directions. IEEE Signal Processing Magazine, 37 (3):50–60, 2020.
2.Peter Kairouz, H Brendan McMahan, Brendan Avent, Aurelien Bellet, Mehdi Bennis, Arjun Nitin Bhagoji, Keith Bonawitz, Zachary Charles, Graham Cormode, Rachel Cummings, et al. Advances and open problems in federated learning. arXiv preprint arXiv:1912.04977, 2019.
3.Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communication-efficient learning of deep networks from decentralized data. In Artificial Intelligence and Statistics, pages 1273–1282. PMLR, 2017.
4.Jakub Konecny, H Brendan McMahan, Felix X Yu, Peter Richtarik, Ananda Theertha Suresh, and Dave Bacon. Federated learning: Strategies for improving communication efficiency. arXiv preprint arXiv:1610.05492, 2016.
5.Anit Kumar Sahu, Tian Li, Maziar Sanjabi, M. Zaheer, Ameet S. Talwalkar, and Virginia Smith. On the convergence of federated optimization in heterogeneous networks. ArXiv, abs/1812.06127, 2018.
6.Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank Reddi, Sebastian Stich, and Ananda Theertha Suresh. Scaffold: Stochastic controlled averaging for federated learning. In International Conference on Machine Learning, pages 5132–5143. PMLR, 2020.
7.Mehryar Mohri, Gary Sivek, and Ananda Theertha Suresh. Agnostic federated learning. In International Conference on Machine Learning, pages 4615–4625. PMLR, 2019.
8.Tian Li, Shengyuan Hu, Ahmad Beirami, and Virginia Smith. Ditto: Fair and robust federated learning through personalization. In International Conference on Machine Learning, pages 6357–6368. PMLR, 2021.
9.Urvashi Khandelwal, Omer Levy, Dan Jurafsky, Luke Zettlemoyer, and Mike Lewis. Generalization through memorization: Nearest neighbor language models. In International Conference on Learning Representations, 2019.
10.Urvashi Khandelwal, Angela Fan, Dan Jurafsky, Luke Zettlemoyer, and Mike Lewis. Nearest neighbor machine translation. arXiv preprint arXiv:2010.00710, 2020.
11.Nicolas Papernot and Patrick McDaniel. Deep k-nearest neighbors: Towards confident, interpretable and robust deep learning. arXiv preprint arXiv:1803.04765, 2018.
12.Emin Orhan. A simple cache model for image recognition. Advances in Neural Information Processing Systems, 31, 2018.
13.Jake Snell, Kevin Swersky, and Richard Zemel. Prototypical networks for few-shot learning. Advances in neural information processing systems, 30, 2017.
14.Yan Wang, Wei-Lun Chao, Kilian Q Weinberger, and Laurens van der Maaten. Simpleshot: Revisiting nearest-neighbor classification for few-shot learning. arXiv preprint arXiv:1911.04623, 2019.
15.Yihan Jiang, Jakub Konecnˇ y, Keith Rush, and Sreeram Kannan. Improving federated learning personalization via model agnostic meta learning. arXiv preprint arXiv:1909.12488, 2019.
16.Tao Yu, Eugene Bagdasaryan, and Vitaly Shmatikov. Salvaging federated learning by local adaptation. arXiv preprint arXiv:2002.04758, 2020.
17.Hong-You Chen and Wei-Lun Chao. On bridging generic and personalized federated learning for image classification. In International Conference on Learning Representations, 2022. URL https://openreview.net/forum?id=I1hQbx10Kxn.
18.Gary Cheng, Karan Chadha, and John Duchi. Fine-tuning is fine in federated learning. arXiv preprint arXiv:2108.07313, 2021.
19.Mikhail Khodak, Maria-Florina F Balcan, and Ameet S Talwalkar. Adaptive gradient-based meta-learning methods. In Advances in Neural Information Processing Systems, pages 5917–5928, 2019.
20.Alireza Fallah, Aryan Mokhtari, and Asuman Ozdaglar. Personalized federated learning with theoretical guarantees: A model-agnostic meta-learning approach. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 3557–3568. Curran Associates, Inc., 2020. URL https://proceedings.neurips.cc/paper/2020/file/24389bfe4fe2eba8bf9aa9203a44cdad-Paper.pdf.
21.Durmus Alp Emre Acar, Yue Zhao, Ruizhao Zhu, Ramon Matas, Matthew Mattina, Paul Whatmough, and Venkatesh Saligrama. Debiasing model updates for improving personalized federated training. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 21–31. PMLR, 7 2021. URL https://proceedings.mlr.press/v139/acar21a.html.
22.Felix Sattler, Klaus-Robert Muller, and Wojciech Samek. Clustered federated learning: Model-agnostic distributed multitask optimization under privacy constraints. IEEE Transactions on Neural Networks and Learning Systems, 2020.
23.Avishek Ghosh, Jichan Chung, Dong Yin, and Kannan Ramchandran. An efficient framework for clustered federated learning. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 19586–19597. Curran Associates, Inc., 2020. URL https://proceedings.neurips.cc/paper/2020/file/e32cc80bf07915058ce90722ee17bb71-Paper.pdf.
24.Yishay Mansour, Mehryar Mohri, Jae Ro, and Ananda Theertha Suresh. Three approaches for personalization with applications to federated learning. arXiv preprint arXiv:2002.10619, 2020.
25.Othmane Marfoq, Giovanni Neglia, Aurelien Bellet, Laetitia Kameni, and Richard Vidal. Federated multi-task learning under a mixture of distributions. arXiv preprint arXiv:2108.10252, 2021.
26.Virginia Smith, Chao-Kai Chiang, Maziar Sanjabi, and Ameet Talwalkar. Federated multi-task learning. In Proceedings of the 31st International Conference on Neural Information Processing Systems, NIPS’17, page 4427–4437, Red Hook, NY, USA, 2017. Curran Associates Inc. ISBN 9781510860964.
27.Paul Vanhaesebrouck, Aurelien Bellet, and Marc Tommasi. Decentralized Collaborative Learning of Personalized Models over Networks. In AISTATS, 2017.
28.Valentina Zantedeschi, Aurelien Bellet, and Marc Tommasi. Fully decentralized joint learning of personalized models and collaboration graphs. volume 108 of Proceedings of Machine Learning Research, pages 864–874, Online, 8 2020. PMLR. URL http://proceedings.mlr.press/v108/zantedeschi20a.html.
29.Filip Hanzely and Peter Richtarik. Federated learning of a mixture of global and local models, 2020.
30.Filip Hanzely, Slavom´ır Hanzely, Samuel Horvath, and Peter Richtarik. Lower bounds and optimal algorithms for personalized federated learning. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 2304–2315. Curran Associates, Inc., 2020. URL https://proceedings.neurips.cc/paper/2020/file/187acf7982f3c169b3075132380986e4-Paper.pdf.
31.Canh T. Dinh, Nguyen Tran, and Josh Nguyen. Personalized federated learning with moreau envelopes. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 21394–21405. Curran Associates, Inc., 2020. URL https://proceedings.neurips.cc/paper/2020/file/f4f1f13c8289ac1b1ee0ff176b56fc60-Paper.pdf.
32.Canh T Dinh, Tung T Vu, Nguyen H Tran, Minh N Dao, and Hongyu Zhang. Fedu: A unified framework for federated multi-task learning with laplacian regularization. arXiv preprint arXiv:2102.07148, 2021.
33.Yutao Huang, Lingyang Chu, Zirui Zhou, Lanjun Wang, Jiangchuan Liu, Jian Pei, and Yong Zhang. Personalized cross-silo federated learning on non-iid data. In AAAI, pages 7865–7873, 2021. URL https://ojs.aaai.org/index.php/AAAI/article/view/16960.
34.Yuyang Deng, Mohammad Mahdi Kamani, and Mehrdad Mahdavi. Adaptive personalized federated learning. arXiv preprint arXiv:2003.13461, 2020.
35.Luca Corinzia and Joachim M. Buhmann. Variational federated multi-task learning, 2019.
36.Michael Zhang, Karan Sapra, Sanja Fidler, Serena Yeung, and Jose M. Alvarez. Personalized federated learning with first order model optimization. In International Conference on Learning Representations, 2021. URL https://openreview.net/forum?id=ehJqJQk9cw.
37.Liam Collins, Hamed Hassani, Aryan Mokhtari, and Sanjay Shakkottai. Exploiting shared representations for personalized federated learning. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 2089–2099. PMLR, 18–24 Jul 2021. URL https://proceedings.mlr.press/v139/collins21a.html.
38.Manoj Ghuhan Arivazhagan, V. Aggarwal, Aaditya Kumar Singh, and Sunav Choudhary. Federated learning with personalization layers. ArXiv, abs/1912.00818, 2019.
39.Idan Achituve, Aviv Shamsian, Aviv Navon, Gal Chechik, and Ethan Fetaya. Personalized federated learning with gaussian processes. Advances in Neural Information Processing Systems, 34, 2021.
40.Paul Pu Liang, Terrance Liu, Liu Ziyin, Nicholas B Allen, Randy P Auerbach, David Brent, Ruslan Salakhutdinov, and Louis-Philippe Morency. Think locally, act globally: Federated learning with local and global representations. arXiv preprint arXiv:2001.01523, 2020.
41.Kang Wei, Jun Li, Ming Ding, Chuan Ma, Howard H. Yang, Farhad Farokhi, Shi Jin, Tony Q. S. Quek, and H. Vincent Poor. Federated learning with differential privacy: Algorithms and performance analysis. IEEE Transactions on Information Forensics and Security, 15:3454–3469, 2020. doi: 10.1109/TIFS.2020.2988575.
42.P.J. Bickel, C.A.J. Klaassen, Y. Ritov, and J.A. Wellner. Efficient and Adaptive Estimation for Semiparametric Models. Johns Hopkins series in the mathematical sciences. Springer New York, 1998. ISBN 9780387984735. URL https://books.google.fr/books?id=lSnTm6SC_SMC.
43.Tao Lin, Lingjing Kong, Sebastian U Stich, and Martin Jaggi. Ensemble distillation for robust model fusion in federated learning. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 2351–2363. Curran Associates, Inc., 2020. URL https://proceedings.neurips.cc/paper/2020/file/18df51b97ccd68128e994804f3eccc87-Paper.pdf.
44.Daliang Li and Junpu Wang. Fedmd: Heterogenous federated learning via model distillation. arXiv preprint arXiv:1910.03581, 2019.
45.Zhuangdi Zhu, Junyuan Hong, and Jiayu Zhou. Data-free knowledge distillation for heterogeneous federated learning. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 12878–12889. PMLR, 18–24 Jul 2021. URL https://proceedings.mlr.press/v139/zhu21b.html.
46.Lan Zhang and Xiaoyong Yuan. Fedzkt: Zero-shot knowledge transfer towards heterogeneous on-device models in federated learning. arXiv preprint arXiv:2109.03775, 2021.
47.Enmao Diao, Jie Ding, and Vahid Tarokh. Heterofl: Computation and communication efficient federated learning for heterogeneous clients. In International Conference on Learning Representations, 2020.
48.Samuel Horvath, Stefanos Laskaridis, Mario Almeida, Ilias Leontiadis, Stylianos Venieris, and Nicholas Donald Lane. FjORD: Fair and accurate federated learning under heterogeneous targets with ordered dropout. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, 2021. URL https://openreview.net/forum?id=4fLr7H5D_eT.
49.Amaury Bouchra Pilet, Davide Frey, and Franc¸ois Ta¨ıani. Simple, efficient and convenient decentralized multi-task learning for neural networks. In IDA, pages 37–49, 2021.
50.Yue Tan, Guodong Long, Lu Liu, Tianyi Zhou, Qinghua Lu, Jing Jiang, and Chengqi Zhang. FedProto: Federated Prototype Learning across Heterogeneous Clients. In AAAI Conference on Artificial Intelligence, 2022.
51.Aviv Shamsian, Aviv Navon, Ethan Fetaya, and Gal Chechik. Personalized federated learning using hypernetworks. In ICML, 2021.
52.Jeff Johnson, Matthijs Douze, and Herve Jegou. Billion-scale similarity search with gpus. IEEE Transactions on Big Data, pages 1–1, 2019.
53.Yu A Malkov and DA Yashunin. Efficient and robust approximate nearest neighbor search using hierarchical navigable small world graphs. IEEE Transactions on Pattern Analysis and Machine Intelligence, 42(4):824–836, 2020.
54.Chirag Gupta, Arun Sai Suggala, Ankit Goyal, Harsha Vardhan Simhadri, Bhargavi Paranjape, Ashish Kumar, Saurabh Goyal, Raghavendra Udupa, Manik Varma, and Prateek Jain. ProtoNN: Compressed and accurate kNN for resource-scarce devices. In Doina Precup and Yee Whye Teh, editors, Proceedings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pages 1331–1340. PMLR, 06–11 Aug 2017. URL https://proceedings.mlr.press/v70/gupta17a.html.
55.Aurelien Bellet, Amaury Habrard, and Marc Sebban. Metric Learning, volume 9 of Synthesis Lectures on Artificial Intelligence and Machine Learning. Morgan & Claypool Publishers (USA), Synthesis Lectures on Artificial Intelligence and Machine Learning, pp 1-151, January 2015. doi: 10.2200/S00626ED1V01Y201501AIM030. URL https://hal.archives-ouvertes.fr/hal-01121733.
56.Shai Shalev-Shwartz and Shai Ben-David. Understanding machine learning: From theory to algorithms. Cambridge university press, 2014.
57.Sebastian Caldas, Sai Meher Karthik Duddu, Peter Wu, Tian Li, Jakub Konecnˇ y, H Brendan McMahan, Virginia Smith, and Ameet Talwalkar. Leaf: A benchmark for federated settings. arXiv preprint arXiv:1812.01097, 2018.
58.Alex Krizhevsky. Learning multiple layers of features from tiny images. Technical report, 2009.
59.Hongyi Wang, Mikhail Yurochkin, Yuekai Sun, Dimitris Papailiopoulos, and Yasaman Khazaeni. Federated learning with matched averaging. In International Conference on Learning Representations, 2020. URL https://openreview.net/forum?id=BkluqlSFDS.
60.Wei Li and Andrew McCallum. Pachinko allocation: Dag-structured mixture models of topic correlations. In Proceedings of the 23rd International Conference on Machine Learning, ICML ’06, page 577–584, New York, NY, USA, 2006. Association for Computing Machinery. ISBN 1595933832. doi: 10.1145/1143844.1143917. URL https://doi.org/10.1145/1143844.1143917.
61.Sashank J. Reddi, Zachary Charles, Manzil Zaheer, Zachary Garrett, Keith Rush, Jakub Konecnˇ y, Sanjiv Kumar, and Hugh Brendan McMahan. Adaptive federated optimization. In International Conference on Learning Representations, 2021. URL https://openreview.net/forum?id=LkFG3lB13U5.
62.Mark Sandler, Andrew Howard, Menglong Zhu, Andrey Zhmoginov, and Liang-Chieh Chen. Mobilenetv2: Inverted residuals and linear bottlenecks. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 4510–4520, 2018.
63.Edward Grefenstette, Karl Moritz Hermann, Mustafa Suleyman, and Phil Blunsom. Learning to transduce with unbounded memory. In NIPS, NIPS’15, pages 1828–1836, Cambridge, MA, USA, 2015. MIT Press. URL http://dl.acm.org/citation.cfm?id=2969442.2969444.
64.Armand Joulin and Tomas Mikolov. Inferring algorithmic patterns with stack-augmented recurrent nets. In NIPS, NIPS’15, pages 190–198, Cambridge, MA, USA, 2015. MIT Press. URL http://dl.acm.org/citation.cfm?id=2969239.2969261.
65.Peter J Bickel, Chris AJ Klaassen, Peter J Bickel, Ya’acov Ritov, J Klaassen, Jon A Wellner, and YA’Acov Ritov. Efficient and adaptive estimation for semiparametric models, volume 4. Johns Hopkins University Press Baltimore, 1993.

Citation

MLA

Marfoq, O., et al. “Personalized Federated Learning Through Local Memorization”. International Conference on Machine Learning, vol. 162, 2022, pp. 15070–92, https://proceedings.mlr.press/v162/marfoq22a.html.

APA

Marfoq, O., Neglia, G., Vidal, R., & Kameni, L. (2022). Personalized Federated Learning through Local Memorization. International Conference on Machine Learning, 162, 15070–15092. https://proceedings.mlr.press/v162/marfoq22a.html

Chicago

Marfoq, O., G. Neglia, R. Vidal, and L. Kameni. 2022. “Personalized Federated Learning Through Local Memorization”. International Conference on Machine Learning 162: 15070–92. https://proceedings.mlr.press/v162/marfoq22a.html.

Harvard

Marfoq, O. et al. (2022) “Personalized Federated Learning through Local Memorization”, International Conference on Machine Learning. PMLR, pp. 15070–15092. Available at: https://proceedings.mlr.press/v162/marfoq22a.html.

Vancouver

1. Marfoq O, Neglia G, Vidal R, Kameni L (2022) Personalized Federated Learning through Local Memorization. In: International Conference on Machine Learning. PMLR, pp 15070–15092

BibTeX

@InProceedings{pmlr-v162-marfoq22a,
  title = 	 {Personalized Federated Learning through Local Memorization},
  author =       {Marfoq, Othmane and Neglia, Giovanni and Vidal, Richard and Kameni, Laetitia},
  booktitle = 	 {Proceedings of the 39th International Conference on Machine Learning},
  pages = 	 {15070--15092},
  year = 	 {2022},
  editor = 	 {Chaudhuri, Kamalika and Jegelka, Stefanie and Song, Le and Szepesvari, Csaba and Niu, Gang and Sabato, Sivan},
  volume = 	 {162},
  series = 	 {Proceedings of Machine Learning Research},
  month = 	 {17--23 Jul},
  publisher =    {PMLR},
  pdf = 	 {https://proceedings.mlr.press/v162/marfoq22a/marfoq22a.pdf},
  url = 	 {https://proceedings.mlr.press/v162/marfoq22a.html},
  abstract = 	 {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.}
}

Metadata:DOI registry

Access the Paper

This paper is available from its original source. Click below to access the PDF.

Open PDF

License: https://creativecommons.org/licenses/by/4.0/

Personalized Federated Learning through Local Memorization

Abstract

Topics

Keywords

Table of Contents

Knowls

Knowl 1 — kNN-Per Algorithm for Personalized Federated Learning

Knowl 2 — Local kNN Estimation and Hypothesis Interpolation in kNN-Per

Knowl 3 — Generalization Error Bound for kNN-Per in Binary Classification

Knowl 4 — Representation-Label Regularity Condition

Knowl 5 — Test Accuracy and Fairness across Federated Benchmarks

Knowl 6 — Generalization to Unseen Clients Joining Post-Training

Knowl 7 — Datastore Adaptation Strategies under Local Distribution Shifts

Knowl 8 — Empirical Relationship between Local Sample Size and Mixing Parameter

Knowl 9 — Datastore Memory Compression with ProtoNN

Knowl 10 — Robustness of kNN-Per to Neighbor Count and Distance Scaling

References

Citation

BibTeX

Access the Paper