Skip to content
#edge computing Preprint

AirMoE: Realizing Over-the-Air Distributed Mixture-of-Experts Inference at the Wireless Edge

Aug 2026 · 0 citations · 43 references
Engineering

TL;DR

An inference-aware AirMoE error metric is constructed to quantify aggregation distortion effects on end-to-end (E2E) inference accuracy via perturbation-based layer-sensitivity calibration, and an activation- and channel-aware expert placement strategy is developed that assigns more important experts to devices with lower channel-power cost.

Abstract

Mixture-of-experts (MoE) architectures enable efficient large language model (LLM) inference at the wireless edge through sparse activation. The wireless distributed MoE (WIDE) architecture addresses edge-resource constraints by distributing experts across devices coordinated by an edge server. However, WIDE suffers from repeated uplink transmissions of high-dimensional expert outputs via orthogonal multiple access. To overcome this bottleneck, we propose AirMoE, an over-the-air computing (AirComp)-enabled framework for simultaneous expert-output aggregation via wireless waveform superposition. Integrating AirComp into MoE inference introduces three challenges: fast-varying aggregation weights, layer-dependent error sensitivity, and channel-aware expert placement. To address these challenges, we construct an inference-aware AirMoE error metric to quantify aggregation distortion effects on end-to-end (E2E) inference accuracy via perturbation-based layer-sensitivity calibration. We then formulate a joint optimization problem to minimize this error and decompose it, without loss of optimality, into a two-timescale framework. At the fast timescale, we derive a globally optimal threshold-based power-control policy that separates devices into coefficient-aligned and full-power groups. At the slow timescale, we develop an activation- and channel-aware expert placement strategy that assigns more important experts to devices with lower channel-power cost. Extensive experiments demonstrate that AirMoE outperforms representative baselines in E2E inference accuracy, especially under strong device heterogeneity.

View source

Similar papers

Preprint Jul 2026

AirMoE: Statistic-Augmented Over-the-Air MoE for Collaborative 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 statistic-augmented 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 single-model competitors. Ablations further confirm the effectiveness of each incorporated component.

Wei-Bin Kou, Jingreng Lei, Guangxu Zhu et al. · 0 citations
Preprint Jul 2026

EdgeCoInfer: Hierarchical Collaborative Inference for On-Device Multimodal Large Models

To deliver ubiquitous intelligence, modern mobile applications increasingly execute concurrent Multimodal Large Language Models (MLLMs) on edge devices, presenting severe challenges under multi-task concurrency and tight resource constraints. To address this, we propose EdgeCoInfer, a hierarchical collaborative inference framework enabling efficient on-device MLLM inference through coarse-to-fine orchestration. Coarsely, EdgeCoInfer decomposes MLLMs into functional modules for inter-task sharing, avoiding redundant model loading. Finely, it partitions models at the neural network layer level and distributes segments across devices and servers. We jointly optimize layer partitioning, module sharing, and resource allocation under tight constraints. To tackle the non-differentiable combinatorial explosion, we propose a Hybrid Evolutionary Hierarchical Reinforcement Learning (HE-HRL) framework. HE-HRL synchronizes a gradient-free genetic algorithm for discrete partitioning and sharing decisions with a gradient-based soft actor-critic agent for continuous resource refinement. We further embed a constructive cut-step decoder with pre-act pruning and a two-phase curriculum to improve feasibility and accelerate convergence. Experimental results show that EdgeCoInfer breaks the edge memory wall and prevents catastrophic out-of-memory and task failures under high concurrency, reducing memory demand by 53.53\% and system cost by 59.86\% compared to existing methods.

Lin Tan, Songtao Guo, Mingyan Li et al. · 0 citations
Conference Jul 2026

Whitening-enhanced and participation-aware over-the-air computation for stable vehicular federated learning

Over-the-air computation enables instantaneous analog aggregation to alleviate uplink latency in vehicular federated learning. However, existing methods often assume perfect channel state information or exact error statistics. This assumption is impractical given the rapid channel changes and dynamic conditions typical of vehicular networks. To address this gap, we propose a participation-aware learning and communication joint design under covariance uncertainty bounded by the Frobenius norm. We introduce a specific participation coefficient for each vehicle, which functions equivalently to an adaptive central server step size for standard federated averaging. This mechanism enables stability control without altering the original algorithmic weights. Furthermore, we prove that the worst-case normalized mean-squared error under our uncertainty model admits an exact deterministic reformulation, solvable via an efficient alternating optimization algorithm. Experiments using the MNIST dataset demonstrate improved aggregation fidelity, faster convergence, and significantly reduced training volatility.

