字节跳动豆包大模型团队提出UltraMem架构：攻克大模型推理瓶颈，速度提升2-6倍

字节跳动豆包大模型团队近日提出了一种名为UltraMem的全新稀疏模型架构，旨在解决大模型推理过程中的高额访存问题和成本问题。

现有的大型语言模型(LLM)在规模扩大后，推理成本和访存效率成为制约其应用的关键瓶颈。虽然混合专家(MoE)架构通过解耦计算和参数来缓解这一问题，但在小批量(batch size)情况下，所有专家都会被激活，导致访存激增，推理速度变慢。

UltraMem架构借鉴了Product Key Memory(PKM)的设计思想，并对其进行了改进。它通过以下三个方面的优化来解决问题：

1. 优化模型结构: 将单层memory layer拆分成多个小型memory layer，并行执行访存操作和Transformer层计算，并增加了跳层连接(skip-layer)以提高效率。
2. 优化value检索方式: 采用更复杂的Tucker Decomposed Query-Key Retrieval(TDQKR)方法，提高value检索的精准度和效率。
3. 隐式扩展稀疏参数: 使用Implicit Value Expansion (IVE)方法，通过虚拟内存和物理内存的结合，在不显著增加显存和部署成本的情况下，隐式扩展稀疏参数数量，提升模型效果。

实验结果表明，UltraMem在保证模型效果甚至超越MoE的同时，推理速度提升了2-6倍，推理成本最高降低了83%。即使在小批量场景下，UltraMem的访存成本也与同计算量的稠密模型相当。该研究已发表在ICLR 2025，论文链接为https://arxiv.org/abs/2411.12364。UltraMem的出现为构建更大规模、更高效的LLM提供了新的可能性，尤其适用于对延迟敏感的应用场景，例如代码补全等。 未来研究方向包括如何更高效地优化稀疏参数、提升稀疏模型推理能力以及更优地激活稀疏参数等。