FineRMoE:细粒度专家的维度扩展及其回收利用方法 FineRMoE: Dimension Expansion for Finer-Grained Expert with Its Upcycling Approach
将细粒度专家设计从中间维度扩展到输出维度,突破现有MoE缩放律的性能瓶颈
前置知识
Mixture-of-Experts (MoE)
MoE是一种稀疏激活架构,通过路由机制为每个输入动态选择并激活一小部分专家网络进行计算,而非激活所有参数。每个专家通常是一个前馈神经网络,路由器根据输入特征决定激活哪些专家。这种设计使模型能够在保持大规模参数的同时,仅激活少量参数进行前向传播,从而实现参数容量与计算效率的平衡。典型的MoE层包括共享专家和稀疏专家,输出通过加权求和或拼接方式融合各专家的输出。
本文提出的FineRMoE是MoE架构的一种扩展,需要理解MoE的基本原理,包括专家激活、路由机制和输出融合方式,才能理解FineRMoE在传统MoE基础上的创新点,特别是双层稀疏计算范式和单路由机制的设计。
细粒度专家设计
细粒度专家设计是指在MoE架构中,将每个专家的中间维度(即FFN的隐藏层大小)设置为远小于传统密集模型FFN的隐藏维度。例如,传统密集模型FFN的中间维度可能是8960,而细粒度专家的中间维度可能只有280(GI=32)。通过减小每个专家的容量,可以增加专家数量,促使每个专家专注于更窄的领域,从而减少专家间的冗余并提高专业化程度。研究发现,在合理范围内,专家粒度越细,模型性能越好,但超过最优阈值后性能会下降。
这是本文的出发点之一。现有的细粒度设计仅应用于中间维度,本文将其扩展到输出维度,打破单维度细粒度设计的性能瓶颈,需要理解中间维度细粒度设计的原理和局限性。
Upcycling(回收利用)
Upcycling是一种高效的MoE训练范式,通过将预训练的密集模型的FFN层转换为MoE层来构建模型,而非从头训练。这避免了从随机初始化开始训练专家,节省大量计算资源和时间。常见的upcycling方法包括复制FFN构建多个专家,或沿中间维度分割FFN生成细粒度专家。由于利用了预训练的权重,upcycling方法通常具有更好的初始性能和更快的收敛速度。
本文提出了适用于FineRMoE架构的通用upcycling方法,理解upcycling的基本原理和现有方法的局限性,有助于理解本文方法如何超越现有的upcycling技术,同时兼容传统MoE架构。
Scaling Law(缩放律)
缩放律描述了模型性能与模型规模、数据量、计算量等资源之间的关系。对于MoE模型,存在关于细粒度设计的缩放律:在中间维度的专家粒度未超过最优阈值时,模型性能随粒度增加而提升;一旦超过阈值,性能反而下降。这种现象为细粒度设计提供了理论指导,但也限制了单维度细粒度设计的潜力。本文的目标正是突破这一限制。
本文的动机正是源于细粒度MoE的缩放律揭示了中间维度细粒度设计的性能瓶颈,需要理解这一缩放律才能领会为何要将细粒度设计扩展到输出维度。
研究动机
现有的细粒度MoE设计仅应用于专家的中间维度,即FFN的隐藏层大小。然而,根据细粒度MoE的缩放律,一旦中间维度的专家粒度超过最优阈值,模型性能就不再提升,甚至出现下降。例如,对于Qwen2.5-1.5B模型,中间维度为8960,当专家的中间维度从8960(GI=1)减小到280(GI=32)时,性能持续提升;但继续减小到140(GI=64)时,性能增益趋于饱和(从56.21提升到56.34)。这一现象限制了单维度细粒度设计的上限,即使投入更多参数也无法获得相应回报。此外,传统的MoE输出融合采用加权求和方式,当专家在输出维度也进行细粒度设计时,会破坏MoE层前后计算的维度一致性。
本文的目标是本文的目标是突破细粒度MoE仅在中间维度设计的性能天花板,将细粒度专家设计扩展到输出维度,从而在超出原缩放律的范畴内为MoE模型解锁进一步性能提升。具体而言,作者提出FineRMoE(FineR-Grained MoE)架构,同时在中间和输出两个维度上进行细粒度设计,通过四个超参数(GI、RI、GO、RO)灵活调整稀疏专家的粒度和模型规模。作者还设计双层稀疏前向计算范式和专门的路由机制来管理激活,并开发通用的upcycling方法以低成本构建FineRMoE模型。
与已有工作不同的是,本文的独特切入角度是重新审视多头注意力机制的设计原理。在多头注意力中,降低QKV变换的输出维度会促使不同注意力头学习到不同的特征提取模式。类似地,作者假设降低专家的输出维度也会鼓励专家学习独立的表示,从而抑制冗余并增强专业化。基于这一洞察,作者首次将细粒度设计从中间维度扩展到输出维度,开创性地提出跨两维度的细粒度专家设计。此外,本文的upcycling方法将主流训练自由的upcycling方法统一到一个协议下,既支持复制FFN构建专家,也支持沿中间维度分割FFN,还能支持同时沿两个维度分割,超越了现有方法的适用范围。
核心方法
FineRMoE架构在传统MoE基础上进行扩展,包含共享专家和稀疏细粒度专家两个部分,两者的输出直接相加作为后续处理的输入。核心创新在于双层稀疏前向计算范式,由稀疏拼接层和稀疏求和层组成。对于每个token,在稀疏拼接层中,通过从RO个候选向量中选择得分最高的向量来拼接GO个分量,恢复完整的输出维度。每个候选向量在稀疏求和层中由一组细粒度专家的输出加权求和得到。整个架构通过四个超参数控制:中间粒度GI(共享专家中间维度H与稀疏专家中间维度He的比值,GI=H/He)、中间扩展率RI(一组稀疏专家的中间维度总和与共享专家中间维度的比值,RI=(Ng·He)/H)、输出粒度GO(模型隐藏维度h与稀疏专家输出维度he的比值,GO=h/he)、输出扩展率RO(每个拼接组件可选择的候选向量数量)。总专家数N=GO·RO·GI·RI。
FineRMoE的核心创新点包括三个方面:首先,首次将细粒度专家设计从中间维度扩展到输出维度,打破单维度细粒度设计的性能瓶颈,使专家能够同时从中间和输出两个层面实现专业化,进一步减少冗余。其次,引入双层稀疏前向计算范式,通过稀疏求和层和稀疏拼接层处理输入token,按照维度缩减再恢复的顺序,解决输出维度细粒度设计带来的维度一致性问题。最后,设计专门的单路由机制,用一个路由网络同时控制稀疏求和层的专家激活和稀疏拼接层的候选向量选择,避免使用两个路由器可能导致的激活冲突,同时减少参数开销。
方法步骤详情
FineRMoE的前向传播过程分为五个步骤。第一步,路由器为每个token计算所有专家的初始得分score ∈ R^(L×N),其中L是序列长度,N是专家总数。第二步,计算稀疏求和层的专家激活掩码:将初始得分reshape为L×(GO·RO)×(GI·RI)的group score,在每个组内对GI·RI个专家选择TI个得分最高的进行激活,得到sum mask。第三步,计算稀疏拼接层的候选向量选择掩码:对group score沿最后一维求和得到每个组的总分cc score ∈ R^(L×(GO·RO)),reshape为L×GO×RO,每个GO分量中选择RO个候选向量中得分最高的一个,得到cc mask并广播到所有专家。第四步,对两个掩码进行逐元素AND操作得到最终掩码final mask,将其应用于初始得分后选择GOTI个得分最高的专家进行实际激活。第五步,对每个token,将其分配给激活的专家进行并行计算(通过permute-dispatch-unpermute机制),每个专家组内同一token的输出通过加权求和形成维度缩减的向量,最后稀疏选择的向量直接拼接恢复维度,得到稀疏专家的最终输出。数学上,给定输入x ∈ R^h,共享专家输出为y_s = xW1_s ⊙ SiLU(xWg_s)W2_s,稀疏专家的输出O = Concat(O_0, ..., O_i, ..., O_{GO-1}),其中O_i = Top1Select(O_{0i}, ..., O_{ji}, ..., O_{(RO-1)i}),每个候选向量O_{ji} = WeightedSum(E_k, x)来自对应的专家组。
技术新颖性
FineRMoE的技术新颖性体现在多个方面。在架构层面,它是首个将细粒度设计从中间维度扩展到输出维度的MoE架构,突破了现有缩放律的单维度限制。引入的双层稀疏前向计算范式(稀疏拼接层+稀疏求和层)是新颖的设计,解决了输出维度细粒化带来的维度一致性挑战,使模型能够在两个维度同时实现专业化。在路由机制层面,单路由器同时控制双层稀疏激活的设计避免了使用两个独立路由器可能导致的激活冲突(即稀疏求和层得分高的专家可能在稀疏拼接层未被选中),同时减少了路由参数。在训练方法层面,提出的通用upcycling方法将主流训练自由的upcycling方法统一到一个协议下,通过四个超参数的配置,既能实现复制FFN(设置GO=RO=GI=1,RI为复制次数),也能实现沿中间维度分割FFN(设置GO=RO=RI=1,GI为分割次数),还能实现沿两个维度同时分割,适用于FineRMoE和传统MoE架构,具有很强的通用性和灵活性。
实验结果
在三个模型规模(Qwen2.5-0.5B、1.5B、7B)和十个标准基准测试(MMLU、BBH、HellaSwag、WinoGrande、ARC-C、ARC-E、AGIEval、MBPP、GSM8K、GPQA)上,FineRMoE在所有规模上都取得了最佳平均性能。对于1.5B模型,FineRMoE平均得分56.62,超过预训练模型(55.39)、持续训练密集模型(55.18)、复制FFN的C32A2(56.50)和沿中间维度分割的S16A4(26.32)。值得注意的是,FineRMoE的总参数仅为5.64B,远少于C32A2的37.62B(少6.7倍),激活参数为1.85B,也少于C32A2的2.94B,但性能仍优于C32A2,表明其高参数效率。消融实验验证了细粒度设计的有效性:仅在输出维度应用细粒度设计就能比预训练模型提升0.73分(55.39→56.12),优于仅在中间维度应用(55.39→53.06)。在两个维度同时应用细粒度设计时性能最佳(56.21),且专家相似度最低(图4),证明专家专业化程度最高。路由机制消融实验显示,单路由器设计(56.21)优于双路由器设计(54.96),避免了激活冲突。架构消融实验显示,添加拼接投影层会损害性能,移除共享专家会导致训练收敛困难,验证了设计的合理性。细粒度配置消融实验(Table 4)表明,当GI≥8时,细粒度MoE能超过预训练模型性能;对于每个GI设置,增加GO从2到8都能带来稳定的性能提升。与最强基线C32A2相比,FineRMoE在7B规模上实现了6倍更高的参数效率(26.65B vs 184.42B总参数,7.94B vs 13.33B激活参数)、281倍更低的prefill延迟(178.3ms vs 50245.9ms TTFT)和136倍更高的解码吞吐量(27.3 tokens/s vs 0.2 tokens/s)。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 综合评估(10个基准平均) | 平均准确率(%) | 56.62 (1.5B), 70.04 (7B) | 56.50 (C32A2, 1.5B), 69.92 (C32A2, 7B) | +0.12 (1.5B), +0.12 (7B) |
| 参数效率 | 总参数量/激活参数量(B) | 5.64/1.85 (1.5B) | 37.62/2.94 (C32A2, 1.5B) | 总参数少6.7倍,激活参数少1.6倍 |
| 推理效率(7B模型) | TTFT/Throughput | 178.3ms / 27.3 tokens/s | 50245.9ms / 0.2 tokens/s (C32A2) | TTFT快281倍,吞吐量高136倍 |
| 细粒度设计有效性 | 平均准确率(%) | 56.21 (Inter&Out FG) | 55.39 (Pre-trained), 53.06 (Inter FG only), 56.12 (Out FG only) | +0.82 vs 预训练, +3.15 vs 仅中间维度 |
| 路由机制 | 平均准确率(%) | 56.21 (Single Router) | 54.96 (Separate Routers) | +1.25 |
局限与改进
作者承认的局限性包括FineRMoE架构的复杂性较高,引入了四个超参数(GI、RI、GO、RO)需要调优,虽然论文提供了消融实验指导,但用户仍需根据具体任务和资源约束进行权衡。此外,虽然FineRMoE在推理效率上表现优异,但其双层稀疏计算范式可能对硬件优化提出更高要求,稀疏concat操作和稀疏求和操作的并行化需要特定的硬件支持。作者还指出,本文的upcycling方法虽然通用,但对于需要训练进行专家归纳的方法(如Branch-Train-Mix)不适用,这限制了某些upcycling技术的兼容性。从自己的观察来看,FineRMoE的优势在较大规模模型上更明显,对于极小模型(如0.5B),性能提升相对有限(41.39 vs 41.28 C32A2),这可能是因为小模型本身的容量有限,细粒度设计难以充分发挥作用。另外,虽然论文验证了10个基准测试,但主要关注知识、推理、代码和数学任务,未在生成任务(如创意写作、对话)上进行充分评估,这些任务可能需要不同的专家分配策略。
独立分析的弱点
FineRMoE的一个潜在弱点是超参数调优的复杂性。四个超参数(GI、RI、GO、RO)的组合会产生不同的模型规模和性能表现,用户需要根据计算资源和任务需求进行权衡。虽然Table 4提供了详细的消融实验,但最优配置可能因任务而异,增加了使用门槛。改进方向可以是开发自动超参数搜索策略或基于启发式的配置推荐。另一个可能的弱点是稀疏concat操作可能带来的内存访问不连续性,影响硬件利用效率。在GPU上,concat操作通常需要额外的内存分配和数据拷贝,可能成为性能瓶颈。改进方向是探索更高效的稀疏融合策略,如使用索引掩码避免显式concat,或开发专门的硬件优化。此外,FineRMoE的单路由机制虽然避免了激活冲突,但可能限制了路由的灵活性,因为稀疏拼接层的选择依赖于稀疏求和层的总分,这可能不总是最优的。改进方向是研究更复杂的路由策略,如软选择或多阶段路由,在保持参数效率的同时提升灵活性。
未来方向
作者提出的未来研究方向包括将FineRMoE扩展到多模态场景,验证其在视觉-语言、音频-语言等多模态任务中的有效性。由于多模态模型通常需要更大的模型容量和更专业的知识处理,FineRMoE的细粒度设计可能带来显著收益。另一个方向是探索更高效的专家激活策略,如动态调整激活专家数量或使用更复杂的路由损失函数,进一步平衡负载和性能。基于论文成果,可以延伸的研究方向包括:1)研究FineRMoE在持续学习场景中的应用,利用稀疏激活特性减少灾难性遗忘;2)探索FineRMoE与参数高效微调(PEFT)方法的结合,如将LoRA应用于特定专家以快速适应新任务;3)分析FineRMoE在不同训练规模(从10B到500B+ tokens)下的缩放行为,验证其是否具有与传统MoE不同的缩放律;4)研究FineRMoE的专家解释性,分析不同维度的细粒度设计如何影响专家的专业化领域,为模型调试和优化提供洞察。
复现评估
FineRMoE的复现性良好。论文提供了开源代码(https://github.com/liaoning97/FineRMoE)和模型权重(https://huggingface.co/collections/NingLiao/finermoe),方便研究者和实践者复现结果。实现基于Megatron-LM框架,这是一个广泛使用的分布式训练框架,具有良好的并行化支持。训练数据来自公开的预训练语料库,包括英文和中文网页数据、知识数据、代码数据、数学数据、The Pile、Wikipedia和书籍数据,数据来源明确且易于获取。训练超参数在Table 5中详细列出,包括优化器(Adam)、学习率(1e-5,余弦衰减到1e-7)、批次大小(2048)、序列长度(8192)、精度(BF16)等,设置合理且标准。主要实验在50B tokens上训练,对于资源有限的研究者,论文还提供了10B tokens的快速验证实验(如Table 2、Table 3、Table 4),降低了复现门槛。算力需求方面,基于Megatron-LM的分布式训练可以在多GPU环境下高效运行,虽然具体GPU数量和训练时间未明确说明,但根据模型规模和训练tokens数估算,训练7B FineRMoE约50B tokens可能在数十张A100/H100 GPU上运行数周,这是大规模语言模型训练的典型成本。总体而言,FineRMoE的复现难度中等,需要一定的硬件资源和分布式训练经验,但由于代码、权重和详细配置的提供,具备条件的实验室应该能够成功复现。
论文图表
Algorithm 1以PyTorch风格的伪代码形式详细描述了FineRMoE路由机制的实现。代码包含24行,分为四个主要部分:第5-6行计算路由得分score;第7-11行计算稀疏求和层中专家激活的掩码,将score resh成L×(GO·RO)×(GI·RI)的group score,每组选择TI个得分最高的专家,构建sum mask;第12-18行计算稀疏拼接层中候选向量选择的掩码,对group score沿最后一维求和得到cc score,reshape后每个GO分量选择RO个候选中得分最高的,构建cc mask并广播;第19-21行对两个掩码进行逐元素AND操作得到final mask;第22-24行将final mask应用于初始得分,选择GOTI个得分最高的专家进行实际激活。
Algorithm 1对于理解FineRMoE的单路由机制至关重要,它精确地展示了如何用一个路由网络同时控制两个稀疏层的激活,避免激活冲突。伪代码的详细程度使得读者可以理解路由机制的具体实现细节,包括掩码的计算、广播和最终激活的选择。这张图对method部分的深入理解不可或缺,特别是对于想要复现或扩展FineRMoE的研究者。
Table 6展示了五种方法(C32A2、S16A4、DU、NVShard、FineRMoE)在推理效率上的对比,所有模型基于Qwen2.5-7B持续训练50B tokens。表格包含两行:第一行是#Param(总参数量)和#A-Param(激活参数量),第二行是#Generated Tokens(生成token数)、TTFT(Time to First Token,首token时间)和Throughput(吞吐量,tokens/s)。FineRMoE的总参数为26.65B,激活参数为7.94B。推理效率方面,FineRMoE的TTFT为178.3ms,吞吐量为27.3 tokens/s。对比C32A2(184.42B总参数,13.33B激活参数),其TTFT为50245.9ms,吞吐量为0.2 tokens/s,FineRMoE快281倍、高136倍。表格还展示了各模型对输入'Give me a short introduction to large language model.'的生成文本,FineRMoE生成了高质量、连贯的介绍性文本,而DU模型产生了重复输出,直到达到最大长度32000 tokens。
Table 6是验证FineRMoE推理效率优势的关键证据,它清楚地展示了FineRMoE在prefill和解码两个阶段的卓越性能。这张表格对理解论文的results部分至关重要,它支持了作者声称的FineRMoE不仅在性能上超过基线,而且在推理效率上显著优于最强基线。表格还展示了不同方法的生成文本质量,间接验证了模型的性能。DU的重复输出问题表明部分重新初始化的upcycling方法在短时间训练(50B tokens)下难以收敛,进一步突显了FineRMoE upcycling方法的有效性。