Coding Practice | YuyaoGe's Website

Coding Practice | Implementing a Tokenizer: Word-based and Byte Pair Encoding

Wed, 03 Jul 2024 00:00:00 +0000

Introduction

A tokenizer is a key tool in natural language processing. Its job is to split a text string into lexical units, or tokens. Tokenizers play a crucial role during text preprocessing, providing the foundation for downstream tasks such as text analysis, information retrieval, and machine translation.

Origins and Development

The history of tokenizers goes back to the early days of computational linguistics and information retrieval. As computer science matured, researchers realized that processing natural language text requires splitting a continuous sequence of characters into basic units that carry meaning (called morphemes in linguistics). Early tokenization methods were mainly rule- and dictionary-based, relying on predefined rules and word lists to identify word boundaries. With the rise of statistical NLP and machine learning, statistical models and data-driven approaches gradually replaced rule-based ones.

In recent years, advances in deep learning have pushed tokenization further. In particular, the introduction of pretrained language models such as BERT (Bidirectional Encoder Representations from Transformers) brought subword-level tokenization methods such as Byte Pair Encoding (BPE), WordPiece, and SentencePiece, which handle out-of-vocabulary words and linguistic diversity far better.

Motivation

To understand tokenizers, we first need to understand why we use them at all.

In NLP tasks, the data we work with is usually raw text.

Take the following sentence as an example:

Jim Henson was a puppeteer

Models, however, can only process numbers, so we need a way to turn raw text into numbers. That is exactly what a tokenizer does.

In short, the goal of a tokenizer is to convert text that humans understand into numbers that machines understand.

Next, we introduce the simplest and most intuitive tokenization method — the Word-based Tokenizer.

The Simplest and Most Intuitive Approach — Word-based Tokenizer

Part of this section draws on the Hugging Face NLP Course ¹

Picking up where we left off: how do we turn the text Jim Henson was a puppeteer into numbers?

One intuitive approach is to split the string on whitespace and assign each word a unique number as its index.

tokenized_text = "Jim Henson was a puppeteer".split()
print(tokenized_text)

['Jim', 'Henson', 'was', 'a', 'puppeteer']

This way each word gets an ID, starting from 0 and running up to the size of the vocabulary. The model uses these IDs to identify each word.

If we wanted a word-based tokenizer to cover an entire language, we would have to assign a unique integer index to every word in that language, which produces an enormous number of indices. English, for example, has more than 500,000 words, so building a mapping from every word to an index would require a dictionary with 500,000 entries. And that is not the only drawback of word-based tokenization.

Furthermore, a word like “dog” is represented differently from a word like “dogs”, and the model has no way of knowing that “dog” and “dogs” are related: it treats them as two unrelated words. The same applies to other similar pairs, such as “run” and “running” — the model will not consider them similar.

We also need a special token to represent words that are not in our vocabulary. This is the “unknown” token, usually written as “[UNK]” or “<unk>”. If you see a tokenizer producing many of these tokens, it is generally a bad sign: it means the tokenizer could not retrieve an index for a word, so information is lost during tokenization — a fatal shortcoming.

To summarize, the word-based tokenizer has several major drawbacks:

Huge vocabulary: because the vocabulary of a language or application domain is very large, a word-based tokenizer needs a very large vocabulary. This not only increases storage and computation costs, but also degrades training and inference efficiency. In practice, many words are low-frequency, and including large numbers of rare words in the vocabulary leads to data sparsity problems.
Difficulty handling morphological variation: many languages have rich morphology (plurals and tenses in English, grammatical gender in French, and so on). A word-based tokenizer has to include every one of these surface forms, inflating the vocabulary even further. Worse, different forms of the same word are treated as distinct tokens, so the model cannot share semantic information across them.
Limited coverage, unable to handle out-of-vocabulary words: a word-based tokenizer requires a fixed vocabulary. New words or misspellings that are not in the vocabulary simply cannot be handled, which hurts generalization.

A Finer-grained Approach — Byte Pair Encoding (BPE)

As discussed above, word-based tokenizers suffer from large vocabularies, difficulty with morphological variation, and limited coverage. Is there a tokenization method that solves all three?

Notice that English words reuse a great many letter combinations. Take “happy”: many words are built on it as a root, such as “happily”, “happiness”, and “unhappy”. These derived words can be decomposed into reusable letter combinations — the sequence “happ”, for instance, appears in all three. Beyond that, combinations like “ily” and “un” occur frequently in other English words, and pairing them with other fragments can express new words.

