ICML 2025 Spotlight｜华为诺亚提出端侧大模型新架构MoLE，内存搬运代价降低1000倍
仅用于站内搜索，没有排版格式，具体信息请跳转上方微信公众号内链接
Mixture-of-Experts（MoE）在推理时仅激活每个token所需的一小部分专家，凭借其稀疏激活的特点，已成为当前LLM中的主流架构。然而，MoE虽然显著降低了推理时的计算量，但整体参数规模依然大于同等性能的Dense模型，因此在显存资源极为受限的端侧部署场景中，仍然面临较大挑战。
现有的主流解决方案是专家卸载（ExpertOffloading），即将专家模块存储在下层存储设备（如CPU内存甚至磁盘）中，在推理时按需加载激活的专家到显存进行计算。但这一方法存在两大主要缺陷：
由于不同token通常激活的专家不同，每一步推理都需要频繁加载不同的专家，导致显著的推理延迟；
在批量解码场景中，各token可能需要不同的专家，在最坏情况下，甚至需要将一整层的所有专家加载到显存中，进一步加剧显存压力并带来额外的推理延迟。
为了解决上述问题，来自北大和华为诺亚的研究人员提出了Mixture-of-Lookup-Experts（MoLE），一种在推理阶段可重参数化的新型MoE架构。
论文链接：https ://arxiv.org/pdf/2503.15798
代码链接：https ://github.com/JieShibo/MoLE
思考
本文的核心思考是，在专家卸载方案中，需要将专家模块加载到显存，主要是为了在GPU上执行高效的矩阵运算。换句话说，如果专家的计算过程能够绕过矩阵运算的需求，就可以避免将专家权重加载到显存，从而根本上规避频繁加载带来的开销。直观来看，专家模块本质上是一个神经网络，用于建模输入到输出的映射。如果能够在推理前预先计算出所有可能的输入-输出对应关系，并将其存储为查找表，那么在推理时即可用简单的查找操作代替矩阵运算。
一般而言，神经网络所建模的映射通常涉及无限多的输入-输出对，因此，要实现查找表替代矩阵运算，首先需要确保专家模块的输入来自一个离散且有限的集合，并且这一离散化方法能够适配大规模预训练任务。其次，由于查找操作发生在GPU之外，还需要保证检索过程本身不依赖密集计算，避免引入新的性能瓶颈。
基于这一思考，作者注意到，大规模语言模型（LLM）中的embeddingtoken（即embedding层的输出）天然具备离散且有限的特性，其数量与词表大小一致，满足了离散有限要求。并且embeddingtoken可以通过tokenID唯一确定，因此查找表的检索可以采用高效的直接寻址。因此，MoLE设计中将专家的输入由中间特征改为embeddingtoken，从而满足了查找表构建的所有要求。
训练阶段
在训练阶段，MoLE相较于传统的MoE架构存在三个主要区别：
输入调整：将所有路由专家（routedexperts）的输入由上一层的输出，改为浅层的embeddingtoken，以确保专家模块可以被预计算并存储为查找表。
激活策略：由于查找表检索在推理时无需额外计算，MoLE无需依赖稀疏激活来控制推理计算量，因此在训练中选择激活所有路由专家。
损失设计：鉴于不再需要通过稀疏激活实现负载均衡，MoLE训练时仅使用语言建模损失，不再引入额外的负载均衡损失项。
除此之外，MoLE的其他设计与标准MoE保持一致，包括路由（router）模块和共享专家（sharedexperts），依然使用上一层的输出作为输入。计算流程如下
推理阶段
在推理前，MoLE通过预先构建查找表来完成专家模块的重参数化。具体而言，embedding层的权重矩阵本身即包含了所有embeddingtoken的向量表示，因此可以直接以该权重矩阵作为专家模块的输入，并通过各个路由专家分别计算对应的输出。这样，便可以高效地获得完整的tokenid到专家输出的映射集合，用于后续的查找操作。具体过程如下所示：
在查找表构建完成后，所有原始的路由专家模块将被删除，查找表则被存储在下层存储设备中。在推理阶段，对于每个token，根据其tokenID直接在查找表中检索对应的专家输出，并将检索到的输出加载到显存中，用于后续的推理计算。整体计算流程如下所示：
复杂度分析
如表所示，在推理阶段，MoLE的计算过程中仅保留了共享专家模块，因此只有共享专家被激活并参与计算，其整体计算量与具有相同激活参数量的Dense模型和传统MoE模型相当。相比之下，MoLE在推理时仅需传输专家输出的结果向量，而传统MoE需要传输中间维度D_r的专家权重矩阵，因此MoLE的传输开销相比MoE减少了数个量级。在存储开销方面，对于端侧部署的模型，词表大小|V|通常在数万左右，与D_r为相同数量级，因此MoLE查找表的存储需求与单个专家模块的大小处于同一数量级，不会带来显著额外的存储负担。
实验结果
本文在Pile数据集的100B-token子集上训练了160M、410M、1B激活参数量的Dense、MoE和MoLE模型。对于MoE和MoLE模型，控制两者的训练阶段参数量相等。由于实验中以及近期OLMoE的结果都发现共享专家会降低MoE的性能，我们对MoE只采用了路由专家。MoLE的专家大小与Dense的FFN保持一致，而MoE由于需要激活两个专家，其专家大小为denseFFN的一半，但专家数量是MoLE的两倍。
实验结果表明MoLE在相同训练参数量和推理激活参数量（即显存使用量）下，具有与MoE相当的性能，相比Dense有显著提升。与专家卸载的MoE相比，MoLE减少了千倍以上的传输开销。
在V100上进行的评测结果表明，在显存用量一定的前提下，MoLE的推理延迟与Dense基本一致，显著优于专家卸载的MoE。在批量解码场景下，随着batchsize的增加，MoE的推理延迟迅速上升，而MoLE与Dense模型的延迟则基本保持稳定，进一步展现了MoLE在高吞吐量推理任务中的优势。
此外，消融实验表明，MoLE的训练确实不需要辅助损失。
在专家数量提升时，模型性能也会提升
然而，如果仅增大专家的隐层维度，由于查找表的大小保持不变，当专家规模增大到一定程度时，推理性能将受限于查找表的固定大小，最终达到饱和。
作者通过将一个MoE模型逐步修改为MoLE模型，系统性地探索了MoLE各组成部分对性能的影响。实验结果表明，使用浅层的embeddingtoken作为专家输入确实会削弱模型的表达能力，这是由于输入中缺乏丰富的上下文信息所致。然而，激活所有专家有效弥补了这一损失，使得MoLE最终能够达到与MoE相当的性能水平。
需要注意的是，路由专家的输入不包含上下文信息，并不意味着专家无法影响模型对上下文的处理。实际上，专家可以通过改变其输出，从而间接影响后续注意力层的输入，实现对上下文的建模。此外，共享专家和路由仍然接收包含上下文信息的输入，进一步保障了模型对上下文理解能力的保留。
最后，作者发现查找表中仍然存在较大程度的冗余。即使将查找表压缩至3-bit精度（例如使用NF3格式），模型性能依然能够基本保持不变。这表明查找表在存储开销上仍具有进一步压缩和优化的潜力。
总结
综上，本文提出了MoLE，一种面向端侧部署优化的新型MoE架构。通过将专家输入改为浅层的embeddingtoken，并采用查找表替代传统的矩阵运算，MoLE有效解决了传统MoE在推理阶段面临的显存开销大、传输延迟高的问题。实验结果表明，MoLE在保持与MoE相当性能的同时，大幅降低了推理延迟，尤其在批量解码场景中展现出显著优势。
©THEEND
转载请联系本公众号获得授权
投稿或寻求报道：liyazhou@jiqizhixin.com
