Source-linked AI summary
Logit-Aware MIMO AirComp for Distributed Mixture-of-Experts LLM Inference over Wireless Edge Networks
Lyutianyang Zhang, Yunjian Jia, Liu Cao, Dengke Wang, Jinke Ren, Shuguang Cui
TL;DR
Conventional AirComp minimizes communication distortion, although hidden-state errors can have unequal effects on next-token logits. The paper weights aggregation blocks by local logit sensitivity and optimizes MIMO AirComp; its closed-loop audit and simulator report advantages over sensitivity-agnostic baselines in their evaluated settings.
Problem
Conventional AirComp minimizes communication distortion, but equal-energy hidden-state errors can differ in their effects on next-token logits and token decisions.
Source: §I · pp. 1–2Method
The design uses block-level sensitivity weights to allocate a shared BS power budget across AirComp blocks.
Source: B. Logit-Sensitivity-Weighted MIMO… · p. 6Results
SW-AirComp achieved the highest accuracy among the compared AirComp methods in the 295-example ARC-Challenge audit and reduced sensitivity-weighted distortion in the modeled MIMO channels.
Source: E. Insights and Limitations · p. 11Takeaways & Limitations
Logit-RMSE and logit-gap diagnostics link hidden-state aggregation error to output-logit perturbation and token-decision risk.
Source: E. Insights and Limitations · p. 11Stale or noisy sensitivity metadata can waste power or leave harmful logit perturbations, so the paper recommends clipping or normalizing weights, refreshing them, and retaining unweighted AirComp as a fallback.
Source: E. Insights and Limitations · p. 11
Abstract
from arXiv · showhide
Distributed mixture-of-experts (MoE) inference is a promising architecture for deploying large language models (LLMs) at wireless edge networks because sparse experts can be placed across coordinated base stations (BSs), while the anchor node and user equipment (UE) can offload LLM inference tasks to BSs. The communication bottleneck is the MoE aggregation, where the anchor BS must recover a weighted sum of selected expert outputs before each decoding step. Over-the-air computation (AirComp) is well matched to this operation because the wireless multiple-access channel naturally superposes simultaneous transmissions. However, conventional AirComp minimizes communication distortion, whereas MoE aggregation errors have unequal impact on LLM outputs. We propose a logit-aware MIMO AirComp framework that estimates local logit sensitivity as block-level weights and jointly optimizes receive combiners and BS precoders under per-BS power constraints. We also develop an alternating algorithm that protects decoding decisions from aggregation perturbations. Using OpenCompass, we evaluate Qwen3-30B-A3B-Instruct-2507-FP8 on GSM8K and ARC-Challenge. At 30 dB aggregation SNR, perturbed Qwen3 retains 99.3% of clean GSM8K accuracy and 94.8% of clean ARC-Challenge accuracy. In a direct ARC-Challenge closed-loop audit over 295 examples, SW-AirComp achieves 43.73% accuracy at 30 dB and 26.78% at 20 dB, outperforming unweighted, matched-filter, and zero-forcing AirComp. In the wireless simulator, SW-AirComp reduces decision-relevant aggregation distortion; at 20 dB, its final weighted-sum mean-squared error is 53.5% lower than unweighted AirComp under the same channels, samples, and sensitivity weights. Logit-RMSE and logit-gap diagnostics further show that the gain comes from reducing output-logit perturbation and lowering the risk of top-token changes.
I. INTRODUCTION
The paper frames distributed MoE inference as weighted expert-output aggregation whose communication errors must be judged by their downstream logit impact, then connects that objective to MIMO AirComp over wireless edge networks.
I. INTRODUCTION
- Equal-energy hidden-state errors can affect next-token logits differently, so the design targets downstream logit distortion rather than average reconstruction error.The paper notes that perturbations exceeding the clean top-two-logit gap can change the selected token.Source: §I · p. 2
A. Literature Review
- Distributed LLM inference sends internal activations through an autoregressive loop, unlike wireless learning that aggregates gradients, updates, or compressed features.The review identifies missing models of how aggregation errors propagate through a pretrained LLM to next-token logits.Source: A. Literature Review · p. 2
B. Main Contributions
- The proposed design derives block-level logit sensitivity and embeds it in MIMO AirComp under per-BS power constraints.Its weighted objective targets next-token preservation rather than average hidden-state reconstruction error.Source: B. Main Contributions · pp. 2–3
- A Jacobian-based perturbation analysis explains why aggregation errors with equal hidden-state energy can have different effects on logits and inference decisions.The analysis also supports weighted transceiver formulation, logit-distortion bounds, token-preservation conditions, and fixed-weight inner-loop updates.Source: B. Main Contributions · p. 3
- The evaluation plan combines Qwen3/OpenCompass replay, a closed-loop ARC-Challenge audit, and wireless simulation using accuracy and aggregation/logit diagnostics.The simulator evaluates 16 aggregation blocks and compares against unweighted, matched-filter, and zero-forcing AirComp.Source: B. Main Contributions · p. 3
II. SYSTEM MODEL
- The system model considers a pretrained LLM deployed across an anchor BS and synchronized expert BSs, with the anchor retaining decoding state and non-expert components.Expert BSs store subsets of the pretrained experts; the anchor maintains the router, output head, and KV cache.Source: §II · p. 3, A. Distributed MoE Inference… · p. 3
A. Distributed MoE Inference Architecture
- At each MoE layer, the anchor routes hidden state and weights to selected expert BSs, which compute weighted contributions and transmit them simultaneously for AirComp aggregation.The anchor recovers the weighted sum without separately decoding each BS contribution, then continues model computation and autoregressive decoding.Source: A. Distributed MoE Inference… · p. 3
B. MIMO AirComp Model
The MIMO AirComp model represents selected experts’ weighted outputs as local BS contributions and uses wireless superposition to recover their sum at the anchor.
B. MIMO AirComp Model
- Each selected expert computes its routed, weighted contribution, and the anchor’s target is the sum of those local outputs.BSs hosting no selected expert contribute zero.Source: B. MIMO AirComp Model · p. 3
- The model uses frequency-domain MIMO channels across information blocks, with quasi-static block fading and a linear, invertible packing map into complex-domain transmissions.The passage also assumes uniform linear arrays for the steering vectors.Source: B. MIMO AirComp Model · pp. 3–4
- Expert BSs linearly precoder their local blocks for simultaneous transmission, and the anchor applies a linear combiner to recover the aggregate.The desired aggregate is achieved when the effective channel–precoder–combiner product approximately aligns with the identity for each active BS.Source: B. MIMO AirComp Model · p. 4
- The model defines aggregation mismatch as the difference between the ideal MoE sum and its AirComp reconstruction, then jointly optimizes BS precoders and the receive combiner.The alignment condition is approximate under per-BS power constraints.Source: B. MIMO AirComp Model · p. 4
III. PROBLEM FORMULATION AND ALGORITHM DESIGN
The design connects AirComp aggregation error to next-token logit perturbation, then uses block-level logit sensitivity to guide wireless aggregation. It also relates logit perturbations to the stability of greedy decoding.
III. PROBLEM FORMULATION AND ALGORITHM DESIGN
- The framework derives block-level logit-sensitivity weights and uses them to optimize MIMO AirComp under per-BS power constraints.Its alternating procedure estimates weights by finite differences and updates receive combiners and transmit precoders.Source: §III · p. 5
A. LLM-Sensitive Objective
- The local mapping Φℓ,t takes the MoE aggregation output at layer ℓ and decoding step t to the next-token logits, whose clean greedy token is their argmax.The formulation uses the local logit Jacobian to connect aggregation perturbations to output changes.Source: A. LLM-Sensitive Objective · p. 5
- Equal-energy hidden-state errors can cause different logit perturbations because the local logit sensitivity is generally direction-dependent.The block error is mapped into hidden space before its effect on next-token logits is measured.Source: A. LLM-Sensitive Objective · p. 5
- The paper’s sensitivity-weighted objective uses block weights to prioritize aggregation errors according to their predicted effect on logits.Fig. 2 depicts construction of scalar AirComp weights from block-level PHY error and local next-token logit sensitivity.Source: Figure 2 · p. 5, A. LLM-Sensitive Objective · p. 5
- A block’s first-order squared logit distortion is r_b^T Ω_b r_b, while a full matrix-weighted design preserves within-block directional sensitivity but complicates the transceiver.The scalar block weight summarizes average sensitivity for a simpler AirComp design.Source: A. LLM-Sensitive Objective · p. 5
- Greedy decoding is unchanged when aggregation-induced logit perturbation is small relative to the clean top-token logit gap.The paper uses this stability condition to interpret the logit-aware objective; it does not add a gap-dependent weight to the optimizer.Source: A. LLM-Sensitive Objective · pp. 5–6
B. Logit-Sensitivity-Weighted MIMO AirComp Design
The design minimizes sensitivity-weighted AirComp error across blocks while jointly accounting for the shared power budgets of base stations. Receive and transmit updates are solved in alternating steps.
B. Logit-Sensitivity-Weighted MIMO AirComp Design
- The signal model assumes zero-mean, mutually uncorrelated sources for each block and receiver noise independent of the sources.The per-block AirComp MSE is the quantity optimized by the design.Source: B. Logit-Sensitivity-Weighted MIMO… · p. 6
- The objective minimizes weighted-sum MSE across AirComp blocks, with the total-power constraint coupling blocks served by the same BS.This shared budget lets sensitivity weights affect how transmit power is allocated across blocks.Source: B. Logit-Sensitivity-Weighted MIMO… · p. 6
- For fixed precoders, the receive-combiner update is a linear-MMSE estimate of the aggregate signal, not of individual user symbols.The physical AirComp MSE supplies the quadratic and linear coefficients for this combiner problem.Source: B. Logit-Sensitivity-Weighted MIMO… · p. 6
- For fixed receive combiners and positive block weights, each BS transmit subproblem is a convex QCQP with a globally optimal update.The update uses a scalar dual variable to enforce the BS power constraint.Source: B. Logit-Sensitivity-Weighted MIMO… · p. 6
- Under a shared BS power budget, reducing a block’s MSE lowers the objective in proportion to its sensitivity weight, so higher-weight blocks receive higher priority.If the unconstrained precoder already satisfies the budget, the weight does not change that precoder; with separate fixed block budgets, scaling a single-block objective also does not change its minimizer.Source: B. Logit-Sensitivity-Weighted MIMO… · p. 6
C. Compact Closed-Loop Algorithm
The compact algorithm estimates block sensitivity weights by finite differences, alternates receive and transmit updates, and refreshes weights in an outer loop. Its monotonicity statement applies only while weights are fixed.
C. Compact Closed-Loop Algorithm
- Finite differences estimate the block sensitivity weight without directly evaluating the full logit Jacobian.Calibration requires S perturbed logit evaluations per calibrated block; the paper notes options to reduce that overhead.Source: C. Compact Closed-Loop Algorithm · p. 7
- When sensitivity is recalibrated over rounds, the stored base weight is updated using a stepsize and clipped to specified bounds.For one-time weight computation, the paper sets α = 1 and uses the clipped estimate directly.Source: C. Compact Closed-Loop Algorithm · p. 7
- With fixed weights, the inner loop’s objective values converge, and any accumulation point is block-stationary if it exists.This statement assumes exact receive and transmit updates and feasible finite-power precoders.Source: C. Compact Closed-Loop Algorithm · p. 7
- The outer loop refreshes sensitivity weights and tracks a changing objective, so monotonicity applies only after the weights are fixed.The per-round computational complexity also depends on the receive and transmit linear solves and the scalar power-budget search.Source: C. Compact Closed-Loop Algorithm · p. 7
IV. SIMULATION RESULTS
The evaluation combines a pretrained-model perturbation replay, a direct ARC-Challenge AirComp audit, and a wireless simulator. The reported results indicate improved ARC-Challenge accuracy over tested AirComp baselines and reduced LLM-relevant wireless distortion.
IV. SIMULATION RESULTS
- The evaluation uses three complementary tests: Qwen3/OpenCompass perturbation replay, a closed-loop ARC-Challenge audit, and a MIMO wireless simulator.The tests examine task accuracy and physical-layer aggregation distortion from distinct angles.Source: §IV · p. 7
- In the direct ARC-Challenge closed-loop audit, SW-AirComp improves task accuracy relative to the tested AirComp baselines.The compared baselines are unweighted, matched-filter, and zero-forcing AirComp.Source: §IV · p. 7
- The wireless simulator results indicate that the proposed transceiver reduces aggregation distortion relevant to LLM inference.The section also tests whether aggregation-output perturbations can change task accuracy in a pretrained Qwen3 replay.Source: §IV · p. 7
A. Benchmark Setup
The benchmark setup uses Qwen3-30B-A3B-Instruct-2507-FP8 on GSM8K and ARC-Challenge, comparing clean inference with controlled aggregation perturbations across SNR levels.
A. Benchmark Setup
- The replay uses Qwen3-30B-A3B-Instruct-2507-FP8, with 48 MoE layers and 128 routed experts in the local implementation.The clean reference uses the original pretrained forward pass; the perturbed replay keeps the checkpoint, prompts, tokenizer, and OpenCompass settings fixed.Source: A. Benchmark Setup · p. 7
- Aggregation noise is injected at the hidden-state interface, with perturbation energy normalized to clean MoE aggregate energy for the target SNR.The passage describes independent Gaussian aggregation noise.Source: A. Benchmark Setup · pp. 7–8
- GSM8K uses zero-shot generation on a fixed 500-example subset, while ARC-Challenge is scored by multiple-choice accuracy across perturbation SNRs of 0–30 dB.Table I lists representative examples from the evaluated subsets.Source: A. Benchmark Setup · pp. 7–8, Table I · p. 8
B. Wireless Simulator
The evaluation combines Qwen3 perturbation replay, a direct ARC-Challenge closed-loop audit, and a wireless simulator to assess task accuracy and sensitivity-weighted aggregation. Across these tests, SW-AirComp outperforms the AirComp baselines in the closed-loop audit and reduces logit-relevant distortion in the simulator.
B. Wireless Simulator
- The wireless simulator uses 16 aggregation blocks and compares sensitivity-weighted MMSE AirComp with unweighted MMSE, matched-filter, and zero-forcing baselines.SW-AirComp weights aggregation error by its impact on final inference output; the baselines do not use logit-sensitivity weights.Source: B. Wireless Simulator · p. 8
C. Qwen3 OpenCompass Replay
- At 30 dB, perturbed Qwen3 retains 99.3% of clean GSM8K accuracy and 94.8% of clean ARC-Challenge accuracy.The replay reports approximately 72.4% versus 72.9% clean accuracy on GSM8K and 45.8% versus 48.3% on ARC-Challenge.Source: C. Qwen3 OpenCompass Replay · p. 8
- At 30 dB, SW-AirComp achieves 43.73% ARC-Challenge accuracy versus 40.34% for UW-AirComp, 37.29% for MF-AirComp, and 33.56% for ZF-AirComp; at 20 dB, it achieves 26.78%.The closed-loop audit inserts calibrated perturbations into all 48 Qwen3 MoE layers during autoregressive decoding and scores 295 examples.Source: Figure 4 · p. 9, C. Qwen3 OpenCompass Replay · p. 9
D. Massive-MIMO Scaling Study
- Residual distortion varies across aggregation blocks, so designs with similar average distortion can differ in error placed on high-sensitivity blocks.The 16 heatmap indices are wireless optimization blocks, not Qwen3 MoE layers.Source: D. Massive-MIMO Scaling Study · p. 9, Figure 5 · p. 9
- Increasing the transmit-antenna count improves weighted aggregation MSE over the tested SNR range.The reported scaling is consistent with increased array gain from a larger transmit aperture.Source: D. Massive-MIMO Scaling Study · p. 9, Figure 6 · p. 9
- At the same channel realization and power budget, SW-AirComp's lower WMSE reflects both improved channel fit and shifting residual error away from blocks that matter more to next-token logits.UW-AirComp isolates the value of sensitivity weighting; MF-AirComp uses channel-matched analog aggregation, while ZF-AirComp uses channel inversion.Source: D. Massive-MIMO Scaling Study · p. 10
- After six transceiver updates, normalized WSMSE falls from 100 to 32.5 for SW-AirComp and to 69.9 for UW-AirComp.Both methods use the same sensitivity weights, channels, source covariances, and samples; most SW-AirComp improvement occurs within the first two iterations.Source: D. Massive-MIMO Scaling Study · p. 10, Figure 8 · p. 10
- In the transition to usable SNR, SW-AirComp reduces logit-RMSE and logit-gap violation risk relative to UW-AirComp, MF-AirComp, and ZF-AirComp.The diagnostics measure output-logit distortion and how often estimated perturbations exceed the clean top-token logit gap; low-SNR saturation is less useful for ranking methods.Source: Figure 9 · p. 10, D. Massive-MIMO Scaling Study · p. 11
E. Insights and Limitations
The paper links block-level logit sensitivity to wireless aggregation design for distributed MoE inference. Its conclusion combines simulator evidence with a direct ARC-Challenge audit, while noting that broader benchmark claims require further evaluation.
E. Insights and Limitations
- Block-level priority metadata lets the wireless optimizer allocate spatial and power resources using inference-side sensitivity while relying on CSI, noise variance, and power constraints.The paper identifies stale or noisy metadata as a practical risk: overestimation can waste resources, while underestimation can leave harmful logit perturbations despite low WMSE.Source: E. Insights and Limitations · p. 11
- Unlike training-time AirComp noise, an inference aggregation error affects the hidden state used for the immediate next-token choice, and a wrong token changes the following context.The paper therefore treats logit-gap violation as a decision-level diagnostic, not merely a communication metric.Source: E. Insights and Limitations · p. 11
V. CONCLUSION
- SW-AirComp gives the highest Qwen3 task accuracy among the tested AirComp methods in the 295-example ARC-Challenge audit, while simulations show lower sensitivity-weighted distortion.The paper calls for more channel seeds, SNRs, tasks, and fully online CSI, routing, and sensitivity calibration before making broader benchmark-level claims.Source: E. Insights and Limitations · p. 11, §V · p. 11
- The conclusion’s framework combines Jacobian-based sensitivity modeling, block-level error weighting, and power-constrained transceiver optimization for distributed MoE inference.Its transmitted symbols use a linear, invertible complex packing of the hidden-state contributions required by MoE aggregation.Source: §V · p. 11
APPENDIX A PROOF OF LEMMA 1
These appendix proofs establish a scalar reduction, bound logit perturbations using a Taylor expansion, and state a condition for a token flip.
APPENDIX A PROOF OF LEMMA 1
- Mutual uncorrelatedness makes the cross terms vanish, yielding the scalar reduction established in Lemma 1.The proof uses this property in reducing the block-error expression.Source: Appendix A · pp. 11–12
APPENDIX B PROOF OF THEOREM 1
- Theorem 1’s proof uses a locally Lipschitz Jacobian and bounds the first-order term and Taylor remainder before substituting them into (20).The cross terms vanish because the block errors are mutually uncorrelated.Source: Appendix B · p. 12
APPENDIX C PROOF OF COROLLARY 1
- A token flip can occur only when the infinity-norm logit perturbation reaches or exceeds the clean top-token logit gap.This is the condition used in the proof of Corollary 1.Source: Appendix C · p. 12
APPENDIX D PROOF OF LEMMA 2
The proof derives the receive-combiner update from the over-the-air signal model, using uncorrelated sources and receiver noise to simplify the objective.
APPENDIX D PROOF OF LEMMA 2
- The proof starts from over-the-air superposition and uses source and receiver-noise independence to eliminate cross terms in the expected squared norm.The resulting expression is equivalent to (26).Source: Appendix D · p. 12
- For fixed transmit precoders, the receive-block derivation uses the source covariance and removes the term independent of the combiner.The proof states that the resulting block is strongly convex because the covariance is positive definite and the block weight is positive.Source: Appendix D · p. 12
- Differentiating the receive-block objective gives the unique receive-combiner solution in (30).Uniqueness follows from the positive-definite covariance and positive block weight stated in the proof.Source: Appendix D · p. 12
APPENDIX E PROOF OF PROPOSITION 1
The appendix formulates each BS’s fixed-combiner precoder update as a convex problem and establishes convergence properties for the alternating updates.
APPENDIX E PROOF OF PROPOSITION 1
- With receive combiners fixed, each BS transmit subproblem is a convex QCQP under its quadratic power constraint.The objective is convex because the map in the precoder is affine, and the constraint is a convex quadratic inequality.Source: Appendix E · p. 12
- The KKT conditions give the optimal precoder form, with the multiplier set to zero when feasible or selected to meet the BS power budget otherwise.The proof states these conditions are sufficient for global optimality because the subproblem is convex.Source: Appendix E · pp. 12–13
APPENDIX F PROOF OF PROPOSITION 2
- Alternating exact receive-combiner and transmit-precoder minimizations yield nonincreasing objective values, and every accumulation point is block-stationary.The objective is bounded below by zero, so its monotone values converge.Source: Appendix F · p. 13