Adaptive Heterogeneous Compression for Resource-Efficient Federated Knowledge Distillation

1

Adaptive Heterogeneous Compression for Resource-Efficient Federated Knowledge Distillation

arXiv:2608.15660v1 [cs.DC] 16 Aug 2026

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

Abstract—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 nonstationary 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. Index Terms—Federated Knowledge Distillation, Heterogeneous, Gradient Compression, Multi-Armed Bandit

I. I NTRODUCTION

F

EDERATED learning (FL) has emerged as a promising distributed learning paradigm for privacy-preserving model training, enabling multiple clients to collaboratively learn a shared model [1]. However, FL inherently suffers from the heterogeneity of client devices and network conditions. In practice, participating clients, such as mobile devices and Internet of Things (IoT) sensors, often exhibit diverse computational resources, memory availability, and communication bandwidths. These disparities make it necessary for clients to adopt heterogeneous model architectures tailored to their local system capabilities, which limits the applicability of traditional parameter-averaging-based FL algorithms that require identical model architectures [2]. Chenwang Liu, Yijun Liu, Chang Liu, and Pengchao Han (corresponding author) are with the School of Information Engineering, Guangdong University of Technology, Guangzhou 510006, China (email: liuchenwang@mails.gdut.edu.cn, 3124002274@mails.gdut.edu.cn, liuchang@gdut.edu.cn, hanpengchao@gdut.edu.cn). Xu Zhang is with the School of Communication and Information Engineering, Chongqing University of Posts and Telecommunications, Chongqing 400065, China (email: zhangxu@cqupt.edu.cn).

Federated knowledge distillation (FedKD) has emerged as a promising approach to address model architecture heterogeneity in FL. In typical FedKD frameworks [3], a prototype denotes a specific model architecture adopted by a group of clients. The server performs hybrid aggregation that incorporates parameter averaging for clients within the same prototype and logit-based knowledge distillation (KD) across different prototypes. However, repeatedly exchanging full gradients introduces significant communication overhead, especially for large-scale models and resource-constrained edge devices. Gradient compression provides an effective solution to alleviate communication overhead in FedKD by transmitting only a subset of gradient entries. Representative methods include Top-K, Random-K, and Periodic-K [4]–[7]. However, most existing works adopt a uniform compression strategy across clients and training rounds, neglecting system heterogeneity. First, clients with different computational capabilities require different optimal strategies. For instance, computationally intensive strategies such as Top-K may benefit powerful devices but overload resource-constrained ones, while lightweight methods such as Random-K may degrade performance. Second, model heterogeneity further affects strategy selection, as Top-K incurs higher sorting cost for large models, whereas it may outperform Random-K for small models due to better update quality. These limitations raise a key question: Key Question: How can compression strategies be selected for clients with diverse model architectures and resource capacities to achieve an effective balance between training efficiency and learning performance? To answer this question, we propose a heterogeneous compression framework for FedKD, where each client adaptively selects a compression strategy from a candidate strategy set according to its model characteristics and resource capacities. Furthermore, we formulate the heterogeneous compression strategy selection as an optimization problem that aims to maximize efficiency-aware performance improvement. However, solving this problem is challenging. First, the performance gain of each compression strategy is revealed after execution, making the corresponding reward unavailable in advance. Second, the optimal compression strategy may change throughout training, where performance-oriented strategies may be preferred in early stages while efficiency-oriented strategies may become more desirable in later stages. To overcome these challenges, we propose an Adaptive heterogeneouS Compression algorithm for fEderated kNowledge Distillation (ASCEND). Specifically, we formulate the client-specific compression strategy selection problem as a non-stationary stochastic multi-armed bandit (MAB) problem,

2

where each arm corresponds to a candidate compression strategy. A reward function is designed to jointly capture model improvement and execution efficiency by considering local training loss, global knowledge alignment, and time cost. An exponential moving average (EMA)-enhanced ϵ-greedy policy is adopted to balance exploration and exploitation during adaptive strategy selection. The main contributions of this paper are as follows: • Heterogeneous Compression Framework for FedKD. We propose a heterogeneous compression framework for FedKD that enables each client to independently select a customized compression strategy from a candidate compression set under heterogeneous model architectures and resource conditions. We provide a convergence anal√ ysis under non-convex loss functions, yielding O(1/ T ) convergence rates. • Compression Strategy Optimization Problem. We formulate a heterogeneous compression strategy optimization problem that maximizes cumulative efficiency-aware performance, explicitly capturing the tradeoff among compression cost, communication cost, and learning performance. • ASCEND Algorithm with Theoretical Guarantee. We model compression strategy selection as a non-stationary stochastic MAB problem and propose ASCEND for adaptive strategy selection under dynamic training conditions. Under a piecewise-stationary approximation, we derive √ a sublinear regret bound of O(1/ M T ), providing a theoretical guarantee for the long-term performance of ASCEND. • Experimental Validation. Extensive experiments on both a real-world platform and simulation settings demonstrate that ASCEND adapts effectively to diverse model and resource settings, significantly reducing communication overhead and training time while maintaining high accuracy. The remainder of this paper is organized as follows. Section II reviews related work. Section III presents the system model and preliminaries, including FL, FedKD, and gradient compression with residual accumulation. Section IV introduces the proposed heterogeneous compression framework and provides a convergence analysis. Section V proposes the ASCEND algorithm and its analysis. Section VI presents the experimental evaluation of the ASCEND algorithm. Finally, Section VII concludes the paper. II. R ELATED W ORK In this section, we review the literature on FedKD and communication-efficient FL via gradient compression. A. Federated Knowledge Distillation FedKD is an effective approach to address model heterogeneity in FL by exchanging knowledge representations, such as soft predictions and intermediate features [3], [8]–[10]. A recent survey [9] revisits KD in collaborative machine learning, highlighting the importance of effective knowledge exchange under heterogeneous learning tasks.

Existing FedKD methods mainly differ in their knowledge sharing mechanisms. FedMD [11] utilizes a shared public dataset to align prediction outputs among heterogeneous client models, while FedGen [12] introduces a generative model to synthesize proxy samples when public datasets are unavailable. Beyond prediction-level knowledge transfer, representationbased FedKD methods exploit intermediate feature information. For example, FedProto [13] introduces class-level feature representations extracted from local models and exchanges these semantic representations among clients. FedTGP [14] constructs and optimizes global class-level feature representations on the server side. Instead of directly aggregating local representations, FedTGP learns more discriminative global semantic representations to improve representation aggregation. Some studies further consider the personalization issue caused by diverse client characteristics. For example, KTpFL [15] incorporates KD into personalized FL by transferring knowledge from the global model to client-specific models. Although these methods effectively mitigate model heterogeneity, the limited amount of shared knowledge may restrict global performance. To improve knowledge aggregation capability, FedDF [16] introduces a hybrid aggregation framework that combines parameter averaging among homogeneous models with KD across heterogeneous models. However, FedDF requires clients to upload gradients for ensemble aggregation, incurring substantial communication overhead in bandwidthconstrained federated environments. B. Communication-Efficient FL via Gradient Compression Communication overhead is one of the major challenges in FL, especially when clients need to frequently exchange highdimensional gradients in bandwidth-constrained networks. Gradient compression has been widely studied to reduce communication cost by transmitting compact representations of gradients while preserving learning performance [7]. Top-K compression is one of the most representative approaches, where only gradients with the largest magnitudes are uploaded. DGC [7] further incorporates momentum correction, local gradient clipping, and residual accumulation to compensate for information loss caused by aggressive compression. These techniques improve communication efficiency while maintaining convergence performance. To reduce the computational overhead of coordinate selection, Random-K compression randomly selects gradient entries for transmission [5]. Although it avoids the sorting operation required by Top-K, it may introduce additional variance during optimization. Periodic-K compression [6] addresses this issue by periodically selecting different coordinates, ensuring that all gradients are updated within a predefined number of iterations. Adaptive compression methods dynamically adjust compression configurations according to training dynamics or system conditions to achieve a better trade-off between communication efficiency and model performance. AdaComp [17] introduces adaptive residual gradient compression, where the compression rate is dynamically adjusted according to gradient residual activity to balance compression efficiency and model accuracy. For a fixed compression method, Han et al. [18] pro-

3

pose an online learning method that adaptively adjusts the TopK compression ratio according to gradient statistics during training, reducing communication overhead while preserving convergence performance. Although existing compression approaches significantly reduce communication overhead, most methods optimize a single compression mechanism, ignoring heterogeneous model architectures, data characteristics, and system resources. They do not consider adaptive selection among multiple heterogeneous compression strategies. In FedKD scenarios, different compression strategies may provide different performanceefficiency trade-offs for different clients. Therefore, adaptive selection among multiple compression strategies remains an open problem, motivating our heterogeneous compression framework and online strategy selection approach. III. S YSTEM M ODEL AND P RELIMINARIES In this section, we first present a brief overview of FL and FedKD. We then introduce gradient compression strategies with residual accumulation, which serve as the building blocks of our heterogeneous compression framework. A. Federated Learning An FL system consists of a central server and a client set N with total N = |N | clients. Each client n ∈ N possesses a local dataset Dn of size |Dn |. Each client trains a local model parameterized by ω to minimize an empirical risk function over its private dataset, defined as fn (ω) =

|Dn | 1 X ℓn (ω; xni ), |Dn | i=1

(1)

aggregate local client models through a hybrid strategy that combines parameter averaging among clients with the same model architecture and logit averaging across heterogeneous models [16]. In FedKD, each training round t proceeds as follows. At the start of each training round, each prototype p with parameter ω pt is broadcast to the corresponding clients to process distributed model training as follows. • Local training: Each client n performs local optimization on its private dataset and computes a stochastic local gradient: g nt = ∇fn (ω nt ; ςtn ), (3) where ω nt denotes the local model parameters of client n at communication round t, and ςtn represents the minibatch sampled from the local dataset of client n. Instead of uploading the updated local model, each client transmits the computed gradient g nt to the server. Upon receiving the local gradients, the server updates the corresponding local models by gradient descent and obtain the updated model ω̃ nt . • Parameter averaging: After obtaining the updated local models {ω̃ nt }N n=1 , the server performs parameter averaging for each model prototype p: ω̄ pt =

