Adaptive Heterogeneous Compression for Resource-Efficient Federated Knowledge Distillation

Chenwang Liu; Yijun Liu; Chang Liu; Xu Zhang; Pengchao Han

Source-linked AI summary

Adaptive Heterogeneous Compression for Resource-Efficient Federated Knowledge Distillation

Chenwang Liu, Yijun Liu, Chang Liu, Xu Zhang, Pengchao Han

arXiv:2608.15660v1Aug 16, 2026cs.DCcs.LG

TL;DR

FedKD still faces substantial gradient-communication overhead, while uniform compression overlooks clients’ heterogeneous models and resources. ASCEND models strategy selection as a non-stationary bandit with an EMA-enhanced ε-greedy policy, achieving a favorable balance between model accuracy and training latency across evaluated settings.

Problem
Existing FedKD compression methods generally use uniform strategies despite heterogeneous client resources, model architectures, and the resulting strategy-selection needs.
Method
ASCEND formulates client compression-strategy selection as a non-stationary stochastic bandit using a time-aware reward and EMA-enhanced ε-greedy policy.
Results
ASCEND consistently achieves a favorable balance between model accuracy and training latency on MNIST and CIFAR-10 under non-IID data and varying bandwidth capacities.
Takeaways & Limitations
Adaptive, client-specific compression can accommodate heterogeneous computation and communication resources in FedKD while preserving competitive accuracy and limiting training latency.

Abstract

from arXiv · show

Federated learning (FL) enables privacy-preserving distributed model training but faces challenges from heterogeneous model architectures and limited communication resources at the network edge. Federated knowledge distillation (FedKD) alleviates model heterogeneity by combining prototype-wise parameter aggregation and knowledge transfer across heterogeneous models. However, transmitting gradients still introduces considerable communication overhead, while existing compression approaches typically apply a uniform strategy across clients and ignore their diverse model characteristics and resource capacities. To address this issue, we propose a heterogeneous compression framework for FedKD that enables each client to select a compression strategy from a candidate strategy set. We formulate the compression strategy selection problem as a non-stationary stochastic multi-armed bandit (MAB), where each arm corresponds to a compression strategy. An efficiency-aware reward is designed by jointly considering local optimization improvement, global knowledge alignment, and execution time. Based on this formulation, we develop an Adaptive heterogeneouS Compression algorithm for fEderated kNowledge Distillation (ASCEND), which employs an exponential moving average (EMA)-enhanced $ε$-greedy policy to balance exploration and exploitation. Experimental results on multiple datasets demonstrate that ASCEND effectively adapts to heterogeneous model and resource settings, reducing communication overhead and training time while maintaining competitive model accuracy.

I. INTRODUCTION … B. Federated Knowledge Distillation

The paper motivates ASCEND by combining FedKD’s heterogeneous-model knowledge transfer with adaptive, client-specific gradient compression. It reviews related compression methods and formalizes FedKD’s training process as the foundation for the proposed resource-efficient framework.

I. INTRODUCTION: FedKD addresses heterogeneous model architectures through prototype-wise parameter averaging and cross-prototype logit-based knowledge distillation, but full-gradient exchange creates substantial communication overhead.A prototype is a specific model architecture adopted by a group of clients.
B. Communication-Efficient FL via Gradient Compression: Gradient compression methods such as Top-K, Random-K, and Periodic-K reduce communication by transmitting compact gradient representations, but most existing approaches use uniform or single-mechanism compression.This uniformity overlooks heterogeneous model architectures, data characteristics, and system resources.
I. INTRODUCTION: ASCEND lets each client select a compression strategy from candidate strategies according to model characteristics and resource capacities.The selection problem maximizes efficiency-aware performance improvement under heterogeneous conditions.
I. INTRODUCTION: The framework models client-specific strategy selection as a non-stationary stochastic multi-armed bandit, with an efficiency-aware reward combining local training loss, global knowledge alignment, and time cost.An EMA-enhanced ϵ-greedy policy balances exploration and exploitation during adaptive selection.
A. Federated Knowledge Distillation: Existing FedKD methods share knowledge through predictions, generated proxy samples, or intermediate representations, while hybrid aggregation combines averaging among homogeneous models with distillation across heterogeneous models.FedMD uses a shared public dataset, FedGen synthesizes proxy samples, and FedDF requires gradient uploads for ensemble aggregation.
A. Federated Learning: The paper’s system model consists of a central server and clients that collaboratively minimize private-data empirical risk, while classical FedAvg averages local updates but generally requires identical architectures.This limitation motivates knowledge transfer across heterogeneous models.
B. Federated Knowledge Distillation: In each FedKD round, clients compute and transmit local gradients; the server averages models within prototypes, constructs global logits on public data, and distills global knowledge into each prototype.Updated prototype parameters are transmitted back to clients for the next training round.

