量子位

电子技术应用专栏作家——量子位。追踪人工智能新趋势,关注科技行业新突破!

Transformer后继有模!MSRA提出全新大模型基础架构:推理速度8倍提升,内存占用减少70%

0
阅读(1343)

  鱼羊 发自 凹非寺

  量子位 | 公众号 QbitAI

  微软大模型新架构,正式向Transformer发起挑战!

  论文标题明晃晃地写道:

  Retentive Network(RetNet):大模型领域Transformer的继任者。

  image.png

  论文提出新的Retention机制来代替Attention。来自微软亚研院和清华的研究人员,毫不讳言“野心”,大胆放话:

  RetNet实现了良好的扩展结果、并行训练、低成本部署和高效推理。

  这些特性使这一基础架构,成为大语言模型中Transformer的有力继承者。

  而实验数据也显示,在语言建模任务上:

  • RetNet可以达到与Transformer相当的困惑度(perplexity)

  • 推理速度达8.4倍

  • 内存占用减少70%

  • 具有良好的扩展性

  并且当模型大小大于一定规模时,RetNet表现会优于Transformer。

 image.png

  Transformer果真“后继有模”了?具体详情,一起来看。

  解决“不可能三角”

  Transformer在大语言模型中的重要性毋庸置疑。无论是OpenAI的GPT系列,还是谷歌的PaLM、Meta的LLaMA,都是基于Transformer打造。

  但Transformer也并非完美无缺:其并行处理机制是以低效推理为代价的,每个步骤的复杂度为O(N);Transformer是内存密集型模型,序列越长,占用的内存越多。

  在此之前,大家也不是没想过继续改进Transformer。但主要的几种研究方向都有些顾此失彼:

  线性attention可以降低推理成本,但性能较差;

  循环神经网络则无法进行并行训练。

  也就是说,这些神经网络架构面前摆着一个“不可能三角”,三个角代表的分别是:并行训练、低成本推理和良好的扩展性能。

image.png

  RetNet的研究人员想做的,就是化不可能为可能。

  具体而言,RetNet在Transformer的基础上,使用多尺度保持(retention)机制替代了标准的自注意力机制。

  与标准自注意力机制相比,保持机制有几大特点:

  引入位置相关的指数衰减项取代softmax,简化了计算,同时使前步的信息以衰减的形式保留下来。

  引入复数空间表达位置信息,取代绝对或相对位置编码,容易转换为递归形式。

  另外,保持机制使用多尺度的衰减率,增加了模型的表达能力,并利用GroupNorm的缩放不变性来提高retention层的数值精度。

 image.png

  每个RetNet块包含两个模块:多尺度保持(MSR)模块和前馈网络(FFN)模块。

  保持机制支持以三种形式表示序列:

  • 并行

  • 递归

  • 分块递归,即并行表示和递归表示的混合形式,将输入序列划分为块,在块内按照并行表示进行计算,在块间遵循递归表示。

  其中,并行表示使RetNet可以像Transformer一样高效地利用GPU进行并行训练。

  递归表示实现了O(1)的推理复杂度,降低了内存占用和延迟。

  分块递归则可以更高效地处理长序列。

  这样一来,RetNet就使得“不可能三角”成为可能。以下为RetNet与其他基础架构的对比结果:

  image.png

  在语言建模任务上的实验结果,进一步证明了RetNet的有效性。

  结果显示,RetNet可以达到与Transformer相似的困惑度(PPL,评价语言模型好坏的指标,越小越好)。

  同时,在模型参数为70亿、输入序列长度为8k的情况下,RetNet的推理速度能达到Transformer的8.4,内存占用减少70%

  在训练过程中,RetNet在内存节省和加速效果方面,也比标准Transformer+FlashAttention表现更好,分别达到25-50%和7倍。

  值得一提的是,RetNet的推理成本与序列长度无关,推理延迟对批量大小不敏感,允许高吞吐量。

  image.png

  另外,当模型参数规模大于20亿时,RetNet的表现会优于Transformer。

  image.png

  研究团队

  RetNet的研究团队,来自微软亚研院和清华大学。

  共同一作为孙宇涛和董力。

  孙宇涛,清华大学计算机系本科,现在在微软亚研院实习。

  董力,微软亚研院研究员。他也是此前引发大量关注的“能记住10亿token的Transformer”的论文作者之一。

  image.png

image.png

  RetNet论文的通讯作者是韦福如。他是微软亚洲研究院全球研究合伙人,10亿token Transformer亦是来自他的研究团队。

  论文地址:

  https://arxiv.org/abs/2307.08621

原文链接:https://mp.weixin.qq.com/s/zJsb2vdpEwqXEUgqHmpInA

量子位.jpg

电子技术应用专栏作家  量子位