n∈Np

min f (ω) = ω

|Dn | fn (ω). PN n=1 n=1 |Dn |

z n (x) = ϕn (ω̃ nt ; x),

FedKD is effective in addressing model architecture heterogeneity in FL by enabling knowledge transfer among clients with diverse model architectures [3], [8], [11], [16]. To enable knowledge transfer among heterogeneous models, KD is performed using a small unlabeled public dataset Dpub . We define a prototype as a specific model architecture. Different clients in FedKD may share the same prototype. Let Np denote the set of clients adopting prototype p. Classical FedKD methods

(4)

∀x ∈ Dpub ,

(5)

where ϕ (ω; x) denotes the logit output of client model n with parameters ω for input sample x. The global logits are computed by averaging the logits of all clients: N

(2)

B. Federated Knowledge Distillation

∀p,

n

z̄(x) =

Classical FL algorithms such as FedAvg [20] perform server-side aggregation by averaging local model updates from participating clients. However, such approaches typically require identical model architectures across clients, limiting their applicability in scenarios with model architecture heterogeneity.

1 ω̃ n , |Np | t

where ω̄ pt represents the aggregated prototype model at round t. • Global knowledge construction (logit averaging): The server evaluates each model on the public dataset Dpub to obtain its logits:

where ℓn (ω; xni ) denotes the sample-wise loss incurred by the model parameters ω on the ith local training sample xni ∈ Dn . In FL, all clients aim to collaboratively learn a global model parameterized by ω by solving the following weighted optimization problem [19]: N X

X

1 X n z (x). N n=1

(6)

The corresponding global soft targets are obtained by applying the softmax function:
ỹ(x) = softmax z̄(x) . (7) •

Global KD: To enable each prototype to learn the global knowledge, for each prototype p, the server distills the global knowledge ỹ(x) into the aggregated prototype model ω̄ pt by minimizing the KL divergence: X

1 LKL softmax ψ p (ω̄ pt ; x) , ỹ(x) , LpKD(ω̄ pt ) = |Dpub | x∈Dpub

(8) where ψ p (ω; x) denotes the logit output of prototype model p with parameters ω for input sample x. After distillation, the prototype parameters are updated to ω pt+1 and transmitted back to the corresponding clients: ω nt+1 ← ω pt+1 ,

∀n ∈ Np .

(9)

4

Top-K

Random-K

In FedKD, clients are required to transmit their local model gradients to the server, which incurs substantial communication overhead, particularly for large-scale models. To alleviate this burden, various gradient compression strategies with residual accumulation have been proposed to reduce communication costs while preserving convergence guarantees [7], [21], [22]. In this section, we first review representative gradient compression techniques, including Top-K, Random-K, and Periodic-K. Then, we introduce the server-side aggregation of the compressed gradients, followed by the analysis to motivate our heterogeneous compression framework. 1) Client-Side Compression with Residual Accumulation: To compensate for information loss introduced by gradient compression, each client maintains a residual vector ent−1 at the end of round t − 1, which accumulates the gradient entries that were not transmitted in previous rounds. At round t, the input to the compression operator C(·; ·) is an augmented update vector unt , defined as the sum of the current gradient and the residual: unt = ent−1 + g nt . (10) For a compression strategy snt ∈ {Top-K, Random-K, Periodic-K} , the compressed gradient on client n is ĝ nt = C(unt ; snt ) .

(11)

The resulting vector ĝ nt is associated with an index set ⊆ {1, . . . , d} of cardinality K, where d denotes the gradient dimension and |I nt | = K. For convenience, we define a reconstruction operator B(·; ·) that maps the compressed representation back to the original gradient dimension. In B(ĝ nt ; I nt ) ∈ Rd , the entries indexed by I nt are filled with the transmitted gradient values and the remaining entries are set to zero. Typical compression strategies are summarized as follows. • Top-K Compression: Each client selects the K entries of unt with the largest magnitudes [7], [22], [23]: I nt

I nt = TopK(unt , K).

(12)

Top-K compression preserves the most informative gradient directions but incurs a relatively high computational cost due to sorting, with computational complexity of O(d log K) [7], [22]. • Random-K Compression: Each client randomly samples K indices from all gradient dimensions [21], [24], [25]: I nt = RandomK(unt , K). (13) This strategy has low computational complexity, i.e., O(K), but it may introduce high gradient variance and low model performance [21], [23], [24]. • Periodic-K Compression: Each client maintains a binary indicator vector bnt ∈ {0, 1}d indicating whether a

0.48 0.40

Accuracy

C. Gradient Compression with Residual Accumulation

Accuracy

0.56

0.55

0.55

0.54

0.54

0.53 0.52

0.32 100120140160180200 Time (s)

(a) K = 105

Periodic-K

Accuracy

Each client uses the received model to start the subsequent training round.

176 184 192 200 208 Time (s)

(b) K = 5 × 105

0.53 0.51 170180190200210220 Time (s)

(c) K = 2 × 106

Fig. 1: Model performance of three compression strategies under different K (dataset: CIFAR-10, α = 0.5).

coordinate has been selected within the current period that contains Tp iterations [23]. The selected index set is I nt = PeriodicK(bnt , K),

(14)

where unvisited coordinates are prioritized to ensure periodic coverage of all coordinates, followed by random completion if necessary. After transmission, the selected coordinates are marked as visited, and the time-flag vector bnt is updated accordingly. Once all coordinates have been selected at least once, the flag vector is reset to initiate a new periodic cycle. Periodic-K compression avoids costly sorting and ensures periodic updates of all gradient coordinates, achieving a favorable balance between computation efficiency and convergence stability. Its computational cost is dominated by index selection, with a time complexity of O(K). After compression, the client transmits the compressed gradients ĝ nt together with its corresponding index set I nt to the server. To ensure convergence of the federated training process, the residual vector ent is updated by preserving the untransmitted gradient entries: ent = unt − B(ĝ nt , I nt ).

(15)

This residual accumulation effectively compensates for information loss due to compression and is crucial for achieving stable convergence. 2) Server-Side Processing of Compressed Gradients: Upon receiving the compressed representation (ĝ nt , I nt ) from participating clients, the server reconstructs the d-dimensional compressed gradient vector using the reconstruction operator: g̃ nt = B(ĝ nt , I nt ).

(16)

The reconstructed gradients are used to update the corresponding models for global aggregation and KD: ω̃ nt = ω nt − ηt g̃ nt ,

(17)

where ηt denotes the learning rate in round t 3) Motivation for Heterogeneous Compression: The three gradient compression strategies, namely Top-K, Random-K, and Periodic-K, exhibit different trade-offs between compression overhead and model performance. To investigate whether a uniform compression strategy is suitable for heterogeneous

5

2.Prototype-wise parameter averaging

0.90

Accuracy

Accuracy

Server

0.96

0.80 0.70 0.450.600.750.901.05 Time (s)

(a) LeNet5Half

Top-K

0.94

Random-K

0.92

1. Gradient reconstruction and model updates

4. Knowledge distillation

Periodic-K

3. Global logit averaging 1 2 3

0.90 0.500.751.001.251.50 Time (s)

Public dataset

(b) LeNet5

Fig. 2: Model performance of three compression strategies under different model prototypes (dataset: MNIST, α = 0.5, K = 1 × 103 ).

Updated prototype models

Download

IV. F ED KD WITH H ETEROGENEOUS C OMPRESSION This section introduces a heterogeneous compression framework that allows clients to adopt different compression strategies based on their model architectures and computational resources. We then provide a convergence analysis for the proposed framework and discuss the motivation for adaptive strategy selection. A. Heterogeneous Compression Framework For each client, we define the candidate compression set S as a collection of three representative strategies: S = {Top-K, Random-K, Periodic-K}.

(18)

Each element in S corresponds to a distinct compression strategy, characterized by different trade-offs among model accuracy, computational overhead, and communication efficiency. Figure 3 illustrates the system architecture of the proposed FedKD framework with heterogeneous gradient compression. As shown, each client independently performs local training on its private dataset, constructs residual-accumulated gradients, and applies a compression strategy from S before transmitting compressed gradients and corresponding indices to the server. The server first reconstructs the compressed gradients and updates the corresponding client models. Then, the updated client models are utilized for prototype-wise parameter aggregation among clients with the same model architecture and

1. Private dataset Data 1

2. Local training

1

2

3

Compressed gradient + indices

Client 2

Client N

3. Residual accumulation

4.Compression strategy Selection

Client 1

FedKD, we evaluate these strategies under different compression levels and model prototypes on CIFAR-10 and MNIST. As shown in Fig. 1, the preferable compression strategy varies with the compression level. Specifically, Top-K tends to achieve better accuracy under aggressive compression, while Random-K becomes more competitive when a larger communication budget is available due to its lower selection overhead. Furthermore, Fig. 2 demonstrates that the optimal strategy also depends on model architectures under the same compression setting, where Top-K is desirable for a smaller model, e.g., LeNet5Half, and Random-K is better for a larger model, e.g., LeNet5. These results indicate that no single compression strategy can consistently achieve the best trade-off across heterogeneous clients. Therefore, enforcing a uniform compression strategy may lead to suboptimal performance in FedKD, motivating the proposed heterogeneous compression framework.

Upload

5. Upload

Output

Input

Data 2

1

Top-K

2

3

Random-K Periodic-K

Data n

Local model training and compression on each client

Compressed gradients

1

2

3

Indices of selected gradients

Different prototypes

Fig. 3: The system architecture of the heterogeneous compression framework for FedKD.

