[Paper Review] Towards Model Agnostic Federated Learning Using Knowledge Distillation

A. N. Afonin, Sai Praneeth Karimireddy|arXiv (Cornell University)|Oct 28, 2021

Privacy-Preserving Technologies in Data35 references4 citations

TL;DR

This paper proposes a model-agnostic federated learning framework using knowledge distillation (KD) to enable collaboration among agents with heterogeneous models and data, without sharing raw data. It introduces a theoretical framework based on Federated Kernel Ridge Regression, revealing that alternating KD degrades performance under data heterogeneity, while averaged KD (AvgKD) avoids degradation and closely matches theoretical predictions in real-world deep learning experiments.

ABSTRACT

Is it possible to design an universal API for federated learning using which an ad-hoc group of data-holders (agents) collaborate with each other and perform federated learning? Such an API would necessarily need to be model-agnostic i.e. make no assumption about the model architecture being used by the agents, and also cannot rely on having representative public data at hand. Knowledge distillation (KD) is the obvious tool of choice to design such protocols. However, surprisingly, we show that most natural KD-based federated learning protocols have poor performance. To investigate this, we propose a new theoretical framework, Federated Kernel ridge regression, which can capture both model heterogeneity as well as data heterogeneity. Our analysis shows that the degradation is largely due to a fundamental limitation of knowledge distillation under data heterogeneity. We further validate our framework by analyzing and designing new protocols based on KD. Their performance on real world experiments using neural networks, though still unsatisfactory, closely matches our theoretical predictions.

Motivation & Objective

• To design a model-agnostic federated learning protocol that enables collaboration without requiring shared model architectures or public data.
• To investigate the limitations of knowledge distillation in federated settings, especially under data and model heterogeneity.
• To develop a theoretical framework to analyze and predict the behavior of KD-based federated learning protocols.
• To design improved KD-based protocols that mitigate performance degradation in heterogeneous settings.
• To validate theoretical insights through empirical evaluation on real-world datasets and models, including MLPs and Random Forests.

Proposed method

• Formulates model-agnostic federated learning as kernel ridge regression on combined datasets, enabling analysis of functional similarity and heterogeneity.
• Introduces a theoretical framework—Federated Kernel Ridge Regression—to model both data and model heterogeneity.
• Analyzes alternating knowledge distillation (AKD) as equivalent to alternating projection on convex sets, revealing sequential information loss.
• Proposes averaged knowledge distillation (AvgKD) as a stable alternative, showing it avoids performance degradation.
• Validates theoretical predictions using real-world experiments with MNIST and CIFAR10, using cross-entropy and MSE loss functions.
• Extends AvgKD to M-agent settings, demonstrating stability and early communication benefits across diverse data heterogeneity levels.

Experimental results

Research questions

• RQ1 Why do standard knowledge distillation-based federated learning protocols degrade in performance under data heterogeneity?
• RQ2 Can a model-agnostic federated learning protocol be designed using only 'fit' and 'predict' as primitives, without requiring shared model architectures?
• RQ3 How does data heterogeneity affect the convergence and performance of knowledge distillation in federated settings?
• RQ4 Can a theoretical framework be developed to predict and analyze the behavior of KD-based federated learning protocols?
• RQ5 Does averaged knowledge distillation (AvgKD) outperform alternating KD in terms of stability and final accuracy across heterogeneous data and model settings?

Key findings

• Alternating knowledge distillation (AKD) degrades performance over rounds, especially under high data heterogeneity, due to sequential information loss in the alternating projection process.
• Averaged knowledge distillation (AvgKD) shows no performance degradation across all tested settings, including with different models like MLPs and Random Forests.
• On MNIST and CIFAR10, AvgKD maintains stable performance across rounds, while AKD consistently degrades, with degradation accelerating under higher regularization, model heterogeneity, or data imbalance.
• Empirical results closely match theoretical predictions from the Federated Kernel Ridge Regression framework, validating its utility for protocol design.
• In M-agent settings with full data heterogeneity (Alpha=0), early stopping improves performance, and all agents benefit from early communication rounds in AvgKD.
• The performance of both agents in AvgKD is highly dependent on data heterogeneity, but no degradation occurs over rounds, even with fundamentally different models such as MLP and Random Forest.