Naturally, this suggests a question: could we design a tokenizer around the idea of finding frequently occurring letter combinations, covering every English word by reusing them?

If so, such a tokenizer would resolve all three drawbacks of the word-based tokenizer, because:

High reuse means a small vocabulary: since words are built from letter combinations that are heavily reused, only a small number of combinations are needed to cover most English words.
Easy handling of morphological variation: in the example above, our vocabulary contains tokens like “ily” and “un”, which also appear frequently in other inflected words, making the method well suited to morphological variation.
Broad word coverage: the method starts from individual letters and searches for letter combinations, so it can cover every English word.

This method is in fact Byte Pair Encoding (BPE).

Byte Pair Encoding (BPE) is a subword segmentation technique that originated in data compression and is now widely used in natural language processing. The core idea of BPE is to iteratively merge the most frequent byte pair, building up subword units step by step and thereby segmenting the text. The method was first proposed by Philip Gage in 1994² for file compression, and was later introduced into machine translation by Sennrich et al. in 2015³ to address oversized vocabularies and sparsity.

Many well-known LLMs use BPE for tokenization, including the GPT series, BERT, RoBERTa, and T5. It is fair to say that BPE is now one of the classic algorithms of natural language processing.

In the next section we will build a BPE tokenizer step by step, following Andrej Karpathy’s project minbpe.

It is worth mentioning that Karpathy has many well-known LLM projects, such as llm.c and llama2.c.

Code Practice: Implementing BPE from Scratch

In this section we build a BPE tokenizer step by step, following Andrej Karpathy’s project minbpe⁴.

The Base Class Tokenizer

We start with the base form of BPE, namely the definition of the Tokenizer class:

class Tokenizer:
 """Base class for Tokenizers"""

 def __init__(self):
 # default: vocab size of 256 (all bytes), no merges, no patterns
 self.merges = {} # (int, int) -> int
 self.vocab = self._build_vocab() # int -> bytes

 def train(self, text, vocab_size, verbose=False):
 # Tokenizer can train a vocabulary of size vocab_size from text
 raise NotImplementedError

 def encode(self, text):
 # Tokenizer can encode a string into a list of integers
 raise NotImplementedError

 def decode(self, ids):
 # Tokenizer can decode a list of integers into a string
 raise NotImplementedError
 def _build_vocab(self):
 # vocab is simply and deterministically derived from merges
 vocab = {idx: bytes([idx]) for idx in range(256)}
 return vocab

The Tokenizer class has four basic functions — init, train, encode, and decode — corresponding to initializing, training, encoding, and decoding. Among them, train, encode, and decode are virtual functions that will be overridden by the subclass.

The minbpe project also supports saving and loading tokenizer models. Since this is only loosely related to our topic, and for reasons of space, we will not go into it here; interested readers can consult the source⁵.

In the init function we define two variables, self.merges = {} and self.vocab. The former records, during tokenization, the mapping from the indices of two adjacent original tokens to the index of the new token they merge into, i.e. (int, int) -> int. The latter represents the vocabulary, recording the mapping from an index to a token, i.e. int -> bytes.

self.vocab is initialized inside init, where we map 256 integers used as indices (idx) to the 256 tokens in hexadecimal representation (bytes([idx])).

If we print vocab inside the _build_vocab function, we get:

{0: b'\x00', 1: b'\x01', 2: b'\x02', 3: b'\x03', ..., 254: b'\xfe', 255: b'\xff'}

BPE

Next, here is the definition of the BPE class:

