arXiv:2503.15798cs.LGcs.CL2025-03ICML被引 13

MoLE让专家模型推理更快更省显存,无需计算直接查表。

Mixture of Lookup Experts

  • 专家在推理时转为查表,输入ID直接查结果,免去计算开销。
  • 同等显存和算力下,推理速度接近稠密模型,远快于需加载的MoE。
  • 适合部署资源受限场景,尤其对低延迟要求高的应用。

Mixture-of-Experts (MoE) 在推理时仅激活部分专家,可保持低计算量和低延迟,但所有专家仍需常驻显存,限制部署。为解决此问题,本文提出 Mixture of Lookup Experts (MoLE),一种通信与显存效率更高的 MoE 架构。训练时,专家为前馈网络(FFN),以嵌入层输出为输入;推理前,这些专家被重参数化为查找表(LUT),根据输入 ID 检索专家输出,并可卸载至存储设备。推理时无需计算,直接按输入 ID 从存储中加载专家输出到显存,通信开销几乎为零。实验表明,在相同 FLOPs 与显存占用下,MoLE 的推理速度与稠密模型相当,显著快于需动态加载的 MoE,且性能与原版 MoE 保持一致。

原文摘要 · Abstract (English)

Mixture-of-Experts (MoE) activates only a subset of experts during inference, allowing the model to maintain low inference FLOPs and latency even as the parameter count scales up. However, since MoE dynamically selects the experts, all the experts need to be loaded into VRAM. Their large parameter size still limits deployment, and offloading, which load experts into VRAM only when needed, significantly increase inference latency. To address this, we propose Mixture of Lookup Experts (MoLE), a new MoE architecture that is efficient in both communication and VRAM usage. In MoLE, the experts are Feed-Forward Networks (FFNs) during training, taking the output of the embedding layer as input. Before inference, these experts can be re-parameterized as lookup tables (LUTs) that retrieves expert outputs based on input ids, and offloaded to storage devices. Therefore, we do not need to perform expert computations during inference. Instead, we directly retrieve the expert's computation results based on input ids and load them into VRAM, and thus the resulting communication overhead is negligible. Experiments show that, with the same FLOPs and VRAM usage, MoLE achieves inference speeds comparable to dense models and significantly faster than MoE with experts offloading, while maintaining performance on par with MoE.

MoE推理优化显存节省查表

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。