据豆包大模型团队消息，字节跳动豆包大模型Foundation团队近期提出UltraMem，一种同样将计算和参数解耦的稀疏模型架构，在保证模型效果的前提下解决了推理的访存问题。据介绍，该架构有效解决了MoE推理时高额的访存问题，推理速度较MoE架构提升2-6倍，推理成本最高可降低83%。
随着模型规模的扩大，推理成本和访存效率已成为限制大模型规模应用的关键瓶颈。近期，字节跳动豆包大模型团队提出了全新的稀疏模型架构UltraMem，该架构有效解决了MoE推理时高额的访存问题，推理速度较MoE架构提升2-6倍，推理成本最高可降低83%。该研究还揭示了新架构的ScalingLaw，证明其不仅具备优异的Scaling特性，更在性能上超越了MoE。
实验结果表明，训练规模达2000万value的UltraMem模型，在同等计算资源下，可同时实现业界领先的推理速度和模型性能，为构建数十亿规模value或expert开辟了新路径。
在Transformer架构下，模型的性能与其参数数量和计算复杂度呈对数关系。随着LLM规模不断增大，推理成本会急剧增加，速度变慢。
尽管MoE架构已经成功将计算和参数解耦，但在推理时，较小的batchsize就会激活全部专家，导致访存急剧上升，进而使推理延迟大幅增加。
为此，字节跳动豆包大模型Foundation团队提出UltraMem，一种同样将计算和参数解耦的稀疏模型架构，在保证模型效果的前提下解决了推理的访存问题。
实验结果表明，在参数和激活条件相同的情况下，UltraMem在模型效果上超越了MoE，并将推理速度提升了2-6倍。此外，在常见batchsize规模下，UltraMem的访存成本几乎与同计算量的Dense模型相当。
目前，该论文已被ICLR2025接收。
Ultra-SparseMemoryNetwork
论文链接：https://arxiv.org/abs/2411.12364
1.MoE与PKM的局限性
LLM的能力增长需要指数级增长的计算资源，这在实时应用等资源有限的环境中颇具挑战。为了解决计算问题，先前的研究者提出了MoE和ProductKeyMemory（PKM）方案，但它们都有各自的局限性。
MoE通过稀疏激活expert解耦了计算和参数，但在推理场景中，速度其实很慢。原因在于，模型在推理时只能一个字一个字的生成，因此batchsize和sequencelength都很小，在这个场景下，MoE的所有专家通常会被全部访问到，极其容易遇到访存瓶颈，导致推理延迟激增。
PKM最早提出largememorylayer，其中包含了数量庞大的稀疏参数value，这里value其实就是一个向量，每个token会根据一个「行路由」和一个「列路由」定位到得分最高的几个value，激活这些value后做weightedsumpooling作为memorylayer的输出。这种方法因为每个token在推理时仅仅只激活极少数的value，所以推理时不会遇到访存瓶颈，但其效果很差，且scaling能力差。
2.UltraMem兼顾访存和效果
UltraMem参考了PKM的设计，但针对PKM的3个缺陷予以补充，以实现更高效的访存、更优质的value检索，同时，降低了显存和部署成本。
(1)优化模型结构
PKM的设计中，memorylayer只有1层，插在整个Transformer的中间层，这对大规模训练并不友好，并且如此庞大的稀疏参数应该尽可能多的参与到每次的残差连接中。
因此，研究团队拆分出多个小memorylayer，以固定的间隔分布在transformerlayer中；并且增加了skip-layer的操作，即当前层的memorylayer的输出会加到后面某层transformerlayer的输出。这使得模型可以并行地执行memorylayer的访存操作和transformerlayer的计算。
(2)优化value检索方式
在检索时，只有score最高的m个value会被激活，PKM的score是通过「行score」+「列score」得到的。团队进一步探索了一种更复杂的乘法方法TuckerDecomposedQuery-KeyRetrieval（TDQKR）。这一方法受启发于TuckerDecomposition。具体来看，给定values，shape为（n，n，h），其中h为hiddensize，那么values的scoreS_grid可以做如下分解：
其中
是可学习的tuckercore。这个结构下，每个value的score由r个行score和r个列score的组合乘加获得，具备更高的复杂度。
(3)隐式扩展稀疏参数
更多的稀疏参数通常会带来更好的效果，但过多的参数又会给显存和部署带来麻烦。为此，研究团队提出了ImplicitValueExpansion（IVE）方法隐式地扩展稀疏参数，并引入了virtualmemory和physicalmemory的概念。
以4倍扩展为例（如下图所示），virtualmemory的数量是physicalmemory的4倍，给定多对（score，index）后，首先按照virtualmemoryaddresstable做查表，4个virtualblock会查询同一个physicalmemorytable，之后各自做weightedsumpooling，并经过不同的线性层，最后再求和输出。
由于最后的Linear和取value之间没有任何非线性操作，因此每个Linear都可以和physicalmemorytable做融合，生成一个全新的memorytable，这个例子下，实际上隐式扩展了4倍的value数量。
3.实验结果：推理速度较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和上文提出的结构改进，最终能拿到C4validationloss-0.092的显著收益，同时稀疏参数和计算量几乎不变。
综上所述，研究团队提出的UltraMem具有极小的访存，因此，相比MoE实现了最高达6倍的速度提升，推理成本最高可降低83%。同时，在性能方面，随着模型容量的增加，在相同的参数和计算量情况下，UltraMem超过了MoE，表明其具有更强的扩展能力。这项工作为开发更高效和可扩展的语言模型提供了一个有希望的方向。
4.写在最后
UltraMem能有效地应用于对延迟要求较高的推理场景（例如代码补全），避免了类似MoE的访存瓶颈。即使在通用场景下，UltraMem相较于MoE都展现出了显著的速度优势，除非batchsize上万的极端场景。
当前，针对UltraMem的技术演进仍存在若干值得探索的方向，包括但不限于：如何高效优化稀疏参数、如何提升稀疏模型推理能力、如何更优地激活稀疏参数等。这些技术方向或将成为后续研究的重要切入点。
本文源自：财联社