class BPE(Tokenizer):

 def __init__(self):
 super().__init__()

 def train(self, text, vocab_size, verbose=False):
 def get_stats(ids, counts=None):
 """
 Given a list of integers, return a dictionary of counts of consecutive pairs
 Example: [1, 2, 3, 1, 2] -> {(1, 2): 2, (2, 3): 1, (3, 1): 1}
 Optionally allows to update an existing dictionary of counts
 """
 counts = {} if counts is None else counts
 for pair in zip(ids, ids[1:]): # iterate consecutive elements
 counts[pair] = counts.get(pair, 0) + 1
 return counts

 def merge(ids, pair, idx):
 """
 In the list of integers (ids), replace all consecutive occurrences
 of pair with the new integer token idx
 Example: ids=[1, 2, 3, 1, 2], pair=(1, 2), idx=4 -> [4, 3, 4]
 """
 newids = []
 i = 0
 while i < len(ids):
 # if not at the very last position AND the pair matches, replace it
 if ids[i] == pair[0] and i < len(ids) - 1 and ids[i+1] == pair[1]:
 newids.append(idx)
 i += 2
 else:
 newids.append(ids[i])
 i += 1
 return newids

 assert vocab_size >= 256
 num_merges = vocab_size - 256
 # input text preprocessing
 text_bytes = text.encode("utf-8") # raw bytes
 ids = list(text_bytes) # list of integers in range 0..255
 # iteratively merge the most common pairs to create new tokens
 merges = {} # (int, int) -> int
 vocab = {idx: bytes([idx]) for idx in range(256)} # int -> bytes
 for i in range(num_merges):
 # count up the number of times every consecutive pair appears
 stats = get_stats(ids)
 # find the pair with the highest count
 pair = max(stats, key=stats.get)
 print(pair)
 # mint a new token: assign it the next available id
 idx = 256 + i
 # replace all occurrences of pair in ids with idx
 ids = merge(ids, pair, idx)
 # save the merge
 merges[pair] = idx
 vocab[idx] = vocab[pair[0]] + vocab[pair[1]]
 if verbose:
 print(f"merge {i+1}/{num_merges}: {pair} -> {idx} ({vocab[idx]}) had {stats[pair]} occurrences")
 # save class variables
 self.merges = merges # used in encode()
 self.vocab = vocab # used in decode()

 def decode(self, ids):
 # given ids (list of integers), return Python string
 text_bytes = b"".join(self.vocab[idx] for idx in ids)
 text = text_bytes.decode("utf-8", errors="replace")
 return text

 def encode(self, text):
 # given a string text, return the token ids
 text_bytes = text.encode("utf-8") # raw bytes
 ids = list(text_bytes) # list of integers in range 0..255
 while len(ids) >= 2:
 # find the pair with the lowest merge index
 stats = get_stats(ids)
 pair = min(stats, key=lambda p: self.merges.get(p, float("inf")))
 # subtle: if there are no more merges available, the key will
 # result in an inf for every single pair, and the min will be
 # just the first pair in the list, arbitrarily
 # we can detect this terminating case by a membership check
 if pair not in self.merges:
 break # nothing else can be merged anymore
 # otherwise let's merge the best pair (lowest merge index)
 idx = self.merges[pair]
 ids = merge(ids, pair, idx)
 return ids

The BPE class overrides train, encode, and decode, while init is inherited from the Tokenizer class.

Let us walk through the functions in the order train, encode, decode.

BPE.train

Inside train there are two helper functions, get_stats and merge.

def get_stats(ids): counts how often each pair of adjacent tokens occurs. For example, when ids = [1, 2, 3, 1, 2], the return value is {(1, 2): 2, (2, 3): 1, (3, 1): 1}⁶.

def merge(ids, pair, idx): merges every occurrence of pair in ids into idx and returns the new list of ids. For example, with ids=[1, 2, 3, 1, 2], pair=(1, 2), idx=4 -> [4, 3, 4], we replace the adjacent 1, 2 in ids with 4, yielding the new ids [4, 3, 4].

In train we first need to convert the characters of the string into indices, which we achieve with text_bytes = text.encode("utf-8") and ids = list(text_bytes) . For instance, when text = 'aaabbc', we get ids = [97, 97, 97, 98, 98, 99].

Next, inside the loop, get_stats first counts the frequency of each adjacent token pair in ids. The astute reader will already have guessed the purpose: it is to merge these high-frequency pairs into new tokens, and that is exactly what merge does. At this point the length of ids has been reduced. Finally, we record the mapping from token pair to index in merges and the mapping from index to new token in the vocabulary vocab.

This process repeats until the loop finishes.

So what is train actually doing?

In essence, train merges frequently co-occurring adjacent token pairs into a single new token. By adding only a small number of mappings to the vocabulary, we can substantially shorten the training text. Note that at the start the vocabulary contains only single letters. As training proceeds, frequent letter pairs are gradually added — first pairs of two letters, then longer letter sequences, and so on.

Taking the training text text = "happily happiness unhappy" as an example, the log produced by each of the three rounds is:

merge 1/3: (104, 97) -> 256 (b'ha') had 3 occurrences
merge 2/3: (256, 112) -> 257 (b'hap') had 3 occurrences
merge 3/3: (257, 112) -> 258 (b'happ') had 3 occurrences