global logit averaging across all client models. Finally, the server performs KD on the public dataset Dpub to align global knowledge and broadcasts the updated prototype models back to the corresponding clients [3], [8]. B. Convergence Analysis of FedKD with Heterogeneous Compression This section analyzes the convergence behavior of the proposed heterogeneous compression framework in FedKD. Unlike conventional FL methods with homogeneous models, FedKD maintains multiple model prototypes to handle model architecture heterogeneity. Therefore, we conduct the convergence analysis from the perspective of each model prototype. The following assumptions are introduced to analyze the convergence of the proposed heterogeneous compression framework. Assumption 1 (Smoothness for Local Training Objective). For each client n, the objective function fn (ω) is Ll smooth [6], [26]. Specifically, there exists a constant Ll > 0 such that ∥∇fn (ω) − ∇fn (ω ′ )∥ ≤ Ll ∥ω − ω ′ ∥ ,

∀ω, ω ′ .

(19)

Assumption 2 (Smoothness for KD Objective). For each prototype p, the objective function LpKD (ω) is Lp -smooth [27]. Specifically, there exists a constant Lp > 0 such that ∥∇LpKD (ω) − ∇LpKD (ω ′ )∥ ≤ Lp ∥ω − ω ′ ∥ ,

∀ω, ω ′ . (20)

Assumption 3 (Unbiased Stochastic Gradient). For each client n ∈ Np and round t, the stochastic gradient g nt satisfies [6] E [g nt | ω pt ] = ∇fn (ω pt ).

(21)

6

Assumption 4 (Bounded Stochastic Gradient Variance). For any client n and round t, the variance of stochastic gradients is bounded by a constant σ 2 > 0 [6], [28], i.e., h i 2 E ∥g nt − ∇fn (ω pt )∥ ≤ σ 2 . (22)

Assumption 5 (Bounded Error-Feedback Residual). For each client n and round t, the accumulated residual satisfies [29], [30]   (23) E ∥ent ∥2 ≤ Γmax , where Γmax denotes the maximum residual amplification factor among all candidate compression strategies. Assumption 5 ensures that these untransmitted errors remain stable and do not diverge after accumulating across multiple rounds of residual propagation. For theoretical analysis, we define the joint optimization objective of prototype p as Fp (ω pt ) = fp (ω pt ) + λLpKD (ω pt ),

(24)

where λ balances the local optimization objective and the global knowledge alignment objective. Thus, we have the following theorem for the convergence of FedKD with heterogeneous compression for non-convex loss functions. Theorem 1 (Convergence of FedKD with Heterogeneous Compression). Under Assumptions 1– 5, for each prototype p, after T communication rounds, the proposed FedKD framework with heterogeneous compression satisfies

V. F ED KD WITH A DAPTIVE H ETEROGENEOUS C OMPRESSION In this section, we first formulate the heterogeneous compression strategy optimization problem in FedKD. Then, we formulate the problem as a MAB and propose the adaptive heterogeneous compression algorithm. Finally, we analyze the regret of the online algorithm. A. Heterogeneous Compression Strategy Optimization Problem The heterogeneous compression strategy optimization problem aims to maximize efficiency-aware model performance by determining the compression strategy snt ∈ S for all clients n ∈ N at round t. To achieve this, we define the model performance gain of client n at round t as G(snt ). In FedKD, the time overhead covers the local training time, the computation time for gradient compression, and the communication time for gradient transmission. Thus, the problem is defined as Problem 1 (Heterogeneous Compression Strategy Optimization Problem). In FedKD with heterogeneous compression, all clients aim to T N X X G(snt ). (27) max n {st }n∈N

n=1 t=1

The proof of the theorem can be found in our online report [*]. Theorem 1 shows that the proposed method converges √ with a convergence rate of O(1/ T ), which equals the convergence rate of existing FL methods.

Solving Problem 1 is challenging. First, the performance gain G(snt ) is stochastic and unknown in advance, which is only revealed after the compression strategy snt is selected and executed by client n. Moreover, the reward distribution may vary over time due to non-stationary factors such as data heterogeneity, changes in loss values, and dynamic communication conditions. This observation motivates us to address Problem 1 in an online decision-making framework. Specifically, we formulate adaptive compression strategy selection as a non-stationary stochastic MAB problem, where each client selects a compression strategy from a finite candidate set at each communication round and observes the corresponding reward as feedback. The objective is to learn an effective strategy-selection policy that maximizes the cumulative expected performance gain over time.

C. Motivation for Adaptive Compression Algorithm

B. MAB Formulation

Although the heterogeneous compression framework enables clients to adopt different compression strategies according to their characteristics, determining the appropriate strategy for each client remains challenging. Specifically, the performance of each compression strategy cannot be known before execution and may vary during training due to dynamic model states and system conditions. Therefore, an online strategy selection mechanism is required to allow each client to adaptively identify suitable compression strategies based on observed training feedback. This motivates the formulation of compression strategy selection as a non-stationary stochastic MAB problem, where each compression strategy corresponds to an arm and the observed efficiency-aware performance gain serves as the reward.

We first formulate the time overhead and reward. Then, we reformulate Problem 1 as an MAB. In this reformulation, each client n decides a sequence of compression strategies {snt } over training rounds. 1) Time overhead: We formulate the client-side training time for each training round t as

T −1  4(Fp (ω p0 ) − Fp∗ + Ck ) 4L 1 X  E ∥∇Fp (ω pt )∥2 ≤ + ηt Cg , T t=0 3ηt T 3 (25) √ By setting ηt = O(1/ T ), the convergence rate becomes   T −1 i 1 1 X h 2 E ∥∇Fp (ω pt )∥ = O √ . (26) T t=0 T

τtn (snt ) = ν n (snt , K) + ζ n (ω nt ) + ϑn (K),

(28)

where ν n (·, ·) denotes the compression time that depends on the selected strategy, model architecture, and compression level K, ζ n (·) represents the local training time of client n, and ϑ(·) represents the transmission time of the compressed gradients and their indices.

7

2) Reward: The reward of client n is defined as an efficiency-aware performance gain: G(snt ) =

∆Ln (t) , τtn (snt )

(29)

where ∆Ln (t) quantifies the model improvement induced by the selected compression strategy. Specifically, for each client n with prototype p at round t, ∆Ln (t) is measured by jointly considering global knowledge alignment that reflects global knowledge improvement and local optimization progress that captures the client-specific loss reduction on private data:   ∆Ln (t) = ρ LpKD (ω̄ pt−1 ) − LpKD (ω̄ pt )   (30) + β fn (ω nt−1 ) − fn (ω nt ) , where ρ, β > 0 denote the weighting coefficients. The global KD loss for prototype p, i.e., LpKD (ω̄ pt ) and the local loss of client n, i.e., fn (ω nt ) are defined in (8) and (1), respectively. The KD loss component is computed at the server and then is fed back to the corresponding clients. By normalizing the combined performance gain with the execution time τtn , the proposed reward encourages compression strategies that achieve faster convergence with lower computational and communication overhead. Furthermore, we apply a clipping function to the reward to avoid numerical instability caused by large reward fluctuations: G(snt ) ← clip(G(snt ), Gmin , Gmax ) .

(31)

In our experiments, the thresholds are set to Gmin = −1 and Gmax = 1. 3) MAB Formulation: The reward of each compression strategy is time-varying because it evolves with dynamic model performance during training. Therefore, we model the proposed strategy selection as a non-stationary stochastic MAB. Unlike adversarial settings where rewards may change arbitrarily, the reward transition in our framework follows the intrinsic optimization dynamics. We define the MAB problem as follows. Problem 2 (Non-stationary Stochastic MAB formulation). We reformulate Problem 1 as a non-stationary stochastic MAB problem. • Arms: each arm corresponds to a compression strategy snt ∈ S; • Policy: In each training round t, client n selects a compression strategy snt to balance exploration and exploitation: ( Random(S), with probability ϵt , n st = (32) arg max Qnt (s), with probability 1 − ϵt . s∈S

where ϵt ∈ (0, 1) denotes the exploration rate and Qnt (s) denotes the local utility estimate of strategy s maintained by client n at round t. • Utility update: Each client maintains a local utility estimate Qnt (s′ ) for each strategy s′ ∈ S. To adapt to non-stationary stochastic reward variation, we employ an EMA-based utility estimator within an ϵ-greedy policy

[31]. After observing the reward G(snt ), client n updates its utility estimate as follows: (
1 − ηq Qnt (s′ ) + ηq G (snt ) if snt = s′ , n ′ Qt+1 (s ) = Qnt (s′ ) if snt ̸= s′ . (33) where ηq ∈ (0, 1] denotes the EMA step size. C. Algorithm Design In this section, we propose the ASCEND algorithm for Problem 2, where each client independently optimizes its compression strategy in an online manner to cope with dynamic model performance and computational efficiency trade-offs in heterogeneous FedKD. Specifically, the proposed algorithm combines an adaptive exploration policy and a rollback-based stability safeguard. 1) Adaptive Exploration Policy: We employ a two-phase exploration policy to adapt to different training stages: • Warm-up Phase: To avoid early-stage bias in utility estimation, each client performs a fixed number of warmup trials for every compression strategy in S. Let S = |S| denote the number of available compression strategies. During the warm-up stage, each strategy snt ∈ S is selected exactly h times by each client, resulting in a total of hS warm-up rounds. The selection order is randomly permuted to prevent systematic ordering bias. Each strategy snt is treated as an independent arm and is associated with a reward estimate G(snt ). The reward of strategy snt is estimated as the empirical mean: h|S|

Qn0 (s′ ) =

1X G(snt )Isnt =s′ . h t=1

(34)

Here, Isnt =s′ indicates whether the selected strategy snt corresponds to the target strategy s′ . The empirical mean operation provides an initial estimate of the average reward for each strategy for the subsequent adaptive EMAenhanced ϵ-greedy updates. As a result, at the end of the warm-up stage, each client maintains S initialized utility values {Qn0 (s)}s∈S , one for each compression strategy. • Adaptive Phase: For rounds t > hS, the compression strategy is selected following (32). Here, Qnt (s) is initialized by Qn0 (s) and updated according to (33). We define the exploration rate in (32) as   c ϵt = min 1, √ , (35) t where c is a fixed hyperparameter. 2) Rollback-Based Stability Safeguard: To prevent unstable convergence caused by overly aggressive or rapidly fluctuating compression strategy selections, we incorporate a rollbackbased stability safeguard inspired by [32], [33]. A stability check is performed by comparing the global loss between two consecutive rounds. For any prototype p, if a sudden loss surge is detected, i.e., LpKD (ω̄ pt ) > γLpKD (ω̄ pt−1 ),

