字節三面：attention中的Q,K,V怎麼來的？

作者：TRiddle@知乎  僅用作學術分享

鏈接：https://www.zhihu.com/question/325839123/answer/3309301644

本質上就是查詢+聚合。而且非常符合直覺，不需要任何公式就能說明這點。我們可以試著「重新發明」一下注意力機制。

現在想像一下，假如你想做一道名菜——佛跳牆，你會怎麼做？你會先去菜市場里找到鮑魚、海參、花膠、瑤柱，然後帶回家將它們煮在一起對吧。

而一個用於文本分類的 BERT 做的事也一樣，也是找到一些東西然後將它們煮在一起。

比方說如果想要識別一篇文章中是否在講佛跳牆，要做的就是查找文章中是否存在相鄰的「佛、跳、牆」（當然，「佛、跳、牆」的上下文也很重要，請允許我做些簡化），然後將這三個存在性信息聚合在一起，最後從 [cls] token 的位置輸出出去。

或者說，用 [cls] 來查詢和聚合「佛、跳、牆」。

1.如何查詢

首先面臨的問題就是 [cls] 怎麼查詢「佛、跳、牆」。

其實有一種辦法很簡單，就是給文章中所有 token 都分別分配一個向量，讓它們滿足 [cls] 的向量同「佛、跳、牆」的三個向量的距離都很小，同除了它們三個之外的其它 token 的向量的距離都很大。

大家通常會用兩個向量的點積來衡量它們之間的距離。具體就是點積大的距離小，點積小的距離大。

正因為距離是可衡量的，我們只需要計算 [cls] 的向量和句子中所有 token 的向量之間的點積，就能知道哪些 token 是「佛、跳、牆」，哪些 token 不是。

具體就是點積大的那些是，點積小的那些不是。這些點積再通過 softmax 歸一化一下，就是我們常說的注意力分數。

我們可以說，注意力分數表示的是查詢的相關性。或者說，表示的是「你有多大概率是我要查的東西」。

怎麼給所有 token 分配向量呢？

先規定一下向量的名字，因為 [cls] 是查詢的主體，所以我們給它分配的向量叫做 Q（query 的縮寫）。

因為「佛、跳、牆」是被查詢的客體，所以我們給它分配的向量叫做 A（answer 的縮寫。等等，為什麼不是 K 和 V？這裏先按下不表）。

向量的數值怎麼分配呢？

我們弄兩個可學習的矩陣 WQ 和 WA，想辦法讓一個 token embedding（加上 position embedding）乘上 WQ 得到 Q ，乘上 WA 得到 A。WQ 和 WA 的參數讓它們自己通過梯度下降學習就行了。

WQ 和 WA 實際上是模型學會的分配 Q 和 A 的邏輯。在我們的例子中的分配邏輯就是，[cls] 的 Q 向量要同「佛、跳、牆」的 A 向量離得近一些，要同其他 token 的 A 向量離得遠一些。

2.如何聚合

接下來需要解決的問題是怎麼將「佛、跳、牆」的存在性信息聚合在一起。一種方法是，直接將「佛、跳、牆」的A向量加起來就行了。

按照前面提到的方法，首先我們計算 [cls] 的 Q 向量和所有 token 的 A 向量之間的距離，然後篩選出距離小的那些，最後把它們的 A 向量加起來。

不過，篩選這個操作太麻煩了。我們明明可以以注意力分數為權重，直接計算序列中所有 token 的 A 向量的加權和。

實際上這和先篩選再加起來是一樣的。由於除「佛跳牆」外其他 token 的注意力分數都很小，這麼做相當於只將「佛、跳、牆」的 A 向量加起來了。

也就是說，以注意力分數為權重對文章中的所有 token 的 A 向量求加權和，就相當於把需要查詢出來的重要 token 給聚合在一起了。

可是，A 向量很忙的啊。既要對查詢的結果負責，又要對聚合的結果負責，我們能不能將它們一分為二呢？

當然可以了，我們用 K 向量（key 的縮寫）和 V 向量（value 的縮寫）替代 A 向量。

前者參與查詢階段的注意力分數計算，後者攜帶著聚合階段需要被加在一起的存在性信息。當然，WA 矩陣也要相應地替換成 WK 矩陣和 WV 矩陣。

讓我們重新表述一下：以 Q 和 K 算出的注意力分數為權重對文章中的所有 token 的 V 向量求加權和，就相當於把需要查詢的重要 token 的給聚合在一起了。

為了更形象地展示這個過程，我畫了一張圖：

以上就是自注意力機制中 Q、K、V 的角色。我們用文本分類場景下的 BERT 舉了個例子。

其它任務或者其它模型的注意力機制也非常好理解，因為本質上都是查詢+聚合。

另外，你可能會關注下面這些問題：

（1）我關注到的 attention 中的 Q、K、V 是矩陣，這裏的怎麼是向量呢？

事實上文章中每個 token 都有自己的 Q、K、V 向量，因為每個 token 都有自己需要查詢的東西，也可能被查詢。把它們的 Q、K、V 向量連接起來就是  Q、K、V 矩陣了。

（2）按照上面說的，注意力機制是不是不能保證「佛、跳、牆」是連續出現的？

其實是能的，只要讓「佛」先查詢「跳、牆」，再讓 [cls] 查詢「佛、跳、牆」就行了。

因為 Q 向量和 K 向量都是攜帶位置信息的，所以查詢的時候是能感知到相對位置的。

所以在「佛」做查詢的時候，就有機會只給它周圍的「跳、牆」一個較高的注意力分數。

而做兩次注意力機制也很容易，只要放在兩個不同的 attention 層中就行了（不然 transformer 為什麼那麼多層）。

（3）一個 token 能不能有多個 Q 向量？

當然能。舉個例子，「佛」可能不一定只用來查詢「跳、牆」，還可能用來查詢「祖」或者「經」等 token。

的確，一個 token 可能是需要多個 Q 的，還可能需要多個 K 和 V，這就是 transformer 中為什麼有 multi-head attention 這種東西。

（4）Q 和 K 用同一個向量行不行？

我的答案是——不行。假如 Q 和 K 用同一個向量，就會出現一些奇怪的相似性傳導問題。

例如，「佛」和「跳、牆」的距離很近，「佛」和「祖」的距離也很近，那豈不是說「跳、牆」和「祖」的距離也很近了？或者從圖的角度來解釋就是：同質圖和異質圖建模的是不同的東西。

鏈接：https://www.zhihu.com/question/325839123/answer/3309301644