Two-Stage Robust Knowledge Distillation

• The paper introduces a two-stage pipeline that robustly filters client models using cosine similarity, clustering, and median-based selection to form a reliable ensemble.
• Stage 1 leverages HDBSCAN clustering on scalar representations of client updates to discard malicious contributions in heterogeneous, non-IID data scenarios.
• Stage 2 distills consensus predictions from the vetted ensemble into a global model, achieving high accuracy (above 80%) with reduced attack success rates (below 17%).

Two-Stage Robust Knowledge Distillation (RKD) is best understood as a server-side robust aggregation pipeline for federated learning in which robust model vetting and ensemble construction precede ensemble-to-global knowledge distillation. In this usage, the phrase aligns most directly with Robust Knowledge Distillation for backdoor-robust federated learning under heterogeneous non-IID client data: client updates are clustered and filtered, a reliable ensemble is formed from representative benign models, and the next global model is distilled from that ensemble on proxy data rather than obtained by direct parameter averaging (Alharbi et al., 1 Feb 2025).

1. Definition, setting, and threat model

In the standard cross-device federated learning architecture, a central server sends a global model to clients in each communication round; clients train locally on private data and return model parameters or updates; the server aggregates them into the next global model. The RKD formulation targets this setting under heterogeneous, non-IID client data and explicitly addresses backdoor attacks in which some clients are malicious and can perform both data poisoning and model poisoning. The malicious clients may collude, poison a fraction of their local data with triggers and target labels, and further manipulate model parameters through scaling, perturbation, or adaptive optimization, but they do not control the server or intercept other clients’ communications (Alharbi et al., 1 Feb 2025).

The method is positioned against several limitations in prior federated defenses. Robust aggregation rules are described as failing under non-IID data or scaling/adaptive attacks; trust-root methods can be bypassed by stealthy updates; differential-privacy or noise-based defenses require delicate tuning and may hurt utility; and prior distillation methods, although often more tolerant of non-IID drift, do not explicitly detect and exclude backdoored client models before distillation. The difficult regimes emphasized are high malicious-client ratios, sophisticated backdoor attacks, and strongly non-IID client distributions generated by Dirichlet partitioning (Alharbi et al., 1 Feb 2025).

The underlying paper does not formally name the method “Two-Stage RKD.” It instead describes three components—Automated Clustering, Model Selection, and Knowledge Distillation. The two-stage interpretation follows from their sequential role: Stage 1 is clustering, model selection, and reliable ensemble construction; Stage 2 is distillation from that vetted ensemble into the new global model. This suggests that “two-stage RKD” is a structural reading of the method rather than its official title (Alharbi et al., 1 Feb 2025).

2. Stage 1: robust filtering and reliable ensemble construction

The first stage converts uploaded client models into a vetted teacher ensemble. Rather than clustering full parameter vectors, the server compresses each client model into a scalar representation: the cosine similarity between the client parameters and the current global model. For client $i$ at round $r$,

$$s_i = \frac{\left(\theta_i^{(r)}\right)^\top M_{\text{global}^{(r)}}}{\|\theta_i^{(r)}\| \, \|M_{\text{global}^{(r)}}\|}, \quad i=1,\dots,N.$$

The rationale given is that malicious backdoor updates often exhibit angular deviation from the benign optimization direction even when their norm is not obviously anomalous. Clustering scalar similarities rather than raw high-dimensional parameters reduces computational burden and avoids the curse of dimensionality (Alharbi et al., 1 Feb 2025).

Clustering is performed with HDBSCAN. Its minimum cluster size is scheduled dynamically by communication round:

$$Q = \max\left(2,\left\lceil 0.2N - r \right\rceil\right),$$

where $N$ is the number of participating clients and $r$ is the current round. HDBSCAN returns labels

$\{L_i\} = \text{HDBSCAN}(\{s_i\}, Q).$

For each cluster $C_k$, the server computes mean cosine similarity

$\mu_k = \frac{1}{|C_k|}\sum_{i\in C_k} s_i.$

The cluster with the largest mean similarity,

$\mu_{\max} = \max_k \mu_k,$

is designated benign, and the remaining clusters are treated as suspicious or malicious. The benign client indices are

$r$0

with corresponding model set

$r$1

The operative heuristic is explicit: the highest-mean-similarity cluster relative to the current global model is retained as benign, and everything else is excluded (Alharbi et al., 1 Feb 2025).