(36)

8

Algorithm 1: ASCEND algorithm Input: Client set N , public dataset Dpub , warm-up length h, exploration rate ϵt , learning rate ηt , ηq . Output: Clients’ models {ω n Tp}n∈N . 1 Initialize prototype models ω 0 , ∀p and distribute them to corresponding clients; n 2 Initialize residual vectors e0 ← 0, ∀n ∈ N and utility estimates Qn (s) ← 0, ∀n ∈ N , ∀s ∈ S; 0 3 for t = 1 to T do Client-side Local Training and Gradient Compression: 4 for each client n ∈ N in parallel do 5 Compute local gradient g n t; n n 6 Accumulate residuals: un t ← et−1 + g t ; 7 Select compression strategy sn t by (37); n n n 8 Compress gradients: (ĝ n t , I t ) ← C(ut ; st ); n − B(ĝ n ; I n ); 9 Update residuals: en ← u t t t t n 10 Upload (ĝ n t , I t ) to the server; 11 12 13 14

15 16 17 18 19

20 21 22

Server-side Prototype-wise Model Aggregation: n n Reconstruct compressed gradients: g̃ n t ← B(ĝ t , I t ), ∀n ∈ N ; for each prototype p do n n Update local models: ω̃ n t ← ω t − ηt g̃ t , ∀n ∈ Np ; Aggregate prototype model by (4); Server-side KD: Evaluate all models on Dpub and compute averaged logits z̄(x) by (6); Update each prototype model by minimizing LpKD (ω̄ pt ); Broadcast updated prototype models and global loss to clients; if LpKD (ω̄ pt ) > γLpKD (ω̄ pt−1 ) then Roll back to round t−1 and enforce unified compression; Client-side Reward and Utility Update: for each client n ∈ N in parallel do Compute reward G(sn t ) by (29); Update utility estimate: n n n n n Qn t+1 (st ) = (1 − ηq ) Qt (st ) + ηq Gt (st ).

where γ > 1 is a predefined stability threshold, the system activates the following stabilization procedure: • Model Rollback: All client models and prototype models are restored to their parameters from the previous round t − 1. • Strategy Unification: To provide a conservative fallback under unstable conditions, all clients adopt the Top-K strategy for the next round. By reverting the system to a previously stable state and temporarily constraining strategy heterogeneity, this rollback mechanism effectively mitigates performance degradation caused by erratic compression decisions and enhances the robustness of the proposed adaptive compression framework. 3) Algorithm Description: The ASCEND algorithm is summarized in Algorithm 1. In each round t, participating clients first perform local training (line 5). Each client then selects a compression strategy according to the EMA-enhanced ϵgreedy policy:  the next in shuffled sequence, t ≤ h|S|,    n st = Random(S), t > h|S| with probability ϵt ,   arg max Qn (s), t > h|S| with probability 1 − ϵt ; s∈S

t

(37) (line 7), compresses its residual-accumulated gradients (line 8), and uploads the resulting compressed gradient representation to the server (line 10).

To aggregate heterogeneous compressed gradients, the server first reconstructs the received compressed gradients (line 11) and updates the corresponding client models (line 13). The updated models within each prototype p are then aggregated through parameter averaging to obtain the prototype models (line 14). Afterwards, the server adopts a logit-based KD mechanism using a small public dataset Dpub . Specifically, the server evaluates all models on Dpub and averages their logits to construct global soft targets, which capture shared knowledge across different prototypes (line 15). Each prototype model further aligns its output distribution with the global soft targets via KL-divergence-based KD (line 16). The distilled prototype models and the global loss are subsequently broadcast to their corresponding clients for the next round of local training (line 17). In addition, clients compute the instantaneous reward (line 21) and update the utilities (line 22) to facilitate the subsequent strategy selection and global optimization. D. Algorithm Analysis In this section, we analyze the regret of the proposed ASCEND algorithm. We first define the regret used to evaluate ASCEND. Since the reward is affected by stochastic gradients, data sampling, and model evolution, it is treated as a random variable. For notational simplicity, we hereinafter use Gtn (s) to denote the reward obtained by client n when selecting compression strategy s in communication round t. The expected reward is defined as µnt (s) = E [Gtn (s)] . (38) Unlike in stationary bandit problems, µnt (s) may vary throughout training because the optimization states and KD outcomes continuously evolve. Accordingly, the optimal compression strategy for client n in round t is defined as sn,∗ = arg max µnt (s). t s∈S

(39)

The regret of ASCEND for client n over T communication rounds is then defined as T X  n n,∗  µt (st ) − µnt (snt ) , (40) RnT = t=1

which measures the cumulative expected-reward loss relative to the round-wise optimal compression strategies. The following assumptions are introduced for analyzing the adaptive strategy selection behavior of ASCEND. Assumption 6 (Piecewise Stationary Reward). The reward process of each compression strategy is assumed to be piecewise stationary during the training process [34], [35]. Specifically, the entire communication process is divided into M stationary segments: {T1 , T2 , . . . , TM },

(41)

where the reward distribution remains approximately unchanged within each segment. That is, for client n, compression strategy s ∈ S, and communication round t ∈ Ti , the expected reward satisfies µnt (s) = µni (s), ∀i,

(42)

9

where µni (s) denotes the stationary expected reward of strategy s within segment Ti . Assumption 7 (Reward Concentration). The stochastic reward obtained from each compression strategy satisfies the subGaussian concentration property [36]. Specifically, for any n ∈ N , s ∈ S, and x > 0,   x2 Pr (|Gtn (s) − µnt (s)| > x) ≤ 2 exp − 2 , ∀t > h|S|, 2σg (43) where σg2 represents the variance of the reward distribution. Assumption 8 (Sufficient Strategy Sampling). For each client n, stationary segment Ti , and strategy s ∈ S, let ri denote the first communication round of segment Ti . We define n Hi,t (s) =

t X

I{snx = s}, ∀s ∈ S

(44)

x=ri

as the number of times that strategy s has been selected from round ri to round t during the segment Ti . Then, there exists a constant πmin ∈ (0, 1] such that n Hi,t (s) ≥ πmin (t − ri + 1), ∀s ∈ S.

(45)

Based on the above assumptions, we introduce the following lemmas to characterize the estimation errors and regret incurred by the EMA, exploration, and exploitation phases of ASCEND. We first characterize how the EMA tracks the expected reward within each stationary segment. Lemma 1 (EMA Error). For any client n and compression strategy s ∈ S, within a stationary segment Ti , the EMA-based utility estimate satisfies |E[Qnt (s)] − µni (s)| ≤ (1 − ηq )m E[Qnri (s)] − µni (s) , (46) where m represents the number of updates of strategy s during this segment. Lemma 1 shows that, within a stationary segment, the expected EMA tracking error decreases geometrically with the number of updates m. Thus, strategies that are sampled more frequently can adapt more rapidly to the reward distribution of the current segment. We then derive separate regret bounds for the exploration and exploitation phases. Lemma 2 (Regret Upper Bound for Exploration ). Under Assumptions 6–7, the cumulative regret incurred by exploration over T communication rounds satisfies √ (47) Rnexplore (T ) ≤ 2c∆max T , where

Lemma 2 indicates that the cumulative cost of exploration grows sublinearly with the number of communication rounds. Hence, the average regret caused by exploration vanishes as T increases. Lemma 3 (Regret Upper Bound for Exploitation). Under Assumptions 6–8, the cumulative regret incurred by exploitation for client n satisfies Rnexploit (T ) ≤ Ce

M p X

|Ti |,

(50)

i=1

where Ce is a constant determined by the EMA update coefficient, reward variation, initialization error, and reward concentration parameter. Consequently, √ (51) Rnexploit (T ) ≤ Ce M T . Lemmas 2 and 3 jointly characterize the regret incurred by the two adaptive selection phases of ASCEND. Under √ a piecewise-stationary reward model, √ we derive an O( T ) exploration-regret bound and an O( M T ) exploitation-regret bound. Based on the preceding lemmas, we establish the overall regret bound of ASCEND. Theorem 2 (Regret Upper Bound of ASCEND). For each client n, the cumulative regret of ASCEND over T communication rounds satisfies √ √ RnASCEND (T ) ≤ h|S|∆max + 2c∆max T + Ce M T , (52) where the three terms correspond to the regret incurred during the warm-up, exploration, and exploitation phases, respectively. Since the warm-up length h|S| is fixed, the overall regret satisfies √ RnASCEND (T ) = O( M T ). (53) The proofs of Lemmas 1–3 and Theorem 2 are provided in [*]. Theorem 2 demonstrates that ASCEND asymptotically approaches the performance of the round-wise optimal compression strategy in a piecewise-stationary training environment. The regret bound further reveals the impact of environmental non-stationarity: a larger number of stationary segments M leads to a higher tracking cost because ASCEND must repeatedly adapt its utility estimates to changes in the reward distributions. Combined with Theorem 1, which establishes convergence under any compression strategy in S, Theorem 2 shows that ASCEND can adaptively select communication-efficient compression strategies without compromising the convergence guarantee of the underlying federated optimization process. VI. E XPERIMENTS

∆max = Gmax − Gmin

(48)

is the maximum instantaneous regret under the clipped reward range, and c is the constant in the exploration rate defined in (35). Consequently, √ Rnexplore (T ) = O( T ). (49)

In this section, we evaluate the effectiveness and robustness of the proposed ASCEND algorithm via experiments. We introduce the experimental settings, present empirical results on a real-world platform and various experimental evaluations, and conduct sensitivity analysis to investigate the impact of critical hyperparameters.

10

TABLE I: Local Training Time per Round (seconds) Dataset MNIST MNIST CIFAR-10 CIFAR-10

Model LeNet5Half LeNet5 WResNet40-2 ResNet18

TABLE II: Gradient Compression Time on MNIST (seconds)