C. Gradient Compression with Residual Accumulation · 1) Client-Side Compression with Residual Accumulation:

This section reviews gradient compression with residual accumulation for reducing FedKD communication overhead while preserving convergence. It describes client-side residual updates and compares Top-K, Random-K, and Periodic-K compression strategies.

C. Gradient Compression with Residual Accumulation: FedKD clients transmit local gradients, creating substantial communication overhead, especially for large-scale models.Gradient compression with residual accumulation is introduced to reduce communication costs while preserving convergence guarantees.
1) Client-Side Compression with Residual Accumulation:: Each client adds its current gradient to a residual vector containing entries not transmitted in previous rounds before compression.This augmented update is passed to the compression operator, and residuals compensate for compression-induced information loss.
1) Client-Side Compression with Residual Accumulation:: Top-K selects the K largest-magnitude entries, preserving informative gradient directions at O(d log K) computational complexity.The compressed representation uses an index set of cardinality K and can be reconstructed in the original gradient dimension with zeros elsewhere.
1) Client-Side Compression with Residual Accumulation:: Random-K randomly samples K gradient indices with O(K) complexity but may cause high gradient variance and low model performance.This strategy trades lower computational cost for less reliable gradient selection.
1) Client-Side Compression with Residual Accumulation:: Periodic-K prioritizes unvisited coordinates within periods, then randomly completes selection when needed to cover all coordinates periodically.Its index-selection cost is O(K), and visited-coordinate flags reset after all coordinates have been selected.
1) Client-Side Compression with Residual Accumulation:: After compression, clients transmit compressed gradients together with their selected index sets to the server.The reconstruction operator maps transmitted values back to the original gradient dimension while setting remaining entries to zero.
1) Client-Side Compression with Residual Accumulation:: Residual accumulation preserves untransmitted gradient entries and is crucial for stable convergence despite compression.The residual vector is updated after transmission using the gradient information that was not sent.

2) Server-Side Processing of Compressed Gradients: … B. Convergence Analysis of FedKD with Heterogeneous Compression

The section describes server-side reconstruction and aggregation in FedKD, motivates heterogeneous compression through strategy-dependent trade-offs, and presents a client-adaptive framework with prototype-wise convergence analysis. The framework supports diverse model architectures and resources while achieving an O(1/T) convergence rate matching existing FL methods.

2) Server-Side Processing of Compressed Gradients:: The server reconstructs compressed gradient vectors before using them to update corresponding models for global aggregation and knowledge distillation.Reconstruction uses a reconstruction operator to recover the d-dimensional gradient vector.
3) Motivation for Heterogeneous Compression:: Top-K, Random-K, and Periodic-K exhibit different trade-offs between compression overhead and model performance across heterogeneous settings.The strategies are evaluated under different compression levels and model prototypes on CIFAR-10 and MNIST.
3) Motivation for Heterogeneous Compression:: Top-K is preferable under aggressive compression and for smaller models, whereas Random-K becomes more competitive with larger communication budgets because of lower selection overhead.The optimal strategy also depends on model architecture under the same compression setting.
3) Motivation for Heterogeneous Compression:: No single compression strategy consistently provides the best trade-off across heterogeneous clients, motivating heterogeneous rather than uniform compression in FedKD.Uniform enforcement may therefore produce suboptimal performance.
IV. FEDKD WITH HETEROGENEOUS COMPRESSION: The proposed heterogeneous compression framework lets clients choose strategies according to their model architectures and computational resources, followed by convergence analysis and adaptive-selection motivation.This section introduces the framework as FedKD with heterogeneous compression.
A. Heterogeneous Compression Framework: Each client independently trains locally, accumulates gradient residuals, applies a strategy from candidate set S, and sends compressed gradients with indices to the server.The server reconstructs gradients, updates client models, performs global logit averaging, applies KD on Dpub, and broadcasts updated prototype models.
B. Convergence Analysis of FedKD with Heterogeneous Compression: Convergence analysis is conducted separately for each model prototype because FedKD maintains multiple prototypes to handle heterogeneous model architectures.The analysis concerns the proposed heterogeneous compression framework.