A second filtering step is then applied within the benign cluster. The server computes an element-wise median parameter vector

$r$2

followed by the $r$3 distance of each benign-cluster model to this median,

$r$4

The reliable ensemble is

$r$5

where $r$6 is chosen using a statistical measure based on the standard deviation of distances from the median. The paper does not provide a fully explicit closed-form expression for $r$7, but the stated criterion is to keep models close to the central tendency of the benign cluster. Clustering therefore removes large directional outliers, while median-distance filtering removes residual subtle outliers inside the presumed-benign group (Alharbi et al., 1 Feb 2025).

3. Stage 2: ensemble-to-global distillation and round-wise server workflow

The second stage distills the vetted ensemble into the new global model. The teachers are the selected models in $r$8, while the student is a server-maintained model $r$9, later promoted to $$s_i = \frac{\left(\theta_i^{(r)}\right)^\top M_{\text{global}^{(r)}}}{\|\theta_i^{(r)}\| \, \|M_{\text{global}^{(r)}}\|}, \quad i=1,\dots,N.$$0. Distillation uses an unlabeled validation dataset $$s_i = \frac{\left(\theta_i^{(r)}\right)^\top M_{\text{global}^{(r)}}}{\|\theta_i^{(r)}\| \, \|M_{\text{global}^{(r)}}\|}, \quad i=1,\dots,N.$$1, specified as $$s_i = \frac{\left(\theta_i^{(r)}\right)^\top M_{\text{global}^{(r)}}}{\|\theta_i^{(r)}\| \, \|M_{\text{global}^{(r)}}\|}, \quad i=1,\dots,N.$$2 of the total training data. For each $$s_i = \frac{\left(\theta_i^{(r)}\right)^\top M_{\text{global}^{(r)}}}{\|\theta_i^{(r)}\| \, \|M_{\text{global}^{(r)}}\|}, \quad i=1,\dots,N.$$3, teacher logits are averaged arithmetically:

$$s_i = \frac{\left(\theta_i^{(r)}\right)^\top M_{\text{global}^{(r)}}}{\|\theta_i^{(r)}\| \, \|M_{\text{global}^{(r)}}\|}, \quad i=1,\dots,N.$$4

Soft pseudo-labels are then formed by temperature-scaled softmax,

$$s_i = \frac{\left(\theta_i^{(r)}\right)^\top M_{\text{global}^{(r)}}}{\|\theta_i^{(r)}\| \, \|M_{\text{global}^{(r)}}\|}, \quad i=1,\dots,N.$$5

and the student is optimized with KL divergence,

$$s_i = \frac{\left(\theta_i^{(r)}\right)^\top M_{\text{global}^{(r)}}}{\|\theta_i^{(r)}\| \, \|M_{\text{global}^{(r)}}\|}, \quad i=1,\dots,N.$$6

Training proceeds for $$s_i = \frac{\left(\theta_i^{(r)}\right)^\top M_{\text{global}^{(r)}}}{\|\theta_i^{(r)}\| \, \|M_{\text{global}^{(r)}}\|}, \quad i=1,\dots,N.$$7 epochs on mini-batches from $$s_i = \frac{\left(\theta_i^{(r)}\right)^\top M_{\text{global}^{(r)}}}{\|\theta_i^{(r)}\| \, \|M_{\text{global}^{(r)}}\|}, \quad i=1,\dots,N.$$8 using SGD with learning rate $$s_i = \frac{\left(\theta_i^{(r)}\right)^\top M_{\text{global}^{(r)}}}{\|\theta_i^{(r)}\| \, \|M_{\text{global}^{(r)}}\|}, \quad i=1,\dots,N.$$9. After each epoch, the server updates a stochastic weight averaging model:

$$Q = \max\left(2,\left\lceil 0.2N - r \right\rceil\right),$$0

and then sets

$$Q = \max\left(2,\left\lceil 0.2N - r \right\rceil\right),$$1

The global update is therefore not a direct weighted average of client parameters but a distilled student trained to imitate a vetted ensemble on proxy data (Alharbi et al., 1 Feb 2025).

This choice is the central robustness mechanism. Any single poisoned model is intended to be excluded before teacher construction; residual malicious behavior is diluted by logit averaging; and the student learns consensus predictions on proxy data rather than inheriting raw parameter-space artifacts. The paper’s interpretation is that this also improves generalization under non-IID client drift, especially when benign models disagree in parameter space but remain broadly aligned in predictive behavior (Alharbi et al., 1 Feb 2025).