Time (s) 0.0025 0.0045 0.120 0.194

Model LeNet5Half LeNet5

Method Top-K Random-K Periodic-K Top-K Random-K Periodic-K

K = 5 × 102 1.70e-7 3.50e-8 1.00e-7 3.05e-7 6.00e-8 1.75e-7

103 3.40e-7 7.00e-8 2.00e-7 6.10e-7 1.20e-7 3.50e-7

5 × 103 1.70e-6 3.50e-7 1.00e-6 3.05e-6 6.00e-7 1.75e-6

A. Setup We introduce the setup, including federated settings, datasets, models, time overheads, and baselines. 1) Federated setting: We build a real-world platform with 10 Raspberry Pi devices, which act as resource-constrained edge clients controlled by a server. For experimental evaluation, we evaluate a total of 10 clients, each of which is assigned a local model architecture from the candidate prototype set. The batch size is set to 32, and the learning rate is 0.02. The weighting parameters ρ and β in (30) are set to 0.6 and 0.4, respectively, and the fixed hyperparameter c in (35) is set to 3 in all experiments. These default values are chosen to achieve favorable empirical performance, and we conduct sensitivity analysis to investigate the impact of hyperparameter variations. 2) Datasets: We conduct experiments on two datasets, namely MNIST [37] and CIFAR-10 [38]. To simulate realistic non-independent and identically distributed (non-IID) settings, datasets are partitioned using a Dirichlet distribution with parameter α that controls the degree of data heterogeneity. A smaller α leads to a more skewed class distribution, while a larger α yields a more balanced data distribution. Unless otherwise specified, we set α = 0.5. 3) Models: For the MNIST dataset, we employ two lightweight convolutional neural networks, namely LeNet5 and LeNet5Half, where LeNet5Half is a shallower and more compact variant of LeNet5. The total number of trainable parameters is 61,706 for LeNet5 and 15,738 for LeNet5Half, respectively. For the CIFAR-10 dataset, we consider two architectures, ResNet18 and WResNet40-2. ResNet18 contains 11,183,562 parameters, while WResNet40-2 contains 2,248,954 parameters. Compared with ResNet18, WResNet402 adopts a wider residual architecture with increased width and a different depth-width trade-off, providing a useful complementary model pair for evaluating the proposed adaptive compression framework under heterogeneous model sizes and architectural characteristics. 4) Time overheads: For the platform, we collect the real computation and communication time and feed them back to the ASCEND algorithm for compression strategy determination. For the experimental evaluation, the computation time, including local training and gradient compression, is instantiated using the real measured execution time of the system on an NVIDIA RTX 3090 GPU and an AMD Ryzen 7 4800H CPU. All computation times are obtained by averaging over multiple runs. The local training time is reported in Table I, while representative compression times for different values of K are summarized in Tables II and III. The communication time for client n is modeled as φn (K) =

K · bv + K · bi , B × 106

(54)

TABLE III: Gradient Compression Time on CIFAR-10 (seconds)

Model WResNet40-2 ResNet18

Method Top-K Random-K Periodic-K Top-K Random-K Periodic-K

K = 105 1.3e-4 2.0e-5 6.5e-5 6.5e-4 4.0e-5 3.2e-4

5 × 105 3.0e-4 3.0e-5 1.5e-4 1.5e-3 5.0e-5 7.5e-4

2 × 106 9.0e-4 7.5e-5 5.0e-4 4.5e-3 2.0e-4 2.55e-3

where K denotes the number of retained gradient coordinates after compression, bv and bi represent the numbers of bits used to encode each gradient value and its corresponding index, respectively. B denotes the bandwidth. In our implementation, gradient values are represented using 32-bit floating point precision (bv = 32), while indices are encoded using log2 d bits, where d denotes the dimensionality of the original gradient vector. The default communication bandwidth B is set to 50 Mbps. 5) Baselines: Since existing adaptive gradient compression mechanisms mainly target compression-ratio adaptation in FL, they are not directly designed for heterogeneous strategy selection in FedKD. Therefore, we compare ASCEND with uniform compression baselines, where all clients adopt the same strategy, and an EXP3-based adaptive compression method. B. Experimental Results and Analysis We verify the effectiveness of the proposed ASCEND algorithm from multiple perspectives, including its performance on a real-world platform, comparison with different baselines at different levels of compression, robustness under varying levels of data heterogeneity and communication bandwidth, as well as the evolution of client-side strategy selection behaviors during training, and sensitivity analysis to investigate the impact of key hyperparameters. 1) Performance on Real-World Platform: To demonstrate the practicality of the proposed ASCEND algorithm, we evaluate its performance on a real-world platform. The experiments are conducted on the MNIST dataset under different compression levels K ∈ {5 × 103 , 104 }, covering relatively aggressive and moderate compression settings. The comparison of different methods under different compression levels on the MNIST dataset is presented in Fig. 4. We can observe that ASCEND achieves competitive performance across different compression levels, demonstrating its ability to effectively balance communication efficiency and model accuracy. In particular, ASCEND achieves the best performance at the compression level K = 104 . Furthermore, at the compression level K = 5 × 103 , ASCEND exhibits faster convergence

11

0.96

0.90

0.93

0.78 0.3 0.6 0.9 1.2 1.5 Time (s)

(a) K = 5 × 102

0.3 0.6 0.9 1.2 1.5 Time (s)

(b) K = 103

0.44

0.50 0.45

40 80 120 160 200 Time (s)

0.40 40 80 120 160 200 Time (s)

0.40

(a) K = 105

(b) K = 5 × 105

(c) K = 2 × 106

0.90

0.90 0.87

0.80 0.70 0.60

0.5 0.8 1.0 1.2 1.5 Time (s)

80 120 160 200 Time (s)

in the early training stage. This behavior can be attributed to its adaptive strategy selection mechanism, which dynamically selects suitable compression strategies based on the observed utility estimates, thereby improving communication efficiency and accelerating empirical risk reduction. 2) Performance Comparison with Baselines: Fig. 5 illustrates the performance on MNIST under different compression levels K ∈ {5 × 102 , 103 , 5 × 103 }, corresponding to compression ratios of approximately 0.81%, 1.62%, and 8.10% for LeNet5 as the reference model. Fig. 6 presents the results on CIFAR-10 with K ∈ {105 , 5 × 105 , 2 × 106 }, corresponding to retained-gradient ratios of approximately 0.89%, 4.47%, and 17.88% for ResNet18 as the reference model. Generally, increasing K improves convergence speed and final accuracy because more gradient information is preserved. However, the optimal compression strategy varies across different compression levels. For example, on MNIST, Top-K performs better when K = 5 × 102 , while Random-K becomes more competitive at K = 5 × 103 ; similarly, on CIFAR-10, Top-K is more effective at K = 105 , whereas Random-K shows advantages at larger budgets such as K = 2×106 . Despite these variations, ASCEND consistently achieves strong performance across all compression levels, demonstrating its capability to adaptively balance the trade-off between communication efficiency and model accuracy. 3) Performance under Varying Data Heterogeneity: To further evaluate the robustness of the proposed method under varying degrees of data heterogeneity, we conduct additional experiments on the MNIST dataset, where α ∈ {0.1, 0.5, 1} correspond to highly heterogeneous and nearly IID data distributions, respectively. In this experiment, the compression level is fixed at K = 103 , and the communication bandwidth is set to 50 Mbps. As shown in Figs. 5 and 7, higher data heterogeneity (e.g.,

0.4

0.8 1.2 Time (s)

Random-K

0.90

Periodic-K

0.85 0.80

1.6

Top-K

0.95

(a) α = 0.1

(c) K = 5 × 103

Fig. 5: Comparison of different methods on the MNIST dataset under different compression levels (α = 0.5).

Accuracy

0.72

0.96

0.84

0.48

ASCEND-K

Accuracy

Accuracy

Accuracy

0.80

0.40

0.52

Fig. 6: Comparison of different methods on the CIFAR-10 dataset under different compression levels (α = 0.5).

Accuracy

Periodic-K EXP3

0.96 0.88

0.48

(b) K = 104

Top-K Random-K

0.55

0.32

Fig. 4: Comparison of different methods on the MNIST dataset deployed on a real-world platform under different compression levels (α = 0.5).

ASCEND-K

0.56

EXP3 0.4

0.8 1.2 Time (s)

1.6

ASCEND-K

(b) α = 1

Fig. 7: Comparison of different methods on the MNIST dataset under different data heterogeneity levels (K = 103 ).

0.96

0.96

Top-K

0.90

0.90

Random-K

0.84

Periodic-K

0.78

EXP3

Accuracy

(a) K = 5 × 103

ASCEND-K

60 90 120 150 180 Time (s)

Periodic-K EXP3

Accuracy

EXP3 0.85

60 90 120 150 180 Time (s)

Periodic-K

Accuracy

0.85

Random-K 0.90

Accuracy

0.90

Top-K Random-K

Top-K

Accuracy

0.97 Accuracy

Accuracy

0.97

0.84 0.78 20.0 30.0 40.0 50.0 Time (s)

(a) 5 Mbps

0.3 0.5 0.6 0.8 Time (s)

ASCEND-K

(b) 5 × 102 Mbps

Fig. 8: Comparison of different methods on the MNIST dataset under different communication bandwidths (α = 0.5, K = 103 ).

α = 0.1) leads to slower convergence and lower final accuracy. Nevertheless, across all settings, ASCEND consistently achieves competitive or superior performance compared to the best fixed compression strategy. 4) Performance under Varying Communication Bandwidths: To evaluate the robustness of the proposed method under varying communication bandwidths, we conduct additional experiments on the MNIST dataset, where the bandwidth is varied over {5 Mbps, 50 Mbps, 500 Mbps}, corresponding to different network conditions. Here, we set K = 103 and α = 0.5. As shown in Figs.5 and 8, we can observe that a small bandwidth (e.g., 5 Mbps) introduces significant communication delays, under which Top-K tends to perform better. In contrast, higher bandwidth (e.g., 500 Mbps) improves convergence efficiency, making Random-K more competitive. Nevertheless, across all settings, ASCEND consistently achieves competitive or superior performance, demonstrating its robustness under varying network conditions. 5) Analysis of Client Compression-Strategy Selection: In this section, we analyze the compression strategy selection results of clients to better understand the selection dynamics of the proposed ASCEND algorithm during the training process.

