A Multi-Prototype-Guided Federated Knowledge Distillation Approach in AI-RAN Enabled Multi-Access Edge Computing System

Luyao Zou; Hayoung Oh; Chu Myaet Thwal; Apurba Adhikary; Seohyeon Hong; Zhu Han

cs.LGMar 10, 2026

A Multi-Prototype-Guided Federated Knowledge Distillation Approach in AI-RAN Enabled Multi-Access Edge Computing System

Authors: Luyao Zou, Hayoung Oh, Chu Myaet Thwal, Apurba Adhikary, Seohyeon Hong, Zhu Han

Organizations: College of Computing and Informatics, Sungkyunkwan University, Republic of Korea · Department of Computer Science and Engineering, Kyung Hee University, Yongin-si, Gyeonggi-do, 17104, Republic of Korea · Department of Information and Communication Engineering, Noakhali Science and Technology University, Noakhali-3814, Bangladesh · Department of Electrical and Computer Engineering, University of Houston, Houston 77004, USA ·

Abstract

With the development of wireless network, Multi-Access Edge Computing (MEC) and Artificial Intelligence (AI)-native Radio Access Network (RAN) have attracted significant attention. Particularly, the integration of AI-RAN and MEC is envisioned to transform network efficiency and responsiveness. Therefore, it is valuable to investigate AI-RAN enabled MEC system. Federated learning (FL) nowadays is emerging as a promising approach for AI-RAN enabled MEC system, in which edge devices are enabled to train a global model cooperatively without revealing their raw data. However, conventional FL encounters the challenge in processing the non-independent and identically distributed (non-IID) data. Single prototype obtained by averaging the embedding vectors per class can be employed in FL to handle the data heterogeneity issue. Nevertheless, this may result in the loss of useful information owing to the average operation. Therefore, in this paper, a multi-prototype-guided federated knowledge distillation (MP-FedKD) approach is proposed. Particularly, self-knowledge distillation is integrated into FL to deal with the non-IID issue. To cope with the problem of information loss caused by single prototype-based strategy, multi-prototype strategy is adopted, where we present a conditional hierarchical agglomerative clustering (CHAC) approach and a prototype alignment scheme. Additionally, we design a novel loss function (called LEMGP loss) for each local client, where the relationship between global prototypes and local embedding will be focused. Extensive experiments over multiple datasets with various non-IID settings showcase that the proposed MP-FedKD approach outperforms the considered state-of-the-art baselines regarding accuracy, average accuracy and errors (RMSE and MAE).

Figures & tables

Fig. 1 : System model of the AI-RAN enabled MEC system, where each edge device can generate local prototypes and local model through training.

Fig. 2 : CHAC for per-class multi-prototype generation on the client (left) and architecture of the proposed MP-FedKD (right).

Dataset	Backbone	ACC	MAE	MFS
CIFAR10	S-CNN	$0.4948$	$1.8022$	$0.4944$
	ResNet-8	$0.6769$	$1.1301$	$0.6758$
	ResNet-10	0.6862	1.1033	0.6823
Fashion-MNIST	S-CNN	$0.8801$	$0.3865$	$0.8781$
	ResNet-8	$0.9007$	$0.3281$	$0.9015$
	ResNet-10	0.9097	0.2982	0.9088

TABLE I : Various backbone comparisons by considering accuracy (ACC), MAE, and Macro F1 Score (MFS).

Dataset	$\zeta_{m}^{c}=2$		$\zeta_{m}^{c}=3$		$\zeta_{m}^{c}=4$
Dataset	ACC	RMSE	ACC	RMSE	ACC	RMSE
CIFAR-10	0.6776	2.4146	$0.6710$	$2.4347$	$0.6390$	$2.5070$
MNIST	$0.9914$	$0.3872$	0.9933	0.3415	$0.9929$	$0.3833$
Fashion-MNIST	$0.9075$	$1.1463$	0.9097	1.1277	$0.9059$	$1.1846$
EuroSAT	$0.8413$	$1.6649$	$0.8390$	$1.7148$	0.8556	1.6456
M+F	$0.9350$	$0.9706$	0.9392	0.9218	$0.9288$	$1.0146$