C. Motivation for Adaptive Compression Algorithm … 3) MAB Formulation:

The paper formulates adaptive heterogeneous compression in FedKD as a non-stationary stochastic multi-armed bandit problem, using efficiency-aware rewards and EMA-enhanced ε-greedy selection to adapt strategies to changing client conditions. The formulation accounts for model improvement, knowledge alignment, optimization progress, execution time, and communication overhead.

C. Motivation for Adaptive Compression Algorithm: The framework selects client-specific compression strategies online because strategy performance is unknown before execution and varies with model states and system conditions.This motivates adaptive rather than uniform strategy selection.
V. FEDKD WITH ADAPTIVE HETEROGENEOUS COMPRESSION: The FedKD framework jointly formulates heterogeneous compression optimization, MAB-based strategy selection, adaptive compression, and online regret analysis.The overall section proceeds from optimization formulation to MAB reformulation and algorithm analysis.
A. Heterogeneous Compression Strategy Optimization Problem: The optimization problem assigns each client a strategy s_n^t from candidate set S to maximize efficiency-aware model performance.The objective is defined over clients and training rounds using strategy-dependent performance gain.
A. Heterogeneous Compression Strategy Optimization Problem: The optimization is difficult because performance gain is stochastic, unknown until execution, and time-varying under data, loss, and communication changes.These properties motivate online decision-making through a non-stationary stochastic MAB.
B. MAB Formulation: Each client chooses a sequence of compression strategies across training rounds, with time overhead combining compression, local training, and gradient transmission time.Compression time depends on strategy, model architecture, and compression level K.
2) Reward:: The reward measures efficiency-aware performance gain by combining global knowledge alignment and local optimization progress, then normalizing by execution time.The reward favors faster convergence with lower computational and communication overhead.
2) Reward:: Reward clipping prevents numerical instability from large fluctuations, using thresholds Gmin = −1 and Gmax = 1.The clipping mechanism limits the effect of extreme reward variations.

C. Algorithm Design · 1) Adaptive Exploration Policy: · 2) Rollback-Based Stability Safeguard:

ASCEND lets each client adapt compression strategies online in heterogeneous FedKD through an adaptive exploration policy and a rollback-based stability safeguard. The policy explores and updates strategy utilities across training stages, while rollback restores stability after sudden global-loss surges.

C. Algorithm Design: ASCEND independently optimizes each client’s compression strategy online to manage dynamic model-performance and computational-efficiency trade-offs in heterogeneous FedKD.The algorithm combines an adaptive exploration policy with a rollback-based stability safeguard.
1) Adaptive Exploration Policy:: ASCEND uses a two-phase exploration policy consisting of warm-up trials followed by adaptive strategy selection across later training rounds.The policy is designed to adapt to different training stages.
1) Adaptive Exploration Policy:: During warm-up, each client selects every compression strategy exactly h times, producing hS rounds with randomly permuted selection order.This avoids early-stage utility-estimation bias and systematic ordering bias.
1) Adaptive Exploration Policy:: Each strategy is treated as an independent arm, and empirical mean rewards initialize one utility estimate per compression strategy for subsequent EMA-enhanced ϵ-greedy updates.At warm-up completion, each client maintains S initialized utility values {Qn0(s)}s∈S.
1) Adaptive Exploration Policy:: After hS rounds, ASCEND selects compression strategies adaptively using utility estimates initialized by Qn0(s) and updated according to (33).The exploration rate includes fixed hyperparameter c.
2) Rollback-Based Stability Safeguard:: The rollback safeguard compares global loss across consecutive rounds and activates when any prototype exhibits a sudden loss surge.The stability threshold is predefined as γ > 1.
2) Rollback-Based Stability Safeguard:: When instability is detected, all client and prototype models roll back to round t −1 parameters, restoring the previously stable state.This temporarily constrains system behavior after unstable compression decisions.
2) Rollback-Based Stability Safeguard:: As a conservative fallback after instability, all clients adopt the Top-K strategy for the next round, reducing strategy heterogeneity and mitigating performance degradation.The safeguard is intended to enhance robustness against erratic compression decisions.