12

2000

(a) Model prototype = LeNet5Half

0

0

500

1000 Round

1500

2000

(b) Model prototype = LeNet5

Fig. 9: Compression strategy selection results of different model prototypes on MNIST dataset (α = 0.5, K = 103 ).

For the three candidate strategies, i.e., Top-K, Random-K, and Periodic-K, we record the average compression strategy selection counts for different prototypes during training. For MNIST with a small K = 103 , as illustrated in Fig. 9, clients equipped with the lightweight LeNet5Half model predominantly select the Top-K strategy, whereas those using the standard LeNet5 model exhibit a more balanced preference between Top-K and Random-K. This discrepancy can be attributed to differences in computational efficiency across model sizes. For smaller models such as LeNet5Half, TopK is more effective since the gradient dimension is relatively low, making the Top-K selection operation computationally efficient while preserving the most important gradient entries. In contrast, for larger models such as LeNet5, the cost of identifying Top-K entries becomes higher, and Random-K provides a more efficient alternative with lower selection overhead while maintaining competitive training performance under the same communication constraints. In contrast, for a large K = 5 × 103 on MNIST, as shown in Fig. 10, both LeNet5Half and LeNet5 clients gradually shift their preference toward the Random-K strategy. At this compression level, more gradient entries can be transmitted, which narrows the performance difference among compression strategies. As a result, the lower selection overhead of Random-K becomes more advantageous, leading to higher computational efficiency during training. The compression strategy selection results for CIFAR-10 under smaller and larger K values, i.e., K = 105 and K = 2 × 106 , are shown in Figs. 11 and 12, respectively. We can observe the same trend as MNIST. Overall, these observations indicate that the proposed ASCEND algorithm is capable of autonomously identifying effective compression strategies and adapting its decisions according to model architecture, resource constraints, and compression strategy features. 6) Sensitivity Analysis: As shown in Fig. 13, the effects of the weighting parameters ρ and β in (30) are evaluated on different datasets. The parameter setting ρ = 0.6 and β = 0.4 achieves stable performance across both datasets. For adaptive strategy selection, the sensitivity of the hyperparameter c in (35) is further investigated. As shown in Fig. 14, different values of c are evaluated to analyze their impacts on the performance of the proposed method. The results demonstrate that c = 3 achieves a good balance between exploration and exploitation, providing relatively favorable performance across different datasets.

0

0

500

1000 Round

1500

2000

Average Usage Count

1500

50

(a) Model prototype = LeNet5Half

Periodic-K

100 50 0

0

500

1000 Round

1500

2000

(b) Model prototype = LeNet5

Fig. 10: Compression strategy selection results of different model prototypes on MNIST dataset (α = 0.5, K = 5 × 103 ).

Top-K

Random-K

400 200 0

0

1500

3000 Round

4500

6000

(a) Model prototype = WResNet40-2

Average Usage Count

1000 Round

100

Periodic-K

200

0

0

1500

3000 Round

4500

6000

(b) Model prototype = ResNet18

Fig. 11: Compression strategy selection results of different model prototypes on CIFAR-10 dataset (α = 0.5, K = 105 ).

Top-K

Random-K

200 100 0

0

1500

3000 Round

4500

6000

(a) Model prototype = WResNet40-2

Average Usage Count

500

Random-K

Periodic-K

200 0

0

1500

3000 Round

4500

6000

(b) Model prototype = ResNet18

Fig. 12: Compression strategy selection results of different model prototypes on CIFAR-10 dataset (α = 0.5, K = 2 × 106 ).

0.96

0.57

0.92

0.54

ρ=0.4, β=0.6

0.51

ρ=0.5, β=0.5

0.48

ρ=0.6, β=0.4

Accuracy

0

Average Usage Count

0

50

Average Usage Count

50

100

Average Usage Count

100

Top-K

Periodic-K

Accuracy

Random-K Average Usage Count

Average Usage Count

Top-K

0.88 0.84 0.80

0.5 0.8 1.0 1.2 1.5 Time (s)

(a) MNIST

0.45

ρ=0.7, β=0.3 45 60 75 90 Time (s)

(b) CIFAR-10

Fig. 13: Impact of the weighting parameters ρ and β on the performance under different datasets.

VII. C ONCLUSION This paper proposes ASCEND, an adaptive heterogeneous compression framework for FedKD, where clients with heterogeneous computation and communication resources can dynamically select suitable compression strategies. The compression strategy selection problem is formulated as a nonstationary stochastic MAB problem. Each client adaptively chooses among Top-K, Random-K, and Periodic-K using an EMA-enhanced ϵ-greedy policy with a time-aware reward. Experimental results on MNIST and CIFAR-10 under non-

0.96

0.57

0.92

0.54

Accuracy

Accuracy

13

0.88 0.84 0.80 0.4 0.6 0.8 1.0 1.2 1.4 Time (s)

(a) MNIST

c=1 c=2

0.51

c=3

0.48 0.45

45 60 75 90 Time (s)

c=4

(b) CIFAR-10

Fig. 14: Impact of the hyperparameter c in the probability ϵt on the performance across different datasets.

IID data and different bandwidth capacities demonstrate that ASCEND consistently achieves a favorable balance between model accuracy and training latency. Future work will investigate the joint optimization of compression strategy selection and compression level adaptation, where the retained gradient budget can be dynamically adjusted according to model performance and system conditions. R EFERENCES [1] P. Kairouz, H. B. McMahan, B. Avent, A. Bellet, M. Bennis, A. N. Bhagoji, K. Bonawitz, Z. Charles, G. Cormode, R. Cummings et al., “Advances and open problems in federated learning,” Foundations and trends® in machine learning, vol. 14, no. 1–2, pp. 1–210, 2021. [2] T. Li, A. K. Sahu, M. Zaheer, M. Sanjabi, A. Talwalkar, and V. Smith, “Federated optimization in heterogeneous networks,” Proceedings of Machine learning and systems, vol. 2, pp. 429–450, 2020. [3] C. Wu, F. Wu, L. Lyu, Y. Huang, and X. Xie, “Communication-efficient federated learning via knowledge distillation,” Nature communications, vol. 13, no. 1, p. 2032, 2022. [4] S. Shi, K. Zhao, Q. Wang, Z. Tang, and X. Chu, “A convergence analysis of distributed sgd with communication-efficient gradient sparsification.” in IJCAI, 2019, pp. 3411–3417. [5] A. Xu and H. Huang, “Detached error feedback for distributed sgd with random sparsification,” in International conference on machine learning. PMLR, 2022, pp. 24 550–24 575. [6] S. U. Stich, “Local sgd converges fast and communicates little,” arXiv preprint arXiv:1805.09767, 2018. [7] Y. Lin, S. Han, H. Mao, Y. Wang, and W. J. Dally, “Deep gradient compression: Reducing the communication bandwidth for distributed training,” arXiv preprint arXiv:1712.01887, 2017. [8] G. Hinton, O. Vinyals, and J. Dean, “Distilling the knowledge in a neural network,” arXiv preprint arXiv:1503.02531, 2015. [9] P. Han, X. Huang, Y. Fang, and G. Han, “Rethinking knowledge distillation in collaborative machine learning: Memory, knowledge, and their interactions,” IEEE Transactions on Network Science and Engineering, vol. 12, no. 6, pp. 4498–4530, 2025. [10] P. Han, Z. Yin, C. Liu, and Y. Fang, “Agentic AI-driven federated feature distillation for adaptive resource–performance tradeoffs in wireless edge networks,” IEEE Transactions on Cognitive Communications and Networking, vol. 12, pp. 6076–6088, 2026. [11] D. Li and J. Wang, “Fedmd: Heterogenous federated learning via model distillation,” arXiv preprint arXiv:1910.03581, 2019. [12] Z. Zhu, J. Hong, and J. Zhou, “Federated learning with heterogeneous data: A generative approach,” in ICML. PMLR, 2021, pp. 12 804– 12 814. [13] Y. Tan, G. Long, L. Liu, T. Zhou, J. Jiang, and C. Zhang, “Fedproto: Federated prototype learning across heterogeneous clients,” Proceedings of the AAAI Conference on Artificial Intelligence, vol. 36, no. 8, pp. 8432–8440, 2022. [14] J. Zhang, Y. Chen, D. Wu, and B. Li, “Fedtgp: Federated learning with global prototypes for heterogeneous models,” IEEE Transactions on Mobile Computing, 2023. [15] J. Zhang, S. Guo, X. Ma, H. Wang, W. Xu, and F. Wu, “Parameterized knowledge transfer for personalized federated learning,” arXiv preprint arXiv:2111.02862, 2021. [16] T. Lin, L. Kong, S. U. Stich, and M. Jaggi, “Ensemble distillation for robust model fusion in federated learning,” Advances in neural information processing systems, vol. 33, pp. 2351–2363, 2020.