Notice that after three iterations, happ has appeared in the vocabulary. This means the frequent letter combination happ can now be used in subsequent decoding and encoding.

BPE.encode

In encode we again start by finding the frequently occurring token pairs in the text. We then look for the token pair with the smallest index in merges. The reason is that during training, a smaller index means a higher frequency in the training set, which is statistically sensible⁷. We then map that token pair to its new token according to the mapping in merges, producing a new token sequence. This process repeats until no token pair can be merged (reduced) any further.

BPE.decode

The decode function takes a list of integer indices ids, converts it into the corresponding byte sequence, and decodes that byte sequence into a UTF-8 string, returning the decoded version of the original text. This function is the inverse of encode.

Conclusion

To summarize, training a BPE amounts to finding the token combinations that occur frequently in the training set; encoding replaces mergeable tokens in the text with new tokens to produce a new token sequence; and decoding is the inverse of encoding.

We can now explain the three advantages of BPE mentioned earlier in this article.

High reuse means a small vocabulary: unlike a word-based tokenizer, BPE does not need a mapping for every single word — it only needs mappings for frequently occurring token sequences.
Easy handling of morphological variation: however a word varies, frequently occurring token sequences such as the past-tense ed or the adverbial ly are already in the vocabulary and can be used efficiently during encoding.
Broad word coverage: even if the current word contains a token sequence that is absent from the vocabulary, we can still represent it with the individual letters that are in the vocabulary. BPE can therefore cover every English word.

https://huggingface.co/learn/nlp-course/zh-CN/chapter2/4 ↩︎
Gage P. A new algorithm for data compression[J]. The C Users Journal, 1994, 12(2): 23-38. ↩︎
Sennrich R, Haddow B, Birch A. Neural machine translation of rare words with subword units[J]. arXiv preprint arXiv:1508.07909, 2015. ↩︎
https://github.com/karpathy/minbpe ↩︎
https://github.com/karpathy/minbpe/blob/master/minbpe/base.py ↩︎
Of course, in a real string the letter a corresponds to index 97 and so on; [1, 2, 3, 1, 2] is used purely for ease of understanding. ↩︎
Note that BPE is a heuristic tokenization method and does not aim for an optimal solution, so the token sequence produced by encoding is not the shortest possible, only a relatively short one. ↩︎

Coding Practice | Implementing Self-Attention from Scratch

Sun, 10 Mar 2024 00:00:00 +0000

Introduction

In contemporary natural language processing (NLP) and deep learning, the Transformer and its core component — the self-attention mechanism — have revolutionized sequence modeling. Since Vaswani et al. first introduced it in the 2017 paper Attention Is All You Need, the Transformer has become the foundation for a wide range of complex tasks, including machine translation, text generation, speech recognition, and image processing.

Within the Transformer, self-attention has drawn wide interest from both academia and industry thanks to its outstanding performance. Attention lets the model access every element of the sequence at each time step. The key idea is selectivity: determining which words matter most in a particular context. It enriches the input embeddings by incorporating information about the surrounding input context. In other words, self-attention lets a model weigh the importance of different elements in the input sequence and dynamically adjust how much each contributes to the output.

In this article we implement self-attention by hand, following Understanding and Coding Self-Attention, Multi-Head Attention, Cross-Attention, and Causal-Attention in LLMs. A Chinese translation of that article was published by Synced as “Still Don’t Understand Self-Attention in the Age of LLMs? This Article Walks You Through Implementing It from Scratch”. Relative to those two pieces, this article adds my own reflections and experience.

Embedding

In this section we use embeddings to turn the discrete symbols of an input sequence (words or characters) into continuous, high-dimensional vector representations. Put simply, this is necessary because deep learning models cannot understand raw text directly; they must learn the semantics and syntax of text from the information carried by these vectors. For a deeper treatment, see the Zhihu article “Understanding Embeddings and How They Relate to Deep Learning” (in Chinese).

Given a sentence, we want to turn it into a continuous vector representation via embedding. We will use the sentence “Life is short, eat dessert first” as our running example.

In the preprocessing stage, we deduplicate the words in the sentence and map each word to an integer index. This is easy to express in Python.

Input:

sentence = 'Life is short, eat dessert first'

dc = {s:i for i,s
 in enumerate(sorted(sentence.replace(',', '').split()))}

