KAIST-AI | 提出Block Transformer架构，大幅提升推理速度和内存效率，20倍增益！

ShuYini

发布于 2024-06-11 19:31:58

980

发布于 2024-06-11 19:31:58

引言

Transformer模型虽然在NLP领域取得了巨大成功，但其Self-Attention机制在处理长序列时会导致计算和内存需求急剧增加，这限制了其在资源受限环境中的实用性。为此，本文作者提出了Block Transformer架构，通过分层的全局到局部建模方法，有效地平衡了全局上下文的捕获和局部依赖关系，减少了推理过程中的内存访问和计算需求，从而实现了高效的语言模型推理。实验结果表明，与一般的Transformer相比，在推理吐量上最高可实现了20倍的增益。

https://arxiv.org/pdf/2406.02657

背景介绍

基于Transformer的自回归语言模型在生成Tokens时会面临内存计算成本问题，其主要原因是Self-Attention机制需要关注所有先前Tokens，这会导致其时间复杂度和内存需求随着序列长度的增加而显著增加。为了能够减少Attention的计算成本，通常会在自回归解码期间缓存所有层上所有 token 的键值 (KV) 状态。然而，尽管每个解码步骤仅计算单个 token 的 KV 状态，但它仍需加载所有先前 token 的 KV 状态才能计算Self-Attention分数。

由此可以看处，KV缓存的输入/输出（IO）在推理成本中占据主要部分，虽然当前的研究已经提出了多种降低Attention组件推理成本的方法，但开发基于Transformer且能够避免Attention开销的语言模型架构，仍旧是一项艰难的挑战。

面对以上挑战，研究者们开始探索新的模型架构，以更有效地处理全局和局部信息。分层的全局到局部（global-to-local）建模方法在其他领域已经显示出其有效性，例如计算机视觉中的层次结构模型。受该架构方法的启发，本文作者提出了Block Transformer 架构，大大提高模型推理速度和内存效率，同时保证了模型性能。

Block Transformer架构

Block Transformer架构（如下图所示），它通过分层范式，将全局和局部注意力机制应用于两个不同的阶段，以分离对完整上下文的理解和详细交互的处理。具体来说：

在较低层，全局上下文以粗粒度的块级粒度被捕捉，每个块由固定数量的标记聚合成单一嵌入。
在上层，通过仅关注来自块解码器的上下文块嵌入，以自回归方式解析多个子词标记的局部依赖关系。

如上图，Block Transformer架构主要由三部分组成，它们分别为：Embedder、Block Decoder、Token Decoder。

「Embedder」 是一个轻量级模块，主要作用是将输入的一系列token（词汇）转换成一个固定大小的输入块（input block），并生成对应的块嵌入（ block embedding）。这些嵌入随后成为Block Decoder（块解码器）的输入单元。Embedder的设计强调了简单性，因为它处理的是小的块长度（在研究中通常是2到8个token）。

「Block Decoder」 块解码器是一种自回归Transformer，其主要负责处理输入嵌入块的全局依赖关系。该模块使用自注意力机制来处理Embedder输出的块嵌入（ block embedding）,并且它只在块级别上进行操作，而不是单个token级别上，进而减少了计算量。通过自注意力它还可以捕捉不同块之间的依赖关系，从而建模全局上下文信息。Block Decoder生成的上下文块嵌入随后被传递给Token Decoder，作为其解码下一个token块的全局上下文信息。

「Token Decoder」 其主要作用是在局部块内解码细粒度的token。Token Decoder使用来自Block Decoder的上下文块嵌入作为全局上下文信息的唯一来源。这种设计显著减少了对全局上下文的计算和存储需求。Token Decoder在解码过程中是自回归的，即每个新token的解码依赖于之前已解码的token序列。Token Decoder的输入是上下文块嵌入和当前块内的token序列，输出是下一个token的预测。由于Token Decoder仅在局部块内应用自注意力，因此避免了对过去所有token的键值（KV）缓存的计算、存储和检索，从而减少了计算和内存访问成本。