TABLE II : Accuracy and RMSE achieved by the proposed method under various cluster numbers (i.e.,

\zeta_{m}^{c}

Fig. 3 : Accuracy comparison between K-Means and CHAC (Ours).

Panel of Figure 6 — Fig. 4 : RMSE and MAE achieved by FedProx, FedProto and the proposed method using various datasets ( $Dir=0.9$ ).

Dataset	Method	Accuracy
Dataset	Method	# 10	# 20	# 50
CIFAR10	FedProx	$0.3733$	$0.3155$	$0.3223$
	FedAS	$0.4816$	$0.4234$	$0.4037$
	E-FPKD	$0.4494$	$0.3501$	$0.3146$
	MOON	$0.6183$	$0.4302$	$0.3700$
	FedProto	$0.4458$	$0.3853$	$0.4109$

TABLE III : Accuracy comparison with considering various number of clients using CIFAR10, MNIST, Fashion-MNIST and EuroSAT datasets.

Dataset	Method	Accuracy
Dataset	Method	# 10	# 20	# 50
M+F	FedProx	$0.8433$	$0.6519$	$0.5615$
	FedAS	$0.9188$	$0.8949$	$0.8838$
	E-FPKD	$0.6117$	$0.5415$	$0.4960$
	MOON	$0.9322$	$0.8197$	$0.7790$
	FedProto	$0.9041$	$0.8342$	$0.8133$

TABLE IV : Accuracy comparison considering various number of clients on heterogeneous datasets.

Panel of Figure 9 — Fig. 5 : (a)-(c) Accuracy comparison among FedAS, MOON, FedProto and the proposed method via various datasets ( $Dir=0.9$ ). (d) illustrates the accuracy achieved by the proposed method with considering different cluster number settings towards M+F dataset.

Fig. 6 : Accuracy versus rounds for M+F dataset. A comparison between K-Means and CHAC based MP-FedKD is conducted.

Dataset	Method	Average Accuracy (# Clients=10)
Dataset	Method	Dir=0.3	Dir=0.5	Dir=0.7	Dir=0.9
MNIST	FedProx	$0.9342$	$0.9417$	$0.9442$	$0.9623$
	MOON	$0.9668$	$0.9695$	$0.9731$	$0.9886$
	FedALA	$0.9719$	$0.9762$	$0.9766$	$0.9859$
	Ours	0.9738	0.9770	0.9786	0.9927
Fashion-MNIST	FedProx	$0.6926$	$0.7856$	$0.7982$	$0.8194$

TABLE V : Average Accuracy achieved by the proposed method under various Dirichlet Distribution parameters.

Parameter	Average Accuracy (# Clients=10)
	FedProx	MOON	FedALA	Ours
$lr=0.1,bs=32$	$0.9880$	$0.8897$	$0.9812$	0.9893
$lr=0.01,bs=32$	$0.9865$	$0.9903$	$0.9854$	0.9920
$lr=0.001,bs=32$	$0.9623$	$0.9886$	$0.9776$	0.9927
$lr=0.001,bs=64$	$0.9632$	$0.9875$	$0.9793$	0.9915
$lr=0.001,bs=128$	$0.9683$	$0.9871$	$0.9786$	0.9881

TABLE VI : Average Accuracy achieved by the proposed method under various learning rate and batch size (MNIST).

Parameter	Average Accuracy (# Clients=10)
	FedProx	MOON	FedALA	Ours
$lr=0.1,bs=32$	$0.1040$	$0.4329$	$0.8191$	0.8387
$lr=0.01,bs=32$	$0.2649$	$0.7750$	$0.8185$	0.8563
$lr=0.001,bs=32$	$0.4821$	$0.7863$	$0.7992$	0.8409
$lr=0.001,bs=64$	$0.6820$	$0.5845$	$0.7078$	0.7672
$lr=0.001,bs=128$	$0.5023$	$0.3981$	$0.6079$	0.6702

TABLE VII : Average Accuracy achieved by the proposed method under various learning rate and batch size (EuroSAT).

Loss	CIFAR-10	MNIST	Fashion-MNIST
w/o PA	$0.6637$	$0.9785$	$0.8890$
w/o LEMGP	$0.6551$	$0.9745$	$0.8916$
Ours	0.6709	0.9927	0.9062

TABLE VIII : Average accuracy achieved by w/o PA, w/o LEMGP and the proposed method.

Federated Knowledge Distillation

Federated knowledge distillation transfers predictive behavior among data-isolated clients to train useful models without sharing raw data.

See all 36 papers

Optimized Federated Knowledge Distillation with Distributed Neural Architecture Search
May 20, 2026Chaimaa Medjadji, Sylvain Kubler, Yves Le Traon +3Communication-Efficient Distributed Training Federated Knowledge Distillation
Interdisciplinary Centre for Security, Reliability and Trust (SnT), University of Luxembourg, Luxembourg · Blekinge Institute of Technology, Sweden · ADSLabs, Umea University, Sweden
Adaptive Heterogeneous Compression for Resource-Efficient Federated Knowledge Distillation
Aug 16, 2026Chenwang Liu, Yijun Liu, Chang Liu +2Federated Knowledge Distillation Federated Learning
School of Information Engineering, Guangdong University of Technology, Guangzhou 510006, China · School of Communication and Information Engineering, Chongqing University of Posts and Telecommunications, Chongqing 400065, China
CRAD: Class-wise Reliability-Aware Distillation for Decentralized Heterogeneous Federated Learning
Aug 31, 2026Baraa Bilbeisi, Mengchen Fan, Baocheng Geng +1Federated Learning Aggregation Federated Knowledge Distillation
Dept. of Electrical & Computer Engineering · Dept. of Computer Science

May 20, 2026J/K move · Enter open · S save

Optimized Federated Knowledge Distillation with Distributed Neural Architecture Search

Chaimaa Medjadji, Sylvain Kubler, Yves Le Traon, Guilain Leduc+2

Interdisciplinary Centre for Security, Reliability and Trust (SnT), University of Luxembourg, Luxembourg · Blekinge Institute of Technology, Sweden · ADSLabs, Umea University, Sweden

Communication-Efficient Distributed Training Federated Knowledge Distillation Non-IID Federated Learning Federated Learning Heterogeneous Federated Learning

Open paper AskPDF arXiv

Federated Learning (FL) enables collaborative model training without centralizing data. However, real-world deployments must simultaneously address statistical heterogeneity across client data (non-IID), system heterogeneity in device capabilities, and communication efficiency. Existing FL approaches mitigate these challenges through improved aggregation, personalization, or knowledge distillation, but they almost universally assume a fixed client architecture, limiting adaptability to heterogeneous data complexity and hardware constraints. This architectural constraint often leads to suboptimal trade-offs between accuracy and efficiency in real-world FL systems. This work introduces FedKDNAS, a distillation-driven FL framework that combines client-side neural architecture selection with distillation of server-coordinated knowledge. Each client autonomously selects a lightweight model under accuracy-resource constraints. It then trains it locally using a hybrid objective combining supervised learning and knowledge distillation and shares only predictions on a public reference set. The server then aggregates and smooths these predictions, optionally combining them with a teacher model, to produce stable distillation targets for the next round. Extensive evaluation on six datasets against six representative FL baselines (FedAvg, Ditto, FedMD, FedDF, FedDistill, Local-KD) demonstrates that FedKDNAS consistently achieves superior Pareto efficiency, improving accuracy by up to 15% under non-IID conditions, reducing client CPU usage by approximately 28%, and decreasing communication overhead by up to 44 times while maintaining lightweight logit-based communication.

Abstract

Federated Knowledge Distillation

Explore similar work

Optimized Federated Knowledge Distillation with Distributed Neural Architecture Search

Adaptive Heterogeneous Compression for Resource-Efficient Federated Knowledge Distillation

CRAD: Class-wise Reliability-Aware Distillation for Decentralized Heterogeneous Federated Learning