给MoE模型瘦身,不再需要校准数据
大模型推理时,专家混合(MoE)模型要频繁从内存里搬权重,这是速度瓶颈。以前的做法是设一个阈值,只加载“重要”的通道,但这个阈值得靠离线数据校准,而且实际读取量会随输入波动——论文发现,按18%预算校准的阈值,实际平均只读了12.9%。他们改用top-k:直接选每个token得分最高的k个通道,完全不需要校准数据。在43组对比中,top-k在40组里不输阈值版,还避免了Gemma 4上阈值版会引发的重复循环崩溃。更妙的是,他们用模型自身的权重构造了一个“阅读器”,用三元草图筛通道、低秩映射补回跳过的部分,在三个模型上只读17.8%的专家权重就保住了98.1%的准确率,推理速度比优化过的vLLM快1.6到2.7倍。这不是你明天能直接用的功能,但它意味着:未来跑MoE大模型可能更快、更省内存,而且不用再为每个模型调一堆校准参数。
📄 原文摘要(英文)
Mixture-of-Experts (MoE) models are bottlenecked by loading routed-expert weights from memory during interactive decoding at small batch sizes. Because most intermediate channels in a routed expert contribute little to any single token, channel selectors predict which channels matter and read only their weights, typically keeping channels above a threshold calibrated on offline data. We show that channel selection needs no calibration data. First, a calibrated threshold does not fix how much a selector reads, because the number of channels passing it varies across inputs; Prox with thresholds calibrated for an 18% read budget actually reads only 12.9% on average across our benchmarks. Second, top-k keeps each token's k highest-scoring channels and needs no calibration data. At the calibrated budget, a selector's top-k version matches or beats its threshold version on the same score in 40 of 43 comparisons at 13% and 18% budgets, including all 23 on generation tasks. The top-k versions also avoid the repetition-loop collapse that calibrated thresholds cause on Gemma 4. Calibrating a threshold thus provides neither of the two things selectors use it for: a known read budget and a consistent accuracy gain at that budget. We furthermore introduce a reader built entirely from the model's weights, which screens channels with a ternary sketch and recovers the skipped ones with a low-rank map. Across three post-trained MoE models, our reader retains 98.1% of full-model accuracy at 17.8% of routed-expert reads, with Triton kernels running MoE blocks 1.60 to 1.90x faster than optimized vLLM at batch one and up to 2.69x at larger batches on the same three models.