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架構如何從簡單縮放轉向更智能的資源分配，以實現更好的長上下文性能。