通信人家园

 找回密码
 注册

只需一步,快速开始

短信验证,便捷登录

搜索

军衔等级:

  少校

注册:2007-10-2912
跳转到指定楼层
1#
发表于 2025-5-7 13:40:12 |只看该作者 |倒序浏览








Mixture-of-Experts(MoE)在推理时仅激活每个 token 所需的一小部分专家,凭借其稀疏激活的特点,已成为当前 LLM 中的主流架构。然而,MoE 虽然显著降低了推理时的计算量,但整体参数规模依然大于同等性能的 Dense 模型,因此在显存资源极为受限的端侧部署场景中,仍然面临较大挑战。

现有的主流解决方案是专家卸载(Expert Offloading),即将专家模块存储在下层存储设备(如 CPU 内存甚至磁盘)中,在推理时按需加载激活的专家到显存进行计算。但这一方法存在两大主要缺陷:

<ol>
  • 由于不同 token 通常激活的专家不同,每一步推理都需要频繁加载不同的专家,导致显著的推理延迟;
  • 在批量解码场景中,各 token 可能需要不同的专家,在最坏情况下,甚至需要将一整层的所有专家加载到显存中,进一步加剧显存压力并带来额外的推理延迟。</ol>为了解决上述问题,来自北大和华为诺亚的研究人员提出了 Mixture-of-Lookup-Experts(MoLE),一种在推理阶段可重参数化的新型 MoE 架构。





    • 论文链接:https://arxiv.org/pdf/2503.15798
    • 代码链接:https://github.com/JieShibo/MoLE
    思考

    本文的核心思考是,在专家卸载方案中,需要将专家模块加载到显存,主要是为了在 GPU 上执行高效的矩阵运算。换句话说,如果专家的计算过程能够绕过矩阵运算的需求,就可以避免将专家权重加载到显存,从而根本上规避频繁加载带来的开销。直观来看,专家模块本质上是一个神经网络,用于建模输入到输出的映射。如果能够在推理前预先计算出所有可能的输入 - 输出对应关系,并将其存储为查找表,那么在推理时即可用简单的查找操作代替矩阵运算。

    一般而言,神经网络所建模的映射通常涉及无限多的输入 - 输出对,因此,要实现查找表替代矩阵运算,首先需要确保专家模块的输入来自一个离散且有限的集合,并且这一离散化方法能够适配大规模预训练任务。其次,由于查找操作发生在 GPU 之外,还需要保证检索过程本身不依赖密集计算,避免引入新的性能瓶颈。

    基于这一思考,作者注意到,大规模语言模型(LLM)中的 embedding token(即 embedding 层的输出)天然具备离散且有限的特性,其数量与词表大小一致,满足了离散有限要求。并且 embedding token 可以通过 token ID 唯一确定,因此查找表的检索可以采用高效的直接寻址。因此,MoLE 设计中将专家的输入由中间特征改为 embedding token,从而满足了查找表构建的所有要求。

    训练阶段

    在训练阶段,MoLE 相较于传统的 MoE 架构存在三个主要区别:

    <ol>
  • 输入调整:将所有路由专家(routed experts)的输入由上一层的输出,改为浅层的 embedding token,以确保专家模块可以被预计算并存储为查找表。
  • 激活策略:由于查找表检索在推理时无需额外计算,MoLE 无需依赖稀疏激活来控制推理计算量,因此在训练中选择激活所有路由专家。
  • 损失设计:鉴于不再需要通过稀疏激活实现负载均衡,MoLE 训练时仅使用语言建模损失,不再引入额外的负载均衡损失项。</ol>除此之外,MoLE 的其他设计与标准 MoE 保持一致,包括路由(router)模块和共享专家(shared experts),依然使用上一层的输出作为输入。计算流程如下









    推理阶段

    在推理前,MoLE 通过预先构建查找表来完成专家模块的重参数化。具体而言,embedding 层的权重矩阵本身即包含了所有 embedding token 的向量表示,因此可以直接以该权重矩阵作为专家模块的输入,并通过各个路由专家分别计算对应的输出。这样,便可以高效地获得完整的 token id 到专家输出的映射集合,用于后续的查找操作。具体过程如下所示:









    在查找表构建完成后,所有原始的路由专家模块将被删除,查找表则被存储在下层存储设备中。在推理阶段,对于每个 token,根据其 token ID 直接在查找表中检索对应的专家输出,并将检索到的输出加载到显存中,用于后续的推理计算。整体计算流程如下所示:





    复杂度分析

    如表所示,在推理阶段,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 由于需要激活两个专家,其专家大小为 dense FFN 的一半,但专家数量是 MoLE 的两倍。





    实验结果表明 MoLE 在相同训练参数量和推理激活参数量(即显存使用量)下,具有与 MoE 相当的性能,相比 Dense 有显著提升。与专家卸载的 MoE 相比,MoLE 减少了千倍以上的传输开销。





    在 V100 上进行的评测结果表明,在显存用量一定的前提下,MoLE 的推理延迟与 Dense 基本一致,显著优于专家卸载的 MoE。在批量解码场景下,随着 batch size 的增加,MoE 的推理延迟迅速上升,而 MoLE 与 Dense 模型的延迟则基本保持稳定,进一步展现了 MoLE 在高吞吐量推理任务中的优势。









    此外,消融实验表明,MoLE 的训练确实不需要辅助损失。

    在专家数量提升时,模型性能也会提升





    然而,如果仅增大专家的隐层维度,由于查找表的大小保持不变,当专家规模增大到一定程度时,推理性能将受限于查找表的固定大小,最终达到饱和。





    作者通过将一个 MoE 模型逐步修改为 MoLE 模型,系统性地探索了 MoLE 各组成部分对性能的影响。实验结果表明,使用浅层的 embedding token 作为专家输入确实会削弱模型的表达能力,这是由于输入中缺乏丰富的上下文信息所致。然而,激活所有专家有效弥补了这一损失,使得 MoLE 最终能够达到与 MoE 相当的性能水平。





    需要注意的是,路由专家的输入不包含上下文信息,并不意味着专家无法影响模型对上下文的处理。实际上,专家可以通过改变其输出,从而间接影响后续注意力层的输入,实现对上下文的建模。此外,共享专家和路由仍然接收包含上下文信息的输入,进一步保障了模型对上下文理解能力的保留。

    最后,作者发现查找表中仍然存在较大程度的冗余。即使将查找表压缩至 3-bit 精度(例如使用 NF3 格式),模型性能依然能够基本保持不变。这表明查找表在存储开销上仍具有进一步压缩和优化的潜力。





    总结

    综上,本文提出了 MoLE,一种面向端侧部署优化的新型 MoE 架构。通过将专家输入改为浅层的 embedding token,并采用查找表替代传统的矩阵运算,MoLE 有效解决了传统 MoE 在推理阶段面临的显存开销大、传输延迟高的问题。实验结果表明,MoLE 在保持与 MoE 相当性能的同时,大幅降低了推理延迟,尤其在批量解码场景中展现出显著优势。


    来源:网易

  • 举报本楼

    本帖有 1 个回帖,您需要登录后才能浏览 登录 | 注册
    您需要登录后才可以回帖 登录 | 注册 |

    版规|手机版|C114 ( 沪ICP备12002291号-1 )|联系我们 |网站地图  

    GMT+8, 2025-5-8 17:37 , Processed in 0.195370 second(s), 16 queries , Gzip On.

    Copyright © 1999-2023 C114 All Rights Reserved

    Discuz Licensed

    回顶部