QKV的重要性

要理解大語(yǔ)言模型效果的底層實(shí)現(xiàn)原理，很大一部分就是理解Transformers Block里面的QKV矩陣。現(xiàn)在前沿的大模型研究工作很大一部分都是圍繞著QKV矩陣去做的，比如注意力、量化、低秩壓縮等等。

其本質(zhì)原因是因?yàn)镼KV權(quán)重占比著大語(yǔ)言模型50%以上的權(quán)重比例，在推理過(guò)程中，QKV存儲(chǔ)量還會(huì)隨著上下文長(zhǎng)度的增長(zhǎng)而線性增長(zhǎng)，計(jì)算量也平方增加。

可以說(shuō)，現(xiàn)在大模型對(duì)外銷(xiāo)售OpenAPI的價(jià)格戰(zhàn)，很大一部分就是對(duì)QKV極致優(yōu)化的技術(shù)戰(zhàn)；圍繞QKV，各大模型廠商使出渾身解數(shù)，在保證效果不變壞的前提下，主要的研究工作就是對(duì)性能和存儲(chǔ)的極致壓縮。

在本文，我們將從純概念的角度，增進(jìn)對(duì)QKV的理解。

QKV在哪里

還是這張圖，QKV矩陣是Transformers Block MHA的重要組成部分，當(dāng)一個(gè)文本Token過(guò)來(lái)的時(shí)候，文本Token的Normalization矩陣會(huì)分別和Q、K、V的權(quán)重矩陣進(jìn)行矩陣乘法操作，得到三個(gè)不同的矩陣，以進(jìn)入后面的計(jì)算過(guò)程。

圖片

以7B模型為例，計(jì)算數(shù)據(jù)的Shape如下：

Token Normalization Shape：[1, 4096]
Q、K、V Shape：[4096, 4096]

如果考慮一次推理的BatchSize和SeqLen（上下文長(zhǎng)度），那就是這個(gè)Shape：

Token Normalization：[BatchSize, SeqLen, 4096]
Q、K、V Shape：[BatchSize, SeqLen, 4096, 4096]
其中，4096就是7B模型對(duì)應(yīng)的Hidden Size。

從Shape就可以看出，Q、K、V的計(jì)算量和存儲(chǔ)量會(huì)隨著SeqLen的增長(zhǎng)而增長(zhǎng)。

QKV的作用

從計(jì)算角度理解QKV始終不容易理解它的算法效果，所以本部分從概念的角度理解QKV的作用。

在大型語(yǔ)言模型（LLM）中，特別是在基于Transformer架構(gòu)的模型中，Q（Query）、K（Key）和V（Value）是自注意力（Self-Attention）機(jī)制中的核心組成部分。它們?cè)谕评磉^(guò)程中的作用如下：

Query (Q)：

角色：Query向量代表當(dāng)前正在處理的token或位置，它表示模型需要“查詢(xún)”的信息。

作用：在自注意力機(jī)制中，Query用于與所有的Key進(jìn)行比較，以確定每個(gè)Key與當(dāng)前token的相關(guān)性。這個(gè)比較的結(jié)果決定了Value的加權(quán)和，從而生成當(dāng)前token的輸出。

變化性：在自回歸推理過(guò)程中，每個(gè)新生成的token都會(huì)有一個(gè)新的Query向量，它只依賴(lài)于當(dāng)前token的信息。

Key (K)：

角色：Key向量代表序列中每個(gè)token的唯一標(biāo)識(shí)，用于與Query進(jìn)行比較。

作用：Key向量用于計(jì)算與Query的相似度或匹配程度，這個(gè)相似度得分決定了相應(yīng)Value在最終輸出中的權(quán)重。

穩(wěn)定性：在自回歸推理中，對(duì)于已經(jīng)生成的token，其Key向量在后續(xù)的推理過(guò)程中會(huì)被重復(fù)使用，因?yàn)樗鼈兇淼氖且呀?jīng)確定的token信息。

Value (V)：

角色：Value向量包含序列中每個(gè)token的實(shí)際內(nèi)容或特征，它對(duì)生成當(dāng)前token的輸出有貢獻(xiàn)。

作用：Value向量根據(jù)與Query的相似度得分（由Key確定）被加權(quán)求和，生成當(dāng)前token的輸出。

穩(wěn)定性：與Key類(lèi)似，對(duì)于已經(jīng)生成的token，其Value向量在后續(xù)的推理過(guò)程中也會(huì)被重復(fù)使用。

在自回歸推理過(guò)程中，模型一次生成一個(gè)token，并且每個(gè)新token都會(huì)基于之前所有token的信息。因此，對(duì)于每個(gè)新生成的token：