The per-round protocol includes an additional communication control. At round $$Q = \max\left(2,\left\lceil 0.2N - r \right\rceil\right),$$2, the server initializes $$Q = \max\left(2,\left\lceil 0.2N - r \right\rceil\right),$$3 and broadcasts it to all clients. In later rounds, the updated model is sent only to clients deemed benign in the previous round. Clients classified as malicious may either be denied the new global model entirely or receive a perturbed global model (PGM) with noise scale around $$Q = \max\left(2,\left\lceil 0.2N - r \right\rceil\right),$$4. This selective dissemination is intended to reduce attacker adaptivity, particularly for A3FL-style attacks that exploit the global model’s trajectory, while still allowing dynamic reclassification of benign clients that may have been temporarily misidentified due to non-IID drift (Alharbi et al., 1 Feb 2025).

4. Experimental protocol and reproducibility parameters

The empirical study uses 30 clients and three vision datasets with task-specific backbones: CIFAR-10 with ResNet-18, EMNIST with a 2-conv CNN plus pooling, dropout, and a fully connected layer, and Fashion-MNIST with a 2-conv CNN plus batch normalization, dropout, and a fully connected layer (Alharbi et al., 1 Feb 2025).

Dataset	Backbone	Client partitioning
CIFAR-10	ResNet-18	Dirichlet non-IID
EMNIST	2-conv CNN + pooling/dropout + FC	Dirichlet non-IID
Fashion-MNIST	2-conv CNN + batchnorm/dropout + FC	Dirichlet non-IID

The non-IID severity is controlled by Dirichlet concentration parameters $$Q = \max\left(2,\left\lceil 0.2N - r \right\rceil\right),$$5, with smaller $$Q = \max\left(2,\left\lceil 0.2N - r \right\rceil\right),$$6 indicating stronger heterogeneity. Malicious clients comprise $$Q = \max\left(2,\left\lceil 0.2N - r \right\rceil\right),$$7, $$Q = \max\left(2,\left\lceil 0.2N - r \right\rceil\right),$$8, or $$Q = \max\left(2,\left\lceil 0.2N - r \right\rceil\right),$$9 of participants, and each malicious client poisons $N$0 of its local training data with a trigger and target label. Evaluated attacks are A3FL, F3BA, DBA, and TSBA, with ADBA reported in the appendix. Baselines are FedAvg, FLAME, FedDF, FedRAD, FedBE, RLR, and FoolsGold (Alharbi et al., 1 Feb 2025).

The reported metrics are main task accuracy (MTA) on clean test data and attack success rate (ASR) on triggered test data:

$N$1

Experiments are repeated five times with confidence intervals at significance level $N$2. Reproduction-critical details include local batch size $N$3, local learning rate $N$4 for CIFAR-10 and $N$5 for EMNIST and Fashion-MNIST, the validation split $N$6 of the total training data, HDBSCAN with $N$7, median-based ensemble selection with an $N$8-distance threshold informed by distance standard deviation, KL-divergence distillation, SGD during KD, SWA over KD epochs, and optional perturbed dissemination with noise scale $N$9 (Alharbi et al., 1 Feb 2025).

Some parameters remain unspecified numerically in the text: the temperature $r$0, the number of KD epochs $r$1, and the exact formula translating the distance standard deviation into $r$2. The paper indicates that these omissions would need to be resolved from code (Alharbi et al., 1 Feb 2025).

5. Empirical findings, ablations, and efficiency profile

The principal empirical claim is that RKD consistently attains lower ASR while preserving higher MTA than the listed baselines across datasets, attack families, attacker ratios, and non-IID regimes. The headline summary is that it keeps accuracy above $r$3 while reducing ASR below $r$4 in difficult scenarios. The strongest advantages are reported under high heterogeneity and high attacker ratios, where competing methods often either maintain clean accuracy while failing to suppress ASR, or suppress ASR at substantial clean-accuracy cost (Alharbi et al., 1 Feb 2025).