[17] C.-Y. Chen, J. Choi, D. Brand, A. Agrawal, W. Zhang, and K. Gopalakrishnan, “Adacomp: Adaptive residual gradient compression for dataparallel distributed training,” in Proceedings of the AAAI Conference on Artificial Intelligence, vol. 32, no. 1, 2018. [18] P. Han, S. Wang, and K. K. Leung, “Adaptive gradient sparsification for efficient federated learning: An online learning approach,” in 2020 IEEE 40th international conference on distributed computing systems (ICDCS). IEEE, 2020, pp. 300–310. [19] B. McMahan, E. Moore, D. Ramage, S. Hampson, and B. A. y Arcas, “Communication-efficient learning of deep networks from decentralized data,” in Artificial intelligence and statistics. PMLR, 2017, pp. 1273– 1282. [20] X. Li, K. Huang, W. Yang, S. Wang, and Z. Zhang, “On the convergence of fedavg on non-iid data,” arXiv preprint arXiv:1907.02189, 2019. [21] W. Wen, C. Xu, F. Yan, C. Wu, Y. Wang, Y. Chen, and H. Li, “Terngrad: Ternary gradients to reduce communication in distributed deep learning,” Advances in neural information processing systems, vol. 30, 2017. [22] A. F. Aji and K. Heafield, “Sparse communication for distributed gradient descent,” arXiv preprint arXiv:1704.05021, 2017. [23] F. Haddadpour, M. M. Kamani, A. Mokhtari, and M. Mahdavi, “Federated learning with compression: Unified analysis and sharp guarantees,” in International Conference on Artificial Intelligence and Statistics. PMLR, 2021, pp. 2350–2358. [24] D. Alistarh, D. Grubic, J. Li, R. Tomioka, and M. Vojnovic, “Qsgd: Communication-efficient sgd via gradient quantization and encoding,” Advances in neural information processing systems, vol. 30, 2017. [25] D. Basu, D. Data, C. Karakus, and S. Diggavi, “Qsparse-local-sgd: Distributed sgd with quantization, sparsification and local computations,” Advances in Neural Information Processing Systems, vol. 32, 2019. [26] A. Koloskova, N. Loizou, S. Boreiri, M. Jaggi, and S. U. Stich, “A unified theory of decentralized SGD with changing topology and local updates,” in Proceedings of the 37th International Conference on Machine Learning, ser. Proceedings of Machine Learning Research, vol. 119. PMLR, 2020, pp. 5381–5393. [27] P. Han, X. Shi, and J. Huang, “Fedal: Black-box federated knowledge distillation enabled by adversarial learning,” IEEE Journal on Selected Areas in Communications, vol. 42, no. 11, pp. 3064–3077, 2024. [28] B. E. Woodworth, K. K. Patel, and N. Srebro, “Minibatch vs local SGD for heterogeneous distributed learning,” in Advances in Neural Information Processing Systems, vol. 33, 2020. [29] S. P. Karimireddy, Q. Rebjock, S. U. Stich, and M. Jaggi, “Error feedback fixes SignSGD and other gradient compression schemes,” in Proceedings of the 36th International Conference on Machine Learning, ser. Proceedings of Machine Learning Research, vol. 97. PMLR, 2019, pp. 3252–3261. [30] P. Richtárik, I. Sokolov, and I. Fatkhullin, “EF21: A new, simpler, theoretically better, and practically faster error feedback,” in Advances in Neural Information Processing Systems, vol. 34, 2021. [31] R. S. Sutton and A. G. Barto, Reinforcement Learning: An Introduction, 2nd ed. Cambridge, MA: MIT Press, 2018. [32] D. Yin, Y. Chen, R. Kannan, and P. Bartlett, “Byzantine-robust distributed learning: Towards optimal statistical rates,” in International conference on machine learning. Pmlr, 2018, pp. 5650–5659. [33] S. P. Karimireddy, S. Kale, M. Mohri, S. Reddi, S. Stich, and A. T. Suresh, “Scaffold: Stochastic controlled averaging for federated learning,” in International conference on machine learning. PMLR, 2020, pp. 5132–5143. [34] A. Garivier and E. Moulines, “On upper-confidence bound policies for switching bandit problems,” in Algorithmic Learning Theory, ser. Lecture Notes in Computer Science, vol. 6925. Springer, 2011, pp. 174–188. [35] F. Liu, J. Lee, and N. Shroff, “A change-detection based framework for piecewise-stationary multi-armed bandit problem,” in Proceedings of the AAAI Conference on Artificial Intelligence, vol. 32, no. 1, 2018. [36] P. Auer, N. Cesa-Bianchi, and P. Fischer, “Finite-time analysis of the multiarmed bandit problem,” Machine learning, vol. 47, no. 2, pp. 235– 256, 2002. [37] Y. LeCun, L. Bottou, Y. Bengio, and P. Haffner, “Gradient-based learning applied to document recognition,” Proceedings of the IEEE, vol. 86, no. 11, pp. 2278–2324, 2002. [38] A. Krizhevsky, “Learning multiple layers of features from tiny images,” Department of Computer Science, University of Toronto, Toronto, ON, Canada, Tech. Rep., 2009.

14

Server

Parameter averaging

Prototype � Client �

p Model Gradient reconstruction updating

�

���

��

Gradient Local Client � training compression � ��� ��� ���

�

��

�

Logit calculation

���

For theoretical analysis, we define the joint optimization objective of prototype p as

Logit KD averaging �

��+1

(61)

where λ balances the local optimization objective and the global knowledge alignment objective. Define ηKD ≜ ληt and substituting (58) into (60) yields

��

�+1

Fp (ω pt ) = fp (ω pt ) + λLpKD (ω pt ),

�+1

���+1

Fig. 15: The pipeline of a single communication round in FedKD with heterogeneous compression.

ω pt+1 = ω pt − ηt (ḡ pt + λ∇LpKD (ω̄ pt )) .

(62)

For convenience, define the aggregated residual and stochastic gradient noise for prototype p as 1 X n et , (63) ept = |Np | n∈Np

A PPENDIX

and

To facilitate understanding in the theoretical analysis, we first present the overall workflow of a single communication round in Figure 15. A. Proof of Theorem 1 We rearrange the training process from the perspective of each prototype. Specifically, at the beginning of training round t, the parameter of prototype p is ω pt . The server broadcasts the updated prototype model after each KD step, all clients belonging to the same prototype start each communication round with identical model parameters, i.e., ω nt = ω pt ,

∀n ∈ Np .

(55)

For local training, we define the local optimization objective for prototype p as 1 X fp (ω) = fn (ω), (56) |Np | n∈Np

where fp (ω) denotes the prototype-wise empirical loss. Based on the received compressed gradients from client n in round t, the server reconstructs the corresponding gradient vector g̃ nt . The reconstructed gradient is used to update the local model of client n by ω̃ nt = ω pt − ηt g̃ nt .

(58)

and E[∥ξ pt ∥2 ] ≤ σ 2 .

where ηKD denotes the KD learning rate.

(60)

(66)

We first propose the following lemma to show the smoothness of Fp (ω). Lemma 4 (Smoothness of the Joint Objective). The joint objective function Fp (ω) defined in (61) is L-smooth. Specifically, there exists a constant L = Ll + λLp > 0 such that ∥∇Fp (ω) − ∇Fp (ω ′ )∥ ≤ L ∥ω − ω ′ ∥ ,

∀ω, ω ′ .

(67)

Proof. According to the definition of the joint optimization objective in (61), therefore, for any two model parameters ω and ω ′ , the gradient difference can be written as ∇Fp (ω) − ∇Fp (ω ′ ) = ∇fp (ω) − ∇fp (ω ′ ) + λ (∇LpKD (ω) − ∇LpKD (ω ′ )) . (68) Taking the norm on both sides and applying the triangle inequality, we have ∥∇Fp (ω) − ∇Fp (ω ′ )∥ ≤ ∥∇fp (ω) − ∇fp (ω ′ )∥

(69)

+ λ ∥∇LpKD (ω) − ∇LpKD (ω ′ )∥ . Applying Assumption 1 and Assumption 2 into the gradient difference bound yields ∥∇Fp (ω) − ∇Fp (ω ′ )∥ ≤ Ll ∥ω − ω ′ ∥ + λLp ∥ω − ω ′ ∥

(70)

′

= (Ll + λLp ) ∥ω − ω ∥ .

n∈Np

ω pt+1 = ω̄ pt − ηKD ∇LpKD (ω̄ pt ),

(64)

Applying Assumption 3 and 4, the stochastic gradient noise satisfies E[ξ pt |ω pt ] = 0, (65)

where the aggregated reconstructed gradient is defined as 1 X n ḡ pt = g̃ t . (59) |Np | After prototype aggregation, the server performs one-step KD on the public dataset Dpub . The KD update is formulated as

1 X n (g t − ∇fn (ω pt )). |Np | n∈Np

(57)

The server performs parameter-averaging for prototype aggregation as formulated in (4). Substituting (57) into (4), the prototype update can be equivalently written as ω̄ pt = ω pt − ηt ḡ pt ,

ξ pt =

By defining the smoothness constant as L = Ll + λLp ,

(71)

∥∇Fp (ω) − ∇Fp (ω ′ )∥ ≤ L ∥ω − ω ′ ∥ ,

(72)

we obtain

which completes the proof.

15

We then propose the following lemma to capture the deviation in each round t. Lemma 5 (Error-feedback Decomposition under Heterogeneous Compression). For each prototype p, the deviation between the aggregated reconstructed gradient ḡ pt and the true prototype gradient satisfies ḡ pt − ∇fp (ω pt ) = ept−1 − ept + ξ pt .

(73)

Proof. According to the error-feedback update, the residual vector at client n satisfies ent = unt − g̃ nt .

(74)

g̃ nt − g nt = ent−1 − ent .

(75)

Using (10), therefore,

For the inner-product term of (83), we have − ηt ⟨∇Fp (ω pt ), ∇Fp (ω pt ) + δ pt ⟩ = − ηt ∥∇Fp (ω pt )∥2 − ηt ∇Fp (ω pt ), ept−1 − ept − ηt ⟨∇Fp (ω pt ), ξpt ⟩ .

(84)

Taking expectation and according to Assumption 3, the stochastic gradient noise term vanishes: E [⟨∇Fp (ω pt ), ξpt ⟩] = 0.

(85)

Therefore, E [−ηt ⟨∇Fp (ω pt ), ∇Fp + δ pt ⟩]     = − ηt E ∥∇Fp (ω pt )∥2 − ηt E ∇Fp (ω pt ), ept−1 − ept . (86) For the second-order term of (83), according to

By adding and subtracting ∇fn (ω pt ), we obtain

∥a + b∥2 ≤ 2∥a∥2 + 2∥b∥2 ,

g̃ nt − ∇fn (ω pt ) =ent−1 − ent

we have Lηt2 ∥∇Fp (ω pt ) + δ pt ∥2 ≤ Lηt2 ∥∇Fp (ω pt )∥2 + Lηt2 ∥δ pt ∥2 . 2 (87)

