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

在短暂的家庭假期之后，我很高兴回来并迎接过去几周忙碌的开源LLM发布。最让我印象深刻的是，新架构越来越注重长上下文效率。

随着推理模型和智能体工作流保留越来越多的令牌（且时间更长），KV缓存大小、内存流量和注意力成本迅速成为主要瓶颈，LLM开发者正在添加越来越多的架构技巧来降低这些成本。

本文主要关注的例子包括：Gemma 4中的KV共享和逐层嵌入、Laguna XS.2中的逐层注意力预算、ZAYA1-8B中的压缩卷积注意力，以及DeepSeek V4中的mHC和压缩注意力。

这些变化中的大多数在我的架构图中看起来像是小调整，但有些是相当复杂的设计变更，值得更详细的讨论。

1. 跨层复用KV张量以缩小缓存（Gemma 4）

在架构进展之旅中，我们首先回顾4月初谷歌发布的新开源Gemma 4模型系列。它们分为三大类：用于移动和小型本地设备的Gemma 4 E2B和E4B模型、针对高效本地推理优化的26B混合专家（MoE）模型，以及追求最高质量和更方便后训练的31B密集模型。

E2B和E4B变体中的第一个小型架构调整是采用了共享KV缓存方案，其中后面的层复用前面层的键值状态，以减少长上下文的内存和计算。

这种KV共享并非Gemma 4首创（例如，见Brandon等人2024年的"Reducing Transformer Key-Value Cache Size with Cross-Layer Attention"），但这是我在流行架构中首次看到此概念应用。

在解释KV共享之前，先简要说明动机。正如我近几个月所写所谈，LLM架构设计的一个主要近期主题是KV缓存大小缩减。而KV缓存大小缩减的动机是减少所需内存，从而允许处理更长的上下文，这在推理模型和智能体时代尤为重要。

实际上，所有流行的注意力变体（如分组查询注意力GQA）都旨在减少KV缓存大小。Gemma 4使用GQA，但除此之外，它还跨不同层共享KV投影，而不是在每个层的注意力模块中计算。这种跨层注意力设计使后面的层不再计算自己的键值投影，而是复用最近一个同类型非共享层的KV张量。

例如，Gemma 4 E2B有35个变换器层，但只有前15个计算自己的KV投影；后20个层复用KV张量。这样大约节省了一半的KV缓存大小。对于最小的E2B模型，在128K长上下文下可节省2.7 GB（bfloat16精度）。

KV共享的缺点是对模型容量的一种"近似"，但根据论文，对小模型的影响可以很小。

2. 逐层嵌入（PLE）与"有效"大小（Gemma 4 E2B/E4B）

Gemma 4 E2B和E4B变体还包括第二个面向效率的设计选择：逐层嵌入（PLE）。这与KV共享方案不同：KV共享减少KV缓存，而PLE是关于参数效率，它让小型Gemma 4模型在不使主变换器堆栈像具有相同总参数量的密集模型那样昂贵的情况下，使用更多令牌特定信息。

例如，Gemma 4 E2B列为2.3B有效参数，若计入嵌入则为5.1B参数。这意味着主变换器堆栈的计算更接近较小的数字，而较大的数字包括额外的嵌入表层。

在PLE路径中，每个变换器块除了正常的注意力和前馈路径外，还会接收一个层特定的令牌向量。该向量通过令牌ID的逐层嵌入查找和正常令牌嵌入的线性投影组合而成，然后经过缩放和重塑，每个层获得自己的切片。在块内部，注意力和前馈分支正常运行，之后隐藏状态用于门控层特定PLE向量，门控后的PLE向量投影回模型隐藏大小，归一化后作为额外的残差更新添加。

PL的好处是：它增加了表示能力，但避免了将整个变换器堆栈扩展到更大参数量的成本。

3. 其他架构进展

文章还讨论了Laguna XS.2的逐层注意力预算，它根据层的重要性动态分配注意力头预算；ZAYA1的压缩卷积注意力，用卷积核近似注意力模式以减少计算；以及DeepSeek V4的mHC（多头压缩）和压缩注意力，进一步优化长上下文效率。

这些设计共同展示了LLM架构如何从简单缩放转向更智能的资源分配，以实现更好的长上下文性能。