print(dc)

Output:

{'Life': 0, 'dessert': 1, 'eat': 2, 'first': 3, 'is': 4, 'short': 5}

With this word-to-index mapping, the sentence can be represented as a sequence of indices.

Input:

import torch

sentence_int = torch.tensor(
 [dc[s] for s in sentence.replace(',', '').split()]
)
print(sentence_int)

Output:

tensor([0, 4, 5, 2, 1, 3])

Now we are ready for the key step, the embedding itself.

In an embedding, each word is represented by a multi-dimensional vector whose dimensionality is determined by the size of the vocabulary. Llama 2, for example, uses an embedding size of 4096. To keep things compact, we use three dimensions in this article.

Input:

vocab_size = 50_000

torch.manual_seed(123)
embed = torch.nn.Embedding(vocab_size, 3)
embedded_sentence = embed(sentence_int).detach()

print(embedded_sentence)
print(embedded_sentence.shape)

Output:

tensor([[ 0.3374, -0.1778, -0.3035],
 [ 0.1794, 1.8951, 0.4954],
 [ 0.2692, -0.0770, -1.0205],
 [-0.2196, -0.3792, 0.7671],
 [-0.5880, 0.3486, 0.6603],
 [-1.1925, 0.6984, -1.4097]])
torch.Size([6, 3])

As the output shows, the six words of our sentence are represented as six row vectors, each of dimension 3. That is, every word in the sentence is represented by a vector of three numbers.

Defining the Weight Matrices

Starting with this section, we introduce the famous “QKV” mechanism of the Transformer.

Self-attention uses three weight matrices, denoted $W_q$, $W_k$, and $W_v$; they are model parameters and are adjusted throughout training. Their role is to project the input into the query, key, and value components of the sequence.

The corresponding query, key, and value sequences are obtained by matrix multiplication between the weight matrix $W$ and the embedded input $x$:

Query sequence: for $i$ in the sequence $1\ldots T$, $q^{(i)}=x^{(i)}W_q$
Key sequence: for $i$ in the sequence $1\ldots T$, $k^{(i)}=x^{(i)}W_k$
Value sequence: for $i$ in the sequence $1\ldots T$, $v^{(i)}=x^{(i)}W_v$
The index $i$ refers to the $token$ index position in the input sequence, whose length is $T$.

Here, both $q^{(i)}$ and $k^{(i)}$ are vectors of dimension $d_k$. The projection matrices $W_q$ and $W_k$ have shape $d × d_k$, while $W_v$ has shape $d × d_v$. Here $d$ denotes the number of dimensions of each word vector $x$, which is 3 in this article.

Because we need to compute the dot product of the query and key vectors, these two vectors must have the same number of elements ($d_q=d_k$). Many LLMs also use value vectors of the same size, that is, $d_q=d_k=d_v$. However, the number of elements in the value vector $v^{(i)}$ can be arbitrary; it determines the size of the resulting context vector.

In the code that follows, we set $d_q=d_k=2$ and $d_v=4$. The projection matrices are initialized as follows:

Input:

torch.manual_seed(123)

d = embedded_sentence.shape[1]

d_q, d_k, d_v = 2, 2, 4

W_query = torch.nn.Parameter(torch.rand(d, d_q))
W_key = torch.nn.Parameter(torch.rand(d, d_k))
W_value = torch.nn.Parameter(torch.rand(d, d_v))

In the original paper Attention Is All You Need, $d_q$, $d_k$, and $d_v$ are typically set to 64, while the total model dimension is 512.

Computing Unnormalized Attention Weights

In this section we walk through the computation using the second token as our example:

As the figure shows, we multiply the input $x$ by $W_q$, $W_k$, and $W_v$ respectively.

The code is as follows:

x_2 = embedded_sentence[1]
query_2 = x_2 @ W_query
key_2 = x_2 @ W_key
value_2 = x_2 @ W_value

print(query_2.shape)
print(key_2.shape)
print(value_2.shape)

Output:

torch.Size([2])
torch.Size([2])
torch.Size([4])

Generalizing, we can multiply embedded_sentence by $W_k$ and $W_v$ respectively; these results will be used in the following steps.

Input:

keys = embedded_sentence @ W_key
values = embedded_sentence @ W_value

print("keys.shape:", keys.shape)
print("values.shape:", values.shape)

Output:

keys.shape: torch.Size([6, 2])
values.shape: torch.Size([6, 4])

