豆包大模型团队 投稿
量子位 | 公众号 QbitAI
字节
出了个全新架构，把
推理成本
给狠狠地打了下去！
有多狠？
推理速度相比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可以做如下分解：
其中S
row
,S
col
∈R
r×n
，C∈R
r×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总参数的稀疏模型上进行了全面的消融实验。