+ (g nt − ∇fn (ω pt )) .

(76)

Taking the average over all clients belonging to prototype p gives ḡ pt − ∇fp (ω pt ) =ept−1 − ept + ξ pt ,

we obtain

Now, we prove Theorem 1. Proof. According to Lemma 4, the joint objective function Fp (ω) is L-smooth. Therefore, we obtain

      E ∥δ pt ∥2 ≤3E ∥ept−1 ∥2 + 3E ∥ept ∥2   + 3E ∥ξ pt ∥2 . By Assumption 4 and Assumption 5, we have,   E ∥δ pt ∥2 ≤ 3(2Γmax + σ 2 ) ≜ Cg .

(78)

According to the prototype update rule after parameter aggregation and KD, the prototype model is updated following (62). The definition of the aggregated gradient deviation, δ pt = ḡ pt − ∇fp (ω pt ).

∥a + b + c∥2 ≤ 3∥a∥2 + 3∥b∥2 + 3∥c∥2 ,

(77)

which completes the proof.

Fp (ω pt+1 ) ≤Fp (ω pt ) + ∇Fp (ω pt ), ω pt+1 − ω pt L + ∥ω pt+1 − ω pt ∥2 . 2

According to Lemma 5, the aggregated gradient deviation can be decomposed as (80). Using the inequality

(79)

According to Lemma 5, the aggregated gradient deviation can be written as δ pt = ept−1 − ept + ξ pt . (80) From (61) and (79), the update direction can be rewritten

+ Lηt2 Cg .

ω pt+1 − ω pt = −ηt (∇Fp (ω pt ) + δ pt ) .

(82)

Substituting (82) into (78), we have Fp (ω pt+1 ) ≤Fp (ω pt ) − ηt ⟨∇Fp (ω pt ), ∇Fp (ω pt ) + δ pt ⟩ Lηt2 ∥∇Fp (ω pt ) + δ pt ∥2 . (83) + 2

(90)

The accumulated residual term can be bounded by the errorfeedback telescoping property. Specifically, T −1 X

(ept−1 − ept ) = ep−1 − epT −1 .

(81)

Therefore, by combining the result with (62), we have

(89)

where Cg denotes the bounded gradient deviation constant induced by heterogeneous compression and stochastic optimization. Combining (86), (87) and (89), (83) becomes   E Fp (ω pt+1 )   ≤E [Fp (ω pt )] − (ηt − Lηt2 )E ∥∇Fp (ω pt )∥2   − ηt E ∇Fp (ω pt ), ept−1 − ept

as ḡ pt + λ∇LpKD (ω̄ pt ) = ∇Fp (ω pt ) + δ pt .

(88)

(91)

t=0

According to the standard error-feedback convergence property [29] and Assumption 5, the residual accumulation term is bounded by a finite constant: T −1 X

  ηt E ⟨∇Fp (ω pt ), ept−1 − ept ⟩ ≤ Ck ,

(92)

t=0

where Ck is a finite constant related to the error-feedback residual bound.

16

Summing (90) from t = 0 to T − 1 gives T −1 X

  (ηt − Lηt2 )E ∥∇Fp (ω pt )∥2

t=0

≤Fp (ω p0 ) − Fp (ω pT ) + Ck + LT ηt2 Cg .

(93)

1 Since ηt ≤ 4L ,

ηt − Lηt 2 ≥

3ηt , 4

(94)

2) Proof of Lemma 2: Proof. During the exploration stage, ASCEND selects a compression strategy randomly with probability ϵt , as defined in (35). Therefore, the expected regret caused by exploration at communication round t is bounded by
n n ϵt µnt (sn,∗ (102) t ) − µt (st ) . Due to the reward clipping operation, the instantaneous regret is bounded by n n µnt (sn,∗ t ) − µt (st ) ≤ ∆max .

therefore, T −1  3ηt X  E ∥∇Fp (ω pt )∥2 4 t=0

Therefore, the cumulative exploration regret after T communication rounds satisfies

≤Fp (ω p0 ) − Fp∗ + Ck + LT ηt2 Cg .

(95)

Rnexplore (T ) ≤ ∆max

  1 E ∥∇Fp (ω pt )∥2 ≤ T t=0

4(Fp (ω p0 ) − Fp∗ + Ck ) 3ηt T

T X

ϵt

t=1 T X

Dividing both sides by 3η4t T , we obtain T −1 X

(103)

4L + ηt Cg . 3 (96)

Therefore, by choosing √ ηt = O(1/ T ), √ the first term and the second term both decrease as O(1/ T ). Therefore, T −1 √  1 X  E ∥∇Fp (ω pt )∥2 = O(1/ T ). T t=0

(97)

≤ c∆max

1 √ . t t=1

Using the integral inequality, Z T T X √ 1 √ ≤1+ x−1/2 dx ≤ 2 T , t 1 t=1

(104)

(105)

we obtain (47). 3) Proof of Lemma 3: Proof. During the exploitation phase, ASCEND selects the compression strategy according to the current utility estimation: snt = arg max Qnt (s). (106) s∈S

Define the utility estimation error as vtn (s) = Qnt (s) − µnt (s).

B. Proof of Theorem 2 1) Proof of Lemma 1: Proof. When strategy s is selected, the EMA utility update follows (33). Within the segment Ti , by Assumption 6 and (38), the expected reward remains constant: E[Gtn (s)] = µni (s).

(98)

Subtracting µni (s) from both sides of (33) yields Qnt (s) − µni (s) =(1 − ηq )(Qnt−1 (s) − µni (s)) + ηq (Gtn (s) − µni (s)).

Since snt maximizes the estimated utility, Qnt (snt ) ≥ Qnt (sn,∗ t ).

n n,∗ µnt (snt ) + vtn (snt ) ≥ µnt (sn,∗ t ) + vt (st ).

(110)

The instantaneous exploitation regret is (99)

n n n n n n,∗ µnt (sn,∗ t ) − µt (st ) ≤ |vt (st )| + |vt (st )|

≤ 2 max |vtn (s)|.

E[Qnt (s) − µni (s)] (100)

Applying the above recursion for m updates gives

Taking the absolute value on both sides gives (46).

(109)

Therefore,

s∈S

E[Qnt (s) − µni (s)] = (1 − ηq )m E[Qnri (s) − µni (s)].

(108)

Substituting Qnt (s) = µnt (s) + vtn (s) gives

n n n n n n,∗ µnt (sn,∗ t ) − µt (st ) ≤ vt (st ) − vt (st ).

Taking expectation gives

=(1 − ηq )E[Qnt−1 (s) − µni (s)].

(107)

(101)

(111)

Next, we analyze the cumulative estimation error within a stationary segment Ti . According to Assumption 6, the expected reward remains unchanged within each segment. The utility estimation error can be decomposed into the EMA bias and the stochastic fluctuation: vtn (s) = Qnt (s) − E[Qnt (s)] + E[Qnt (s)] − µni (s) . | {z } | {z } stochastic fluctuation

EMA bias

(112)

17

n (s) According to Assumption 7 and Lemma 1, after Hi,t updates of strategy s, its utility estimation error satisfies

E [|vtn (s)|] ≤ q

C1 n (s) ∨ 1 Hi,t

+ C2 (1 − ηq )

n Hi,t (s)

,

(113)

where a ∨ b = max{a, b}, and C1 and C2 are constants determined by the reward variance, EMA coefficient, and initialization error. Combining (113) with (111), the expected instantaneous exploitation regret at round t ∈ Ti satisfies   E µnt (sn,∗ ) − µnt (snt )  t 1 1  +q ≤ C1  q n n n Hi,t (st ) ∨ 1 ) ∨ 1 Hi,t (sn,∗ t h i n,∗ n n Hi,t (sn ) t + C2 (1 − ηq ) + (1 − ηq )Hi,t (st ) . (114) Under Assumption 8, for every s ∈ S, (114) yields   n n E µnt (sn,∗ t ) − µt (st ) 2C1 + 2C2 (1 − ηq )πmin (t−ri +1) . ≤p πmin (t − ri + 1)

(115)

Summing (115) over t ∈ Ti , we obtain 4C1 p 2C2 Rnexploit,i ≤ √ |Ti | + . πmin 1 − (1 − ηq )πmin Let

4C1 2C2 Ce = √ + , πmin 1 − (1 − ηq )πmin

(116)

where Ce is independent of |Ti |. Since |Ti | ≥ 1, we have p Rnexploit,i ≤ Ce |Ti |. (117) Finally, summing over all stationary segments gives Rnexploit (T ) =

M X

Rnexploit,i

i=1

≤ Ce

M p X

|Ti |

i=1

v u M u X ≤ Ce tM |Ti | √

i=1

= Ce M T ,

(118)

where the second inequality follows from the Cauchy–Schwarz inequality. 4) Proof of Theorem 2: Proof. We first analyze the regret incurred during the warmup phase. During warm-up, each compression strategy is selected exactly h times to initialize the utility estimation. Therefore, the number of warm-up rounds is h|S|. Since the reward is clipped into [Gmin , Gmax ], the maximum instantaneous regret is bounded by (48). Therefore, the warmup regret is bounded by Rnwarm ≤ h|S|∆max .

(119)

Combining (119), Lemmas 2 and 3, we obtain RnASCEND (T ) = Rnwarm + Rnexplore (T ) + Rnexploit (T ), (120) √ √ ≤ h|S|∆max + 2c∆max T + Ce M T . (121) Therefore,

√ RnASCEND (T ) = O( M T ).

(122)

Related documents

• Adaptive Heterogeneous Compression for Resource-Efficient Federated Knowledge Distillation#616809
• Adaptive Heterogeneous Compression for Resource-Efficient Federated Knowledge Distillation#616309
• Knowledge Distillation in Wide Neural Networks: Risk Bound, Data Efficiency and Imperfect Teacher#974116
• Data Unlearning via Inverse Distillation#1130941
• A Unified Per-Token Gating Family for On-Policy Distillation: FKL/RKL Mixing with Multi-Channel and Bias Coefficients#999455