DraftExpert: Expansion-Aware Self-Speculative Decoding for End-Device MoE Inference

2026-07-27Machine Learning

Machine LearningArtificial Intelligence
AI summary

The authors address how to make large language models with many 'experts' work faster on devices like phones or laptops, where memory is limited. They note that picking experts to run for each word can slow things down when experts are loaded on demand. To fix this, the authors create DraftExpert, a method that trains small 'draft' experts to guess outputs quickly and decides when to fully check with the big model only if needed. Their approach improves speed and accuracy when running these models across different hardware setups.

Mixture-of-Expertslanguage modelsself-speculative decodingexpert routingmodel offloadinginference accelerationprefetchingself-distillationlatency-critical applications
Authors
Dengke Han
Abstract
Large Mixture-of-Experts (MoE) language models are attractive for end-device deployment because only a small subset of experts is active per token, but their routed expert weights often exceed accelerator memory. We target latency-critical single-user settings where routed experts are staged on demand from CPU memory to a GPU or from Flash to a mobile NPU. In this setting, self-speculative decoding faces a new bottleneck: increasing the draft expert set improves accuracy but triggers extra expert loading, while cheap small-footprint drafts have low acceptance; moreover, verifying a multi-token block activates the union of target experts and is no longer close to one target step. We propose DraftExpert, an expansion-aware self-speculative decoding framework for expert-offloaded MoE inference. DraftExpert trains one lightweight accelerator-resident draft expert per layer by self-distilling residual, logit/token, and router-agreement signals from the frozen target MoE. At inference time, it uses a fixed-footprint shared+top-1+draft-expert drafter together with confidence--expansion truncation and target-expert prefetching, while final tokens are still exactly verified by the target model. On DeepSeek-V2-Lite and Moonlight-16B-A3B across CPU-GPU and Flash-NPU offload, DraftExpert improves decode throughput by 1.45x on average, raises draft acceptance to 84~87%, and achieves 86~88% prefetch hit rates.