Q：需要重新計(jì)算，因?yàn)樗蕾?lài)于當(dāng)前token的信息。

K 和 V：可以被緩存（即KV Cache），因?yàn)樗鼈兇碇耙呀?jīng)生成的token的信息，這些信息在生成后續(xù)token時(shí)不需要重新計(jì)算。

總結(jié)來(lái)說(shuō)，Q、K、V在推理過(guò)程中共同工作，通過(guò)自注意力機(jī)制允許模型在生成每個(gè)新token時(shí)動(dòng)態(tài)地聚焦于序列中相關(guān)的信息。

Q代表了當(dāng)前token的查詢(xún)需求，而K和V則提供了序列中每個(gè)token的標(biāo)識(shí)和內(nèi)容，使得模型能夠根據(jù)當(dāng)前token的需求加權(quán)組合之前的信息，生成連貫和相關(guān)的輸出。

為什么需要QKV，不能直接從Q得到輸出呢？

在Transformer模型中，區(qū)分Q（Query）、K（Key）和V（Value）的原因主要有以下幾點(diǎn)：

角色分離：

Q、K、V的設(shè)計(jì)允許模型在處理序列數(shù)據(jù)時(shí)，更有效地捕捉不同位置之間的關(guān)系。Q代表當(dāng)前位置希望獲得的信息，K代表序列中各位置能提供的信息，而V則代表當(dāng)查詢(xún)和鍵匹配時(shí)，應(yīng)該從各位置獲取的實(shí)際內(nèi)容。

通過(guò)這種方式，模型可以靈活地捕捉不同位置之間的依賴(lài)關(guān)系。

增強(qiáng)模型能力：

使用獨(dú)立的Q、K、V允許模型學(xué)習(xí)到更復(fù)雜的關(guān)系。每個(gè)位置不僅可以決定與其他位置的關(guān)聯(lián)程度（通過(guò)Q和K），還可以決定從每個(gè)位置獲取什么樣的信息（通過(guò)V）。

靈活性和表達(dá)能力：

獨(dú)立的Q、K、V矩陣增加了模型的靈活性和表達(dá)能力。模型可以學(xué)習(xí)到如何根據(jù)上下文將重點(diǎn)放在不同的信息上。

并行處理：

Transformer模型的設(shè)計(jì)允許在處理序列時(shí)進(jìn)行高效的并行計(jì)算。Q、K、V的獨(dú)立使得模型可以同時(shí)計(jì)算整個(gè)序列中所有位置的注意力分?jǐn)?shù)，這大大提高了計(jì)算效率。

信息檢索類(lèi)比：

可以把Q、K、V機(jī)制類(lèi)比為一個(gè)信息檢索過(guò)程。

Q類(lèi)似我們提出的搜索問(wèn)題，目的是從大量信息中尋找相關(guān)答案；K類(lèi)似信息庫(kù)中的索引，它們決定哪些信息與查詢(xún)相關(guān)；V類(lèi)似實(shí)際的內(nèi)容，是查詢(xún)找到相關(guān)信息后的返回結(jié)果。

Q用來(lái)提出問(wèn)題，K用來(lái)匹配相關(guān)性，V則是我們最終希望獲取的信息。

綜上所述，QKV的區(qū)分使得Transformer模型能夠更加靈活和有效地處理序列數(shù)據(jù)，捕捉復(fù)雜的依賴(lài)關(guān)系，并適應(yīng)不同的任務(wù)需求。

直接從Q得到V會(huì)限制模型的表達(dá)能力和靈活性，因?yàn)樗雎粤送ㄟ^(guò)K來(lái)確定相關(guān)性的重要性，并且減少了模型處理信息的靈活性。

在LLM推理時(shí)，為什么KV可以Cache？

在LLM（Large Language Model，大型語(yǔ)言模型）推理過(guò)程中，KV（Key-Value）Cache可以被緩存的原因主要基于以下幾個(gè)方面：

減少重復(fù)計(jì)算：

在自注意力機(jī)制中，如果沒(méi)有KV Cache，每次生成新token時(shí)，模型需要重新計(jì)算整個(gè)歷史序列的Key和Value向量，并參與注意力計(jì)算，這導(dǎo)致了大量的重復(fù)計(jì)算。通過(guò)緩存歷史序列的K和V，可以避免這種重復(fù)計(jì)算，顯著降低推理的計(jì)算復(fù)雜度。

提升推理速度：

KV Cache通過(guò)緩存Key和Value向量，使得模型在生成新token時(shí)只需計(jì)算當(dāng)前token的Query向量，并與緩存的Key和Value進(jìn)行注意力計(jì)算，這樣可以加快推理速度。

降低計(jì)算復(fù)雜度：

