新闻资讯

股票软件定制：字节豆包大模型团队提出稀疏模型架构 推理成本最高可降低83%

作者：金策略小编 发布时间：2025-02-12 浏览量：

股票软件定制：根据字节跳动豆包大模型 Foundation 团队的消息，近期他们提出了 UltraMem，这是一种将计算和参数解耦的稀疏模型架构。在保证模型效果的同时，它成功解决了推理过程中的访存问题。据悉，该架构有效攻克了 MoE 推理时高额的访存难题，相比 MoE 架构，推理速度提升了 2-6 倍，推理成本最高能降低 83%。
随着模型规模不断扩大，推理成本和访存效率成为了限制大模型广泛应用的关键阻碍。字节跳动豆包大模型团队适时推出全新的稀疏模型架构 UltraMem，成功解决了 MoE 推理时高额的访存问题，让推理速度比 MoE 架构快了 2-6 倍，推理成本最高可降低 83%。该研究还揭示了新架构的 Scaling Law，证实它不仅具备出色的 Scaling 特性，性能上也超过了 MoE。
实验结果显示，训练规模达 2000 万 value 的 UltraMem 模型，在相同计算资源条件下，能同时实现业界领先的推理速度和模型性能，为构建数十亿规模 value 或 expert 开辟了新途径。
在 Transformer 架构下，模型性能与参数数量和计算复杂度成对数关系。随着 LLM 规模持续增大，推理成本会急剧上升，速度也会变慢。
虽然 MoE 架构已成功将计算和参数解耦，但在推理时，较小的 batch size 就会激活全部专家，导致访存急剧增加，进而大幅延长推理延迟。
为此，字节跳动豆包大模型 Foundation 团队提出 UltraMem，这一同样将计算和参数解耦的稀疏模型架构，在保证模型效果的基础上解决了推理的访存问题。
实验表明，在参数和激活条件相同的情况下，UltraMem 在模型效果上超越了 MoE，推理速度提升了 2-6 倍。此外，在常见 batch size 规模下，UltraMem 的访存成本几乎和同计算量的 Dense 模型一样。
目前，相关论文已被 ICLR 2025 接收，论文题目为 “Ultra-Sparse Memory Network”，论文链接：https://arxiv.org/abs/2411.12364。

1. MoE 与 PKM 的局限性

LLM 能力的增长需要指数级增长的计算资源，这在实时应用等资源有限的环境中面临很大挑战。为解决计算问题，之前的研究者提出 MoE 和 Product Key Memory（PKM）方案，但它们都存在各自的局限。
MoE 通过稀疏激活 expert 解耦了计算和参数，但在推理场景中速度很慢。因为模型推理时只能逐字生成，batch size 和 sequence length 都很小，在这种情况下，MoE 的所有专家通常都会被访问到，极易遭遇访存瓶颈，致使推理延迟大幅增加。
PKM 最早提出 large memory layer，其中包含大量稀疏参数 value（value 是一个向量），每个 token 会依据一个 “行路由” 和一个 “列路由” 定位到得分最高的几个 value，激活这些 value 后进行 weighted sum pooling 作为 memory layer 的输出。这种方法在推理时每个 token 仅激活极少数 value，所以不会遇到访存瓶颈，但其效果不佳，且 scaling 能力差。

2. UltraMem 兼顾访存和效果

UltraMem 参考了 PKM 的设计，并针对 PKM 的 3 个缺陷进行改进，以实现更高效的访存、更优质的 value 检索，同时降低显存和部署成本。
  • 优化模型结构：PKM 的 memory layer 只有 1 层，插在整个 Transformer 的中间层，这不利于大规模训练，而且庞大的稀疏参数应尽可能多地参与每次残差连接。因此，研究团队拆分出多个小 memory layer，以固定间隔分布在 transformer layer 中，还增加了 skip-layer 操作，即当前层 memory layer 的输出会加到后面某层 transformer layer 的输出。这使得模型能并行执行 memory layer 的访存操作和 transformer layer 的计算。
  • 优化 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 可以进行如下分解：（此处插入公式部分，因格式问题暂未呈现，原文中有详细公式）其中（此处插入公式部分，因格式问题暂未呈现，原文中有详细公式）是可学习的 tucker core。在这个结构下，每个 value 的 score 由 r 个行 score 和 r 个列 score 的组合乘加获得，复杂度更高。
  • 隐式扩展稀疏参数：更多的稀疏参数通常能带来更好的效果，但过多参数会给显存和部署造成麻烦。为此，研究团队提出 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 数量。

3. 实验结果：推理速度较 MoE 最高提升 6 倍

  • 模型性能评估：研究团队在 151M、680M、1.6B 三个尺寸的激活参数上进行了大量实验，MoE、PKM 和 UltraMem 的总稀疏参数保持在激活参数的 12 倍。结果表明，UltraMem 在 680M、1.6B 上具有明显的效果优势。研究还探究了随着稀疏参数增加，UltraMem 的效果和推理速度的变化。结果显示，持续增加稀疏参数和 loss 的下降呈对数关系；稀疏度越小，模型效果越好，但稀疏度持续降低带来的收益逐渐饱和。而 UltraMem 在持续增加稀疏参数时推理时间几乎不变，MoE 却有显著增长的趋势。
  • 消融实验：研究团队在 151M 激活、1.5B 总参数的稀疏模型上进行了全面的消融实验。从最原始的 PKM 开始，逐步增加一些技巧和上述结构改进，最终能获得 C4 validation loss -0.092 的显著收益，同时稀疏参数和计算量几乎不变。
综上所述，研究团队提出的 UltraMem 访存极小，相比 MoE 实现了最高达 6 倍的速度提升，推理成本最高可降低 83%。在性能方面，随着模型容量增加，在相同的参数和计算量情况下，UltraMem 超过了 MoE，展现出更强的扩展能力。这项工作为开发更高效和可扩展的语言模型提供了一个有前景的方向。

4. 写在最后

UltraMem 能有效应用于对延迟要求较高的推理场景（如代码补全），避免了类似 MoE 的访存瓶颈。即便在通用场景下，UltraMem 相较于 MoE 也展现出显著的速度优势，除非是 batch size 上万的极端场景。
目前，针对 UltraMem 的技术演进仍有多个值得探索的方向，比如如何高效优化稀疏参数、如何提升稀疏模型推理能力、如何更优地激活稀疏参数等。这些技术方向可能会成为后续研究的重要切入点。