Now that we have $query^{(2)}$ along with all the keys and values, let us compute the unnormalized attention weights $\omega$, as shown below:

As the figure shows, $\omega (i,j)$ is the dot product between the query and key sequences: $\omega (i,j) = q^{(i)}k^{(j)}$.

For example, we can compute the unnormalized attention between the 2nd token’s query and the 5th token as follows:

Input:

omega_24 = query_2.dot(keys[4])
print(omega_24)

Output:

tensor(1.2903)

Generalizing, we can multiply query_2 by keys to obtain the unnormalized attention scores between the 2nd token and every other token.

Input:

omega_2 = query_2 @ keys.T
print(omega_2)

Output:

tensor([-0.6004, 3.4707, -1.5023, 0.4991, 1.2903, -1.3374])

Computing the Attention Weights

In the previous section we computed the unnormalized attention scores between the 2nd token and every other token. In practice, we also need to normalize these scores.

The purpose of normalization is to put the attention weight at each sequence position between 0 and 1, with all positions summing to 1. This probabilistic interpretation lets the model express, in probabilistic terms, how much attention it pays to different parts of the sequence: a higher weight means the model attends more to that position.

As the figure shows, self-attention first scales $\omega$ by $1/√{d_k} $ and then normalizes it with the softmax function.

Scaling by $d_k$ ensures that the Euclidean lengths of the weight vectors are all roughly on the same scale. This helps prevent the attention weights from becoming too small or too large — which could cause numerical instability or hurt the model’s ability to converge during training.

We can implement the attention weight computation in code like this:

Input:

import torch.nn.functional as F

attention_weights_2 = F.softmax(omega_2 / d_k**0.5, dim=0)
print(attention_weights_2)

Output:

tensor([0.0386, 0.6870, 0.0204, 0.0840, 0.1470, 0.0229])

The final step is to compute the context vector $z^{(2)}$ — an attention-weighted version of the original query input $x^{(2)}$ that incorporates all the other input elements as context through the attention weights:

Input:

context_vector_2 = attention_weights_2 @ values

print(context_vector_2.shape)
print(context_vector_2)

Output:

torch.Size([4])
tensor([0.5313, 1.3607, 0.7891, 1.3110])

Note that this output vector has more dimensions ($d_v=4$) than the input vector ($d=3$), because we set $d_v > d$ earlier. The embedding size $d_v$, however, can be chosen arbitrarily.

Self-Attention

Now let us pull together the code for the self-attention mechanism from the previous sections.

We can condense everything above into a compact Self-Attention class:

import torch.nn as nn

class SelfAttention(nn.Module):

 def __init__(self, d_in, d_out_kq, d_out_v):
 super().__init__()
 self.d_out_kq = d_out_kq
 self.W_query = nn.Parameter(torch.rand(d_in, d_out_kq))
 self.W_key = nn.Parameter(torch.rand(d_in, d_out_kq))
 self.W_value = nn.Parameter(torch.rand(d_in, d_out_v))

 def forward(self, x):
 keys = x @ self.W_key
 queries = x @ self.W_query
 values = x @ self.W_value

 attn_scores = queries @ keys.T # unnormalized attention weights 
 attn_weights = torch.softmax(
 attn_scores / self.d_out_kq**0.5, dim=-1
 )

 context_vec = attn_weights @ values
 return context_vec

Input:

torch.manual_seed(123)

# reduce d_out_v from 4 to 1, because we have 4 heads
d_in, d_out_kq, d_out_v = 3, 2, 4

sa = SelfAttention(d_in, d_out_kq, d_out_v)
print(sa(embedded_sentence))

Output:

tensor([[-0.1564, 0.1028, -0.0763, -0.0764],
 [ 0.5313, 1.3607, 0.7891, 1.3110],
 [-0.3542, -0.1234, -0.2627, -0.3706],
 [ 0.0071, 0.3345, 0.0969, 0.1998],
 [ 0.1008, 0.4780, 0.2021, 0.3674],
 [-0.5296, -0.2799, -0.4107, -0.6006]], grad_fn=<MmBackward0>)

The second row is exactly the value of context_vector_2 from the previous section: tensor([0.5313, 1.3607, 0.7891, 1.3110]).

Conclusion

In this article we implemented self-attention by hand, with code and explanations, following Understanding and Coding Self-Attention, Multi-Head Attention, Cross-Attention, and Causal-Attention in LLMs.