3) Algorithm Description: · D. Algorithm Analysis

ASCEND adaptively compresses client gradients, aggregates heterogeneous models, and transfers knowledge through prototype averaging and logit-based distillation. Its regret analysis models evolving rewards and establishes sublinear adaptive-selection regret, asymptotic optimality, and preserved convergence guarantees.

3) Algorithm Description:: In each round, clients train locally, select compression strategies via EMA-enhanced ε-greedy, compress residual-accumulated gradients, and upload compressed representations.The server reconstructs gradients, updates client models, averages models within each prototype, and uses public-data logit averaging for knowledge distillation.
D. Algorithm Analysis: ASCEND evaluates client-specific strategy selection using stochastic rewards affected by gradients, data sampling, and model evolution.The optimal strategy is defined per client and round, while regret measures cumulative expected-reward loss against round-wise optimal strategies.
D. Algorithm Analysis: The analysis assumes piecewise-stationary rewards, sub-Gaussian reward concentration, and sufficient sampling of every strategy within each stationary segment.These assumptions support analysis under evolving reward distributions during training.
D. Algorithm Analysis: Within a stationary segment, EMA tracking error decreases geometrically with the number of strategy updates, enabling frequently sampled strategies to adapt faster.Lemma 1 characterizes the EMA-based utility estimate error using the number of updates m.
D. Algorithm Analysis: Exploration regret grows sublinearly with communication rounds, so its average regret vanishes as T increases.Lemma 2 derives the exploration upper bound under the piecewise-stationary and reward-concentration assumptions.
D. Algorithm Analysis: The adaptive phases incur O(T) exploration regret and O(MT) exploitation regret under a piecewise-stationary reward model.Lemmas 2 and 3 jointly characterize regret from exploration and exploitation, with exploitation depending on the number of stationary segments M.
D. Algorithm Analysis: Theorem 2 bounds ASCEND’s cumulative regret through warm-up, exploration, and exploitation terms, while fixed warm-up length yields the resulting asymptotic bound.The bound approaches the performance of the round-wise optimal compression strategy, with larger M increasing tracking cost.
D. Algorithm Analysis: Together with convergence under any strategy in S, the regret result shows ASCEND can select communication-efficient compression strategies without compromising federated optimization convergence.This conclusion combines Theorem 1’s convergence guarantee with Theorem 2’s adaptive-selection analysis.

VI. EXPERIMENTS … 4) Time overheads:

The experiments evaluate ASCEND’s effectiveness and robustness across heterogeneous federated settings, using real-world edge devices, multiple datasets and model architectures, and measured computation and communication overheads. The setup varies data heterogeneity and models while feeding measured execution times into compression-strategy selection.

VI. EXPERIMENTS: The experiments assess ASCEND’s effectiveness and robustness through empirical evaluations and sensitivity analysis of critical hyperparameters.The evaluation includes a real-world platform and various experimental studies.
A. Setup: The setup specifies federated settings, datasets, models, time overheads, and baselines.These components define the experimental evaluation framework.
1) Federated setting:: 10 Raspberry Pi devices serve as resource-constrained edge clients, with each client assigned a local architecture from the candidate prototype set.The batch size is 32, and the learning rate is 0.02.
3) Models:: MNIST uses LeNet5 and LeNet5Half, while CIFAR-10 uses ResNet18 and WResNet40-2 to represent heterogeneous model architectures.The corresponding parameter counts are 61,706 and 15,738 for MNIST models, and 11,183,562 and 2,248,954 for CIFAR-10 models.
4) Time overheads:: ASCEND receives real computation and communication times to determine compression strategies, with computation measured on an NVIDIA RTX 3090 GPU and AMD Ryzen 7 4800H CPU.Computation time includes local training and gradient compression, and measurements are averaged over multiple runs.
4) Time overheads:: The default communication bandwidth is 50 Mbps, with 32-bit gradient values and indices encoded using log2 d bits.K is the number of retained gradient coordinates, bv = 32, and d is the original gradient-vector dimensionality.

