字节出了个全新架构，把推理成本给狠狠地打了下去！
有多狠？
推理速度相比 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，表明其具有更强的扩展能力。
这项工作可以说是为开发更高效和可扩展的语言模型提供了一个有希望的方向。
论文地址：
arxiv.org/abs/2411.12…
— 完 —
