LLM架构的最新发展：KV共享、mHC与压缩注意力 [P]

# LLM架构的最新进展：KV共享、mHC与压缩注意力 来源：https://magazine.sebastianraschka.com/p/recent-developments-in-llm-architectures 短暂的家庭假期之后，我很兴奋地回归，并梳理了过去几周忙碌的开源权重LLM发布动态。让我印象最深刻的是，新的架构多么注重长上下文效率。 随着推理模型和智能体工作流保留更多Token（且时间更久），KV缓存大小、内存流量和注意力成本迅速成为主要瓶颈，LLM开发者正在增加越来越多的架构技巧来降低这些成本。 我想重点研究的主要例子包括：Gemma 4中的KV共享和逐层嵌入、Laguna XS.2中的逐层注意力预算分配、ZAYA1-8B中的压缩卷积注意力，以及DeepSeek V4中的mHC加压缩注意力。 在我的架构图中，这些变化大多看起来像是小调整，但其中一些是相当复杂的设计变更，值得深入讨论。 [图片说明：图1. 近期主要开源权重LLM发布的架构图（4月至5月）。图片及更多详情请访问我的LLM架构画廊。图中未显示所有模型大小；Qwen3.6包含27B和35B-A3B变体，ZAYA1以8B模型为代表（省略了ZAYA1-base和ZAYA1-reasoning-base）。虚线框中的架构将在本文中详细讨论。] 请注意，本文主要关注架构设计，因此我会跳过数据集混合、训练计划、后训练细节、RL配方、基准测试表和产品比较等内容。即便范围如此缩小，仍有大量内容需要覆盖。而且，和往常一样，文章最终比预期更长，所以我会重点讨论Transformer块、残差流、KV缓存或注意力计算中发生的变化。 另外请注意，我只涉及那些有趣（新颖）的设计选择，且尚未在其他地方介绍过的内容。包括： 1. Gemma 4中的KV共享和逐层嵌入 2. ZAYA1中的压缩卷积注意力 3. Laguna XS.2中的注意力预算分配 4. DeepSeek V4中的mHC和压缩注意力 在进入新部分之前，这里有两篇我会引用的早期文章。第一篇提供了关于近期MoE模型、路由专家、激活参数和模型大小比较的更广泛架构背景。第二篇涵盖了下面反复出现的注意力背景知识，包括MHA、MQA、GQA、MLA、滑动窗口注意力、稀疏注意力和混合注意力设计。 ## 现代LLM注意力变体的视觉指南 （链接：https://magazine.sebastianraschka.com/p/visual-attention-variants） 我原本计划撰写关于DeepSeek V4的文章。由于它尚未发布，我便利用时间做了件一直想做的事：收集、整理并细化我过去几年介绍过的各种LLM架构。 我还将其中一些解释整理成简短的独立教程页面，放在LLM架构画廊中。例如，读者可以找到关于GQA、MLA、滑动窗口注意力、DeepSeek稀疏注意力、MoE路由等概念的简洁解释，这些解释从对应的模型卡和概念标签处链接。 在这趟架构进步与调整之旅中，我们将回到四月初，当时谷歌发布了全新的开源权重Gemma 4模型系列。它们大致分为三类： - 用于移动设备和小型本地（嵌入式）设备（即物联网）的Gemma 4 E2B和E4B模型， - 针对高效本地推理优化的Gemma 4 26B混合专家（MoE）模型， - 以及追求最高质量和更方便后训练（因为MoE更难处理）的Gemma 4 31B密集模型。 [图片说明：图2. Gemma 4架构图。] E2B和E4B变体中的第一个小的架构调整是采用了共享KV缓存方案，其中靠后的层重用早期层的键值状态，以减少长上下文内存和计算。 这种KV共享并非Gemma 4首创。例如，参见Brandon等人的论文“Reducing Transformer Key-Value Cache Size with Cross-Layer Attention”（NeurIPS 2024）。但这是我看到的第一个应用此概念的流行架构。（注意：跨层注意力不要与交叉注意力混淆。） 在进一步解释KV共享之前，我们简单谈一下动机。正如我在最近几个月撰写和讨论的，LLM架构设计的最新主题之一就是KV缓存大小缩减。而KV缓存大小缩减的动机是减少所需内存，从而能处理更长的上下文，这在推理模型和智能体时代尤为重要。关于KV缓存的更多背景知识，请参见我的文章《从零理解并编码LLM中的KV缓存》： 从零理解并编码LLM中的KV缓存 （链接：https://magazine.sebastianraschka.com/p/coding-the-kv-cache-in-llms） 实际上，我在之前的文章《现代LLM注意力变体的视觉指南》中描述的所有流行注意力变体，都是为了减少KV缓存大小： 现代LLM注意力变体的视觉指南 （链接：https://magazine.sebastianraschka.com/p/visual-attention-variants） 举一个经典例子（Gemma 4仍在使用的）：分组查询注意力（GQA）已经在不同查询头之间共享键值（KV）头以减少KV缓存大小，如下图所示。 [图片说明：图3. 分组查询注意力（GQA）在多个查询头（Q）之间共享相同的键（K）和值（V）头。] 如前所述，Gemma 4使用了GQA。然而，除了作为GQA一部分的查询间KV共享外，Gemma 4还跨不同层共享KV投影，而不是在每个层的注意力模块中计算它们。这种KV共享方案，也称为跨层注意力，如下图所示。 [图片说明：图4. 常规Transformer块在每个注意力模块中分别计算Q、K、V投影（左）。跨层注意力设计（右）在多个层之间共享相同的K和V投影。] 如图2的架构概览中简要提到的，Gemma 4 E2B以4:1的模式使用常规GQA和滑动窗口注意力。（更准确地说，Gemma 4 E2B使用MQA，即GQA的单个KV头特例。） 在GQA（或MQA）的情况下，KV共享的工作原理如下：靠后的层不再计算自己的键和值投影，而是重用最近一个相同注意力类型的非共享较早层的KV张量。换句话说，滑动窗口层与之前的滑动窗口层共享KV；全注意力层与之前的全注意力层共享KV。各层仍然计算自己的查询投影，因此每个层可以形成自己的注意力模式，但昂贵且占用大量内存的KV缓存在多个层之间重用。 例如，Gemma 4 E2B有35个Transformer层，但只有前15层计算自己的KV投影；最后20层重用最近一个相同注意力类型的非共享较早层的KV张量。类似地，Gemma 4 E4B有42层，其中24层计算自己的KV，最后18层共享它们。 这实际上能节省多少？由于我们跨层共享了大约一半的KV，因此节省了大约一半的KV缓存大小。对于最小的E2B模型，在长128K上下文的情况下，这节省了2.7 GB（以bfloat16精度），如下所示。（对于E4B变体，在128K时节省约6 GB。） [图片说明：图5. Gemma 4 E2B类似设置中，来自GQA和跨层KV共享的KV缓存内存节省。为简化，未显示滑动窗口注意力的额外节省。] 当然，KV共享的缺点是它是对真实情况的“近似”。或者更准确地说，它降低了模型容量。然而，根据跨层注意力论文，影响可以很小（对于所测试的小模型而言）。 Gemma 4 E2B和E4B变体还包含第二个面向效率的设计选择，称为逐层嵌入（PLE）。这与上述KV共享方案是分开的。 KV共享减少KV缓存。而PLE则关乎参数效率，它让小型Gemma 4模型能够使用更多的令牌特定信息，而无需将主Transformer堆栈变得和具有相同总参数量的密集模型一样昂贵。 例如，Gemma 4 E2B和E4B中的“E”代表“effective”（有效）。具体来说，Gemma 4 E2B被列为23亿有效参数，或者当嵌入被计入时是51亿参数。（类似地，Gemma 4 E4B被列为45亿有效参数，或者含嵌入时是80亿参数。） 简而言之，在“E”模型中，主Transformer堆栈的计算量接近较小的数字，而较大的数字包括额外的嵌入表层的参数。（有关嵌入层如何工作的说明，请参阅我的代码笔记本《理解嵌入层与线性层的区别》。） 从概念上讲，新的PLE路径如下所示： [图片说明：图6. 简化后的Gemma 4块，包含PLE残差路径。常规块首先计算注意力和前馈残差更新。得到的隐藏状态对层特定的PLE向量进行门控，然后投影后的PLE更新在块末尾作为额外的残差更新添加。] PLE向量本身是在重复的Transformer块之外准备的。简化来说，PLE构建有两个输入。首先，令牌ID通过一个逐层嵌入查找。其次，常规令牌嵌入通过一个线性投影进入相同压缩的PLE空间。这两部分相加、缩放并重塑为一个每层一个切片的张量。注意，每个块随后接收自己的切片。 [图片说明：图7. 简化的PLE构建。令牌ID提供逐层嵌入查找，而常规令牌嵌入被投影到相同空间。两个贡献组合在一起并重塑，使每个Transformer块收到自己特定层的PLE切片。] 重要的细节是，PLE并没有给每个Transformer块一个完整的、独立的常规令牌嵌入层副本。相反，逐层嵌入查找只计算一次。然后，如前所述，它给每一层一个小的令牌特定嵌入切片（通过“reshape / select layer l”）。 因此，对于每个输入令牌，Gemma 4准备一个压缩的PLE张量，其中包含每个解码器层的一个小向量。然后，在正向传播期间，第l层只接收它自己的切片（图6中Gemma4WithPLEBlock中的ple_l）。 在Transformer块内部，常规的注意力和前馈分支照常运行。首先，块计算注意力残差更新。然后计算前馈残差更新。在第二次残差相加之后，得到的隐藏状态（我在图6的伪代码中记为*z*）用于门控层特定的PLE向量。门控后的PLE向量被投影回模型隐藏大小、归一化，并作为额外的残差更新添加。 因此，有用的心智模型是：Transformer块仍然具有相同的主要注意力和前馈路径，但Gemma 4在前馈分支之后添加了一个小的层特定令牌向量。这通过嵌入参数和小型投影增加了表示能力。这增加了计算开销，但避免了将整个Transformer堆栈扩展到更大参数量的成本。 但为什么选择PLE？更简单的替代方案是让密集模型更小，使用更少的层、更窄的隐藏状态或更小的前馈网络。那会减少内存和延迟，但它也移除了模型执行主要计算部分的容量。 PLE设计将昂贵的Transformer块保持在较小的“有效”大小附近，同时将额外容量存储在逐层嵌入表中。这些使用起来比添加更多注意力或FFN权重便宜得多，因为它们主要是可以缓存的查找式参数。 而且，我们只能相信谷歌的说法，即这是一个有效且值得的设计选择。如果能有一些比较研究，看看这种E2B设计如何与常规的Gemma 4 2.3B模型和常规的Gemma 4 5.1B模型进行比较，那将会很有趣。 另外，原则上，PLE并非固有地仅限于小模型。我们也可以将逐层嵌入切片附加到更大的模型上。然而，较大的模型已经具有足够的容量，这些额外的嵌入可能帮助不大。此外，对于较大的模型，我们已经使用MoE设计作为一种技巧，在保持较小计算占用的同时增加容量。 顺便提一下，如果你对相对简单且可读性强的代码实现感兴趣，我在这里实现了Gemma 4 E2B和E4B模型的从零实现。 Laguna是poolside公司发布的第一个开源权重模型，poolside是一家专注于训练编码应用LLM的欧洲公司。我以前的几位同事近年加入了poolside，他们拥有一支非常有才华的团队。很高兴看到更多公司也以开源权重变体形式发布他们的一些模型。 总之，下面描绘的Laguna XS.2架构乍一看非常标准。然而，有一个我没有展示（/试图硬塞进去）的细节，我们可以称其为“逐层注意力预算分配”。