沒人警告你的那個快取
模型生成文字時,一次產生一個 token,而每產生一個新 token,都必須回頭注意每一個先前的 token。每步重算全部會無望地慢,所以模型保留一個 KV 快取:每個過去 token 的鍵與值存在記憶體裡重複使用。這讓生成變快——但快取隨每個 token 成長,在長對話裡,塞不進 GPU 的東西變成是它、而非模型權重。
KV 缓存占用的内存随每个因子同时增长——层数、注意力头数、每头维度、序列长度和批量——所以削减头数才是关键。
多查詢:所有頭共用一組鍵
標準的多頭注意力給每個頭自己的查詢、鍵、值——比方 32 個頭,就有 32 組獨立的鍵/值要快取。多查詢注意力(MQA)保留 32 個查詢頭,但讓它們全部共用單一鍵頭與單一值頭。快取一夜縮小 32 倍,生成也變快,因為每步要從記憶體讀的少得多。
多头注意力示意图:多个并行的头分别产生输出,再拼接并混合。
代價是:把所有鍵與值塌縮成一個頭,丟掉了表達空間,純 MQA 模型品質會掉一點,訓練也可能變得不穩。這是個正確、但稍微走過頭的想法。
分組查詢:甜蜜點
分組查詢注意力(GQA)是如今人人都採用的折衷。不是 32 個鍵/值頭(完整多頭)也不是 1 個(MQA),而是用少數幾個——比方 8 個——讓每個鍵/值頭由一組 4 個查詢頭共用。你幾乎拿回完整注意力的全部品質,同時把快取砍 4 倍。它是 Llama、Mistral、Qwen 與多數現行開放模型的預設,因為它正好坐在品質對記憶體的甜蜜點上。
注意力示意图:查询通过 softmax 与各个键匹配,从而对值加权。
# config knobs you will actually see num_attention_heads: 32 # query heads num_key_value_heads: 8 # GQA: 32/8 = 4 query heads share each KV head # num_key_value_heads == 32 -> plain multi-head # num_key_value_heads == 1 -> multi-query
滑動視窗:別什麼都看
GQA 縮小每個 token 的快取,但快取仍隨序列長度成長。滑動視窗注意力直接攻擊長度:不讓每個 token 回頭注意整段歷史,而是限制在最近的 W 個 token——一個固定視窗,比方 4096。一旦超過視窗,快取就停止成長,因為舊的鍵可以丟掉。
交互式自注意力:选中一个词元会高亮它所关注的其他词元。
這不會丟掉長距離資訊嗎?比你想的少。疊很多視窗層,有效觸及範圍會累積:一個 token 可以透過一連串重疊視窗,影響到 4096 個之外的 token,就像深層卷積網路用小濾波器看見大圖。許多模型把視窗層與完整注意力層交錯,以取得便宜的局部性又偶爾的全域觸及——這是最後一篇長上下文架構的初嘗。