Block Transformer：通过全局到局部的语言建模加速LLM推理

在基于transformer的自回归语言模型（LMs）中，生成令牌的成本很高，这是因为自注意力机制需要关注所有之前的令牌，通常通过在自回归解码过程中缓存所有令牌的键值（KV）状态来解决这个问题。但是，加载所有先前令牌的KV状态以计算自注意力分数则占据了LMs的推理的大部分成本。

在这篇论文中，作者提出了Block Transformer架构，该架构通过在较低层次之间的粗糙块（每个块代表多个令牌）的自注意力来模拟全局依赖性，并在较高层次的每个局部块内解码细粒度的令牌，如下图所示。

论文的主要贡献包括：

发现了在自回归变换器中，全局和局部建模在推理时的核心作用和好处，特别是局部模块的重要性。
利用这些见解可以优化架构中的推理吞吐量，与普通transformers相比，显著提高了性能与吞吐量

Block Transformer

Block Transformer包括三个组成部分：

嵌入器：嵌入器将每个LB令牌的块聚合成一个输入块嵌入。
块解码器：块解码器对整个块序列应用自注意力以模拟全局依赖关系。
令牌解码器：令牌解码器在每个块内应用自注意力以处理细粒度的局部依赖性并解码个别令牌。

为什么Block Transformer高效？

全局到局部的方法可以通过将全局建模的昂贵瓶颈隔离到较低层并在上层的独立块内进行局部建模，这样可以减轻检索先前KV缓存的延迟和内存开销。
粗粒度的全局建模（块级解码）通过块长度因子缓解了KV缓存的瓶颈，同时保持了考虑完整上下文的能力。局部解码几乎没有预填充的成本，并且几乎消除了KV缓存开销，因此在推理硬件上的利用率中受益。
令牌解码器可以用更多的FLOPs进行细粒度的语言建模，对推理吞吐量的影响最小。
虽然block transformer需要比普通transformer更多的参数以保持可比的性能，但实际的吞吐量瓶颈是KV缓存开销，并且仍然可以实现更高的速度提升。

嵌入器 Embedder

嵌入器优先考虑简单性，主要处理小块长度（2-8），使用查找表Eemb∈RV×Demb来检索和连接可训练的令牌嵌入，其中令牌嵌入维度Demb设置为D/LB，D是整个网络中使用的块表示维度。

块解码器 Block decoder

块解码器的目标是通过关注前面的块来使块表示具有上下文性，利用嵌入器的输出作为输入。这种自回归transformer在块级别操作，产生输出块嵌入（也称为上下文嵌入），使得令牌解码器能够自回归地解码后续块的令牌内容。从嵌入器得到的输入块嵌入，源自输入令牌x0:(i×LB−1)，块解码器输出一个上下文嵌入，包含预测x(i×LB)😦(i+1)×LB−1)所需的信息。这种方法通过使用粗粒度块输入而不是单个令牌，减轻了自注意力的二次成本，从而减少了给定序列的上下文长度，同时保持了全局建模能力和硬件加速密集注意力的便利性。

令牌解码器 Token decoder

令牌解码器使用来自上下文块嵌入的全局上下文信息局部解码下一个块的个别令牌。令牌解码器也是一个标准的自回归transformer，具有自己的嵌入表Etok∈RV×Dtok和分类器。令牌解码器消除了预填充（仅在块解码器中必需），因为上下文信息由输出块嵌入提供，因此称之为上下文嵌入。KV缓存IO，批量解码期间的一个主要瓶颈，几乎被消除。与普通transformer相比，因为与完整上下文长度的成本是线性的，而普通注意力的KV缓存IO与完整上下文长度是二次的，因此计算单元的利用率更高。