自注意力機(jī)制的計(jì)算復(fù)雜度為O(n^2?d)，其中n是序列長(zhǎng)度，d是向量維度。使用KV Cache后，計(jì)算復(fù)雜度可以降低到O(n?d)，顯著減少了計(jì)算量。

跨請(qǐng)求復(fù)用：

在某些場(chǎng)景下，多次請(qǐng)求的Prompt可能會(huì)共享同一個(gè)前綴（Prefix），這些情況下，很多請(qǐng)求的前綴的KV Cache計(jì)算結(jié)果是相同的，可以被緩存起來(lái)，給下一個(gè)請(qǐng)求復(fù)用。

綜上所述，KV Cache在LLM推理中通過(guò)緩存Key和Value向量，有效減少了重復(fù)計(jì)算，降低了計(jì)算復(fù)雜度，提升了推理速度，并且優(yōu)化了顯存資源的使用，從而提高了模型的推理效率和吞吐量。

那為什么Q不可以Cache？

因?yàn)镼不需要Cache…

在LLM（Large Language Model，大型語(yǔ)言模型）推理過(guò)程中，不緩存Q（Query）的原因主要有以下幾點(diǎn)：

依賴(lài)性差異：

在自回歸Transformer模型中，每個(gè)新生成的token的輸出（即Q）只依賴(lài)于當(dāng)前token的Q以及之前所有token的K和V。

因此，對(duì)于計(jì)算下一個(gè)token的輸出，不需要重復(fù)使用之前的Q，而K和V則需要被重復(fù)使用。

計(jì)算效率：

由于每次推理只會(huì)用到當(dāng)前的Q，而這個(gè)Q在下次推理時(shí)不會(huì)被再次使用，因此緩存Q不會(huì)帶來(lái)效率上的提升。

相反，K和V在每次推理中都需要被使用，緩存KV可以避免重復(fù)計(jì)算，從而加速推理過(guò)程。

自回歸特性：

在自回歸模型中，每個(gè)Token的生成僅依賴(lài)于它之前的所有Token，這意味著每個(gè)新Token的生成只需要當(dāng)前的Q和之前所有Token的KV。由于每個(gè)Q都是基于前面序列來(lái)生成的，緩存Q對(duì)于計(jì)算Attention沒(méi)有意義。

綜上所述，由于Q在自回歸Transformer模型中的使用特性和計(jì)算過(guò)程中的不對(duì)稱(chēng)性，緩存Q不會(huì)帶來(lái)推理效率的提升，因此LLM推理過(guò)程中通常不緩存Q。

為什么區(qū)分MHA和MLP？

在Transformer模型中，Multi-Head Attention (MHA) 和 Multi-Layer Perceptron (MLP) 是兩個(gè)核心組件，它們各自承擔(dān)著不同的功能，共同協(xié)作以提升模型的性能。

MHA 的作用：

捕捉上下文信息：MHA通過(guò)多個(gè)頭的方式，可以同時(shí)關(guān)注輸入序列的不同部分，從而增強(qiáng)模型對(duì)上下文信息的捕捉能力。

提高模型的表達(dá)能力：MHA允許模型在不同的表示子空間中學(xué)習(xí)信息，這有助于模型學(xué)習(xí)到更豐富的特征表示。

**處理長(zhǎng)序列數(shù)據(jù)：**雖然MHA的計(jì)算量較大，特別是對(duì)于長(zhǎng)序列，但它通過(guò)并行計(jì)算和優(yōu)化策略（如稀疏注意力）來(lái)提高效率1。

MLP 的作用：

非線性映射：MLP通過(guò)多層感知機(jī)對(duì)MHA的輸出進(jìn)行非線性變換，這有助于模型學(xué)習(xí)到更復(fù)雜的特征表示。

提升模型的表達(dá)能力：MLP的引入使得模型能夠捕捉到輸入數(shù)據(jù)中的非線性關(guān)系，從而提高模型的表達(dá)能力。

整合特征：MLP將MHA提取的特征進(jìn)行整合和進(jìn)一步處理，為模型的輸出提供必要的特征表示。

MHA和MLP的協(xié)同作用：

特征提取與整合：MHA負(fù)責(zé)提取輸入序列的上下文特征，而MLP則負(fù)責(zé)對(duì)這些特征進(jìn)行進(jìn)一步的整合和轉(zhuǎn)換，兩者共同工作以提高模型的性能2。

通過(guò)區(qū)分MHA和MLP兩個(gè)部分，Transformer模型能夠更有效地捕捉輸入序列的上下文信息，并學(xué)習(xí)到更豐富的特征表示，從而在自然語(yǔ)言處理任務(wù)中取得優(yōu)異的性能。