Xiaoyu Zhang, Jiasheng Zhu, Lei Wang et al. · 0 citations
Conference Jul 2026

Distribution-Level AirComp for Bayesian FL

Conventional federated learning (FL) methods face critical challenges in realistic wireless edge networks, where training data are often limited and heterogeneous, which can lead to unstable training and poor generalization. To address these challenges in a principled manner, we propose a novel Bayesian wireless FL framework grounded in Bayesian inference. By explicitly modeling uncertainty, the proposed framework mitigates local overfitting and client drift, thereby enabling more reliable inference. Nevertheless, adopting Bayesian FL increases communication overhead due to the need to transmit richer model information and fundamentally changes the aggregation process beyond simple averaging. To overcome this limitation, we design a dedicated over-the-air computation (AirComp) scheme tailored to Bayesian FL, which efficiently aggregates local posterior distributions at the distribution level by exploiting the superposition property of wireless channels. Simulations validate the proposed framework, demonstrating significant gains in test accuracy over conventional FL methods, particularly in data-scarce and heterogeneous environments.

Jun-Pyo Hong, Hyowoon Seo, Kisong Lee · 0 citations
2026

Communication-Efficient Over-the-Air Federated Learning via Lightweight Gradient Compression

Integrating over-the-air computations into the model aggregation process of federated learning (FL) offers a promising solution to mitigate the communication bottleneck in FL model training. In this approach, all the clients modulate their intermediate parameters, such as gradients, onto the same set of orthogonal waveforms and transmit the resulting signals to the edge server simultaneously. Capitalizing on the superposition property of the radio channel, the server can extract an automatically aggregated global gradient from the received radio signal. However, the limited number of orthogonal waveforms imposes a constraint on the dimensionality of transmittable updates, hindering the adoption of more advanced, but high-dimensional models. In light of this challenge, we propose OFLight, a lightweight, yet effective, gradient compression algorithm tailored for OTA-FL systems. Specifically, in each communication round, the edge server constructs a low-rank projection matrix based on the received gradient matrix from the previous round (initialized with an independent and identically distributed standard normal matrix in the first round) and broadcasts it, along with the global model, to all clients in the system. Based on this matrix, every client projects its locally updated gradient matrix into a low-dimensional subspace through a linear operation. The clients upload only their compressed gradients via OTA computations, and the edge server can perform a linear decompression on the received signal, retrieving the original gradient dimension. Moreover, an error feedback mechanism is incorporated to compensate for the approximation error under aggressive compression. We derive analytical expressions for the convergence rate of both convex and non-convex loss functions, quantitatively demonstrating the effect of OFLight on the OTA-FL training efficiency. We also conduct extensive experiments to corroborate the efficacy of the proposed method.

Jiaqi Zhu, H. Yang, Nikolaos Pappas et al. · 0 citations
Open access Aug 2026

Adaptive Unequal Error Protection for Semantic Split Learning Over Wireless Channels

We propose a task-aware semantic split learning (SL) framework for wireless edge–cloud inference, in which the reliability of transmitted latent representations is dynamically adapted to their relevance for the downstream task. An autoencoder (AE)-based physical (PHY) layer enables end-to-end learning of the communication interface, while unequal error protection (UEP) is realized via mutual information (MI)-driven prioritization of latent components during training. The gradient of the estimated MI with respect to each latent component serves as a sensitivity-based proxy for task relevance, providing a fully learning-driven prioritization that adapts to both the data distribution and the downstream task. We further show that this prioritization translates into measurable physical-layer effects: MI-guided UEP assigns significantly higher transmit power to the most task-critical latent components compared to the equal error protection (EEP) baseline. Experiments on real-world IoT sensing data demonstrate consistent gains over equal and fixed-UEP baselines across SNR regimes. Additional analysis confirms ranking stability, estimator robustness and generalization across datasets and task types, indicating broad applicability of the proposed framework.

Vukan Ninkovic, D. Vukobratović, D. Mišković et al. · 0 citations

Related blog posts

Microsoft Research Blog Aug 31, 2026

GigaPath-Flash and GigaTIME-Flash: Toward population-scale discovery with efficient pathology foundation models

What if pathology foundation models could do more with less? GigaPath-Flash and GigaTIME-Flash cut computational demands while maintaining strong performance, opening the door to larger studies and broader exploration. The post GigaPath-Flash and GigaTIME-Flash: Toward population-scale discovery with efficient pathology foundation models appeared first on Microsoft Research.

MIT News · Artificial Intelligence Aug 27, 2026

Looking beyond natural sequences

A new machine-learning framework aims to improve the success rate of computational protein design while moving away from results that reproduce sequences found in nature.