JOURNAL OF LATEX CLASS FILES, VOL. 18, NO. 9, SEPTEMBER 2020
1
AirMoE: Statistic-Augmented Over-the-Air MoE for Collaborative Intelligence
arXiv:2607.16562v1 [cs.DC] 18 Jul 2026
Wei-Bin Kou, Jingreng Lei, Guangxu Zhu∗ , Yujiu Yang∗
Abstract—In modern edge intelligence, Mixture of Experts (MoE) are increasingly deployed over wireless cloud-edge networks, as a single edge device lacks sufficient resources to host large-scale models locally. In this distributed architecture, a cloud-hosted pretrained Large Model (LM) acts as a shared backbone for latent feature extraction, while heterogeneous experts deployed across distributed, wirelessly-connected clients collaboratively form the task head. However, deploying MoE over wireless links exposes two coupled bottlenecks. On the one hand, routing which clients to activate generally overloads bandwidth-limited uplinks due to required raw feature transmission. On the other hand, aggregating the activated experts’ outputs over wireless links is hindered by channel noise and poor scalability. To break these bottlenecks, we propose a statisticaugmented over-the-air MoE (AirMoE) paradigm. Specifically, on the routing side, each client queries its local Feature Retrieval Library (FRL) with a cloud-broadcast compact query, retrieves a prototype-induced statistic, and reports it digitally to the cloud, drastically reducing uplink traffic; the cloud then selects the most relevant clients by aligning these statistics with the LM-extracted features via Jensen–Shannon (JS) divergence. On the aggregating side, selected experts simultaneously transmit their outputs over the multiple-access channel, which physically computes the reweighted sum via waveform superposition, with reweighting coefficients realized through channel-aware power control. The two mechanisms are thus decoupled both algorithmically and physically. We further provide theoretical analyses on convergence and iteration complexity. Taking semantic segmentation task as an example, extensive experiments demonstrate that AirMoE outperforms MoE baselines and singlemodel competitors. Ablations further confirm the effectiveness of each incorporated component. Index Terms—Mixture of Experts (MoE), Statistic-Augmented Routing, Over-the-Air Aggregating, Edge Intelligence, CloudEdge Distributed Inference, Large Models (LMs).
I. I NTRODUCTION The relentless growth of Mixture of Experts (MoE) models has pushed their resource demands far beyond the capability of a single edge device. Consequently, modern edge intelligence increasingly deploys MoE in a cloud–edge collaborative setting to split the computation [1]–[6]. In this distributed architecture, a powerful cloud server hosts a heavyweight, pretrained Large Model (LM) that acts as a fixed, shared backbone responsible for latent feature extraction. Meanwhile, lightweight, specialized experts are deployed across Wei-Bin Kou and Yujiu Yang are with Tsinghua Shenzhen International Graduate School, Tsinghua University, Shenzhen, China. Jingreng Lei is with the Department of Electrical and Electronic Engineering, The University of Hong Kong, Hong Kong, China. Guangxu Zhu is with Shenzhen Research Institute of Big Data, Shenzhen, China. Corresponding authors: Guangxu Zhu and Yujiu Yang
Fig. 1. Illustration of the two coupled bottlenecks in wireless MoE.
distributed, wirelessly connected clients to collaboratively form the task head. Because each client may have acquired expertise on a specific input distribution (e.g., a particular environment, sensor, or operating condition), the population of clients naturally forms a heterogeneous expert pool. As in any MoE system, harnessing this expert pool hinges on two core operations: routing, which selects the most relevant experts for each input’s LM-extracted features, and aggregating, which fuses the selected experts’ outputs into the final prediction. However, once the experts are scattered behind wireless links, these two routine operations turn into two coupled bottlenecks that do not exist in a single-machine MoE. First, routing becomes a communication-dominant problem rather than a purely algorithmic one: naively transporting the highdimensional latent features required for routing decisions would saturate the bandwidth-limited uplink channel [7]. Second, aggregating becomes a multiple-access scalability problem: instead of fusing expert outputs by a weighted sum at one location as in single-machine MoE, the cloud has to collect each selected client’s full output over orthogonal uplinks (e.g., TDMA) and fuse them digitally, thus both the uplink cost and the scheduling latency scale linearly with the number of activated experts. The two bottlenecks of the wireless MoE are illustrated in Fig. 1. Fortunately, the two bottlenecks exhibit fundamentally different natures, which in turn suggests different remedies. Routing, on the one hand, is inherently a discrete selection problem. It suffices to select most relevant experts merely based on the compact statistics rather than costly raw features. This works because the routing function is invariant to most of the variation in the raw features. For example, two inputs from the same road condition should be routed identically even if their pixellevel features differ. The statistic deliberately eliminates the variations that are irrelevant to routing decisions in the input.
JOURNAL OF LATEX CLASS FILES, VOL. 18, NO. 9, SEPTEMBER 2020
2
Fig. 2. Overview of the proposed AirMoE.
On the other hand, aggregating admits an elegant physicallayer solution via Over-the-air computing (AirComp) [8]–[10], which leverages the waveform-superposition property of the shared multiple-access channel to sum signals. This property is particularly suited to the wireless MoE aggregating, since, when routed experts transmit their analog-modulated symbols concurrently, the channel can inherently superimpose them into a channel-weighted sum. Instead of transmitting experts’ output over orthogonal uplink and fusing them at the cloud, over-the-air aggregation allows all selected experts to transmit output simultaneously and computes their sum by the channel itself. Consequently, both the cost and the latency of the fusion become independent of the number of activated experts. This fundamental dichotomy between wireless MoE’s routing and aggregating motivates decoupled design. We address them with individual mechanisms tailored to each characteristic. Specifically, we propose AirMoE, a statisticaugmented, routing-aggregating-decoupled over-the-air MoE paradigm. The framework comprises the following integral components. (i) Cloud-side Feature Extraction and Compact Query Generation: the cloud deploys a pretrained LM backbone to extract latent features for each input and summarizes it into a compact query. Then, the LM-extracted latent features and the summarized query are broadcast to all clients. At client side, the received latent features are used to compute client-wise intermediate features, while the received compact query is used to retrieve clients’ expertise. (ii) Client-Wise Feature Retrieval Library (FRL): each client equips the hosted expert with a FRL which stores prototypical feature patterns that encode the client’s expertise obtained from its historical experiences. The FRL is maintained by a read-then-update fashion using attention rule based on client-side hidden features. (iii) Statistic-Augmented Routing Mechanism (MoERM): as mentioned before, the cloud broadcasts the sum-
marized compact query to all clients and each client locally retrieves a prototype-induced statistic from its equipped FRL based on the cloud query. Subsequently, each client reports the retrieved statistic over a small digital uplink to the cloud. After receiving statistics from all clients, the cloud aligns these clients’ statistics with the summarized compact query via Jensen–Shannon (JS) divergence, which is used by the cloud to select the most aligned clients. Finally, the cloud feeds back which clients are routed. (iv) Over-the-Air Aggregating Mechanism (MoE-AM): when receiving the activated feedback from the cloud, each routed client compresses the cloud-broadcast latent features to obtain its own hidden features. Subsequently, it calculates and sends the JS divergence of its hidden features’ statistics against the cloud query to the cloud. After receiving all routed clients’ JS, the cloud calculates each routed client’s aggregation weight based on those JSs, and then sends back the calculated weight to each routed client. Upon receiving their aggregation weights, the routed clients simultaneously transmit power-scaled expert outputs over the multiple-access channel, which physically computes the weighted MoE aggregation via AirComp. In a nutshell, AirMoE decouples routing from aggregating both algorithmically and physically: routing is digital, low-rate, and robust, whereas aggregating is analog, one-shot, and bandwidth efficient. In summary, the proposed AirMoE is illustrated in Fig. 2, and the main contributions are highlighted as follows: •
We formulate MoE over a cloud–edge wireless network, and identify routing and aggregating as the two communication bottlenecks. We propose AirMoE to address these bottlenecks from a statistical perspective by integrating the following integral components: cloud-side feature extraction and compact query generation, client-wise FRL, a statistic-augmented routing mechanism (MoE-RM), and an over-the-air aggregating mechanism (MoE-AM).
JOURNAL OF LATEX CLASS FILES, VOL. 18, NO. 9, SEPTEMBER 2020
Client-wise FRL is updated in a read-then-update manner to store each client’s expertise as compact prototypes. MoE-RM uses prototype-induced statistics instead of raw features to select clients to reduce uplink saturation, and MoE-AM exploits physical waveform superposition to compute the expert fusion directly over the channel via power control, achieving latency and bandwidth that are invariant to the number of activated clients. • We provide theoretical guarantees for AirMoE, including convergence analysis and ϵ-stationarity complexity, in over-the-air setting based on reasonable assumptions. • Taking semantic segmentation task as an example, extensive experiments show that AirMoE outperforms MoE and single-model baselines, and ablations further confirm the effectiveness of the incorporated components. The remainder of this paper proceeds as follows. Section II reviews related works. Section III formulates the wireless MoE over cloud–edge setting. Section IV details the proposed AirMoE paradigm. Section V provides theoretical guarantees for the proposed AirMoE. Section VI presents experiments and ablations. Section VII concludes this paper. •
II. R ELATED W ORK A. Distributed and Edge Foundation Models Splitting foundation models over a cloud–edge setting has been studied under split inference and collaborative edge intelligence [2], [11]–[14], where an early stage of a deep network runs on the cloud server and a later stage runs on the edge device, or vice versa. Recent works refined this paradigm along several axes. For example, split computing frameworks introduce learnable bottleneck layers to compress the intermediate features transmitted across the split point [15], [16]; early-exit and adaptive-depth architectures allow inference to terminate at intermediate stages to trade accuracy for latency under fluctuating channel conditions [17]; and device–edge co-inference schemes jointly optimize the partition point, feature compression, and resource allocation to meet end-to-end latency and energy budgets [18], [19]. These studies, however, predominantly assume a monolithic network that is merely partitioned across nodes, leaving the deployment of heterogeneous, specialized sub-models on distributed devices largely unexplored. Furthermore, these studies largely overlook the communication constraints. In parallel, recent efforts on edge foundation models pursue model compression, quantization, and parameter-efficient adaptation to fit foundation-scale backbones within constrained hardware [20]–[22]. Meanwhile, other parameter-efficient fine-tuning techniques such as low-rank adapters [23] make it feasible to specialize a frozen backbone with a small number of trainable parameters. Our work differs from the above in three ways: (i) we keep the foundation-scale backbone fixed at the cloud and place the adaptation on parallel, distributed MoE experts across devices; (ii) we explicitly account for the wireless links that connect the cloud backbone to the parallel, distributed edge experts, treating communication not as an afterthought but as a core design constraint; and (iii) we decouple the MoE routing and aggregating mechanisms from the underlying wireless channel.
3
B. MoE and Its Distributed Variant MoE has become a pivotal method for scaling models and enhancing task performance [3], [24], and its strength lies in harnessing the specialization of individual experts across diverse data [4], [25]. The modern sparsely-gated MoE was popularized by [26], which demonstrated that conditional computation could increase model capacity while keeping pertoken compute nearly constant. Subsequent efforts simplified and stabilized the routing mechanism. For instance, Switch Transformers [27] route each token to a single expert to reduce communication and computation; GShard [28] introduced automatic sharding primitives to enable MoE models to scale to billions of parameters across thousands of devices. In addition, GLaM [29] showed that MoE language models can match or exceed dense counterparts at a fraction of the training and inference cost, and recent open MoE systems such as Mixtral [30] and DeepSeekMoE [31] have further refined expert granularity and load balancing for practical deployment. A primary difficulty in MoE training is balancing expert utilization to avoid load collapse, where a small subset of experts dominates routing [32], [33]. To this end, BASE Layers [32] cast tokento-expert assignment as a balanced linear assignment problem, while Expert Choice routing [33] inverts the selection so that experts choose tokens, guaranteeing balanced loads. Emerging works place experts on different machines or accelerators for capacity scaling [28], [34], [35]. Systems such as Tutel [34] provide optimized expert-parallel runtimes, and dedicated communication-scheduling techniques [36] have been proposed to mitigate the all-to-all bottleneck that dominates distributed MoE execution. However, most of these works assume high-bandwidth, reliable interconnects, and the allto-all dispatch and aggregation operations are designed under the assumption of essentially lossless, symmetric links. Our considered setting, where experts are hosted at heterogeneous wireless clients and connected to the cloud through noisy and bandwidth-limited links, introduces routing and aggregating challenges that are absent under above the ideal assumptions, and that are the focus of this paper. C. Over-the-Air Computing (AirComp) AirComp leverages the signal-superposition property of the multiple-access channel to compute the weighted sum of distributed data directly during transmission. The informationtheoretic foundations of this idea trace back to the study of reliable computation over multiple-access channels [37], where it was shown that uncoded analog transmission can be optimal for estimating functions of correlated sources, in stark contrast to the separation-based design philosophy of digital communication. AirComp has been extensively applied to fast wireless data aggregation and to federated learning (FL) [38]–[41], where the global model update is a sum over clients [8]–[10]. In the FL setting, over-the-air aggregation has been shown to dramatically reduce communication latency relative to orthogonal multiple access [8], and convergence guarantees have been established even in the presence of channel fading and additive noise [42]. Subsequent works have refined this paradigm through joint device-selection and
JOURNAL OF LATEX CLASS FILES, VOL. 18, NO. 9, SEPTEMBER 2020
beamforming [43]. Power control and signal misalignment have likewise been identified as central design governing the accuracy of the aggregated estimate [44]. We are, to our knowledge, the first to map AirComp onto the wireless MoE aggregation, i.e., the expert fusion weights become transmit-power-control coefficients, and the channel performs the MoE fusion. This view also reframes expert sparsity as a communication budget, where the TopK routing not only preserves specialization but also bounds the number of simultaneous transmitters. D. MoE Routing and Aggregating MoE routing is generally achieved by a gating function that assigns inputs to preferred experts. The dominant paradigm remains sparse TopK gating, first introduced in the sparselygated MoE layer [26] and later streamlined to single-expert routing in Switch Transformers [27] to minimize routing overhead. Linear (e.g., softmax) gating [4], [25] is simple and common, while nonlinear alternatives such as cosine-similarity gating [45] project inputs onto a hypersphere and compare to expert embeddings. A recurring concern with sparse gating is training instability and load imbalance, which has motivated auxiliary load-balancing losses [27], assignment-based formulations that enforce balanced routing by construction [32], [33], and stochastic or noise-injected routers that improve exploration over experts [46]. MoE aggregation typically combines selected experts by weighted sum based on gatingderived weights. Beyond simple weighted sums, recent studies revisit the aggregation mechanism. For instance, hierarchical and multi-level routing structures organize experts into groups before fusion [47]; analyses of expert specialization show that the quality of the aggregated output depends critically on how routing decisions and combination weights interact [31]. The statistical behavior of these gating and aggregation functions has also been studied for expert estimation [45]. In this work, we separate routing and aggregating from a statistical perspective, and further decouple their communication modalities, aiming to enhance the MoE overall performance under wireless constraints. III. OVER - THE -A IR M O E F ORMULATION We formulate the wirelessly connected MoE in this section. This formulation makes explicit the two coupled constraints (i.e., uplink routing cost and wireless aggregation scalability) that motivate the proposed AirMoE in Section IV. Related key notations are summarized in Table I.
4
TABLE I K EY N OTATIONS IN THE PROPOSED A IR M O E Notation
Explanation
D(i)
The i-th input image of dataset D. Parameters of the cloud LM backbone. LM-extracted latent feature of D(i) . (i) Cloud-side compact query of FLM and its normalization. Number of clients/experts. Expert at client j and its parameters. Intermediate rep. and output of expert j. Feature Retrieval Library (FRL) of client j. k-th prototype and its importance weight in Mj . Attention weight on pj,k for input i.
ωLM (i) FLM (i) q , Q(i) N Ej , θj (i) (i) hj , yj Mj pj,k , wj,k (i) αj,k p̃j , Pj η (i) Rj
(i)
(i)
Retrieved prototype of client j and its normalization. FRL memory update rate. Output-induced distributions.
(i)
(i)
Routing score and probability of client j. Index set of TopK selected clients. (i) JS divergence between Q(i) and Rj .
sj , π j S (i) (i) δj (i)
βj γj (i) bj σc2 ŷ (i)
Aggregation weight of client j. Uplink channel coefficient of client j. Input-dependent transmit pre-scaling (power control) of client j. Complex receiver-noise power at the cloud. Over-the-air aggregated output for input i.
a particular SNweather regime, illumination, or road type), so that D = j=1 Dj can reflect the significantly heterogeneous working scenarios of real-world applications. This heterogeneity is precisely what a single centralized model fails to capture, and what the distributed MoE heads are intended to exploit. B. Distributed MoE Routing and Aggregating For the i-th input D(i) , the cloud first extracts a shared latent (i) (i) feature FLM . Based on FLM , a routing operator R selects a subset of clients via (i) (i) S (i) = R FLM ; {Mj }N |S (i) | = K, j=1 ⊆ {1, . . . , N }, (1) where K ≪ N is the number of routed clients by TopK (i) and Mj denotes any client-side state of input D(i) used by routing. Meanwhile, an aggregation operator A fuses the (i) selected outputs {yj } (where j ∈ S (i) ) into the final prediction ŷ (i) through X X (i) (i) (i) (i) βj yj , βj = 1, (2) ŷ (i) = A {yj }j∈S (i) = j∈S (i)
j∈S (i)
(i)
A. Network and Computation Setting We consider a cloud–edge MoE system consisting of one cloud server and N distributed clients indexed by j ∈ {1, . . . , N }, connected over a shared wireless multiple-access channel. The cloud hosts a pretrained LM backbone with parameters ωLM , while client j hosts an expert Ej with parameters θj = θjenc ∪θjdec . The collection of clients {Ej }N j=1 forms a heterogeneous expert pool, where each Ej is specialized on a subset Dj ⊆ D of the data distribution (e.g.,
where {βj }j∈S (i) are the nonnegative fusion weights. Eq. (2) exposes the structural property of MoE aggregation, which is a weighted sum and computable over a multiple-access channel. C. Communication Model Realizing Eq. (1)–Eq. (2) over the wireless network incurs two distinct communication challenges. Prohibitive routing uplink. To evaluate the routing operator R, for each client j (j ∈ {1, · · · , N }), it must report some (i) descriptor Mj to the cloud over a bandwidth-constrained
JOURNAL OF LATEX CLASS FILES, VOL. 18, NO. 9, SEPTEMBER 2020
5
(i)
digital uplink. Let b(Mj ) denote its size. The total routing cost per input is XN (i) (i) Croute = b Mj . (3) j=1
(i)
A naive choice of Mj is the client-side latent features, which (i) makes Croute prohibitive for bandwidth-limited uplink. Poor scalability of digitally wireless aggregation. Given S (i) , the selected clients transmit their outputs simultaneously (i) with input-dependent pre-scaling bj ∈ C under the per(i) (i) client power budget E ∥bj yj ∥22 ≤ P0 . Through the uplink channels γj ∈ C and receiver noise n ∼ CN (0, σc2 I), the cloud-aggregated signal is X (i) (i) r(i) = γj bj yj + n. (4) (i) j∈S
For the conventional orthogonal digital transmission (e.g., TDMA), the uplink cost and latency scale linearly with the number of activated experts, and becomes non-negligible as the number of the activated clients increases. In addition, realizing Eq. (4) digitally introduces aggregation distortion. The aggregation distortion can be formulated as (i)
2
(i) Eagg = E ŷota − ŷ (i) 2 ,
(5)
(i)
where ŷota is an estimate of the target aggregation ŷ (i) . This (i) distortion depends jointly on {bj }, {γj }, and σc2 . D. Optimization Objective of Wireless MoE The optimization goal is to jointly design the routing opera(i) tor R, the fusion weights {βj }, the input-dependent transmit (i) pre-scalings {bj }, the positive receive scalings {ρ(i) }, and the expert parameters {θj } so as to minimize the expected task loss L(·) subject to the communication budget, physical-layer alignment, and power constraints, i.e., h i (i) min Ei L ŷota , Y (i) (6) (i)
(i)
{θj }, R, {βj }, {bj }, {ρ(i) }
s.t.
(i) S (i) = R FLM ; {Mj }N j=1 , |S
(i)
| = K,
(i) Croute ≤ Cmax , (i) (i) 2 E bj yj 2 ≤ P0 , j ∈ S (i) , p (i) (i) γj bj = ρ(i) βj , j ∈ S (i) , P (i) (i) j∈S (i) βj = 1, βj ≥ 0.
(6a) (6b) (6c) (6d) (6e)
must operate on highly compressed descriptors; and (iii) the aggregation in the objective is computed through the noisy channel (Eq. (4)), coupling the learning problem with the physical-layer design via the distortion (Eq. (5)). These three difficulties motivate our decoupled design AirMoE in Section IV: a statistic-augmented digital routing mechanism that satisfies (6a)–(6c) by reporting only compact prototype statistics, and a statistic-augmented over-theair aggregating mechanism that realizes the weighted fusion (Eq. (2)) directly through the channel (Eq. (4)), while controlling the distortion (Eq. (5)) under the power constraint (6d). IV. M ETHODOLOGY The proposed AirMoE consists of (i) cloud-side feature extraction and compact query generation, (ii) client-wise FRL, (iii) statistic-augmented routing mechanism (MoE-RM), and (iv) over-the-air aggregating mechanism (MoE-AM). These components are detailed in Section IV-A to Section IV-D. We then present the training objective of AirMoE in Section IV-E. A. Cloud-side Feature Extraction and Query Generation As discussed previously, we consider a cloud server connected to N distributed clients {1, . . . , N } over a wireless network, where the cloud hosts a pretrained LM backbone with parameters ωLM , and client j hosts an expert Ej with parameters θj . For the i-th input D(i) of dataset D, the cloud (i) extracts the latent features FLM in a zero-shot manner via (i) FLM = ωLM D(i) , (7) (i)
where FLM is a tensor. To enable both routing and aggregating in a common, communication-friendly space, the cloud (i) summarizes FLM into a compact prototype-space query by (i) q (i) = ϕ FLM ∈ Rd , Q(i) = Norm q (i) , (8) where ϕ(·) is a lightweight feature-to-vector map and Norm(·) converts a vector into a normalized distribution. (i) When the extracted features FLM and the compact query (i) Q are prepared, the cloud broadcasts them to all clients in a one-to-many manner via the downlink channel. After each (i) client receives them, the extracted features FLM is adopted (i) to compute the instant features for input D at each client, and the compact query Q(i) is used to retrieve each client’s expertise from client-wise FRL.
(6f)
Here, Cmax is the routing uplink budget and P0 is the perclient power budget. Constraint (6e) explicitly couples the fusion weight and transmit pre-scaling so that, after receiver p (i) normalization by 1/ ρ(i) , the channel coefficient of yj (i) (i) equals the desired weight βj . Since ŷota is the argument of the task loss in Eq. (6), the aggregation distortion in Eq. (5) affects the objective implicitly rather than being imposed as a separate constraint. Problem Eq. (6) is challenging for three reasons: (i) the routing constraint (6a) and (6b) are combinatorial due to the discrete TopK selection; (ii) the budget (6c) forbids transporting raw features, therefore, R
B. Client-Wise Feature Retrieval Library (FRL) Recall that wireless MoE’s routing requires the cloudqueried, client-side state information as formalized in Eq. (1). This client-side state can be its intermediate representation. However, transmitting intermediate representation from clients to the cloud leads to channel saturation and inference latency for bandwidth-constrained uplink owing to their costly size. To surmount such uplink challenges, instead of transmitting the raw states, we propose to transmit feature prototype-based expertise from each client to the cloud. To this end, each client j is proposed to equip its hosted expert Ej with a
JOURNAL OF LATEX CLASS FILES, VOL. 18, NO. 9, SEPTEMBER 2020
6
Fig. 3. Detailed illustration of the proposed AirMoE.
FRL. The FRL stores prototypical patterns summarizing the client’s specialization, which are obtained from its experiencing history in a read-then-update fashion. Specifically, FRL termed as Mj contains Kj entries of client-specific expertise, Kj , where pj,k ∈ Rd are prototype i.e., Mj = {(pj,k , wj,k )}k=1 vectors and wj,k ≥ 0 are importance weights. Given the cloud-broadcast query Q(i) , client j reads its FRL through cross attention, i.e., exp sim(Q(i) , pj,k ) (i) (9) αj,k = PKj , (i) ℓ=1 exp sim(Q , pj,ℓ ) XKj (i) (i) p̃j = αj,k pj,k , (10) k=1
where sim(·, ·) is a similarity metric in the shared ddimensional space, thus Q(i) and pj,k are directly comparable (i) and p̃j summarizes the most relevant FRL content for input (i) D . When client j is activated and computes a forward pass (Section IV-D), its FRL is renewed in a read-then-update manner through (i) (i) (i) pj,k ← (1 − ηαj,k ) pj,k + ηαj,k ψj hj , (11) (i)
(i)
wj,k ← (1 − ηαj,k ) wj,k + ηαj,k , (i) where hj (i)
(12)
is the client-side intermediate features for input D (as shown in Eq. (17)), η ∈ (0, 1] is the memory update rate, and ψj : Rdh → Rd is a learnable projection into the prototype space. The read-then-update order ensures that the current sample is routed against the pre-updated memory before its information is absorbed, avoiding trivial self-rewriting. Because the FRL resides and operates locally, this update just involves computation but incurs no communication.
C. Statistic-Augmented Digital Routing (MoE-RM) To select the most relevant clients without uploading raw features through uplink, we route on clients’ compact statistics. (i) Specifically, after each client j retrieves its expertise p̃j based on the cloud query Q(i) , it forms a prototype-induced distribution by normalizing its locally retrieved prototype via (i) (i) Pj = Norm p̃j , (13) (i)
and uploads it digitally to the cloud via uplink. Since Pj is a low-dimensional normalized vector, this uplink report is (i) much smaller than the size of client-side hidden feature hj . (i) At the cloud, we align Pj with the input-induced cloud-side distribution Q(i) using the Jensen–Shannon (JS) divergence, i.e., (i) (i) (i) (i) DJS Q(i) ∥ Pj = 12 DKL (Q(i) ∥Ij )+DKL (Pj ∥Ij ) , (14) (i) (i) where Ij = 21 (Q(i) + Pj ). We prefer JS divergence over raw KL divergence or cosine similarity, because it not only is symmetric and bounded, but also compares two normalized distributions through a shared midpoint, making the routing score less sensitive to scale mismatch across heterogeneous clients. Routing scores and probabilities are computed through (i)
sj =
1 (i)
ϵ + DJS (Q(i) ∥Pj )
,
(15)
(i)
exp(τ sj ) (i) πj = PN , (i) ℓ=1 exp(τ sℓ )
(16)
with stability constant ϵ > 0 and temperature τ > 0. The (i) cloud applies TopK to {πj }N j=1 , producing the activated
JOURNAL OF LATEX CLASS FILES, VOL. 18, NO. 9, SEPTEMBER 2020
7
client set S (i) ⊆ {1, . . . , N }. Only clients in S (i) proceed to the expert forward pass and following over-the-air MoE aggregation, therefore, TopK controls both expert sparsity and the number of simultaneous uplink transmitters (i.e., the involved communication overheads). D. Over-the-Air Statistic-Augmented Aggregating (MoE-AM) (i)
(i)
In general, for input D , only each activated client j ∈ S (i) produces its client-side hidden representation hj and task (i) output yj via (i) (i) hj = Ejenc FLM ; θjenc , (17) (i) (i) yj = Ejdec hj ; θjdec , (18) (i)
where FLM is the cloud LM-extracted latent features that are broadcast by the cloud to all clients, and Ej = Ejdec ◦ Ejenc represents the expert that resides at client j. In contrast, clients that are not in S (i) remain dummy in processing input D(i) . Statistical fusion weights. To weigh experts by how well their local features agree with the LM-extracted features, each client computes a client-side latent feature-induced distribution (i) Rj and its JS distance to Q(i) through (i) (i) Rj = Norm ψj (hj ) , j ∈ S (i) , (19) (i) (i) δj = DJS Q(i) ∥ Rj , j ∈ S (i) . (20) (i)
This process generates a small number of statistics δj that each activated client needs to report to the cloud to calculate its aggregation weight. Specifically, at the cloud side, the aggregation weights are reciprocal-distance, normalized over the activated set S (i) , i.e., (i) w̃j =
1 (i)
ϵ + δj
,
(i) w̃j (i) βj = P , (i) ℓ∈S (i) w̃ℓ
j ∈ S (i) . (21)
Re{r(i) } X (i) (i) (i) ŷota = p = β yj + j∈S (i) j ρ(i) | {z } ŷ (i)
Re{n} p ρ(i) | {z }
j
2
.
(23) This indicates that the channel computes the statistically reweighted fusion in a single transmission, perturbed by zeromean real-equivalent noise e(i) with covariance (σc2 /(2ρ(i) ))I. Moreover, maximizing ρ(i) under (6d) minimizes the resulting noise-induced distortion in Eq. (5). Unlike orthogonal digital transmission, the over-the-air scheme in Eq. (23) uses a single channel regardless of |S (i) |, its latency and bandwidth therefore are invariant to the number of activated experts. To handle clients in deep fade (where small |γj | would force ρ(i) → 0 and amplify noise), we adopt truncated channel inversion. Specifically, clients with |γj |2 < γth are pruned from S (i) before ρ(i) is computed. This truncation couples cleanly with MoE-RM, because weak experts are already unlikely to be routed, the truncation therefore rarely removes high-relevance experts. E. Overall Training Objective of AirMoE The proposed AirMoE is trained in an end-to-end manner (i) with the over-the-air estimated output ŷota (i.e., the channel is considered in the forward graph). Taking semantic segmentation as the example task, the total objective L combines the cross-entropy loss LCE , the MoE load-balancing term LLB , and the FRL regularizer LFRL , i.e., (24)
where λLB , λFRL ≥ 0 control regularization strengths. Among these terms, LLB follows the standard auxiliary load-balancing regularization widely used in MoE training to encourage more uniform expert utilization [26], [27], whereas LFRL is a taskspecific regularizer introduced in this paper to keep the FRL memory compact and stable. In the proposed AirMoE, term LLB is introduced to encourage balanced expert usage during training, and is defined as XB XN (i) πj , (25) LLB = uj log uj −log N, uj = B1 i=1
j=1
PN
where B is the batch size. In this formulation, j=1 uj log uj (minimized when {uj }N j=1 is uniform) pushes inputs to be distributed across experts and thereby reduces the risk of expert collapse. Adding − log N zero-centers the objective at the uniform optimum. The FRL regularizer keeps client memories compact and stable, i.e.,
(i)
where P0 is the per-client power budget and ρ is the largest receive scaling that satisfies the vector power constraint (6d). The first equality is exactly the physical-layer
(i)
aggregation noise e(i)
L = LCE + λLB LLB + λFRL LFRL ,
After these aggregation weights are already calculated, the (i) cloud feeds back the scalar βj to each activated client with a small downlink overhead. Each activated client j uses its (i) individual weight βj to control its transmission power for over-the-air computing the MoE aggregation. Notably, using the same normalized comparison space for (i) (i) Q(i) , Pj , and Rj keeps routing and aggregating decoupled while remaining comparable at the distribution level. Over-the-air computing the MoE aggregation. The target fusion of MoE is the weighted sum ŷ (i) through Eq. (2). Each (i) activated client transmits its analog-modulated output yj pre(i) scaled by bj via channel γj . We realize the MoE summation r(i) physically via waveform superposition as Eq. (4), and (i) combine channel-inversion power control (through bj and γj ) (i) with the fusion weight βj via p (i) ρ(i) βj P0 |γj |2 (i) , ρ(i) = min bj = , (22) γj j∈S (i) (β (i) )2 E∥y (i) ∥2 j
(i)
alignment constraint (6e); hence, it prevents bj and βj from being optimized independently. Substituting Eq. (22) into Eq. (4), coherently projecting the received signal p onto its real component, and applying receiver scaling 1/ ρ(i) yield the over-the-air estimate of the MoE aggregation as
LFRL =
Kj N X X j=1 k=1
XB (i) 2 ∥pj,k ∥22 + wj,k + |αj,k | , {z } | i=1 | {z } norm/weight decay
sparse attention
(26)
JOURNAL OF LATEX CLASS FILES, VOL. 18, NO. 9, SEPTEMBER 2020
8
Algorithm 1: AirMoE
A. Assumptions
Input : Dataset D; cloud LM ωLM ; experts {Ej }N j=1 with {θj }N j=1 ; K, τ, ϵ, η, γth , P0 , λLB , λFRL . Kj Init : FRLs Mj = {(pj,k , wj,k )}k=1 , ∀j. Output: Trained {θj }, updated FRLs.
e Assumption 1 (Smooth surrogate). Let L(Θ, M) = (i) (i) Ei ℓ(ŷ (Θ), Y ) be the noiseless, relaxed objective, where Θ collects all trainable parameters, ℓ(·, Y (i) ) is the persample loss, and M = {Mj }N j=1 is the FRL state. Assume ϕ(·), ψj (·), Ejenc , Ejdec , Norm(·) are continuously differe is lower bounded by Leinf and L-smooth in entiable, and L(·) Θ, uniformly over bounded M.
for each mini-batch {D(i) }B i=1 do 2 foreach input D(i) do // Cloud feature extraction & query generation (i) 3 Extract FLM , form Q(i) by Eqs. (7) and (8); (i) broadcast Q(i) , FLM
1
4 5 6
7 8
9
10 11 12
// Client-wise FRL read foreach client j ∈ {1, . . . , N } do (i) (i) Get αj,k , p̃j by Eqs. (9) and (10) (i) Form Pj by Eq. (13); upload digitally // Statistic-augmented routing (i) (i) Compute sj , πj by Eqs. (14) to (16) (i) Select S (i) ← TopK({πj }N j=1 )
Assumption 2 (Bounded states and channel). There is C > 0 (i) (0) with ∥q (i) ∥2 ≤ C, ∥ψj (hj )∥2 ≤ C, ∥pj,k ∥2 ≤ C, and 0 ≤ (0) wj,k ≤ 1. The estimated channel satisfies that the truncation |γj |2 ≥ γth for all activated clients, so by Eq. (22) the perinput scaling obeys ρ(i) ≥ ρmin > 0. Assumption 3 (Loss and forward-map regularity). The persample loss ℓ(·, Y ) is twice continuously differentiable in the output, with Ly -Lipschitz gradient ( ∇ŷ2 ℓ(ŷ, Y ) 2 ≤ Ly ) and ρH -Lipschitz Hessian ( ∇ŷ2 ℓ(ŷ, Y )−∇ŷ2 ℓ(ŷ ′ , Y ) 2 ≤ ρH ∥ŷ− ŷ ′ ∥2 ). The forward-map Jacobian is bounded, ∥JΘ (ŷ (i) )∥2 ≤ G, uniformly over the region of analysis.
// Expert forward pass, j ∈ S (i) (i) (i) Compute hj , yj by Eqs. (17) and (18); form (i) (i) (i) Rj , δj by Eqs. (19) and (20); upload δj digitally
Assumption 4 (Over-the-air noise). Conditioned p on the data and Θ, the aggregation noise e(i) = Re{n}/ ρ(i) in Eq. (23) is zero-mean, independent of the data, with a symmetric distribution and covariance Σ(i) ⪯ (σc2 /(2ρmin ))I, and its dimension (also output dimension) is dy .
// Over-the-air aggregation (i) Compute βj by Eq. (21); feed back to clients (i) (i) Set bj , ρ by Eq. (22) Simultaneous transmission ⇒ r(i) by Eq. (4); (i) estimate ŷota by Eq. (23)
Assumption 5 (Stochastic, channel-perturbed update). Define the noise-smoothed objective L̄(Θ, M) := Ei Ee ℓ(ŷ (i) (Θ) + e(i) , Y (i) ) . (27)
(i)
13
// FRL renewal (local, j ∈ S ) Update pj,k , wj,k by Eqs. (11) and (12)
14
// End-to-end optimization Form total loss L by Eq. (24); update {θj } via ∇L
where norm/weight decay term prevents unbounded growth of prototype vectors and their importance scalars, controlling scale and avoiding trivial wins by increasing magnitude. Sparse attention term over prototypes encourages retrieving a few relevant prototypes rather than averaging many, which (i) keeps retrieved statistic Pj informative and interpretable. In summary, AirMoE is illustrated in Fig. 3 and outlined in Algorithm 1. V. T HEORETICAL G UARANTEES UNDER C HANNEL N OISE This section provides theoretical guarantees of AirMoE under channel noise. Because TopK routing makes the system piecewise smooth, we analyze a smooth surrogate following classical nonconvex stochastic approximation. The analysis applies either to a soft-routing relaxation of TopK or to any local region where the activated set S (i) is fixed. The key new ingredient relative to a single-machine MoE is the over-the-air aggregation noise e(i) in Eq. (23).
For fixed bounded M, parameters are updated by Θt+1 = Θt − αt Ht gt , where αt is the update step, gt is the overthe-air stochastic gradient, and Ht is a symmetric positivedefinite preconditioning aggregation-scaling matrix applied to the gradient before the parameter step. The preconditioner satisfies mI ⪯ P PHt 2⪯ M I, and the stepsizes satisfy αt > 0, α = ∞, t t t αt < ∞. B. FRL Stability and Gradient Statistics Theorem 1 (FRL stability). Under Assumption 2 and η ∈ (0, 1], each read-then-update step t keeps every prototype and (t) (t) weight bounded: ∥pj,k ∥2 ≤ C and 0 ≤ wj,k ≤ 1 for all t, j, k. Proof. Refer to Appendix A. Theorem 1 guarantees M as bounded. By certifying that the read-then-update dynamics keep prototypes inside a ball of radius C and weights inside [0, 1] for all rounds t, it ensures that the memory cannot diverge or collapse as online updates accumulate. Notably, Theorem 1 holds under the mild and practical update rate η ∈ (0, 1]. Lemma 1 (Unbiasedness for the smoothed objective and inflated variance). Under Assumptions 3–5, the over-the-air gradient is an unbiased estimator of the gradient of the smoothed objective, E[gt | Θt , M] = ∇Θ L̄(Θt , M),
(28)
JOURNAL OF LATEX CLASS FILES, VOL. 18, NO. 9, SEPTEMBER 2020
and its variance is bounded by E ∥gt − ∇Θ L̄(Θt , M)∥22 ≤ σ 2 + κ
σc2 2ρmin
2 ≜ σeff , (29)
where σ 2 is the data-sampling variance and κ := G2 L2y dy . Proof. Refer to Appendix B. Lemma 1 is the bridge between the physical-layer channel model and the optimization analysis. It converts the effect of over-the-air aggregation into the two statistical quantities, 2 i.e., an unbiased gradient and a bounded variance σeff . It conceptualizes that channel noise, once viewed through the smoothed objective L̄, contributes no bias to the descent direction and instead manifests purely as inflated variance. Specifically, the standard data-sampling noise variance σ 2 is additively augmented by the channel term κ σc2 /(2ρmin ), which cleanly exposes how the complex channel-noise power σc2 , the truncation floor ρmin , and the problem geometry κ = G2 L2y dy each scale the effective noise after coherent real projection. This decomposition lets the convergence theorem 2 invoke off-the-shelf nonconvex-SGD with σeff in place of 2 2 the usual variance σ , and drives the O(σeff /ϵ2 ) iteration complexity. Lemma 2 (Bias gap to the noiseless objective). Under Assumptions 3–4, for every Θ, e ∇Θ L̄(Θ, M)−∇Θ L(Θ, M) 2 ≤
σc2 G ρH dy · ≜ Bσ . 2 2ρmin (30)
Proof. Refer to Appendix C. Lemma 2 quantifies the price paid for optimizing the tractable smoothed surrogate L̄ instead of the true objective e It certifies that the two gradients never differ by more L. than the constant Bσ , uniformly over all Θ. This complements Lemma 1. Specifically, Lemma 1 guarantees that gt tracks ∇L̄ exactly, and this lemma bounds how far ∇Θ L̄ itself sits from e Therefore, they together ensure that the SGD converges ∇Θ L. to a stationary point of Le bounded by Bσ , where the bound scales with the constant ρH , G, dy , and σc2 /ρmin . C. Main Result Theorem 2 (Convergence to a stationary point of the smoothed objective). Fix any bounded M satisfying Theorem 1, and suppose Assumptions 1–5 hold with L̄(·) L-smooth and lower bounded by L̄inf (both inherited from Assumptions 1, 3). If αt ≤ m/(LM 2 ), then PT 2 2(L̄1 − L̄inf ) t=1 αt E ∥∇Θ L̄(Θt , M)∥2 ≤ + PT PT m t=1 αt t=1 αt PT 2 2 LM 2 σeff t=1 αt , (31) PT m t=1 αt where L̄1 = L̄(Θ1 , M). Consequently lim inf T →∞ E[∥∇Θ L̄(ΘT , M)∥22 ] = 0. Proof. Refer to Appendix D. Theorem 2 certifies that SGD update based on Assumption 5 actually converges to a stationary point of L̄(·) despite
9
over-the-air noise. The bound cleanly separates twoPerror sources, where an optimization term 2(L̄1 − L̄inf )/(m αt ) that decays as the accumulated and a noise P 2 stepsize grows, P 2 term proportional to σ α that conditions ( αt = ∞, t eff P 2 αt < ∞) force to vanish. Working together with Lemma 2, Theorem 2 upgrades to approximate stationarity of the true objective Le whose radius is governed by the channel bias Bσ . Remark 1 (On the fixed-M assumption). Theorem 2 fixes the FRL state M, whereas in Algorithm 1 the prototypes are updated online via Eq. (11)–Eq. (12). Theorem 1 guarantees M stays in a bounded set for all t. Given a memory rate η small relative to αt , the FRL evolves on a slow timescale, and we can treat M as quasi-static for the fast Θ-updates. Lemma 3 (Rate, SNR dependence, and the noiseless-objective √ floor). With αt = α/ T and α ≤ m/(LM 2 ), Theorem 2 gives κσc2 1 2 2 min E ∥∇Θ L̄(Θt , M)∥2 = O √ σ + 1≤t≤T 2ρmin T 2 σeff =O √ . (32) T e 2 ≤ 2∥∇L̄∥2 +2Bσ2 with the bias Bσ of Moreover, using ∥∇L∥ 2 2 Lemma 2, the iterates approach a neighborhood of a stationary e point of the noiseless objective L: 2 σeff e t , M)∥22 ≤ O √ min E ∥∇Θ L(Θ + 1≤t≤T T G2 ρ2H d2y σc4 . (33) · 2 8 ρ {z min} | 2 channel-induced floor = 2Bσ
Proof. Refer to Appendix E. Two effects of the channel are now explicit and qualitatively distinct. (i) The transient term in Eq. (33) is inflated by the variance κσc2 /(2ρmin ). A better channel (larger ρmin , i.e. higher receive SNR) reduces the noise that the optimizer must average out, but this term still decays as T −1/2 . (ii) The steady-state floor 2Bσ2 = O(σc4 /ρ2min ) does not vanish with T . It is the reason why AirMoE converges to a neighborhood of a noiseless stationary point rather than to an exact one. Both terms shrink monotonically with the receive SNR. D. ϵ-Stationarity and Iteration Complexity The following theorem reports the number of iterations needed to reach an ϵ-stationary point. Theorem 3 (ϵ-stationarity of the smoothed objective). Fix any bounded M satisfying Theorem 1 and suppose Assumptions 1– 5 hold. Run the update with the constant stepsize s ( ) m 2∆ , , (34) αt ≡ α = min 2 T LM 2 LM 2 σeff where ∆ := L̄1 − L̄inf denotes the initial-optimality gap of the smoothed objective. Then r 2 2LM 2 ∆ 2M 2L∆ σeff 2 min E ∥∇Θ L̄(Θt , M)∥2 ≤ + . 2 1≤t≤T m T m T (35)
JOURNAL OF LATEX CLASS FILES, VOL. 18, NO. 9, SEPTEMBER 2020
10
Consequently, the iterate with the smallest expected gradient norm is an ϵ-stationary point of L̄ after at most 2 LM 2 ∆ σeff 1 κ σc2 2 T (ϵ) = O =O 2 σ + (36) m2 ϵ2 ϵ 2ρmin
between LM-extracted features and expert-wise FRL prototypes in Section VI-D. Finally, we conduct ablation studies to investigate various hyperparameters’ effect on AirMoE’s performance in Section VI-E.
iterations, where the second equality isolates the dependence 2 on the channel through σeff = σ 2 + κσc2 /(2ρmin ).
A. Datasets, Evaluation Metrics and Implementation
Proof. Refer to Appendix F. Theorem 3 converts the asymptotic guarantee of Theorem 2 into a complexity result. With the stepsize in Eq. (34), it certifies the best iteration that reaches an ϵ-stationary point of L̄(·) within a finite number of rounds. Through the two terms of Eq. (35), it also recovers the canonical nonconvexSGD behavior √with an O(1/T ) transient followed by the dominant O(1/ T ) term. In addition, as Eq. (36), by isolating the channel’s contribution as an additive κσc2 /(2ρmin ) term 2 inside σeff , it makes the communication–computation tradeoff quantitative and actionable. This indicates that the training cost degrades linearly with channel power σc2 and improves with the receive-SNR floor ρmin , cleanly separating the datasampling cost σ 2 from the communication-induced overhead. Lemma 4 (ϵ-stationarity of the noiseless objective, above the floor). Let Bσ = 41 GρH dy σc2 /ρmin be the bias of Lemma 2. e 2 ≤ Under the setting of Lemma 3, the relation ∥∇L∥ 2 2∥∇L̄∥22 + 2Bσ2 implies that an ϵ-stationary point of the noiseless objective Le is reachable iff ϵ > ϵfloor := 2Bσ2 =
G2 ρ2H d2y σc4 · 2 , 8 ρmin
in which case it is attained after at most 2 LM 2 ∆ σeff T (ϵ) = O m2 (ϵ − ϵfloor )2
(37)
(38)
iterations. For ϵ ≤ ϵfloor , no finite T suffices because the e channel-induced bias prevents exact stationarity of L. Proof. Refer to Appendix G. Lemma 4 completes the analytical loop by transferring the finite-time guarantee from the smoothed surrogate L̄ back e and further exposes a fundamental to the true objective L, accuracy floor intrinsic to over-the-air training. Its central message is a sharp dichotomy: (i) any target accuracy strictly above the floor ϵfloor = 2Bσ2 is reachable in finite rounds, 2 which shares the same O(σeff /(ϵ − ϵfloor )2 ) complexity as Eq. (38) but is merely re-centered by the floor. (ii) in contrast, no finite budget can drive Le below ϵfloor = 2Bσ2 , because the channel bias Bσ irreducibly separates stationary points of the surrogate from those of the true loss. VI. E XPERIMENTS AND A NALYSES In this section, we carry out extensive simulations to verify the proposed AirMoE, taking semantic segmentation task as an example. Specifically, we first introduce the simulation setup in Section VI-A. We then compare and analyze experimental results in Section VI-B. Subsequently, we carry out evaluations about ϵ-stationarity iteration complexity under over-the-air setting in Section VI-C. We then visualize the relationship
1) Datasets: The Cityscapes dataset [48] includes 2,975 training and 500 validation images annotated with masks for 19 semantic classes. The CamVid dataset [49] contains 701 images across 11 semantic classes, with 600 used for training and the other 101 for testing. A subset of the Apolloscapes dataset [50], featuring 854 training and 400 test images, provides pixel-level labels for 23 classes. The CARLA ADV dataset, generated via the CARLA simulator (version 0.9.13) [51], focuses on adverse weather conditions (e.g., fog and rain) and includes 2,764 training and 1,921 test images, annotated for 23 pixel-level classes. These benchmarks include both real-world and simulation scenarios, covering various domain biases. 2) Evaluation Metrics: We evaluate the proposed AirMoE on semantic segmentation task using four metrics: mIoU, which quantifies the overlap between prediction and ground truth; mPre, which measures the accuracy of positive prediction; mRec, which evaluates the model’s ability to identify relevant instances; and mF1, which balances mPre and mRec. 3) Implementation Details: We take ViT [52] as the cloudside LM backbone to extract shared features and adopt ASPP architecture [53] for all experts within AirMoE, where the ViT backbone keeps frozen and client-side experts are trained from scratch. The Adam optimizer is chosen for expert optimization with Betas values of (0.9, 0.999), a weight decay of 1e-4, and a learning rate of 3e-4. The default hyperparameters of the proposed AirMoE include: 10 experts, a TopK value of 5, 16 FRL prototypes per expert, a load-balancing weight (λLB ) of 0.01, and a FRL regularization weight (λFRL ) of 1e-4. For evaluation, on the one hand, we compare the proposed AirMoE with other MoE routing strategies, such as LinearMoE [25], NonlinearMoE [54], and SoftMoE [55]. For this MoEto-MoE comparison, all methods share the same frozen ViT backbone, the same expert architecture, and the same training protocol, yet only the routing/aggregation rule is changed. This isolates the effect of the proposed MoE-RM and MoEAM. On the other hand, we also compare AirMoE with other single-model methods, such as BiSecNetV2 [56], SegNet [57], SegFormer [58], AttaNet [59], HRDA [60], TopFormer [61], and SeaFormer [62]. This comparison is intended to position AirMoE against representative semantic segmentation models. Notably, the proposed AirMoE, all MoE baselines, and all single-model competitors are implemented using the PyTorch framework and trained using NVIDIA GeForce 4090 GPUs. For MoE-based models, including AirMoE, LinearMoE, NonlinearMoE, and SoftMoE, a pretrained ViT backbone is used and frozen, with only the experts (i.e., ASPPs) within the MoE being trained. In contrast, all single-model competitors are trained from scratch using the aforementioned datasets. Relative to prior MoE baselines, AirMoE introduces extra FRL retrieval and JS-divergence computations, but these operations
JOURNAL OF LATEX CLASS FILES, VOL. 18, NO. 9, SEPTEMBER 2020
11
TABLE II P ERFORMANCE COMPARISON OF THE PROPOSED A IR M O E AGAINST OTHER M O E BASELINES AND SINGLE - MODEL APPROACHES FOR ALL ADOPTED METRICS ACROSS MULTIPLE SEMANTIC SEGMENTATION DATASETS
mIoU
Apolloscapes mF1 mPre
mRec
mIoU
CamVid mF1 mPre
mRec
mIoU
CARLA ADV mF1 mPre
mRec
mIoU
Cityscapes mF1 mPre
mRec
BiSecNetV2 SegNet SegFormer AttaNet HRDA TopFormer SeaFormer
22.92 21.01 20.89 22.19 20.84 20.58
27.12 24.60 24.85 27.27 24.71 24.53
26.64 34.39 26.37 25.54
25.67 26.94 25.70 25.28
47.89 46.60 39.37 51.12 64.42 48.50 47.85
53.33 50.18 46.23 58.89 75.65 57.02 56.62
58.83 83.66 59.34 56.22
60.96 71.80 57.52 58.65
28.80 31.67 28.97 34.98 32.32 28.68
33.59 37.15 34.46 43.10 38.22 34.24
35.63 52.13 41.61 37.76
34.54 40.74 37.20 33.20
33.63 43.13 39.37 22.96 38.89 22.15 20.51
43.32 53.87 46.23 27.28 49.38 26.70 24.02
26.11 64.30 25.48 23.36
30.87 45.80 30.24 26.85
ViT+ASPP
17.11
21.18
-
-
68.12
77.01
-
-
31.64
37.58
-
-
26.31
30.37
-
-
LinearMoE NonLinearMoE SoftMoE
21.92 22.32 22.22
26.76 27.30 27.11
37.23 37.18 39.02
26.79 27.03 26.95
71.07 72.09 71.54
81.36 82.28 81.71
84.96 85.66 85.20
79.74 80.40 79.54
34.97 35.41 36.21
42.95 43.50 44.37
52.48 54.08 53.39
40.57 41.38 42.14
39.80 41.74 42.08
51.53 54.07 54.45
66.07 65.90 66.93
47.62 50.15 49.79
AirMoE (Ours)
24.75
28.47
38.35
29.48
72.43
82.89
85.31
80.92
37.97
45.56
55.19
43.39
44.74
55.99
68.30
53.47
Method
(a) mIoU
(b) mF1
Fig. 4. Convergence comparison of the proposed AirMoE against all baselines on Cityscapes dataset.
are performed on compact feature summaries instead of rerunning the backbone. B. Main Results and Empirical Analyses 1) Performance Evaluation: Table II compares the main results of the proposed AirMoE against the aforementioned baselines. To interpret the results in Table II, instead of reporting the numbers cell-by-cell, we organize the analyses around four research questions (RQs) that progressively isolate the source and nature of AirMoE’s advantage across all adopted datasets. RQ1: Does AirMoE deliver consistent gains, or are its improvements confined to particular datasets or metrics? AirMoE attains the best mIoU and mF1 on every dataset without exception, improving mIoU over the strongest prior baseline by +2.43 on Apolloscapes (24.75 vs. NonLinearMoE’s 22.32), +0.34 on CamVid (72.43 vs. NonLinearMoE’s 72.09), +1.76 on CARLA ADV (37.97 vs. SoftMoE’s 36.21), and +2.66 on Cityscapes (44.74 vs. SoftMoE’s 42.08). Because mIoU and mF1 are the two threshold-independent, classbalanced measures of segmentation quality, their simultaneous first-place ranking across all four benchmarks answers RQ1 affirmatively, i.e., the gains are systematic rather than a datasetspecific or metric-specific artifact.
RQ2: Do the improvements stem from the MoE paradigm in general, or to AirMoE’s specific design? A clear three-tier hierarchy is visible across the method families. The MoE-based models (LinearMoE, NonLinearMoE, SoftMoE, AirMoE) uniformly dominate both the single-model segmentors and the ViT+ASPP baseline, confirming that MoEbased models are inherently better suited to the heterogeneous, multi-domain nature of semantic segmentation scenes. Crucially, within this MoE family, AirMoE still advances the frontier on the primary metrics, which isolates the contribution of its over-the-air aggregation and prototype-guided FRL design from the generic MoE mechanism. RQ2 is thus answered on two levels: MoE certainly helps, and AirMoE’s particular design helps beyond that. RQ3: What is the preference of AirMoE’s advantage in the precision–recall space? The advantage is concentrated in recall. Specifically, AirMoE achieves the highest mRec on all four datasets (e.g. +3.32 on Cityscapes over SoftMoE and +2.45 on Apolloscapes over NonLinearMoE), while its precision is marginally edged out by two cases (i.e., SoftMoE’s 39.02 vs. 38.35 on Apolloscapes, and NonLinearMoE’s 85.66 vs. 85.31 on CamVid). This profile is not a weakness but a favorable trade. The two competing methods raise small precision at the cost of lower recall, whereas AirMoE’s bal-
JOURNAL OF LATEX CLASS FILES, VOL. 18, NO. 9, SEPTEMBER 2020
109
12
1012
109
σc2 = 0.05, ǫfloor = 0.08
10
σc2 = 0.05
ρmin = 2
σc2 = 0.06, ǫfloor = 0.115
8
1011
σc2 = 0.07, ǫfloor = 0.157 σc2 = 0.08, ǫfloor = 0.205
10
ρmin = 8
1010
107
σc2 = 0.06
ρmin = 4
4
10
107
Iterations T
Iterations T
Iterations T
105
108 107 106
103
σc2 = 0.08
ρmin = 16
109 106
σc2 = 0.07
8
106
105
105 104
102
104
101
0.1
0.15
0.2
0.25
0.3
103
0.35 0.4 0.45 0.50.550.6
0
0.05
0.1
0.15
0.2
0.25
0.3
103
0.35
108 σc2 = 0.06
dy = 16 dy = 32 dy = 64 target ǫ = 0.05
102
σc2 = 0.07 σc2 = 0.08
10
4
5
6
7
Best reachable ǫ of Le
ǫfloor = 2Bσ2
105
9
10
σc2 = 0.05, floor= 0.08
0.45
σc2 = 0.07, floor= 0.157
σc2 = 0.06, floor= 0.115 σc2 = 0.08, floor= 0.205
0.35
100
106
8
0.55 0.5 0.4
101
Iterations T
3
Truncation / receive-SNR floor ρmin (c) T vs ρmin
103 σc2 = 0.05
7
2
Channel noise power σc2 (b) T vs σc2
Target stationarity ǫ (a) T vs ǫ
10−1
10−2
0.3 0.25 0.2
0.15
10−3 104 10−4
3
10
0
5
10
15
20
25
Output dimension dy (κ ∝ dy , ǫfloor ∝ d2y ) (d) T vs dy
30
10−5
0.1
0
5
10
15
20
25
30
102
Receive SNR 10 log10 (ρmin /σc2 ) (dB) (e) Irreducible floor ǫfloor ∝ (σc2 /ρmin )2
103
104
105
106
Iteration budget T (f) ǫ(T ) → ǫfloor
Fig. 5. Illustration of the relationship of ϵ-stationarity complexity iteration with respect to various involved factors.
anced expert utilization broadens semantic coverage so that its overall mIoU and mF1 remain superior despite the tiny precision gap. RQ3 therefore reveals that AirMoE’s mechanism improves coverage of under-represented regions rather than merely shifting a decision threshold. RQ4: Where does AirMoE help most, and does this align with practical difficulty? Ordering datasets’ difficulty by the best baseline mIoU yields Apolloscapes (22.32) > CARLA ADV (36.21) > Cityscapes (42.08) > CamVid (72.09). AirMoE’s relative mIoU improvement tracks this ordering closely: +10.9% on Apolloscapes (22.32 → 24.75), +6.3% on Cityscapes (42.08 → 44.74), +4.9% on CARLA ADV (36.21 → 37.97), and near-saturated +0.5% on the easiest benchmark CamVid (72.09 → 72.43). mRec gains exhibit the same monotone pattern, rising from +0.52 on CamVid to +2.45 on Apolloscapes and +3.32 on Cityscapes over the respective strongest baseline. These cues are the key evidence for RQ4: AirMoE contributes the most precisely on the low-score, high-confusion datasets where competing models saturate or collapse on under-represented classes, and the most on the recall axis that governs coverage of those rare classes. Consequently, AirMoE’s benefit scales with task difficulty rather than generating gains on already-easy scenes, which underscores the practical robustness of the proposed AirMoE for real-world semantic segmentation. 2) Convergence Comparison: We also compare the convergence of the proposed AirMoE against other MoE routing baselines and single-model competitors, and the results can be viewed in Fig. 4. From Fig. 4, we can derive the following insights: (i) ViT backbone-based methods (including ViT+ASPP, LinearMoE, NonlinearMoE, SoftMoE, and our AirMoE) generally converge faster than single-model approaches. This can be attributed to the pretraining of ViT on a vast number of datasets. (ii) The proposed AirMoE converges faster than
other MoE routing baselines, thanks to the statistic-augmented, decoupled routing and aggregating strategies. C. AirMoE’s ϵ-Stationarity Iteration Complexity This section explores how involved factors impact the ϵstationarity iteration complexity of the noiseless objective e The specific feature of ϵ-stationarity is that over-the-air L. aggregation injects a systematic bias whose squared magnitude is ϵfloor , and the SGD can certify progress only until the expected squared gradient norm reaches this residual level. Fig. 5 dissects how each physical and algorithmic factor acts. 1) Effect of the target accuracy ϵ: Panel (a) of Fig. 5 reports the iteration complexity T as a function of the target stationarity accuracy ϵ, with one curve per channel-noise level σc2 . The following regimes are visible. When ϵ ≫ ϵfloor , the term (ϵ − ϵfloor )2 in the denominator of Eq. (38) is dominated by ϵ2 , thus the bound reduces to the classical non-convex rate T = O(ϵ−2 ) and the channel merely rescales the constant. As ϵ decreases toward the floor, the correction ϵfloor becomes non-negligible, the denominator shrinks faster than ϵ2 , and T grows super-quadratically. Finally, as ϵ approaches downward to ϵfloor , the denominator vanishes and T → ∞, producing the vertical asymptote at ϵfloor . Furthermore, when ϵ ≤ ϵfloor , no finite T satisfies the bound and the target is infeasible. This behavior is the algorithmic counterpart of the dichotomy in Lemma 4. Over-the-air aggregation contributes a systematic bias to the gradient of the true objective whose squared norm is exactly ϵfloor = 2Bσ2 . The SGD can only certify progress while the expected squared gradient norm exceeds this bias. Once it falls to the level of the bias, the true gradient is no longer distinguishable from the channel-induced perturbation, and further decrease cannot be guaranteed. The floor is therefore an irreducible accuracy barrier rather than a
JOURNAL OF LATEX CLASS FILES, VOL. 18, NO. 9, SEPTEMBER 2020
slow-convergence artifact, and it cannot be overcome by any additional iterations, only by reducing the bias itself. 2) Effect of the channel-noise power σc2 : Panel (b) of Fig. 5 fixes the target accuracy ϵ and sweeps the channel-noise power σc2 . This is the decisive factor in Eq. (38) because it is the only 2 quantity that enters both the numerator σeff and the floor ϵfloor , and it enters them with different orders with compounding effects. Specifically, for the numerator, σc2 appears linearly 2 through the effective variance σeff = σ 2 + κσc2 /(2ρmin ), reflecting the real-component receiver noise that inflates the variance of the aggregated stochastic gradient. Acting alone, this term would raise T only linearly, exactly as the intrinsic sampling variance σ 2 does. On the other hand, for the floor, σc2 appears with a strictly higher order ϵfloor = 2Bσ2 ∝ (σc2 )2 , i.e. quadratic in the noise power and quartic in the noise amplitude. The interaction of these two orders explains the shape of each curve. For small σc2 , the floor is negligible relative to ϵ, the denominator (ϵ−ϵfloor )2 ≈ ϵ2 is essentially constant, and T 2 rises only linearly through σeff . As σc2 increases, the quadratic floor grows far faster than the linear numerator; the denominator begins to collapse, and T turns sharply super-linear. Once ϵfloor (σc2 ) reaches the fixed target ϵ, the denominator vanishes, T → ∞ and the curve terminates. This monotonethen-divergent profile quantifies the compounding penalty of noisy analog aggregation, i.e., channel noise degrades not only the convergence rate (via the numerator) but also the best attainable accuracy (via the floor). 3) Effect of the SNR floor ρmin : Panel (c) of Fig. 5 sweeps the receive-SNR floor ρmin that lower-bounds the effective SNR ratio of the aggregated update. Whereas the channelnoise power σc2 is an environmental quantity that we cannot control, ρmin is a design parameter, and it enters Eq. (38) only through the ratio σc2 /ρmin . Raising ρmin is thus equivalent to improving the effective SNR and acts favorably in both terms of the bound but with different orders. In the numerator, ρmin suppresses the effective variance linearly, since 2 σeff = σ 2 + κ σc2 /(2ρmin ). It accelerates the T −1/2 approach to stationarity but saturates once the channel contribution κσc2 /(2ρmin ) falls below the intrinsic sampling variance σ 2 . In the floor, as ϵfloor ∝ ρ−2 min , ρmin acts more strongly, because the aggregation bias Bσ scales with σc2 /ρmin and the floor is its square. Increasing ρmin therefore lowers the irreducible accuracy barrier quadratically. The two effects combine monotonically, i.e., raising ρmin both shortens the trajectory (via the numerator) and lowers the floor it approaches (via the denominator), thus T decreases throughout. The benefit is most pronounced in the high-noise regime, where the floor ϵfloor ∝ (σc2 /ρmin )2 dominates the denominator. Conversely, in the low-noise regime, the floor is already negligible against ϵ, only the linear numerator term remains active, and the marginal value of increasing ρmin diminishes. This asymmetry identifies ρmin as the primary design lever under a hostile channel. 4) Effect of the output dimension dy : Panel (d) of Fig. 5 varies the output dimension dy , which enters Eq. (38) at two distinct orders. In the numerator, it appears linearly through 2 κ = G2 L2y dy in the effective variance σeff . In the floor, it 2 appears quadratically (ϵfloor ∝ dy ), because the aggregation
13
bias Bσ itself grows linearly with dy and the floor is the squared bias 2Bσ2 . These two orders produce competition. For small dy , the floor is negligible relative to the target ϵ and the denominator (ϵ − ϵfloor )2 ≈ ϵ2 is nearly constant, thus T rises 2 only linearly through σeff . As dy grows, the quadratic floor overtakes the linear numerator and the denominator collapses, thus, T diverges once ϵfloor (dy ) reaches the fixed target ϵ. This behavior exposes a curse of dimensionality intrinsic to over-the-air aggregation. Specifically, transmitting higherdimensional updates over the analog channel narrows the feasible-accuracy region quadratically in dy , which motivates dimensionality-reducing measures (e.g. sparsification or subspace projection) as means of lowering the floor rather than merely the convergence rate. 5) The irreducible floor ϵfloor vs. receive SNR: Panel (e) of Fig. 5 isolates the irreducible floor in Eq. (37) and plots it against the receive SNR. Because the floor depends on the channel solely through σc2 /ρmin , with ϵfloor ∝ (σc2 /ρmin )2 = 10−SNR/5 , its logarithm is affine in the SNR, i.e., 10 log10 ϵfloor = const − 2 SNR.
(39)
Therefore, on the log–log axes, each curve is a straight line of slope −20 dB per decade of SNR. Crucially, this slope is universal, and it is set by the squaring of the bias and is 2 independent of the numerator constants L, ∆, M/m, and σeff . The dashed horizontal line marks a representative target accuracy ϵ. Its intersection with a given floor curve determines the minimum receive SNR at which that target is attainable: to the right of the intersection ϵfloor < ϵ and the target is feasible (finite T ), whereas to the left ϵfloor ≥ ϵ and it is unattainable for any iteration budget. This intersection is precisely the SNR at which the vertical asymptote of panel (a) coincides with the chosen ϵ, therefore, panels (a) and (e) describe the same feasibility boundary from complementary viewpoints, i.e., T diverging as ϵ approaches downward to ϵfloor in (a), and ϵfloor crossing ϵ as the SNR falls in (e). In addition, increasing the output dimension dy shifts every curve upward by 20 log10 dy dB, since ϵfloor ∝ d2y without changing the −20 dB/decade slope. Higher-dimensional updates therefore demand a correspondingly higher receive SNR to reach the same target accuracy, which restates the curse of dimensionality of panel (d) directly in terms of the SNR budget required for feasibility. 6) Best reachable accuracy vs. iteration budget T : Panel (f) of Fig. 5 takes the complementary view to the preceding panels, i.e., instead of asking how many iterations a target ϵ requires, it asks for the tightest accuracy attainable given a fixed budget T . Inverting Eq. (38) for ϵ gives r 2 C L M 2 ∆ σeff . (40) ϵ(T ) = ϵfloor + 2 | {z } m | {z T } irreducible
transient, O(T −1/2 )
The decomposition separates the two roles. The transient term decays at the classical stochastic rate T −1/2 and carries all of the numerator dependence, including the smoothness L, 2 the initial gap ∆, the effective variance σeff , and the squared 2 2 preconditioner conditioning M /m . Each of these enlarges
JOURNAL OF LATEX CLASS FILES, VOL. 18, NO. 9, SEPTEMBER 2020
14
TABLE III C OMPARISON AMONG CASES WITH DIFFERENT EXPERT NUMBERS Expert Number
Other Fixed Settings
mIoU
Cityscapes mF1 mPre
mRec
7 8 9 10
TopK = 5, FRL prototypes = 16, λLB = 0.01, λFRL = 10−4
45.45 45.02 44.99 44.74
56.76 56.08 55.82 55.99
53.22 53.68 53.87 53.47
71.65 70.19 69.07 68.30
TABLE IV C OMPARISON AMONG CASES WITH DIFFERENT T OP K VALUES
Fig. 6. Relationship visualization between ViT-extracted features and expertwise FRL prototypes.
the constant and hence slows the descent, but none changes the value the curve approaches. The irreducible term ϵfloor is independent of T and sets that limiting value. The decisive feature is therefore the asymptote (dotted lines). As T → ∞, the transient term vanishes and ϵ(T ) approaches downward to ϵfloor rather than 0. No iteration budget, no matter how large, can drive the noiseless objective below the channel-induced floor. The curve is eventually flat and further iterations are wasted. This panel thus consolidates the central message of Fig. 5. Two families of parameters play qualitatively distinct and non-interchangeable roles, where the optimization constants 2 (L, ∆, σeff , M/m) govern only the speed at which the floor is approached, whereas the channel parameters (σc2 , ρmin , dy ) set the floor itself and hence the fundamental limit of achievable stationarity. Improving the former accelerates convergence toward a fixed barrier, and improving the latter is the only way to lower the barrier. Panels (a) and (e) locate this barrier from the accuracy and SNR axes, respectively, and panel (f) shows its consequence along the iteration axis, giving a single coherent account of feasibility for over-the-air aggregation. D. Visualization of the Relationship between ViT-extracted Features and Expert-wise FRL Prototypes This section visualizes the relationship between cloud-side ViT-extracted features and the most relevant expert’s FRL prototypes by t-SNE. The result is presented in Fig. 6. For clarity, we merely visualize sampled cloud-side shared tokens and the most aligned client’s prototypes. Specifically, each plot overlays randomly sampled ViT token features, represented by smaller points, together with the prototypes stored in the single top-routed expert, depicted as larger “X” markers. Because each expert FRL contains specific expertise, these “X” markers appear quite close in the plot. As indicated by Fig. 6, the learned prototypes of the most relevant expert align closely with the majority of the ViT-extracted token features, demonstrating the effectiveness of the proposed expert-wise FRLs.
TopK Value
Other Fixed Settings
mIoU
Cityscapes mF1 mPre
mRec
3 4 5 6
Experts = 10, FRL prototypes = 16, λLB = 0.01, λFRL = 10−4
45.85 45.46 44.74 44.33
56.75 55.47 55.99 55.65
53.69 53.41 53.47 52.48
69.91 68.82 68.30 67.44
E. Ablation Study This section investigates how AirMoE-embedded hyperparameters (including the number of experts, the TopK value, the FRL prototype number, λLB , and λFRL ) affect the overall performance of AirMoE. These ablations should be interpreted jointly, where the number of experts and the TopK value together determine routing sparsity, the FRL size controls how sharp or diffuse each expert memory becomes, and λLB along with λFRL controls the regularization strength. Notably, we keep the FRL update rate η fixed throughout this paper. 1) The Influence of the Number of Experts on AirMoE Performance: Table III reports the influence of the number of experts on AirMoE. From this table, we can observe the following patterns: (i) As the expert pool grows from 7 to 10 under a fixed TopK= 5, mIoU declines steadily from 45.45 to 44.74, mF1 from 56.76 to 55.99, and mPre from 71.65 to 68.30. This behavior can be interpreted through the activation ratio K/N . With K = 5, enlarging N lowers the activation ratio K/N from 5/7 to 5/10, therefore, each expert is selected and updated less frequently, resulting in declined performance. (ii) It is also notable that the degradation is not uniform across metrics, where mRec stays essentially flat and even peaks at N = 9 (53.87), whereas mPre drops the most sharply (−3.35). This asymmetry indicates that a larger, more diffuse expert population still covers the relevant semantic regions (preserving recall) but weakens per-region confidence (eroding precision). Because N = 7 sits at the lower boundary of the search range, we refrain from claiming global optimality. Instead, the consistent decline suggests that AirMoE favors a compact expert pool whose activation ratio remains high. 2) The Effect of the TopK Value on AirMoE Performance: Table IV examines the routing sparsity by sweeping TopK from 3 to 6 while fixing the expert pool at N = 10. The sparsest configuration of TopK= 3 dominates every metric (i.e., 45.85 in mIoU, 56.75 in mF1, 69.91 in mPre, 53.69 in mRec), and increasing K degrades performance steadily until TopK= 6. The reason is that a larger K forces the router to admit lower-affinity experts beyond the well-matched top ones. These marginal experts contribute noisy logits into the fused output, which is exactly why mPre falls most sharply (−2.47
JOURNAL OF LATEX CLASS FILES, VOL. 18, NO. 9, SEPTEMBER 2020
15
TABLE V C OMPARISON AMONG CASES WITH VARIOUS FRL PROTOTYPE NUMBERS
TABLE VII C OMPARISON AMONG CASES WITH DIFFERENT λFRL VALUES
FRL Prototype Number
Other Fixed Settings
mIoU
Cityscapes mF1 mPre
mRec
λFRL
4 8 12 16
Experts = 10, TopK = 5, λLB = 0.01, λFRL = 10−4
45.59 45.27 44.47 44.74
57.35 56.92 55.25 55.99
53.84 53.68 54.04 53.47
0.00 1e-4 2e-4 4e-4 8e-4
70.76 69.16 67.99 68.30
TABLE VI C OMPARISON AMONG CASES WITH DIFFERENT λLB VALUES λLB
Other Fixed Settings
mIoU
Cityscapes mF1 mPre
mRec
0.00 0.01 0.02 0.04 0.08
Experts = 10, TopK = 5, FRL prototypes = 16, λFRL = 10−4
44.23 44.74 43.97 45.44 44.04
55.48 55.99 55.19 57.17 55.46
71.13 68.30 68.87 70.86 69.30
50.82 53.47 52.06 53.09 51.78
from K=3 to K=6) while mRec also erodes (−1.21). This confirms that AirMoE benefits from decisive, low-K routing rather than broad ensembling. This is further reinforced from the over-the-air perspective. Each additionally activated expert must transmit its contribution over the shared channel, hence, a larger K raises aggregation noise and power contention under the truncated channel-inversion scaling in Eq. (22), compounding the model-side dilution with a communicationside penalty. In summary, AirMoE prefers a small number of routing experts. 3) The Impact of the Number of FRL Prototypes on AirMoE Performance: Table V probes the capacity of FRL module by varying its prototype count from 4 to 16. The dominant trend here reveals a preference for a small prototype set. The smallest configuration of 4 prototypes achieves the best mIoU (45.59), mF1 (57.35), and mPre (70.76), and enlarging the prototype bank generally erodes performance. This is consistent with a representation-granularity argument. That is, each prototype acts as a learnable anchor to partition the feature space, and when the number of prototypes exceeds the number of genuinely required modes in the data, the surplus anchors fragment the feature assignment and dilute the discriminative signal that the FRL regularizer is meant to sharpen. The precision column again reflects this most clearly (−2.77 from 4 to 12 prototypes), since redundant prototypes blur the boundaries between confidently separated classes. Importantly, the degradation is not strictly monotonic. Concretely, mIoU recovers slightly from 44.47 at 12 prototypes to 44.74 at 16, and mRec is in fact maximized at 12 (54.04). We attribute this mild rebound to a two-regime behavior, wherein an oversized bank first introduces harmful redundancy but a sufficiently large bank (16) begins to re-specialize its anchors into finer sub-clusters, partially restoring assignment stability. The recall peak at 12 likewise indicates that a denser prototype set can marginally broaden region coverage even as it sacrifices precision. In conclusion, a compact prototype set best matches the intrinsic semantic complexity of the task, and we therefore adopt a small prototype number to keep the FRL representation both discriminative and computationally economical.
Other Fixed Settings
mIoU
Cityscapes mF1 mPre
mRec
Experts = 10, TopK = 5, FRL prototypes = 16, λLB = 0.01
44.23 44.74 44.16 45.44 45.03
55.71 55.99 55.47 57.17 56.44
52.18 53.47 52.19 53.60 52.29
68.62 68.30 69.65 71.17 69.93
4) The Effect of the Load-Balancing Weight λLB on AirMoE Performance: Table VI studies how strongly the loadbalancing regularizer λLB pushes the router toward an even utilization of experts. In contrast to the previous ablations, the response here is distinctly non-monotonic and interior-optimal. The performance instead is maximized at an intermediate value λLB = 0.04 (45.44 mIoU, 57.17 mF1), forming an invertedU profile. The two ends of the sweep are diagnostic. At the end of λLB = 0, the router is free to collapse onto a small subset of favored experts, and the metric confirms this expert collapse precisely, where mPre attains its global maximum (71.13) while mRec goes to its global minimum (50.82). This is the classic overconfident-specialist regime, in which a few dominant experts fire decisively on the classes (high precision) but the starved, rarely-routed experts leave large portions of the label space uncovered (low recall). On the other end, an overly strong constraint (λLB = 0.08) begins to force uniform routing even when a token clearly belongs to one expert. This overregularization homogenizes the experts and drags all metrics back down (44.04 mIoU, 51.78 mRec). Overall, the inverted-U confirms that load balancing is beneficial only in moderation. The chosen regularization strength must be strong enough to prevent expert collapse yet weak enough to preserve the input-conditioned specialization on which AirMoE’s accuracy depends. 5) The Impact of the FRL Loss Weight λFRL on AirMoE Performance: Table VII investigates the strength of the FRL regularizer by sweeping λFRL from 0 to 8e-4. The effect is non-monotonic and interior-optimal, peaking cleanly at λFRL = 4e-4, which uniquely attains the best value on all metrics (i.e., 45.44 in mIoU, 57.17 in mF1, 71.17 in mPre, 53.60 in mRec). This indicates that the optimal FRL strength improves precision and recall jointly rather than trading one against the other, i.e., it genuinely sharpens the feature space instead of merely shifting the decision threshold. Specifically, at λFRL = 0, the FRL objective is inactive, the prototypes are unconstrained, and the features remain comparatively diffuse, yielding the lowest mIoU (44.23). As the weight increases toward 4e-4, this regularization progressively pulls sameclass features toward their prototypes and enforces inter-class separation, lifting every metric to its maximum. However, an overly large λFRL = 8e-4 lets the auxiliary regularizer begin to dominate the primary segmentation loss and over-constrain features toward the prototypes, at the expense of the taskrequired pixel-level discrimination. Consequently, mRec falls to its lowest value (52.29) and mIoU retreats to 45.03. Taken together, the inverted-U confirms that the FRL regularizer is most effective as a moderate auxiliary signal. It should be
JOURNAL OF LATEX CLASS FILES, VOL. 18, NO. 9, SEPTEMBER 2020
strong enough to impose prototype-guided feature clustering, yet subordinate to the main segmentation objective so as not to distort task-critical discriminative cues. VII. C ONCLUSION We studied MoE over a cloud–edge wireless network, where a pretrained LM backbone resides at the cloud and specialized experts are distributed across clients. We identified routing and aggregating as two communication bottlenecks with fundamentally different needs. To solve such bottlenecks, we proposed AirMoE, which decouples them both algorithmically and physically. Specifically, MoE-RM performs reliable, lowrate digital routing using compact prototype statistics, while MoE-AM realizes the statistically reweighted expert fusion directly over the multiple-access channel via channel-aware power control, achieving aggregation latency and bandwidth that are invariant to the number of activated experts. Furthermore, theoretical guarantees make AirMoE’s dependence on channel quality explicit. Finally, we conducted extensive experiments on semantic segmentation that show the proposed AirMoE obtains consistent gains over MoE and single-model baselines. Current limitations include the mere focus on the idealized analog front-end, and extending AirMoE to incorporate imperfect channel-state information and synchronization errors is a promising direction for future work. R EFERENCES [1] X. Ye, H. Lin, X. Song, Y. Wu, and L. Fu, “Multi-objective isac for low-altitude economy based on multi-task deep reinforcement learning with mixture of experts,” IEEE Transactions on Mobile Computing, pp. 1–17, 2026. [2] W. Shi, J. Cao, Q. Zhang, Y. Li, and L. Xu, “Edge computing: Vision and challenges,” IEEE Internet of Things Journal, vol. 3, no. 5, pp. 637–646, 2016. [3] S. Mu and S. Lin, “A comprehensive survey of mixture-of-experts: Algorithms, theory, and applications,” arXiv preprint arXiv:2503.07137, 2025. [4] C. Riquelme, J. Puigcerver, B. Mustafa, M. Neumann, R. Jenatton, A. Susano Pinto, D. Keysers, and N. Houlsby, “Scaling vision with sparse mixture of experts,” Advances in Neural Information Processing Systems (NeurIPS), vol. 34, pp. 8583–8595, 2021. [5] N. Xue, Y. Sun, Z. Chen, M. Tao, X. Xu, L. Qian, S. Cui, and P. Zhang, “Wdmoe: Wireless distributed large language models with mixture of experts,” in GLOBECOM 2024-2024 IEEE Global Communications Conference. IEEE, 2024, pp. 2707–2712. [6] Q. Song, S. Jing, S. Zhang, S. Zhang, and C. Huang, “Mixture-of-experts for distributed edge computing with channel-aware gating function,” in 2025 IEEE International Conference on Communications Workshops (ICC Workshops). IEEE, 2025, pp. 1353–1358. [7] J. Lei, Y. Li, Z. Wang, Q. Lin, Y.-F. Liu, and Y.-C. Wu, “A unified distributed algorithm for hybrid near-far field activity detection in cellfree massive mimo,” IEEE Transactions on Wireless Communications, vol. 25, pp. 18 259–18 275, 2026. [8] G. Zhu, Y. Wang, and K. Huang, “Broadband analog aggregation for low-latency federated edge learning,” IEEE Transactions on Wireless Communications, vol. 19, no. 1, pp. 491–506, 2020. [9] M. M. Amiri and D. Gündüz, “Machine learning at the wireless edge: Distributed stochastic gradient descent over-the-air,” IEEE Transactions on Signal Processing, vol. 68, pp. 2155–2169, 2020. [10] G. Zhu, J. Xu, K. Huang, and S. Cui, “Over-the-air computing for wireless data aggregation in massive IoT,” IEEE Wireless Communications, vol. 28, no. 4, pp. 57–65, 2021. [11] Z. Hu, J. Yan, and Y.-J. A. Zhang, “Computation-efficient federated prompt-tuning with vision-language foundation model compression over resource-constrained edge networks,” IEEE Transactions on Mobile Computing, pp. 1–17, 2026.
16
[12] Z. Li, W. Wu, S. Wu, and X. Shen, “Fast ai model partition for split learning over edge networks,” IEEE Transactions on Mobile Computing, pp. 1–15, 2026. [13] Z. Li, Z. Tang, J. Guo, W. Jia, and W. Zhao, “Efficient layer-granularity unloading for llms in edge computing,” IEEE Transactions on Mobile Computing, pp. 1–13, 2026. [14] W.-B. Kou, Q. Lin, M. Tang, S. Xu, R. Ye, Y. Leng, S. Wang, G. Li, Z. Chen, G. Zhu, and Y.-C. Wu, “pfedlvm: A large vision model (lvm)-driven and latent feature-based personalized federated learning framework in autonomous driving,” IEEE Transactions on Intelligent Transportation Systems, vol. 26, no. 10, pp. 15 915–15 931, 2025. [15] A. E. Eshratifar, A. Esmaili, and M. Pedram, “BottleNet: A deep learning architecture for intelligent mobile cloud computing services,” in Proc. IEEE/ACM Int. Symp. Low Power Electronics and Design (ISLPED), 2019, pp. 1–6. [16] Y. Matsubara, M. Levorato, and F. Restuccia, “Split computing and early exiting for deep learning applications: Survey and research challenges,” ACM Computing Surveys, vol. 55, no. 5, pp. 1–30, 2022. [17] S. Teerapittayanon, B. McDanel, and H. T. Kung, “BranchyNet: Fast inference via early exiting from deep neural networks,” in Proc. Int. Conf. Pattern Recognition (ICPR), 2016, pp. 2464–2469. [18] J. Shao and J. Zhang, “Communication-computation trade-off in resource-constrained edge inference,” IEEE Communications Magazine, vol. 58, no. 12, pp. 20–26, 2020. [19] A. E. Eshratifar, M. S. Abrishami, and M. Pedram, “JointDNN: An efficient training and inference engine for intelligent mobile cloud computing services,” in IEEE Transactions on Mobile Computing, vol. 20, no. 2, 2021, pp. 565–576. [20] R. Yi, L. Guo, S. Wei, A. Zhou, S. Wang, and M. Xu, “EdgeMoE: Fast on-device inference of moe-based large language models,” arXiv preprint arXiv:2308.14352, 2023. [21] Z. Liu, B. Oguz, C. Zhao, E. Chang, P. Stock, Y. Mehdad, Y. Shi, R. Krishnamoorthi, and V. Chandra, “LLM-QAT: Data-free quantization aware training for large language models,” in Findings of the Association for Computational Linguistics (ACL), 2024, pp. 467–484. [22] J. Lei, Y. Li, L.-Y. Yung, Y. Leng, Q. Lin, and Y.-C. Wu, “Understanding complex-valued transformer for modulation recognition,” IEEE Wireless Communications Letters, vol. 13, no. 12, pp. 3523–3527, 2024. [23] E. J. Hu, Y. Shen, P. Wallis, Z. Allen-Zhu, Y. Li, S. Wang, L. Wang, and W. Chen, “LoRA: Low-rank adaptation of large language models,” in Proc. Int. Conf. Learning Representations (ICLR), 2022. [24] Y. Xiao, Y. Lu, F. Tang, R. Wang, and G. Wang, “A federated recommendation system framework based on variational autoencoder with mixture of experts,” IEEE Transactions on Mobile Computing, vol. 25, no. 6, pp. 8597–8610, 2026. [25] Z. Fan, R. Sarkar, Z. Jiang, T. Chen, K. Zou, Y. Cheng, C. Hao, Z. Wang et al., “M3 vit: Mixture-of-experts vision transformer for efficient multitask learning with model-accelerator co-design,” Advances in Neural Information Processing Systems, vol. 35, pp. 28 441–28 457, 2022. [26] N. Shazeer, A. Mirhoseini, K. Maziarz, A. Davis, Q. Le, G. Hinton, and J. Dean, “Outrageously large neural networks: The sparsely-gated mixture-of-experts layer,” arXiv preprint arXiv:1701.06538, 2017. [27] W. Fedus, B. Zoph, and N. Shazeer, “Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity,” Journal of Machine Learning Research (JMLR), vol. 23, no. 120, pp. 1–39, 2022. [28] D. Lepikhin, H. Lee, Y. Xu, D. Chen, O. Firat, Y. Huang, M. Krikun, N. Shazeer, and Z. Chen, “GShard: Scaling giant models with conditional computation and automatic sharding,” in Proc. Int. Conf. Learning Representations (ICLR), 2021. [29] N. Du, Y. Huang, A. M. Dai, S. Tong, D. Lepikhin, Y. Xu, M. Krikun, Y. Zhou, A. W. Yu, O. Firat, B. Zoph, L. Fedus, M. Bosma, Z. Zhou, T. Wang, Y. E. Wang, K. Webster, M. Pellat, K. Robinson, K. MeierHellstern, T. Duke, L. Dixon, K. Zhang, Q. V. Le, Y. Wu, Z. Chen, and C. Cui, “GLaM: Efficient scaling of language models with mixtureof-experts,” in Proc. Int. Conf. Machine Learning (ICML), 2022, pp. 5547–5569. [30] A. Q. Jiang, A. Sablayrolles, A. Roux, A. Mensch, B. Savary, C. Bamford, D. S. Chaplot, D. de las Casas, E. B. Hanna, F. Bressand, G. Lengyel, G. Bour, G. Lample, L. R. Lavaud, L. Saulnier, M.-A. Lachaux, P. Stock, S. Subramanian, S. Yang, S. Antoniak, T. Le Scao, T. Gervet, T. Lavril, T. Wang, T. Lacroix, and W. El Sayed, “Mixtral of experts,” arXiv preprint arXiv:2401.04088, 2024. [31] D. Dai, C. Deng, C. Zhao, R. X. Xu, H. Gao, D. Chen, J. Li, W. Zeng, X. Yu, Y. Wu, Z. Xie, Y. K. Li, P. Huang, F. Luo, C. Ruan, Z. Sui, and W. Liang, “DeepSeekMoE: Towards ultimate expert specialization in mixture-of-experts language models,” arXiv preprint arXiv:2401.06066, 2024.
JOURNAL OF LATEX CLASS FILES, VOL. 18, NO. 9, SEPTEMBER 2020
[32] M. Lewis, S. Bhosale, T. Dettmers, N. Goyal, and L. Zettlemoyer, “BASE layers: Simplifying training of large, sparse models,” in Proc. Int. Conf. Machine Learning (ICML), 2021, pp. 6265–6274. [33] Y. Zhou, T. Lei, H. Liu, N. Du, Y. Huang, V. Y. Zhao, A. M. Dai, Z. Chen, Q. V. Le, and J. Laudon, “Mixture-of-experts with expert choice routing,” in Advances in Neural Information Processing Systems (NeurIPS), vol. 35, 2022, pp. 7103–7114. [34] C. Hwang, W. Cui, Y. Xiong, Z. Yang, Z. Liu, H. Hu, Z. Wang, R. Salas, J. Jose, P. Ram, J. Chau, P. Cheng, F. Yang, M. Yang, and Y. Xiong, “Tutel: Adaptive mixture-of-experts at scale,” in Proc. Machine Learning and Systems (MLSys), vol. 5, 2023. [35] H. Li and L. Duan, “Theory of mixture-of-experts for mobile edge computing,” in IEEE INFOCOM 2025-IEEE Conference on Computer Communications. IEEE, 2025, pp. 1–10. [36] J. Li, Y. Jiang, Y. Zhu, C. Wang, and H. Xu, “Accelerating distributed MoE training and inference with Lina,” in Proc. USENIX Annual Technical Conference (ATC), 2023, pp. 945–959. [37] B. Nazer and M. Gastpar, “Computation over multiple-access channels,” IEEE Transactions on Information Theory, vol. 53, no. 10, pp. 3498– 3516, 2007. [38] W.-B. Kou, Q. Lin, M. Tang, R. Ye, S. Wang, G. Zhu, and Y.-C. Wu, “Fast-convergent and communication-alleviated heterogeneous hierarchical federated learning in autonomous driving,” IEEE Transactions on Intelligent Transportation Systems, vol. 26, no. 7, pp. 10 496–10 511, 2025. [39] W.-B. Kou, S. Wang, G. Zhu, B. Luo, Y. Chen, D. W. Kwan Ng, and Y.-C. Wu, “Communication resources constrained hierarchical federated learning for end-to-end autonomous driving,” in 2023 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS), 2023, pp. 9383–9390. [40] W.-B. Kou, Q. Lin, M. Tang, S. Wang, G. Zhu, and Y.-C. Wu, “Fedrc: A rapid-converged hierarchical federated learning framework in street scene semantic understanding,” in 2024 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS), 2024, pp. 2578– 2585. [41] W.-B. Kou, G. Zhu, B. Cheng, S. Wang, M. Tang, and Y.-C. Wu, “Fedema: Federated exponential moving averaging with negative entropy regularizer in autonomous driving,” in 2025 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS), 2025, pp. 10 190– 10 197. [42] M. M. Amiri and D. Gündüz, “Machine learning at the wireless edge: Distributed stochastic gradient descent over-the-air,” IEEE Trans. Signal Process., vol. 68, pp. 2155–2169, 2020. [43] K. Yang, T. Jiang, Y. Shi, and Z. Ding, “Federated learning via overthe-air computation,” IEEE Transactions on Wireless Communications, vol. 19, no. 3, pp. 2022–2035, 2020. [44] X. Cao, G. Zhu, J. Xu, and S. Cui, “Optimized power control design for over-the-air federated edge learning,” IEEE J. Sel. Areas Commun., vol. 40, no. 1, pp. 342–358, 2022. [45] H. Nguyen, P. Akbarian, T. Pham, T. Nguyen, S. Zhang, and N. Ho, “Statistical advantages of perturbing cosine router in sparse mixture of experts,” arXiv preprint arXiv:2405.14131, 2024. [46] D. Dai, L. Dong, S. Ma, B. Zheng, Z. Sui, B. Chang, and F. Wei, “StableMoE: Stable routing strategy for mixture of experts,” in Proc. Annu. Meeting Assoc. Comput. Linguistics (ACL), 2022, pp. 7085–7095. [47] S. Roller, S. Sukhbaatar, A. Szlam, and J. Weston, “Hash layers for large sparse models,” in Advances in Neural Information Processing Systems (NeurIPS), vol. 34, 2021, pp. 17 555–17 566. [48] M. Cordts, M. Omran, S. Ramos, T. Rehfeld, M. Enzweiler, R. Benenson, U. Franke, S. Roth, and B. Schiele, “The cityscapes dataset for semantic urban scene understanding,” in Proc. of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2016. [49] G. J. Brostow, J. Shotton, J. Fauqueur, and R. Cipolla, “Segmentation and recognition using structure from motion point clouds,” in Computer Vision–ECCV 2008: 10th European Conference on Computer Vision, Marseille, France, October 12-18, 2008, Proceedings, Part I 10. Springer, 2008, pp. 44–57. [50] P. Wang, X. Huang, X. Cheng, D. Zhou, Q. Geng, and R. Yang, “The apolloscape open dataset for autonomous driving and its application,” IEEE transactions on pattern analysis and machine intelligence, 2019. [51] A. Dosovitskiy, G. Ros, F. Codevilla, A. Lopez, and V. Koltun, “Carla: An open urban driving simulator,” in Proceedings of The 1st Annual Conference on Robot Learning, 2017, pp. 1–16. [52] A. Dosovitskiy, L. Beyer, A. Kolesnikov, D. Weissenborn, X. Zhai, T. Unterthiner, M. Dehghani, M. Minderer, G. Heigold, S. Gelly et al., “An image is worth 16x16 words: Transformers for image recognition at scale,” arXiv preprint arXiv:2010.11929, 2020.
17
[53] L.-C. Chen, G. Papandreou, F. Schroff, and H. Adam, “Rethinking atrous convolution for semantic image segmentation,” arXiv preprint arXiv:1706.05587, 2017. [54] H. Nguyen, P. Akbarian, T. Pham, T. Nguyen, S. Zhang, and N. Ho, “Statistical advantages of perturbing cosine router in sparse mixture of experts,” arXiv preprint arXiv:2405.14131, 2024. [55] J. Puigcerver, C. Riquelme Ruiz, B. Mustafa, and N. Houlsby, “From sparse to soft mixtures of experts,” in International Conference on Learning Representations, vol. 2024, 2024, pp. 28 435–28 445. [56] C. Yu, C. Gao, J. Wang, G. Yu, C. Shen, and N. Sang, “Bisenet v2: Bilateral network with guided aggregation for real-time semantic segmentation,” International Journal of Computer Vision, vol. 129, pp. 3051–3068, 2021. [57] V. Badrinarayanan, A. Kendall, and R. Cipolla, “Segnet: A deep convolutional encoder-decoder architecture for image segmentation,” IEEE transactions on pattern analysis and machine intelligence, vol. 39, no. 12, pp. 2481–2495, 2017. [58] E. Xie, W. Wang, Z. Yu, A. Anandkumar, J. M. Alvarez, and P. Luo, “Segformer: Simple and efficient design for semantic segmentation with transformers,” in Neural Information Processing Systems (NeurIPS), 2021. [59] Q. Song, K. Mei, and R. Huang, “Attanet: Attention-augmented network for fast and accurate scene parsing,” in Proceedings of the AAAI conference on artificial intelligence, vol. 35, no. 3, 2021, pp. 2567– 2575. [60] L. Hoyer, D.-X. Dai, and L. Van Gool, “Domain adaptive and generalizable network architectures and training strategies for semantic image segmentation,” IEEE Transactions on Pattern Analysis and Machine Intelligence, 2023. [61] W. Zhang, Z. Huang, G. Luo, T. Chen, X. Wang, W. Liu, G. Yu, and C. Shen, “Topformer: Token pyramid transformer for mobile semantic segmentation,” in Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2022, pp. 12 083–12 093. [62] Q. Wan, Z. Huang, J. Lu, Y. Gang, and L. Zhang, “Seaformer: Squeezeenhanced axial transformer for mobile semantic segmentation,” in The eleventh international conference on learning representations, 2023.