TrimMoE A communication aware and adaptive depth framework for distributed edge inference

Technical Summary: TrimMoE – A Communication-Aware and Adaptive-Depth Framework for Distributed Edge Inference

Problem Statement
The paper addresses the challenge of serving Mixture-of-Experts (MoE) Large Language Models (LLMs) across distributed edge servers. While existing distributed inference approaches focus on accelerating the transmission of tokens to remote experts (e.g., via better scheduling or caching), they fail to address a fundamental inefficiency: the unnecessary execution of layers and the subsequent cross-server transmissions they trigger. In edge environments, characterized by heterogeneous resources, limited bandwidth, and time-varying network conditions, forcing every token to traverse the full depth of an MoE model—even when the representation has stabilized or specific layers are of low importance—creates a significant bottleneck. The core problem is to determine whether a given layer, and the layers following it, need to be executed at all, rather than simply optimizing how to reach the required remote expert faster.

Methodology
The authors propose TrimMoE, a communication-aware, adaptive-depth framework that jointly optimizes layer skipping, confidence-based early exit, substitute execution, and server-expert selection under a unified quality degradation budget. The framework operates in two stages:

1. Offline Stage (Calibration and Deployment):

   • Exit-Head Training: Without fine-tuning the backbone, lightweight per-layer exit heads are trained on a calibration set to predict a confidence score indicating whether the current token representation is sufficient for the task.
   • Threshold Calibration: Per-layer importance thresholds are calibrated based on the empirical relationship between layer importance (derived from router gating weights) and the quality degradation caused by skipping.
   • Communication-Aware Redundancy: The framework refines expert replication by calculating a "skip/exit-aware redundancy benefit." Instead of replicating frequently activated experts, it prioritizes replicating experts that are both frequently executed (i.e., not skipped or exited early) and currently expensive to reach across the network. This concentrates memory on experts that genuinely reduce cross-server traffic.

2. Online Stage (Runtime Decision Making):

   • Joint Decision Making: For each token at each layer, TrimMoE decides whether to skip the layer (bypassing the expert via the residual path), exit early (terminating inference), or execute (selecting a specific server and expert).
   • Look-Ahead Mechanism: The decision process utilizes a transition-aware look-ahead that anticipates future token movement based on expert transition statistics. This allows the system to avoid decisions (like skipping a local layer) that might force a more expensive cross-server transmission in a subsequent layer.
   • Adaptive Feedback: The framework employs dual feedback rules to dynamically adapt the delay-quality trade-off weights and the exit confidence threshold based on real-time network conditions and workload, ensuring robustness under varying loads.
   • Quality Guarantees: The system enforces a hard constraint where the total quality degradation (from substitution and skipping) never exceeds a configured per-user budget. Early exits are only permitted when the confidence score meets a calibrated gate, ensuring the representation is sufficiently refined.

Key Contributions

• Novel Problem Formulation: The paper formulates distributed edge MoE inference as a joint optimization of layer skipping, early exit, and server-expert selection. Unlike prior works that treat every layer as mandatory, this formulation allows tokens to skip low-importance layers and terminate early, fundamentally changing the question from "how to transmit" to "whether to transmit."
• Communication-Aware Deployment: A novel offline scheme is designed to refine expert redundancy. By accounting for skip and exit behaviors, it ensures that memory replication targets experts that actually reduce cross-server transmission, rather than those that are frequently bypassed.
• Online Skip-Exit Algorithm: A communication-aware online algorithm is developed that jointly weighs transmission delay, computation delay, and quality loss. It incorporates a look-ahead mechanism to anticipate future costs and self-calibrates its parameters based on runtime observations.
• Theoretical Guarantees: The authors prove that the substitution-and-skipping proxy degradation is strictly bounded by the configured budget and that early exits are admitted only under a calibrated confidence gate, ensuring quality preservation by construction.

Experimental Results
The framework was evaluated on a heterogeneous 10-server edge testbed using three MoE models: Switch-Base-8E, Qwen-MoE-A2.7B, and Mixtral-8x7B.

• Latency Reduction: TrimMoE reduced average latency by up to 62.8% compared to baseline methods (e.g., EdgeShard) and significantly outperformed prior substitute-only schemes.
• Traffic Reduction: Cross-server traffic was reduced by 77.2% compared to baselines, and the remote execution ratio dropped to 16% (from 71% in baselines).
• Quality Preservation: Despite aggressive depth reduction, the task-quality degradation was kept within a 2% bound across various datasets (WikiText-103, SQuAD, GSM8K).
• Robustness: Under heavy load (up to 80 requests per second), TrimMoE maintained a higher SLA satisfaction ratio (71% vs. <50% for baselines) and demonstrated superior tail latency performance.

Significance and Claims
The paper claims to be the first work to jointly exploit communication-aware layer skipping and confidence-based early exit, coordinated with server-expert selection, to suppress cross-server transmission in distributed edge MoE inference. The significance lies in reframing the optimization objective: instead of merely accelerating necessary transmissions, TrimMoE eliminates unnecessary layer executions and the transmissions they would trigger. The authors emphasize that this approach achieves substantial latency and bandwidth savings while maintaining strict, provable bounds on inference quality, making it a viable solution for deploying large models in resource-constrained and bandwidth-limited edge environments.