5) Baselines: … VII. CONCLUSION

ASCEND is evaluated against uniform and adaptive compression baselines across real-world, heterogeneous-data, bandwidth, strategy-selection, and sensitivity settings. Results show competitive performance, adaptive client decisions, and a favorable accuracy–latency balance on MNIST and CIFAR-10.

5) Baselines; B. Experimental Results and Analysis: ASCEND is compared with uniform compression baselines and an EXP3-based adaptive method across compression, heterogeneity, bandwidth, selection, and sensitivity experiments.The evaluation covers real-world deployment, baseline comparisons, robustness, client strategy evolution, and hyperparameter sensitivity.
1) Performance on Real-World Platform:: On real-world MNIST experiments, ASCEND performs competitively at K ∈{5 × 103, 104}, achieves best performance at K = 104, and converges faster early at K = 5×103.The results demonstrate a balance between communication efficiency and model accuracy, attributed to adaptive strategy selection.
2) Performance Comparison with Baselines:: For MNIST and CIFAR-10, tested compression levels correspond to LeNet5 retained-gradient ratios of approximately 0.81%, 1.62%, and 8.10%, and ResNet18 ratios of 0.89%, 4.47%, and 17.88%.The experiments use MNIST K ∈{5 × 102, 103, 5 × 103} and CIFAR-10 K ∈{105, 5 × 105, 2 × 106}.
3) Performance under Varying Data Heterogeneity:: Under α ∈{0.1, 0.5, 1} with K = 103 and 50 Mbps, higher heterogeneity slows convergence and lowers final accuracy, while ASCEND remains competitive or superior to the best fixed strategy.The settings span highly heterogeneous to nearly IID data distributions.
4) Performance under Varying Communication Bandwidths:: Across bandwidths of 5 Mbps, 50 Mbps, and 500 Mbps with K = 103 and α = 0.5, ASCEND remains competitive or superior despite Top-K and Random-K advantages in specific conditions.Top-K performs better at 5 Mbps, whereas Random-K becomes more competitive at 500 Mbps.
5) Analysis of Client Compression-Strategy Selection:: ASCEND selects among Top-K, Random-K, and Periodic-K according to model architecture and compression level, favoring Top-K for LeNet5Half at K = 103 and Random-K at K = 5 × 103.LeNet5 shows a more balanced Top-K and Random-K preference at K = 103, and CIFAR-10 exhibits the same trend.
6) Sensitivity Analysis:: The sensitivity analysis identifies ρ = 0.6 and β = 0.4 as stable across datasets, while c = 3 balances exploration and exploitation with relatively favorable performance.These parameters are evaluated for weighting and adaptive strategy selection, respectively.
VII. CONCLUSION: ASCEND formulates heterogeneous compression selection as a non-stationary stochastic MAB and uses EMA-enhanced ϵ-greedy selection with a time-aware reward, achieving a favorable accuracy–training-latency balance.The conclusion reports results on MNIST and CIFAR-10 under non-IID data and different bandwidth capacities; future work will adapt compression levels jointly with strategy selection.

APPENDIX · A. Proof of Theorem 1

The appendix formalizes one communication round and proves Theorem 1 for prototype-wise federated knowledge distillation with heterogeneous compression. The proof establishes smoothness, decomposes compression error through error feedback, and derives an O(1/...) convergence decrease under bounded-noise and residual assumptions.

