推理成本比MoE直降83%！字节最新大模型架构入围ICLR 2025

888974111

字节出了个全新架构，把推理成本给狠狠地打了下去！
有多狠？
推理速度相比MoE架构提升2-6倍，推理成本最高可降低83%。

这个全新的稀疏模型架构叫做UltraMem，有效地解决了目前主流的MoE架构和PKM架构所存在的局限性。
例如MoE在做推理时，较小的batch size会激活全部专家，导致访存急剧上升，推理延迟增加；而PKM虽然减少了访存开销，但效果较差且扩展能力有限。
实验结果表明，训练规模达2000 万value的UltraMem模型，在同等计算资源下，可同时实现业界领先的推理速度和模型性能，为构建数十亿规模value或expert开辟了新路径。
这项研究目前已经被ICLR 2025接收。

那么UltraMem架构具体是如何做到这点的呢？我们继续往下看。
兼顾访存和模型效果

随着大语言模型能力的提升，所需的计算资源呈指数级增长，这在实时应用等资源受限的环境中带来了巨大挑战。
为了应对计算资源的瓶颈，研究者们提出了多种解决方案，其中包括MoE和Product Key Memory (PKM)。
然而，这些方法各自存在一定的局限性。
MoE 通过稀疏激活专家（expert）的方式，将计算与参数解耦，从而在训练时能够有效减少计算量；但在推理场景中，MoE 的表现却不尽如人意。
由于推理时模型需要逐字生成输出，batch size和sequence length通常较小，而通常情况下，数量较少的token就可以激活几乎所有的专家。
这种全专家激活的模式使得访存需求急剧增加，进而引发严重的访存瓶颈，最终导致推理延迟大幅上升。
另一方面，PKM提出了 large memory layer的概念，其中包含了大量的稀疏参数 value（每个 value 实际上是一个向量）。
在推理时，每个 token 会通过“行路由”和“列路由”定位到得分最高的几个 value，并对其进行加权求和（weighted sum pooling），作为 memory layer 的输出。
由于每个token在推理时仅激活极少数的value，PKM能够有效避免访存瓶颈。然而，尽管PKM在访存效率上表现优异，其模型效果却较差，且扩展能力（scaling）有限，难以应对更大规模的模型需求。
总的来说，MoE和PKM虽然在一定程度上缓解了计算资源的问题，但它们在推理效率、模型效果和扩展能力等方面仍存在明显不足，亟需新的解决方案来进一步优化大模型的推理性能。

而UltraMem则是参考了PKM的设计，但针对PKM的3个缺陷予以补充，以实现更高效的访存、更优质的value检索，同时，降低了显存和部署成本。
1、优化模型结构

PKM的设计中，memory layer只有1层，插在整个Transformer的中间层，这对大规模训练并不友好，并且如此庞大的稀疏参数应该尽可能多的参与到每次的残差连接中。
因此，研究团队拆分出多个小memory layer，以固定的间隔分布在 Transformer layer中；并且增加了skip-layer的操作，即当前层的memory layer的输出会加到后面某层Transformer layer 的输出。
这使得模型可以并行地执行memory layer的访存操作和Transformer layer的计算。

2、优化value检索方式

在检索时，只有score最高的m个value会被激活，PKM的score是通过“行score”+“列score”得到的。
团队进一步探索了一种更复杂的乘法方法Tucker Decomposed Query-Key Retrieval（TDQKR）。
这一方法受启发于Tucker Decomposition。具体来看，给定values，shape为（n，n，h），其中h为hidden size，那么values的score S_grid可以做如下分解：

其中Srow,Scol∈Rr×n，C∈Rr×r是可学习的tucker core。这个结构下，每个value的score由r个行score和r个列score的组合乘加获得，具备更高的复杂度。

3、隐式扩展稀疏参数

更多的稀疏参数通常会带来更好的效果，但过多的参数又会给显存和部署带来麻烦。
为此，研究团队提出了Implicit Value Expansion （IVE）方法隐式地扩展稀疏参数，并引入了virtual memory和physical memory的概念。
以4倍扩展为例（如下图所示），virtual memory的数量是physical memory的4倍，给定多对（score， index）后，首先按照virtual memory address table做查表，4个virtual block会查询同一个physical memory table，之后各自做 weighted sum pooling，并经过不同的线性层，最后再求和输出。
由于最后的Linear和取value之间没有任何非线性操作，因此每个Linear都可以和physical memory table做融合，生成一个全新的memory table，这个例子下，实际上隐式扩展了4倍的value 数量。

较MoE最高提速6倍

1、模型性能评估

研究团队在151M、680M、1.6B三个尺寸的激活参数上做了广泛的实验，其中MoE、PKM和UltraMem的总稀疏参数保持在激活参数的12倍。
如下表所示，可以发现 UltraMem在680M、1.6B上具有显著的效果优势。

随着稀疏参数的增加，UltraMem的效果和推理速度如何变化？

• 下图（b）展示了UltraMem的效果变化，横轴为稀疏参数和稠密参数的比值，每个颜色的线代表了一种稀疏度。稀疏度定义为value的数量 / 每个token激活的value数量。观察发现，持续增加稀疏参数和loss的下降呈对数关系；且稀疏度越小，模型效果越好；但是稀疏度持续降低带来的收益在逐渐饱和。
  
• 下图（c）展示了UltraMem的推理时间变化，横轴为稀疏参数和稠密参数的比值。观察发现，UltraMem在持续增加稀疏参数时推理时间几乎不变，反观MoE有显著增长的趋势。
  

2、消融实验

研究团队在151M激活、1.5B总参数的稀疏模型上进行了全面的消融实验。
从最原始的PKM开始，逐渐增加一些 trick 和上文提出的结构改进，最终能拿到C4 validation loss -0.092的显著收益，同时稀疏参数和计算量几乎不变。

综上所述，研究团队提出的UltraMem具有极小的访存，因此，相比 MoE 实现了最高达6倍的速度提升，推理成本最高可降低83%。
同时，在性能方面，随着模型容量的增加，在相同的参数和计算量情况下，UltraMem超过了MoE，表明其具有更强的扩展能力。
这项工作可以说是为开发更高效和可扩展的语言模型提供了一个有希望的方向。
论文地址：https://arxiv.org/abs/2411.12364