Client-Aware Security Adaptive Early Exit Vision Transformers for Heterogeneous Federated Learning

Sahar Refaei Ghalibafan; Samira Nazari; Christian Herglotz; Maksim Jenihhin; Mahdi Taheri

cs.LGOct 16, 2026

Client-Aware Security Adaptive Early Exit Vision Transformers for Heterogeneous Federated Learning

Authors: Sahar Refaei Ghalibafan, Samira Nazari, Christian Herglotz, Maksim Jenihhin, Mahdi Taheri

Organizations: Brandenburg University of Technology, Cottbus, Germany · University of Zanjan, Zanjan, Iran · Tallinn University of Technology, Tallinn, Estonia · Humboldt University of Berlin, Berlin, Germany

Abstract

Federated learning (FL) trains a shared model across data-owning clients without centralizing raw data, but client data difficulty and compute budgets can differ. We study local early-exit Vision Transformer (ViT) inference, where an intermediate classifier may return a prediction before the final Transformer stage. The proposed method replaces one global acceptance threshold with a client-specific base threshold from label entropy and shallow-exit confidence, then adjusts each input using predictive entropy, confidence, top-two class margin, and exit depth. Aggregation uses only clients that updated a parameter. Security-aware routing reaches 88.70% CIFAR-10 accuracy with 32.0% depth saving and 73.7% ImageNet-1K top-1 accuracy with 30% saving. On matched six-client non-IID CIFAR-100, full-depth accuracy is 72.28% versus 68.20--71.80% for four heterogeneous-FL baselines; robustness-aware routing reaches 70.01% with 32.55% saving.

Figures & tables

Policy	$a$	$b$	$c$	$d$
Fixed	0	0	0	0
Global resource-aware	0	.25	0	.10
Client-specific	.25	0	0	0
Robustness-aware	0	.25	.20	.10
Combined client+robust.	.25	.25	.20	.10

TABLE I: Instantiated Exit Policies

Panel of Figure 2 — Fig. 1: Federated training (a) and local client inference (b).

Setting	Global acc. (%)	Mean client (%)	Saved (%)
3 clients, IID	90.62	91.78	–
6 clients, IID	90.17	91.30	–
9 clients, IID	89.94	91.02	–
Full depth	90.17	–	0.0
Fixed threshold	87.00	–	37.5
Global resource-aware	87.60	–	36.3

TABLE II: CIFAR-10 Evaluation Across Client Counts and Exit Policies

Policy	Top-1 (%)	Avg. depth	Saved (%)
Full depth	75.4	6.00	0
Fixed threshold	72.1	3.96	34
Global resource-aware	72.6	4.02	33
Robustness-aware	73.7	4.20	30

TABLE III: Six-Client ImageNet-1K Accuracy–Depth Trade-off

Method	Acc. (%)	$\Delta$ vs. ours (pp)	Compute control
HeteroFL [ 12 ]	68.20	-4.08	model width
DepthFL [ 13 ]	70.10	-2.18	trainable depth
ScaleFL [ 14 ]	71.30	-0.98	width + depth
ReeFL [ 15 ]	71.80	-0.48	recurrent exits
Ours, full depth	72.28	0.00	6 stages
Ours, robustness-aware	70.01	-2.27	avg. depth 4.05

TABLE IV: Matched CIFAR-100 Comparison: 6 Clients, Dirichlet

\alpha=0.1

Fig. 2: CIFAR-100 analysis: client disparity and communication (a–b), non-IID degradation and clean accuracy–depth trade-off (c–d), and FGSM/PGD routing comparisons (e–f).

Policy	Acc. (%)	Avg. depth	Saved (%)
Full depth	72.28	6.00	0.00
Fixed threshold	68.80	3.85	35.90
Global resource-aware	68.95	3.83	36.25
Robustness-aware	70.01	4.05	32.55

TABLE V: CIFAR-100 Global-Test Accuracy–Depth Trade-off

Early-Exit Neural Networks

Early-exit neural networks stop processing inputs at intermediate layers when predictions are sufficiently reliable, trading accuracy for lower latency and energy.

See all 44 papers

Heterogeneity-Aware Client Selection Methodology For Efficient Federated Learning
Feb 24, 2026Nihal Balivada, Shrey Gupta, Shashank Shreedhar Bhatt +1Client Selection in Federated Learning Heterogeneous Federated Learning
University of Oregon, US · Boston College, US · Microsoft Research, India
Byzantine-Robust Federated Representation Learning
Sep 29, 2026Leonardo F. Toso, James Anderson, Rafael Pinot +1Federated Learning Byzantine-Robust Federated Learning
Columbia University, USA · Sorbonne Universit´e and Universit´e Paris Cit´e, CNRS, LPSM, France · University of Copenhagen, Denmark
FoggyTrust: Robust Federated Learning with Hierarchical Trust Networks
Jun 26, 2026Emmanuel Rassou, Tomas GonzalezHierarchical Federated Learning Federated Learning Aggregation
School of Engineering and Applied Sciences (SEAS) Harvard Cambridge, MA, USA

Feb 24, 2026J/K move · Enter open · S save

Heterogeneity-Aware Client Selection Methodology For Efficient Federated Learning

Nihal Balivada, Shrey Gupta, Shashank Shreedhar Bhatt, Suyash Gupta

University of Oregon, US · Boston College, US · Microsoft Research, India

Client Selection in Federated Learning Heterogeneous Federated Learning Non-IID Federated Learning

Open paper AskPDF arXiv

Federated Learning (FL) enables a distributed client-server architecture where multiple clients collaboratively train a global Machine Learning (ML) model without sharing sensitive local data. However, FL often results in lower accuracy than traditional ML algorithms due to statistical heterogeneity across clients. Prior works attempt to address this by using model updates, such as loss and bias, from client models to select participants that can improve the global model's accuracy. However, these updates neither accurately represent a client's heterogeneity nor are their selection methods deterministic. We mitigate these limitations by introducing Terraform, a novel client selection methodology that uses gradient updates and a deterministic selection algorithm to select heterogeneous clients for retraining. This bi-pronged approach allows Terraform to achieve up to 47 percent higher accuracy over prior works. We further demonstrate its efficiency through comprehensive ablation studies and training time analyses, providing strong justification for the robustness of Terraform.

Abstract

Early-Exit Neural Networks

Explore similar work

Heterogeneity-Aware Client Selection Methodology For Efficient Federated Learning

Byzantine-Robust Federated Representation Learning

FoggyTrust: Robust Federated Learning with Hierarchical Trust Networks