据豆包大模子团队讯息,字节开首豆包大模子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 的终端和推理速率奈何变化?
下图(b)展示了 UltraMem 的终端变化,横轴为寥落参数和粘稠参数的比值,每个神采的线代表了一种寥落度。寥落度界说为 value 的数目 / 每个 token 激活的 value 数目。不雅察发现,抓续加多寥落参数和 loss 的下跌呈对数磋磨;且寥落度越小,模子终端越好;可是寥落度抓续缩小带来的收益在慢慢足够。
下图(c)展示了 UltraMem 的推理时刻变化,横轴为寥落参数和粘稠参数的比值。不雅察发现,UltraMem 在抓续加多寥落参数时推理时刻的确不变,反不雅 MoE 有权贵增长的趋势。

② 消融实验
究诘团队在 151M 激活、1.5B 总参数的寥落模子上进行了全面的消融实验。从最原始的 PKM 驱动,慢慢加多一些 trick 和上文提议的结构纠正,最终能拿到 C4 validation loss -0.092 的权贵收益,同期寥落参数和蓄意量的确不变。

要而论之,究诘团队提议的 UltraMem 具有极小的访存,因此,比拟 MoE 完毕了最高达 6 倍的速率扶助,推理本钱最高可缩小 83%。同期,在性能方面,跟着模子容量的加多,在调换的参数和蓄意量情况下,UltraMem 超越了 MoE,标明其具有更强的扩展材干。这项使命为建造更高效和可扩展的道话模子提供了一个有但愿的标的。
4. 写在终末
UltraMem 能有用地诓骗于对延长要求较高的推理场景(举例代码补全),幸免了肖似 MoE 的访存瓶颈。即使在通用场景下,UltraMem 相较于 MoE 都展现出了权贵的速率上风,除非 batch size 上万的极点场景。
现时开云(中国)kaiyun网页版登录入口,针对 UltraMem 的工夫演进仍存在些许值得探索的标的,包括但不限于:奈何高效优化寥落参数、奈何扶助寥落模子推理材干、奈何更优地激活寥落参数等。这些工夫标的或将成为后续究诘的紧迫切入点。