Introduction<\/mdspan><\/h2>\nThe Attention Mechanism<\/a> is often associated with the transformer architecture, but it was already used in RNNs. In Machine Translation or MT (e.g., English-Italian) tasks, when you want to predict the next Italian word, you need your model to focus, or pay attention, on the most important English words that are useful to make a good translation. <\/p>\n![\"Attention](\"https:\/\/i0.wp.com\/contributor.insightmediagroup.io\/wp-content\/uploads\/2025\/03\/1DtqI-w0CXu29NdXK28TPxQ.png?ssl=1\")
<\/figure>\nI will not go into details of RNNs, but attention helped these models to mitigate the vanishing gradient problem and to capture more long-range dependencies among words.<\/p>\n
At a certain point, we understood that the only important thing was the attention mechanism, and the entire RNN architecture was overkill. Hence, Attention is All You Need!<\/a> <\/p>\nSelf-Attention in Transformers<\/h2>\n
Classical attention indicates where words in the output sequence should focus attention in relation to the words in input sequence. This is important in sequence-to-sequence tasks like MT.<\/p>\n
The self-attention<\/strong> is a specific type of attention. It operates between any two elements in the same sequence. It provides information on how \u201ccorrelated\u201d the words are in the same sentence.<\/p>\nFor a given token (or word) in a sequence, self-attention generates a list of attention weights corresponding to all other tokens in the sequence. This process is applied to each token in the sentence, obtaining a matrix of attention weights (as in the picture).<\/p>\n![\"\"](\"https:\/\/www.researchgate.net\/publication\/356242242\/figure\/fig3\/AS:1095657444122624@1638236499150\/Self-attention-weights-from-the-product-MSC-model-for-a-product-titled-Handmade-Vintage.ppm\")
<\/figure>\nThis is the general idea, in practice things are a bit more complicated because we want to add many learnable parameters to our neural network, let\u2019s see how.<\/p>\n
K, V, Q representations<\/h2>\n
Our model input is a sentence like \u201cmy name is Marcello Politi\u201d<\/em>. With the process of tokenization<\/strong>, a sentence is converted into a list of numbers like [2, 6, 8, 3, 1]. <\/p>\nBefore feeding the sentence into the transformer we need to create a dense representation for each token.<\/p>\n
How to create this representation? We multiply each token by a matrix. The matrix is learned during training. <\/p>\n
Let\u2019s add some complexity now.<\/p>\n
For each token, we create 3 vectors instead of one, we call these vectors: key, value<\/em> and query<\/em>. (We see later how we create these 3 vectors).<\/p>\n![\"\"](\"https:\/\/i0.wp.com\/contributor.insightmediagroup.io\/wp-content\/uploads\/2025\/03\/image-195.png?ssl=1\")
<\/figure>\nConceptually these 3 tokens have a particular meaning:<\/p>\n
\n- The vector key represents the core information captured by the token<\/li>\n
- The vector value captures the full information of a token <\/li>\n
- The vector query, it\u2019s a question about the token relevance for the current task. <\/li>\n<\/ul>\n
So the idea is that we focus on a particular token i , and we want to ask what is the importance of the other tokens in the sentence regarding the token i we are taking into consideration. <\/p>\n
This means that we take the vector q_i (we ask a question regarding i) for token i, and we do some mathematical operations with all the other tokens k_j (j!=i). This is like wondering at first glance what are the other tokens in the sequence that look really important to understand the meaning of token i.<\/p>\n
What is this magical mathematical operation?<\/p>\n![\"\"](\"https:\/\/i0.wp.com\/substackcdn.com\/image\/fetch\/f_auto%2Cq_auto%3Agood%2Cfl_progressive%3Asteep\/https%253A%252F%252Fsubstack-post-media.s3.amazonaws.com%252Fpublic%252Fimages%252F8e197d7d-a2c7-4253-992b-256520cdafec_893x349.png?ssl=1\")
<\/figure>\nWe need to multiply (dot-product) the query vector by the key vectors and divide by a scaling factor. We do this for each k_j token. <\/p>\n
In this way, we obtain a score for each pair (q_i, k_j). We make this list become a probability distribution by applying a softmax operation on it. Great now we have obtained the attention weights<\/strong>!<\/p>\nWith the attention weights, we know what is the importance of each token k_j to for undestandin the token i. So now we multiply the value vector v_j associated with each token per its weight and we sum the vectors. In this way we obtain the final context-aware vector<\/strong> of token_i<\/strong>. <\/p>\nIf we are computing the contextual dense vector of token_1 we calculate:<\/p>\n
z1 = a11*v1 + a12*v2 + \u2026 + a15*v5<\/p>\n
Where a1j are the computer attention weights, and v_j are the value vectors. <\/p>\n
Done! Almost\u2026<\/p>\n
I didn\u2019t cover how we obtained the vectors k, v and q of each token. We need to define some matrices w_k, w_v and w_q so that when we multiply:<\/p>\n
\n- token * w_k -> k<\/li>\n
- token * w_q -> q<\/li>\n
- token * w_v -> v <\/li>\n<\/ul>\n
These 3 matrices are set at random and are learned during training, this is why we have many parameters in modern models such as LLMs.<\/p>\n
Multi-head Self-Attention in Transformers (MHSA)<\/h2>\n
Are we sure that the previous self-attention mechanism is able to capture all important relationships among tokens (words) and create dense vectors of those tokens that really make sense?<\/p>\n
It could actually not work always perfectly. What if to mitigate the error we re-run the entire thing 2 times with new w_q, w_k and w_v matrices and somehow merge the 2 dense vectors obtained? In this way maybe one self-attention managed to capture some relationship and the other managed to capture some other relationship.<\/p>\n
Well, this is what exactly happens in MHSA. The case we just discussed contains two heads because it has two sets of w_q, w_k and w_v matrices. We can have even more heads: 4, 8, 16 etc.<\/p>\n
The only complicated thing is that all these heads are managed in parallel, we process the all in the same computation using tensors.<\/p>\n![\"\"](\"https:\/\/i0.wp.com\/contributor.insightmediagroup.io\/wp-content\/uploads\/2025\/03\/image-196.png?ssl=1\")
<\/figure>\nThe way we merge the dense vectors of each head is simple, we concatenate them (hence the dimension of each vector shall be smaller so that when concat them we obtain the original dimension we wanted), and we pass the obtained vector through another w_o learnable matrix. <\/p>\n
Hands-on <\/h2>\nPython\">import torch<\/code><\/pre>\nSuppose you have a sentence. After tokenization, each token (word for simplicity) corresponds to an index (number):<\/p>\n
tokenized_sentence = torch.tensor([\n 2, #my\n 6, #name\n 8, #is\n 3, #marcello\n 1 #politi\n])\ntokenized_sentence<\/code><\/pre>\nBefore feeding the sentence into the transofrmer we need to create a dense representation for each token.<\/p>\n
How to create these representation? We multiply each token per a matrix. This matrix is learned during training.<\/p>\n
Let\u2019s build this embedding matrix.<\/p>\n
torch.manual_seed(0) # set a fixed seed for reproducibility\nembed = torch.nn.Embedding(10, 16)<\/code><\/pre>\nIf we multiply our tokenized sentence with the embeddings, we obtain a dense representation of dimension 16 for each token<\/p>\n
sentence_embed = embed(tokenized_sentence).detach()\nsentence_embed<\/code><\/pre>\nIn order to use the attention mechanism we need to create 3 new We define 3 matrixes w_q, w_k and w_v. When we multiply one input token time the w_q we obtain the vector q. Same with w_k and w_v.<\/p>\n
d = sentence_embed.shape[1] # let's base our matrix on a shape (16,16)\n\nw_key = torch.rand(d,d)\nw_query = torch.rand(d,d)\nw_value = torch.rand(d,d)<\/code><\/pre>\nCompute attention weights<\/strong><\/p>\nLet\u2019s now compute the attention weights for only the first input token of the sentence.<\/p>\n
token1_embed = sentence_embed[0]\n\n#compute the tre vector associated to token1 vector : q,k,v\nkey_1 = w_key.matmul(token1_embed)\nquery_1 = w_query.matmul(token1_embed)\nvalue_1 = w_value.matmul(token1_embed)\n\nprint(\"key vector for token1: n\", key_1) \nprint(\"query vector for token1: n\", query_1)\nprint(\"value vector for token1: n\", value_1)<\/code><\/pre>\nWe need to multiply the query vector associated to token1 (query_1) with all the keys of the other vectors.<\/p>\n
So now we need to compute all the keys (key_2, key_2, key_4, key_5). But wait, we can compute all of these in one time by multiplying the sentence_embed times the w_k matrix.<\/p>\n
keys = sentence_embed.matmul(w_key.T)\nkeys[0] #contains the key vector of the first token and so on<\/code><\/pre>\nLet\u2019s do the same thing with the values<\/p>\n
values = sentence_embed.matmul(w_value.T)\nvalues[0] #contains the value vector of the first token and so on<\/code><\/pre>\nLet\u2019s compute the first part of the attions formula.<\/p>\n<\/figure>\nimport torch.nn.functional as F<\/code><\/pre>\n# the following are the attention weights of the first tokens to all the others\na1 = F.softmax(query_1.matmul(keys.T)\/d**0.5, dim = 0)\na1<\/code><\/pre>\nWith the attention weights we know what is the importance of each token. So now we multiply the value vector associated to each token per its weight.<\/p>\n
To obtain the final context aware vector of token_1.<\/p>\n
z1 = a1.matmul(values)\nz1<\/code><\/pre>\nIn the same way we could compute the context aware dense vectors of all the other tokens. Now we are always using the same matrices w_k, w_q, w_v. We say that we use one head.<\/p>\n
But we can have multiple triplets of matrices, so multi-head. That\u2019s why it is called multi-head attention.<\/p>\n
The dense vectors of an input tokens, given in oputut from each head are at then end concatenated and linearly transformed to get the final dense vector.<\/p>\n
Implementing MultiheadSelf-Attention<\/h4>\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\ntorch.manual_seed(0) # fixed seed for reproducibility<\/code><\/pre>\nSame steps as before\u2026<\/p>\n
# Tokenized sentence (same as yours)\ntokenized_sentence = torch.tensor([2, 6, 8, 3, 1]) # [my, name, is, marcello, politi]\n\n# Embedding layer: vocab size = 10, embedding dim = 16\nembed = nn.Embedding(10, 16)\nsentence_embed = embed(tokenized_sentence).detach() # Shape: [5, 16] (seq_len, embed_dim)<\/code><\/pre>\nWe\u2019ll define a multi-head attention mechanism with h heads (let\u2019s say 4 heads for this example). Each head will have its own w_q, w_k, and w_v matrices, and the output of each head will be concatenated and passed through a final linear layer.<\/p>\n
Since the output of the head will be concatenated, and we want a final dimension of d, the dimension of each head needs to be d\/h. Additionally each concatenated vector will go though a linear transformation, so we need another matrix w_ouptut as you can see in the formula.<\/p>\n
d = sentence_embed.shape[1] # embed dimension 16\nh = 4 # Number of heads\nd_k = d \/\/ h # Dimension per head (16 \/ 4 = 4)<\/code><\/pre>\nSince we have 4 heads, we want 4 copies for each matrix. Instead of copies, we add a dimension, which is the same thing, but we only do one operation. (Imagine stacking matrices on top of each other, its the same thing).<\/p>\n
# Define weight matrices for each head\nw_query = torch.rand(h, d, d_k) # Shape: [4, 16, 4] (one d x d_k matrix per head)\nw_key = torch.rand(h, d, d_k) # Shape: [4, 16, 4]\nw_value = torch.rand(h, d, d_k) # Shape: [4, 16, 4]\nw_output = torch.rand(d, d) # Final linear layer: [16, 16]<\/code><\/pre>\nI\u2019m using for simplicity torch\u2019s einsum. If you\u2019re not familiar with it check out my\u00a0blog post<\/a>.<\/p>\nThe einsum operation torch.einsum('sd,hde->hse', sentence_embed, w_query)<\/code> in PyTorch uses letters to define how to multiply and rearrange numbers. Here\u2019s what each part means:<\/p>\n\n- \nInput Tensors:<\/strong><\/p>\n
\n- \n
sentence_embed<\/code> with the notation 'sd'<\/code>:<\/p>\n\n- \n
s<\/code> represents the number of words (sequence length), which is 5.<\/li>\n- \n
d<\/code> represents the number of numbers per word (embedding size), which is 16.<\/li>\n- The shape of this tensor is
[5, 16]<\/code>.<\/li>\n<\/ul>\n<\/li>\n- \n
w_query<\/code> with the notation 'hde'<\/code>:<\/p>\n\n- \n
h<\/code> represents the number of heads, which is 4.<\/li>\n- \n
d<\/code> represents the embedding size, which again is 16.<\/li>\n- \n
e<\/code> represents the new number size per head (d_k), which is 4.<\/li>\n- The shape of this tensor is
[4, 16, 4]<\/code>.<\/li>\n<\/ul>\n<\/li>\n<\/ul>\n<\/li>\n- \nOutput Tensor:<\/strong><\/p>\n
\n- The output has the notation
'hse'<\/code>:\n\n- \n
h<\/code> represents 4 heads.<\/li>\n- \n
s<\/code> represents 5 words.<\/li>\n- \n
e<\/code> represents 4 numbers per head.<\/li>\n- The shape of the output tensor is
[4, 5, 4]<\/code>.<\/li>\n<\/ul>\n<\/li>\n<\/ul>\n<\/li>\n<\/ol>\n# Compute Q, K, V for all tokens and all heads\n# sentence_embed: [5, 16] -> Q: [4, 5, 4] (h, seq_len, d_k)\nqueries = torch.einsum('sd,hde->hse', sentence_embed, w_query) # h heads, seq_len tokens, d dim\nkeys = torch.einsum('sd,hde->hse', sentence_embed, w_key) # h heads, seq_len tokens, d dim\nvalues = torch.einsum('sd,hde->hse', sentence_embed, w_value) # h heads, seq_len tokens, d dim<\/code><\/pre>\nThis einsum equation performs a dot product between the queries (hse) and the transposed keys (hek) to obtain scores of shape [h, seq_len, seq_len], where:<\/p>\n
\n- h -> Number of heads.<\/li>\n
- s and k -> Sequence length (number of tokens).<\/li>\n
- e -> Dimension of each head (d_k).<\/li>\n<\/ul>\n
The division by (d_k ** 0.5) scales the scores to stabilize gradients. Softmax is then applied to obtain attention weights:<\/p>\n
# Compute attention scores\nscores = torch.einsum('hse,hek->hsk', queries, keys.transpose(-2, -1)) \/ (d_k ** 0.5) # [4, 5, 5]\nattention_weights = F.softmax(scores, dim=-1) # [4, 5, 5]<\/code><\/pre>\n# Apply attention weights\nhead_outputs = torch.einsum('hij,hjk->hik', attention_weights, values) # [4, 5, 4]\nhead_outputs.shape<\/code><\/pre>\nNow we concatenate all the heads of token 1<\/p>\n
# Concatenate heads\nconcat_heads = head_outputs.permute(1, 0, 2).reshape(sentence_embed.shape[0], -1) # [5, 16]\nconcat_heads.shape<\/code><\/pre>\nLet\u2019s finally multiply per the last w_output matrix as in the formula above<\/p>\n
multihead_output = concat_heads.matmul(w_output) # [5, 16] @ [16, 16] -> [5, 16]\nprint(\"Multi-head attention output for token1:n\", multihead_output[0])<\/code><\/pre>\nFinal Thoughts<\/h2>\n
In this blog post I\u2019ve implemented a simple version of the attention mechanism. This is not how it is really implemented in modern frameworks, but my scope is to provide some insights to allow anyone an understanding of how this works. In future articles I\u2019ll go through the entire implementation of a transformer architecture.<\/p>\n
Follow me on\u00a0TDS<\/a>\u00a0if you like this article! ![\"\ud83d\ude01\"](\"https:\/\/i0.wp.com\/s.w.org\/images\/core\/emoji\/15.0.3\/72x72\/1f601.png?ssl=1\")
<\/p>\n![\"\ud83d\udcbc\"](\"https:\/\/i0.wp.com\/s.w.org\/images\/core\/emoji\/15.0.3\/72x72\/1f4bc.png?ssl=1\")
\u00a0Linkedin<\/a>\u00a0| ![\"\ud83d\udc26\"](\"https:\/\/i0.wp.com\/s.w.org\/images\/core\/emoji\/15.0.3\/72x72\/1f426.png?ssl=1\")
\u00a0X (Twitter)<\/a>\u00a0|\u00a0![\"\ud83d\udcbb\"](\"https:\/\/i0.wp.com\/s.w.org\/images\/core\/emoji\/15.0.3\/72x72\/1f4bb.png?ssl=1\")
<\/a>\u00a0Website<\/a><\/p>\n
\nUnless otherwise noted, images are by the author<\/p>\n
The post A Simple Implementation of the Attention Mechanism from Scratch<\/a> appeared first on Towards Data Science<\/a>.<\/p>\n<\/div>\n
The Attention Mechanism<\/a> is often associated with the transformer architecture, but it was already used in RNNs. In Machine Translation or MT (e.g., English-Italian) tasks, when you want to predict the next Italian word, you need your model to focus, or pay attention, on the most important English words that are useful to make a good translation. <\/p>\n I will not go into details of RNNs, but attention helped these models to mitigate the vanishing gradient problem and to capture more long-range dependencies among words.<\/p>\n At a certain point, we understood that the only important thing was the attention mechanism, and the entire RNN architecture was overkill. Hence, Attention is All You Need!<\/a> <\/p>\n Classical attention indicates where words in the output sequence should focus attention in relation to the words in input sequence. This is important in sequence-to-sequence tasks like MT.<\/p>\n The self-attention<\/strong> is a specific type of attention. It operates between any two elements in the same sequence. It provides information on how \u201ccorrelated\u201d the words are in the same sentence.<\/p>\n For a given token (or word) in a sequence, self-attention generates a list of attention weights corresponding to all other tokens in the sequence. This process is applied to each token in the sentence, obtaining a matrix of attention weights (as in the picture).<\/p>\n This is the general idea, in practice things are a bit more complicated because we want to add many learnable parameters to our neural network, let\u2019s see how.<\/p>\n Our model input is a sentence like \u201cmy name is Marcello Politi\u201d<\/em>. With the process of tokenization<\/strong>, a sentence is converted into a list of numbers like [2, 6, 8, 3, 1]. <\/p>\n Before feeding the sentence into the transformer we need to create a dense representation for each token.<\/p>\n How to create this representation? We multiply each token by a matrix. The matrix is learned during training. <\/p>\n Let\u2019s add some complexity now.<\/p>\n For each token, we create 3 vectors instead of one, we call these vectors: key, value<\/em> and query<\/em>. (We see later how we create these 3 vectors).<\/p>\n Conceptually these 3 tokens have a particular meaning:<\/p>\n So the idea is that we focus on a particular token i , and we want to ask what is the importance of the other tokens in the sentence regarding the token i we are taking into consideration. <\/p>\n This means that we take the vector q_i (we ask a question regarding i) for token i, and we do some mathematical operations with all the other tokens k_j (j!=i). This is like wondering at first glance what are the other tokens in the sequence that look really important to understand the meaning of token i.<\/p>\n What is this magical mathematical operation?<\/p>\n We need to multiply (dot-product) the query vector by the key vectors and divide by a scaling factor. We do this for each k_j token. <\/p>\n In this way, we obtain a score for each pair (q_i, k_j). We make this list become a probability distribution by applying a softmax operation on it. Great now we have obtained the attention weights<\/strong>!<\/p>\n With the attention weights, we know what is the importance of each token k_j to for undestandin the token i. So now we multiply the value vector v_j associated with each token per its weight and we sum the vectors. In this way we obtain the final context-aware vector<\/strong> of token_i<\/strong>. <\/p>\n If we are computing the contextual dense vector of token_1 we calculate:<\/p>\n z1 = a11*v1 + a12*v2 + \u2026 + a15*v5<\/p>\n Where a1j are the computer attention weights, and v_j are the value vectors. <\/p>\n Done! Almost\u2026<\/p>\n I didn\u2019t cover how we obtained the vectors k, v and q of each token. We need to define some matrices w_k, w_v and w_q so that when we multiply:<\/p>\n These 3 matrices are set at random and are learned during training, this is why we have many parameters in modern models such as LLMs.<\/p>\n Are we sure that the previous self-attention mechanism is able to capture all important relationships among tokens (words) and create dense vectors of those tokens that really make sense?<\/p>\n It could actually not work always perfectly. What if to mitigate the error we re-run the entire thing 2 times with new w_q, w_k and w_v matrices and somehow merge the 2 dense vectors obtained? In this way maybe one self-attention managed to capture some relationship and the other managed to capture some other relationship.<\/p>\n Well, this is what exactly happens in MHSA. The case we just discussed contains two heads because it has two sets of w_q, w_k and w_v matrices. We can have even more heads: 4, 8, 16 etc.<\/p>\n The only complicated thing is that all these heads are managed in parallel, we process the all in the same computation using tensors.<\/p>\n The way we merge the dense vectors of each head is simple, we concatenate them (hence the dimension of each vector shall be smaller so that when concat them we obtain the original dimension we wanted), and we pass the obtained vector through another w_o learnable matrix. <\/p>\n Suppose you have a sentence. After tokenization, each token (word for simplicity) corresponds to an index (number):<\/p>\n Before feeding the sentence into the transofrmer we need to create a dense representation for each token.<\/p>\n How to create these representation? We multiply each token per a matrix. This matrix is learned during training.<\/p>\n Let\u2019s build this embedding matrix.<\/p>\n If we multiply our tokenized sentence with the embeddings, we obtain a dense representation of dimension 16 for each token<\/p>\n In order to use the attention mechanism we need to create 3 new We define 3 matrixes w_q, w_k and w_v. When we multiply one input token time the w_q we obtain the vector q. Same with w_k and w_v.<\/p>\n Compute attention weights<\/strong><\/p>\n Let\u2019s now compute the attention weights for only the first input token of the sentence.<\/p>\n We need to multiply the query vector associated to token1 (query_1) with all the keys of the other vectors.<\/p>\n So now we need to compute all the keys (key_2, key_2, key_4, key_5). But wait, we can compute all of these in one time by multiplying the sentence_embed times the w_k matrix.<\/p>\n Let\u2019s do the same thing with the values<\/p>\n Let\u2019s compute the first part of the attions formula.<\/p>\n With the attention weights we know what is the importance of each token. So now we multiply the value vector associated to each token per its weight.<\/p>\n To obtain the final context aware vector of token_1.<\/p>\n In the same way we could compute the context aware dense vectors of all the other tokens. Now we are always using the same matrices w_k, w_q, w_v. We say that we use one head.<\/p>\n But we can have multiple triplets of matrices, so multi-head. That\u2019s why it is called multi-head attention.<\/p>\n The dense vectors of an input tokens, given in oputut from each head are at then end concatenated and linearly transformed to get the final dense vector.<\/p>\n Same steps as before\u2026<\/p>\n We\u2019ll define a multi-head attention mechanism with h heads (let\u2019s say 4 heads for this example). Each head will have its own w_q, w_k, and w_v matrices, and the output of each head will be concatenated and passed through a final linear layer.<\/p>\n Since the output of the head will be concatenated, and we want a final dimension of d, the dimension of each head needs to be d\/h. Additionally each concatenated vector will go though a linear transformation, so we need another matrix w_ouptut as you can see in the formula.<\/p>\n Since we have 4 heads, we want 4 copies for each matrix. Instead of copies, we add a dimension, which is the same thing, but we only do one operation. (Imagine stacking matrices on top of each other, its the same thing).<\/p>\n I\u2019m using for simplicity torch\u2019s einsum. If you\u2019re not familiar with it check out my\u00a0blog post<\/a>.<\/p>\n The einsum operation This einsum equation performs a dot product between the queries (hse) and the transposed keys (hek) to obtain scores of shape [h, seq_len, seq_len], where:<\/p>\n The division by (d_k ** 0.5) scales the scores to stabilize gradients. Softmax is then applied to obtain attention weights:<\/p>\n Now we concatenate all the heads of token 1<\/p>\n Let\u2019s finally multiply per the last w_output matrix as in the formula above<\/p>\n In this blog post I\u2019ve implemented a simple version of the attention mechanism. This is not how it is really implemented in modern frameworks, but my scope is to provide some insights to allow anyone an understanding of how this works. In future articles I\u2019ll go through the entire implementation of a transformer architecture.<\/p>\n Follow me on\u00a0TDS<\/a>\u00a0if you like this article! Unless otherwise noted, images are by the author<\/p>\n The post A Simple Implementation of the Attention Mechanism from Scratch<\/a> appeared first on Towards Data Science<\/a>.<\/p>\n<\/div>\n<\/figure>\nSelf-Attention in Transformers<\/h2>\n
<\/figure>\n
K, V, Q representations<\/h2>\n
<\/figure>\n\n
<\/figure>\n\n
Multi-head Self-Attention in Transformers (MHSA)<\/h2>\n
<\/figure>\nHands-on <\/h2>\n
Python\">import torch<\/code><\/pre>\ntokenized_sentence = torch.tensor([\n 2, #my\n 6, #name\n 8, #is\n 3, #marcello\n 1 #politi\n])\ntokenized_sentence<\/code><\/pre>\ntorch.manual_seed(0) # set a fixed seed for reproducibility\nembed = torch.nn.Embedding(10, 16)<\/code><\/pre>\nsentence_embed = embed(tokenized_sentence).detach()\nsentence_embed<\/code><\/pre>\nd = sentence_embed.shape[1] # let's base our matrix on a shape (16,16)\n\nw_key = torch.rand(d,d)\nw_query = torch.rand(d,d)\nw_value = torch.rand(d,d)<\/code><\/pre>\ntoken1_embed = sentence_embed[0]\n\n#compute the tre vector associated to token1 vector : q,k,v\nkey_1 = w_key.matmul(token1_embed)\nquery_1 = w_query.matmul(token1_embed)\nvalue_1 = w_value.matmul(token1_embed)\n\nprint(\"key vector for token1: n\", key_1) \nprint(\"query vector for token1: n\", query_1)\nprint(\"value vector for token1: n\", value_1)<\/code><\/pre>\nkeys = sentence_embed.matmul(w_key.T)\nkeys[0] #contains the key vector of the first token and so on<\/code><\/pre>\nvalues = sentence_embed.matmul(w_value.T)\nvalues[0] #contains the value vector of the first token and so on<\/code><\/pre>\n<\/figure>\nimport torch.nn.functional as F<\/code><\/pre>\n# the following are the attention weights of the first tokens to all the others\na1 = F.softmax(query_1.matmul(keys.T)\/d**0.5, dim = 0)\na1<\/code><\/pre>\nz1 = a1.matmul(values)\nz1<\/code><\/pre>\nImplementing MultiheadSelf-Attention<\/h4>\n
import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\ntorch.manual_seed(0) # fixed seed for reproducibility<\/code><\/pre>\n# Tokenized sentence (same as yours)\ntokenized_sentence = torch.tensor([2, 6, 8, 3, 1]) # [my, name, is, marcello, politi]\n\n# Embedding layer: vocab size = 10, embedding dim = 16\nembed = nn.Embedding(10, 16)\nsentence_embed = embed(tokenized_sentence).detach() # Shape: [5, 16] (seq_len, embed_dim)<\/code><\/pre>\nd = sentence_embed.shape[1] # embed dimension 16\nh = 4 # Number of heads\nd_k = d \/\/ h # Dimension per head (16 \/ 4 = 4)<\/code><\/pre>\n# Define weight matrices for each head\nw_query = torch.rand(h, d, d_k) # Shape: [4, 16, 4] (one d x d_k matrix per head)\nw_key = torch.rand(h, d, d_k) # Shape: [4, 16, 4]\nw_value = torch.rand(h, d, d_k) # Shape: [4, 16, 4]\nw_output = torch.rand(d, d) # Final linear layer: [16, 16]<\/code><\/pre>\ntorch.einsum('sd,hde->hse', sentence_embed, w_query)<\/code> in PyTorch uses letters to define how to multiply and rearrange numbers. Here\u2019s what each part means:<\/p>\n\n
\n
sentence_embed<\/code> with the notation 'sd'<\/code>:<\/p>\n\n
s<\/code> represents the number of words (sequence length), which is 5.<\/li>\nd<\/code> represents the number of numbers per word (embedding size), which is 16.<\/li>\n[5, 16]<\/code>.<\/li>\n<\/ul>\n<\/li>\nw_query<\/code> with the notation 'hde'<\/code>:<\/p>\n\n
h<\/code> represents the number of heads, which is 4.<\/li>\nd<\/code> represents the embedding size, which again is 16.<\/li>\ne<\/code> represents the new number size per head (d_k), which is 4.<\/li>\n[4, 16, 4]<\/code>.<\/li>\n<\/ul>\n<\/li>\n<\/ul>\n<\/li>\n\n
'hse'<\/code>:\n\n
h<\/code> represents 4 heads.<\/li>\ns<\/code> represents 5 words.<\/li>\ne<\/code> represents 4 numbers per head.<\/li>\n[4, 5, 4]<\/code>.<\/li>\n<\/ul>\n<\/li>\n<\/ul>\n<\/li>\n<\/ol>\n# Compute Q, K, V for all tokens and all heads\n# sentence_embed: [5, 16] -> Q: [4, 5, 4] (h, seq_len, d_k)\nqueries = torch.einsum('sd,hde->hse', sentence_embed, w_query) # h heads, seq_len tokens, d dim\nkeys = torch.einsum('sd,hde->hse', sentence_embed, w_key) # h heads, seq_len tokens, d dim\nvalues = torch.einsum('sd,hde->hse', sentence_embed, w_value) # h heads, seq_len tokens, d dim<\/code><\/pre>\n\n
# Compute attention scores\nscores = torch.einsum('hse,hek->hsk', queries, keys.transpose(-2, -1)) \/ (d_k ** 0.5) # [4, 5, 5]\nattention_weights = F.softmax(scores, dim=-1) # [4, 5, 5]<\/code><\/pre>\n# Apply attention weights\nhead_outputs = torch.einsum('hij,hjk->hik', attention_weights, values) # [4, 5, 4]\nhead_outputs.shape<\/code><\/pre>\n# Concatenate heads\nconcat_heads = head_outputs.permute(1, 0, 2).reshape(sentence_embed.shape[0], -1) # [5, 16]\nconcat_heads.shape<\/code><\/pre>\nmultihead_output = concat_heads.matmul(w_output) # [5, 16] @ [16, 16] -> [5, 16]\nprint(\"Multi-head attention output for token1:n\", multihead_output[0])<\/code><\/pre>\nFinal Thoughts<\/h2>\n
<\/p>\n\u00a0Linkedin<\/a>\u00a0| \u00a0X (Twitter)<\/a>\u00a0|\u00a0<\/a>\u00a0Website<\/a><\/p>\n
\n