Under A3FL with $r$5, the reported explanation for RKD’s advantage combines directional clustering, HDBSCAN outlier detection, median-based model selection, consensus transfer by knowledge distillation, SWA smoothing, and restricting suspected malicious clients’ access to the latest global model. Under F3BA and TSBA with $r$6, robustness is maintained even at $r$7 and $r$8 attacker ratios. Under DBA with $r$9, the advantage over prior distillation-based methods is attributed to the fact that RKD does not distill indiscriminately from all client models (Alharbi et al., 1 Feb 2025).

The ablation study strongly supports the two-stage interpretation. Removing clustering and model selection while retaining only distillation causes severe robustness collapse as attacker fraction rises, showing that KD alone is insufficient if poisoned teachers remain in the ensemble. Removing model selection but retaining clustering still degrades performance at higher attack ratios, indicating that HDBSCAN filtering alone is not enough because residual outliers inside the benign cluster still matter. Removing KD while retaining clustering and model selection lowers ASR but causes about a $\{L_i\} = \text{HDBSCAN}(\{s_i\}, Q).$0 drop in accuracy relative to full RKD, especially under non-IID data. The empirical implication is that robust filtering constructs the right teachers, while distillation converts them into a global model that remains useful under heterogeneity (Alharbi et al., 1 Feb 2025).

The method also occupies a middle point in the reported defense-time spectrum. Comparative defense times are RLR $\{L_i\} = \text{HDBSCAN}(\{s_i\}, Q).$1 s, FLAME $\{L_i\} = \text{HDBSCAN}(\{s_i\}, Q).$2 s, RKD $\{L_i\} = \text{HDBSCAN}(\{s_i\}, Q).$3 s, FedDF $\{L_i\} = \text{HDBSCAN}(\{s_i\}, Q).$4 s, and FedBE $\{L_i\} = \text{HDBSCAN}(\{s_i\}, Q).$5 s. This places RKD slightly slower than FLAME but substantially faster than prior distillation-heavy baselines, consistent with the design choice to cluster one scalar similarity score per client rather than full parameter vectors (Alharbi et al., 1 Feb 2025).

6. Limitations, interpretive status, and terminology in the broader KD literature

The method is empirical rather than theorem-driven. The paper does not provide a formal guarantee, but instead offers several mechanistic intuitions: cosine similarity captures angular deviation commonly associated with poisoned updates; HDBSCAN isolates dense groups without fixing the number of clusters; the element-wise median and $\{L_i\} = \text{HDBSCAN}(\{s_i\}, Q).$6-distance pruning suppress residual anomalous models; ensemble distillation transfers consensus predictive behavior; SWA encourages flatter minima; and selective withholding or perturbation of the global model limits adaptive attacker feedback (Alharbi et al., 1 Feb 2025).

The limitations are also explicit. In an extreme threat model where malicious clients somehow know all benign updates in a round, a coordinated adversary can craft one update to shift the average and another to align perfectly in cosine similarity with that manipulated average, thereby infiltrating the benign cluster. More generally, the approach may degrade when benign updates are so heterogeneous that the highest-similarity cluster is no longer a reliable proxy for benignness, or when attackers match benign directional statistics closely enough to survive both HDBSCAN and median-distance filtering. The method also depends on the availability and quality of a server-side distillation dataset; the paper fixes this at $\{L_i\} = \text{HDBSCAN}(\{s_i\}, Q).$7 of total training data and does not deeply study proxy-data quality. Selective withholding or perturbing of the global model could additionally slow or complicate convergence for benign clients that are temporarily misclassified, although dynamic reclassification is intended to mitigate this (Alharbi et al., 1 Feb 2025).

The phrase “RKD” is not unique in arXiv usage, and this affects interpretation. “Residual Knowledge Distillation” uses an assistant model that learns the residual error between teacher and student representations rather than robustly filtering teachers (Gao et al., 2020). “RobustKD” against a backdoored teacher purifies teacher feature maps by feature variance before distillation, which is a different robustness problem from federated backdoor defense (Chen et al., 2024). In 3D detection, “RKD” can mean “Reflectance Knowledge Distillation” inside the broader RPKD framework (Jing et al., 23 May 2025). A separate semantic communication paper explicitly uses the phrase “two-stage robust knowledge distillation” for a stage-1 semantic transfer and stage-2 robustness-enhancement procedure tied to a channel codec, again in a different task domain (Ding et al., 4 Aug 2025). This suggests that “Two-Stage Robust Knowledge Distillation” is better treated as a descriptive structural label than as a single canonical algorithmic name across the literature.