[Paper Review] Self-Aware Personalized Federated Learning

Huili Chen, Jie Ding+5 more|arXiv (Cornell University)|Apr 17, 2022

Privacy-Preserving Technologies in Data4 citations

TL;DR

This paper proposes Self-Aware Personalized Federated Learning (Self-FL), a novel framework that uses uncertainty quantification from Bayesian hierarchical models to automatically balance local personalization and global model aggregation. By deriving client-specific training steps and aggregation rules from inter-client and intra-client uncertainty, Self-FL achieves superior personalization performance without hyperparameter tuning, outperforming FedAvg, DITTO, and PerFedAvg on synthetic, MNIST, FEMNIST, CIFAR10, Sent140, and real-world Alexa audio data.

ABSTRACT

In the context of personalized federated learning (FL), the critical challenge is to balance local model improvement and global model tuning when the personal and global objectives may not be exactly aligned. Inspired by Bayesian hierarchical models, we develop a self-aware personalized FL method where each client can automatically balance the training of its local personal model and the global model that implicitly contributes to other clients' training. Such a balance is derived from the inter-client and intra-client uncertainty quantification. A larger inter-client variation implies more personalization is needed. Correspondingly, our method uses uncertainty-driven local training steps and aggregation rule instead of conventional local fine-tuning and sample size-based aggregation. With experimental studies on synthetic data, Amazon Alexa audio data, and public datasets such as MNIST, FEMNIST, CIFAR10, and Sent140, we show that our proposed method can achieve significantly improved personalization performance compared with the existing counterparts.

Motivation & Objective

• To address the challenge of balancing local model personalization and global model aggregation in personalized federated learning when local and global objectives are misaligned.
• To develop a method that automatically adapts training and aggregation based on statistical uncertainty rather than heuristic or size-based rules.
• To establish a principled connection between personalized FL and Bayesian hierarchical modeling for improved theoretical and empirical performance.
• To reduce reliance on hyperparameter tuning by enabling built-in auto-tuning per client, improving usability for FL developers.
• To evaluate the method on diverse data distributions, including real-world non-IID audio data from Amazon Alexa devices.

Proposed method

• The method is grounded in a two-level Bayesian hierarchical model, where inter-client uncertainty captures distributional differences across clients, and intra-client uncertainty models local data variability.
• Self-FL uses uncertainty-driven local training steps that adapt based on the client’s data heterogeneity and uncertainty estimates.
• The aggregation rule is derived from uncertainty quantification, replacing conventional sample size-based weighting with a statistically principled approach.
• A practical deep learning version of Self-FL is proposed, maintaining the same communication and computation overhead as FedAvg while improving personalization.
• The framework automatically tunes client-specific training without requiring manual hyperparameter adjustment, reducing developer burden.
• Theoretical convergence is justified for the hierarchical model setting, providing analytical grounding for the adaptive mechanisms.

Experimental results

Research questions

• RQ1 What is the theoretical upper bound of personalization performance in federated learning, and how can it be approached via uncertainty quantification?
• RQ2 How can the global model be used as an effective intermediary to transfer useful information across clients without sharing raw data?
• RQ3 How can local training and model aggregation be adaptively calibrated based on data distribution discrepancies across clients?
• RQ4 Can uncertainty quantification replace heuristic or size-based aggregation rules to improve personalization in non-IID settings?
• RQ5 How can personalized FL be auto-tuned per client without requiring manual hyperparameter configuration?

Key findings

• Self-FL achieved the lowest relative false alarm rate (0.92) on device type A and consistently outperformed all baselines across all device types in global model evaluation.
• On personalized model performance, Self-FL achieved relative FA values of 0.90–0.99 across device types, significantly outperforming FedAvg-e (0.95–0.98), DITTO-e (0.96–0.96), and PerFedAvg (1.02–1.11).
• Sample size-based aggregation (FedAvg-w and DITTO-w) degraded global model performance, with relative FA values as high as 12.80, indicating that size-based weighting fails under high data heterogeneity.
• The ablation study confirmed that Self-FL maintains strong performance even under client sampling, demonstrating robustness in realistic FL settings.
• Self-FL achieved superior personalization on real-world Alexa audio data, where data distributions are highly non-IID due to accents, noise, and device-specific characteristics.
• To our knowledge, Self-FL is the first method to formally connect personalized FL with hierarchical modeling and use uncertainty quantification to drive adaptive training and aggregation.