A key is the representation an attention mechanism compares with a query to decide how much weight to give an associated value. Think of the query as what a token is looking for, the key as the features available for matching, and the value as the content to be weighted. A key determines relevance; it is not itself the content summed into the output. section 2.2.3
In scaled dot-product attention:
- Project input embeddings into separate query, key, and value matrices: Q, K, and V. section 2.2.3
- Compare queries with keys using Q · Kᵀ, divide by the square root of the key dimension, and apply softmax to obtain attention weights. section 2.2.3
- Multiply those weights by V to form the attention output. Multihead attention projects keys separately for each head. section 2.2.3
Matching “not” with “bad”
In “The movie was not bad,” a query for “not” can match the key for “bad” strongly. That match gives the value associated with “bad” greater weight when forming the output for “not.” section 2.2.3
Listing 2.3 Simplified scaled dot-attention implementation
embed_1 = np.array([0, 1, 0]) #1 embed_2 = np.array([1, 0, 1]) embed_3 = np.array([0, 1, 1]) embed_4 = np.array([1, 1, 0]) embeddings = np.array([embed_1, embed_2, embed_3, embed_4]) #2 Wq = rand(3, 3) #3 Wk = rand(3, 3) Wv = rand(3, 3) Q = embeddings.dot(Wq) #4 K = embeddings.dot(Wk) V = embeddings.dot(Wv) attention_scores = softmax(Q.dot(K.T) / sqrt(K.shape[1]), axis=1) #5 attention_output = attention_scores.dot(V) #6
The NumPy listing shows where keys enter the calculation: K = embeddings.dot(Wk), followed by softmax(Q.dot(K.T) / sqrt(K.shape[1])). The resulting weights, rather than K, are multiplied by V. section 2.2.3