APPENDIX: The appendix first presents the overall workflow of a single communication round in Figure 15.
A. Proof of Theorem 1: Each prototype round begins from a server-broadcast model, with clients assigned to the same prototype sharing identical parameters.
A. Proof of Theorem 1: The server reconstructs compressed client gradients, updates local models, aggregates parameters prototype-wise, and applies one-step KD on Dpub.
A. Proof of Theorem 1: The joint objective Fp(ω) combines local optimization and global knowledge alignment, with λ balancing the two objectives.
A. Proof of Theorem 1: Lemma 4 proves that Fp(ω) is L-smooth with L = Ll + λLp > 0.
A. Proof of Theorem 1: Lemma 5 decomposes the deviation between the aggregated reconstructed gradient and the true prototype gradient under heterogeneous compression.
A. Proof of Theorem 1: Theorem 1’s proof uses vanishing expected stochastic-gradient noise, bounded gradient deviation Cg, and bounded residual accumulation Ck.
A. Proof of Theorem 1: With the prescribed parameter choice, both convergence-bound terms decrease as O(1/...).

B. Proof of Theorem 2 · 1) Proof of Lemma 1: · 2) Proof of Lemma 2:

Theorem 2 is established through two lemmas: one analyzes EMA utility updates under piecewise-constant expected rewards, while the other bounds exploration regret over communication rounds. The proofs use recursion, reward clipping, and an integral inequality to derive the stated bounds.

1) Proof of Lemma 1:: Lemma 1 begins from the EMA utility update for a selected compression strategy.The update is given by equation (33).
1) Proof of Lemma 1:: Within segment T_i, Assumption 6 and equation (38) imply that the expected reward remains constant.The constancy holds throughout the segment T_i.
1) Proof of Lemma 1:: Rearranging equation (33) yields the utility recursion used in the proof.The recursion is obtained by manipulating both sides of the update equation.
1) Proof of Lemma 1:: Applying the recursion for m updates produces the multi-step expression required for Lemma 1.The proof iterates the update over m steps.
1) Proof of Lemma 1:: Taking absolute values on both sides yields equation (46), completing the bound used in Lemma 1.This step converts the preceding recursion into the stated absolute-value result.
2) Proof of Lemma 2:: Lemma 2 bounds exploration regret using ASCEND’s random strategy selection during the exploration stage.At communication round t, a compression strategy is selected randomly with probability ϵ_t according to equation (35).
2) Proof of Lemma 2:: Reward clipping bounds the instantaneous regret in each communication round.The clipping operation provides the per-round regret bound used in the proof.
2) Proof of Lemma 2:: The resulting per-round bound is accumulated to establish the cumulative exploration-regret bound after T communication rounds.The final derivation applies an integral inequality.

3) Proof of Lemma 3: · 4) Proof of Theorem 2:

Lemma 3 bounds ASCEND’s exploitation regret by decomposing utility-estimation error into EMA bias and stochastic fluctuation over stationary segments. Theorem 2 then combines warm-up analysis with Lemmas 2 and 3 to obtain the overall regret bound.

3) Proof of Lemma 3:: During exploitation, ASCEND selects the compression strategy with the highest current estimated utility.This selection rule initiates the instantaneous exploitation-regret analysis.
3) Proof of Lemma 3:: The instantaneous exploitation regret is controlled through the utility-estimation error of the selected strategy.The proof defines the estimation error and relates it to the regret incurred when the estimated maximizer is chosen.
3) Proof of Lemma 3:: Within each stationary segment Ti, the utility-estimation error decomposes into EMA bias and stochastic fluctuation.The decomposition uses the unchanged expected reward within a stationary segment.
3) Proof of Lemma 3:: After Hn updates of strategy s, Lemma 1 and Assumption 7 bound its utility-estimation error using constants C1 and C2.The constants depend on reward variance, the EMA coefficient, and initialization error, with a ∨ b defined as max{a, b}.
3) Proof of Lemma 3:: Combining the estimation-error bound with the exploitation analysis yields an expected instantaneous-regret bound for each round t ∈ Ti.The bound is then summed across stationary segments, with the second inequality obtained using Cauchy–Schwarz.
4) Proof of Theorem 2:: During warm-up, each compression strategy is selected exactly h times, producing h|S| warm-up rounds.This phase initializes the utility estimates before exploitation begins.
4) Proof of Theorem 2:: Because rewards are clipped to [Gmin, Gmax], the warm-up regret is bounded by the maximum instantaneous regret, and Theorem 2 follows by combining this bound with Lemmas 2 and 3.The final result combines the warm-up contribution with the previously established regret bounds.