Big picture

Basically Transformer builds a graph network where a node is a position-encoded token in a sequence.

During training:

1. Get un-connected tokens as a sequence (e.g. sentence).
2. Wires connections among tokens by having looked at the co-occurrences of them in billions of sequences.

What roles Q and K will play to build this graph network? You could be Q in your society trying to build the social graph network with other people. Each person in the people is K and you will build the connection with them. Eventually by having billions of interactions with other people, the connections become dependent on the contexts even with the same person K.

You may be superior to a person K at work, but you may be a master of martial art to K. As you remember such connections/relations with others based on the contexts, Transformer model (trained on a specific dataset) figures out such context dependent connections from Q to K (or from you to other person(s)), which is a memory that it offers.

If the layers go up higher, your individual identity as K will be blended into larger parts via going through the BoW process which plays the role.

With regard to the Markov Chain (MC), there is only one static connection from Q to K as P(K|Q) in MC as MC does not have the context memory that Transformer model offers.

First, understand Q and K

First, focus on the objective of First MatMul in the Scaled dot product attention using Q and K.

Intuition on what is Attention

For the sentence "jane visits africa".

When your eyes see jane, your brain looks for the most related word in the rest of the sentence to understand what jane is about (query). Your brain focuses or attends to the word visit (key).

This process happens for each word in the sentence as your eyes progress through the sentence.

First MatMul as Inquiry System using Vector Similarity

The first MatMul implements an inquiry system or question-answer system that imitates this brain function, using Vector Similarity Calculation. Watch CS480/680 Lecture 19: Attention and Transformer Networks by professor Pascal Poupart to understand further.

Think about the attention essentially being some form of approximation of SELECT that you would do in the database.

Think of the MatMul as an inquiry system that processes the inquiry: "For the word q that your eyes see in the given sentence, what is the most related word k in the sentence to understand what q is about?" The inquiry system provides the answer as the probability.

Note that the softmax is used to normalize (in yellow) values into probabilities so that their sum becomes 1.0.

There are multiple ways to calculate the similarity between vectors such as cosine similarity. Transformer attention uses simple dot product.

Where are Q and K from

The transformer encoder training builds the weight parameter matrices WQ and Wk in the way Q and K builds the Inquiry System that answers the inquiry "What is k for the word q".

The calculation goes like below where x is a sequence of position-encoded word embedding vectors that represents an input sentence.

1. Picks up a word vector (position encoded) from the input sentence sequence, and transfer it to a vector space Q. This becomes the query.
   $Q = X \cdot W_{Q}^T$

2. Pick all the words in the sentence and transfer them to the vector space K. They become keys and each of them is used as key.
   $K = X \cdot W_K^T$

3. For each (q, k) pair, their relation strength is calculated using dot product.
   $q\_to\_k\_similarity\_scores = matmul(Q, K^T)$

4. Weight matrices $W_Q$ and $W_K$ are trained via the back propagations during the Transformer training.

We first needs to understand this part that involves Q and K before moving to V.

Borrowing the code from Let's build GPT: from scratch, in code, spelled out. by Andrej Karpathy.

import torch
from torch import nn
from torch.nn import functional as F


# let's see a single Head perform self-attention
torch.manual_seed(1337)

# B: batch size
# T: time steps or number of tokens to iterate or sequence size (512 for BERT)
# C: channels or embedding vector dimension or features
B,T,C = 4,8,32 # batch, time, channels
head_size = 16
    
Wk = nn.Linear(C, head_size, bias=False)
Wq = nn.Linear(C, head_size, bias=False)

def calculate_similarity_score(x):
    k = Wk(x)   # (B, T, head_size)
    q = Wq(x)   # (B, T, head_size)
    
    # First MatMul: (B, T, head_size) @ (B, head_size, T) ---> (B, T, T)
    score =  q @ k.transpose(-2, -1) 
    
    tril = torch.tril(torch.ones(T, T))
    score = score.masked_fill(tril == 0, float('-inf'))
    score = F.softmax(score, dim=-1)

    return score    # shape:(B, T, T)

Then, understand how V is created using Q and K

Second Matmul

Self Attention then generates the embedding vector called attention value as a bag of words (BoW) where each word contributes proportionally according to its relationship strength to q. This occurs for each q from the sentence sequence. The embedding vector is encoding the relations from q to all the words in the sentence.

Citing the words from Andrej Karpathy:

What is the easiest way for tokens to communicate. The easiest way is just average.

He makes it simple for the sake of tutorial but the essence is BoW.

Wv = nn.Linear(C, head_size, bias=False)

def calculate_attension_value(score, x):
    v = Wv(x)            # (B,T,C) @ (C,head_size) -> (B,T,head_size)
    value = score @ v    # (B,T,T) @ (B,T,head_size) -> (B,T,head_size)

    return value

Example to get Self Attention value.

x = torch.randn(B,T,C)

similarity_score = calculate_similarity_score(x)
attention_value = calculate_attension_value(similarity_score, x)

attention_value.shape    # (B,T,head_size)
-----
torch.Size([4, 8, 16])

References

There are multiple concepts that will help understand how the self attention in transformer works, e.g. embedding to group similars in a vector space, data retrieval to answer query Q using the neural network and vector similarity.

• Transformers Explained Visually (Part 2): How it works, step-by-step give in-detail explanation of what the Transformer is doing.

• CS480/680 Lecture 19: Attention and Transformer Networks - This is probably the best explanation I found that actually explains the attention mechanism from the database perspective.

• Illustrated Guide to Transformers Neural Network: A step by step explanation
  

• Distributed Representations of Words and Phrases and their Compositionality - It helps understand how word2vec works to group/categorize words in a vector space by pulling similar words together, and pushing away non-similar words using negative sampling.

• Generalized End-to-End Loss for Speaker Verification - Continuation to understand embedding to pull together siimilars and pushing away non-similars in a vector space.

• Transformer model for language understanding - TensorFlow implementation of transformer

• The Annotated Transformer - PyTorch implementation of Transformer

Update

Getting meaning from text: self-attention step-by-step video has visual representation of query, key, value.

Update 2

Andrej Karpathy explained by regarding a sentence as a graph as in CS25 I Stanford Seminar - Transformers United 2023: Introduction to Transformers w/ Andrej Karpathy.

Each token (position encoded word embedding) in a sentence is a node having edges to other nodes that represent attentions. When we focus on one node, it is a query that communicates with other nodes which are keys.