专栏名称: 字节跳动技术团队
字节跳动的技术实践分享
目录
相关文章推荐
字节跳动技术团队  ·  火山引擎:单机部署 DeepSeek-R1 ... ·  23 小时前  
美团技术团队  ·  可信实验白皮书系列04:随机轮转实验 ·  4 天前  
美团技术团队  ·  可信实验白皮书系列03:随机对照实验 ·  4 天前  
51好读  ›  专栏  ›  字节跳动技术团队

Seed Research | 全新超稀疏架构,推理成本较 MoE 最高可降 83%!

字节跳动技术团队  · 公众号  · 架构  · 2025-02-14 18:00

正文

请到「今天看啥」查看全文


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







请到「今天看啥」查看全文