Understanding Single-Head Attention in TransformersBy Markelic / July 31, 2026 Table of contents Introduction Attention What problem does attention solve? A conceptual overview Embeddings are vectors A quick review on neural network training How can the model transform a static embedding into its contextualized version without knowing that version beforehand? How do we take context into account? The database analogy How to compute the attention weights for a single query? How to apply this to embeddings? Why scaling? How to compute the attention weights for all queries using matrix notation? Scaled dot-product attention (Single-Head) Where do the vectors come from? Abstract Queries vs. Human Questions So, how does a filtered vector become a “question”? How to code single-head attention? Summary What’s Next? References Introduction To understand transformers, you must understand attention. Many explanations focus on the how rather than the why, making terms like queries, keys, and the underlying math feel arbitrary. You can memorize the attention formula and still have no idea why it actually works. But once you grasp the core intuition behind each step, attention becomes easy to remember and reason about. In this article, we’ll explore both the how and the why. Specifically, we will look at: The role and intuition behind queries, keys, and values. Why we “scale” in scaled dot-product attention. This article focuses on single-head attention. Its extension, multi-head attention, will be covered in a follow-up piece. Note: English is not my native language, so I used AI for proofreading and polish. The concepts and explanations are entirely my own. Attention Vaswani et al. [1] established modern attention as the core of the transformer architecture. While attention itself was not new (see [2]), their variant—Scaled Dot-Product Self-Attention—changed how models process text. Prior architectures like RNNs and LSTMs processed text sequentially word by word, creating two critical bottlenecks: No parallelization: Step-by-step processing slowed down training. Vanishing long-range context: Connecting distant words (e.g., a subject at the start of a long sentence to a verb 20 words later) was difficult. Self-attention eliminated these bottlenecks by allowing every token (a word or sub-word unit) to directly attend to every other token in a single operation. This direct connection enabled massive parallelization and solved context loss, paving the way for modern large-scale models. In this article, we will examine the self-attention mechanism as presented in the original paper. A basic familiarity with vector embeddings, elementary Python, and linear algebra fundamentals—specifically matrix multiplication and dot products—will help, though key concepts are explained along the way. What problem does attention solve? A conceptual overview Let’s first clarify the problem we want to solve. Understanding language is hard, partly because the meaning of a word depends on its context. Take these two sentences: “You should duck when the ceiling is low.” “The hunter saw a duck in the pond.” In the first sentence, “duck” functions as a verb (duck: action), which we deduce because it appears alongside “ceiling” and “low.” The second sentence refers to the animal (duck: animal), which we infer from “hunter” and because the duck is the direct object being “seen.” You can’t physically see a verb. This context dependency posed a major challenge for natural language processing (NLP). But before tackling context, researchers had to solve an even more fundamental issue: how to represent words so computers could process them at all. Solving this was an incremental process, evolving from one-hot encodings to count-based approaches and eventually generative methods ([3], [4]) like Word2Vec [5]. The unifying idea was mapping each word to a vector of numbers—a word embedding. A word embedding captures semantic meaning, placing words with similar meanings close together in vector space. For example, the embeddings for “man” and “boy” should be close to each other, but far from an unrelated word like “banana.” The text processing pipeline works roughly as follows: Tokenization: A sequence of text is split into pieces (words, subwords, or characters), and each unique piece is assigned a numerical identifier called a token. (For simplicity in this text we’ll assume that one word is mapped to a token.) In frameworks like Hugging Face, the AutoTokenizer class handles this step. Embedding Lookup: Each token is mapped to a vector, historically using pre-trained lookup tables like Word2Vec or GloVe [6]. Today, models typically start with randomly initialized embeddings that are learned end-to-end during training. The important point here is that regardless of how these embeddings are initialized, the same token always maps to the same vector. Because of this, we call them static embeddings. Static embeddings are completely oblivious to context. They provide a solid baseline representation of a word’s default meaning, but they must be refined to fit the sentence they appear in. And this is precisely what the attention mechanism solves: it takes static embeddings as input and updates them into contextual embeddings based on the surrounding tokens. This transformation is illustrated in Figure 1. Notice that the input and output vectors share the same dimensionality (\(d_{\text{model}}\)), while \(seq_{\text{len}}\) denotes the number of tokens processed together. Figure 1: Attention creates contextualized versions of static input embeddings. Input and output vectors share the same dimension, \(d_{\text{model}}\). Here, vectors are represented as row vectors, and \(seq_{\text{len}}\) denotes sequence length. To see how contextual embeddings are computed, we can break the mechanism down into three core ideas: Vector Transformations: Embeddings are vectors and vectors can be transformed into other vectors via matrix multiplication. Weighted Averaging: To incorporate context, attention uses a form of weighted averaging. This was originally introduced via a database analogy. This analogy is where the terms queries, keys, and values originate: in a traditional database, you send a query to match an exact key and retrieve its associated value. Attention extends this by replacing exact matching with soft matching. Every key receives a relevance score relative to the query, yielding a weighted sum of all values. (Statisticians will recognize this as a vector-valued generalization of Nadaraya-Watson kernel regression; see [7], [8].) Learned Projections: Queries, keys, and values are vectors generated by multiplying static input embeddings by three projection matrices. Learning these matrices during training is how the transformer adapts to language. Let’s examine each of these ideas in detail. Embeddings are vectors It might sound trivial, but this is a fundamental concept to keep in mind: embeddings are vectors. As such, they point to specific positions in an \(n\)-dimensional space (see Figure 2). Consider an ambiguous word whose meaning depends on context—like “duck” from our previous example. It is helpful to think of the initial, static embedding of “duck” as a baseline vector that lies somewhere between its possible specific meanings. This idea is visualized in Figure 2, where the static vector for “duck” is shown in green, alongside the target vectors for duck:animal and duck:action. If we encounter a sentence where the context demands the “action” meaning, the attention mechanism must do two things: 1. Identify the target: Determine from the surrounding tokens that we need to head toward duck:action (the black vector). 2. Apply the transformation: Shift the green static vector to its new, contextual position in space. Figure 2: Visualization of word embeddings in a 2D space. The green vector represents the static baseline embedding for “duck.” Context reveals that the sentence requires the specific meaning “duck:action” (black vector). Attention determines this target from the surrounding tokens and transforms the static vector into its contextualized representation. Now, how can a model transform the static embedding into a contextualized version when it doesn’t even know what that final version should look like? This seems like an impossible task. To see how this works, let’s take a quick step back to review the core idea of neural networks and how they are trained. (Feel free to skip the next section if you are already familiar with the basics.) A quick review on neural network training Remember that the attention mechanism is a core component of the transformer—a neural network composed of layers with trainable parameters. The central premise of neural networks is that an optimal set of parameter values exists to produce our desired output. In image classification, for example, we input images and expect correct labels. Initially, random parameters yield poor predictions. Training is the process of discovering parameters that make those predictions accurate. Crucially, we never calculate these ideal values explicitly. Instead, we frame learning as an optimization problem. We measure performance using a loss function—which quantifies how far our predictions are from the correct answers. Minimizing this loss directly yields a well-performing model. Optimization algorithms (typically variants of gradient descent) dictate how to adjust parameters at each step to steadily reduce error. We repeat this process across many examples until performance stabilizes. The key takeaway: we don’t manually program the parameters. Instead, we define the architecture, the loss function, and the optimization method—and the model learns the ideal parameters through training. With this reminder in place, we can return to our original question. How can the model transform a static embedding into its contextualized version without knowing that version beforehand? The short answer: The transformation that “finds” the right contextual embedding is learned. To understand this, recall that our goal is to map an input vector to a specific output vector. Mathematically, this is a linear transformation: we multiply the input vector by a matrix (see Figure 3). Figure 3: An input vector is multiplied by a matrix, resulting in a transformed output vector. The matrix parameters act as trainable weights that are adjusted during neural network training. Initially, this matrix is filled with random values, producing poor, unhelpful embeddings that lead to a high loss (see Figure 4, top row). During training, the optimization algorithm continuously nudges these parameters toward values that minimize the loss (see Figure 4, middle row). Crucially, the matrix entries are what we actually learn. Once training concludes and the transformer achieves high performance, the matrix has successfully encoded the exact transformation required. We discover both the transformation and the final embeddings without ever having to explicitly define or know those target embeddings beforehand (see Figure 4, bottom row). Figure 4: A vector \(\mathbf{x}\) is multiplied by a transformation matrix \(\mathbf{T}\), resulting in the projected vector \(\mathbf{y}\). Initially, the matrix entries are random, leading to embeddings that cause a high loss (top row). The learning mechanism in the neural network adapts the matrix parameters to lower the loss, producing better projected vectors (middle and bottom rows). This is the foundational principle behind contextual embeddings. However, a single standard linear transformation isn’t enough to capture complex language relationships. To truly incorporate context, the attention mechanism uses not just one, but several interacting matrices, each playing a distinct role. That is what we will explore next. How do we take context into account? Our goal is to understand how the meaning of a word is shaped by its surrounding words—its context. Every word in the sequence contributes to creating the final contextual embedding, i.e. every word in the sequence is part of the context. The core idea of attention is: each contextual embedding is a weighted sum of all static input embeddings. The weights reflect how relevant each surrounding word (represented by its static input embedding) is to the target word for which we want to compute the contextual embedding. As already mentioned, if you are familiar with non-parametric statistics, you can view this as an instance of Nadaraya-Watson kernel regression (see [7]). If not, don’t worry - the following database analogy is an intuitive way to picture it. The database analogy Imagine a simple database of colors, which we can represent as a Dictionary in Python as follows: import numpy as np database = { 'Red': np.array([255, 0, 0 ]), 'Blue': np.array([0, 0, 255]), 'Yellow': np.array([255, 255, 0 ]), 'Green': np.array([0, 128, 0 ]) } The keys are color names and the values are their RGB vectors. If you query the database for an exact key, like database['Red'], you get [255, 0, 0]. This is equivalent to a weighted sum where the weight for ‘Red’ is 1 and all others are 0. But what if the query were something more abstract? For example: “What is the color of a sunset?” There is no exact match for that in our database. Instead, we match the keys against this query in a “soft,” probabilistic way. Imagine the keys as individuals scoring their own relevance to the question: ‘Red’ says: “A sunset contains a large part of me.” It claims a high relevance weight of 0.6. ‘Yellow’ feels relevant but a bit less so, responding with 0.3. ‘Blue’ senses a faint connection at twilight, giving a 0.1. ‘Green’ bows out entirely with a 0.0. This gives a clean set of weights that are between 0 and 1 and sum to 1: Key Attention Weight Red 0.6 Yellow 0.3 Blue 0.1 Green 0.0 In deep learning literature, these are called attention weights. They dictate exactly how much attention to pay to each key. The final result of our query is a weighted sum of the color vectors: result = 0.6*database['Red'] + 0.3*database['Yellow'] + 0.1*database['Blue'] \ + 0.0*database['Green'] print(result) # = [229.5 76.5 25.5], a warm orange-red [229.5 76.5 25.5] This might seem like a silly example, but it captures exactly the idea behind attention: the result is not a single retrieved item, but a sum of all values, weighted by their relevance to the query. This is how context gets taken into account — the more relevant a key within the sequence is to the query, the more its value contributes to the final result. And that final result is, of course, the contextual embedding. Now that we have the intuition down, the remaining question is practical: how do we actually calculate these weights mathematically, and where do the queries and keys come from? That is what we will look at next. How to compute the attention weights for a single query? Computing attention weights is straightforward — if queries, keys, and values are vectors. Imagine a query \(\mathbf{q}\) is not a sentence in human language, but a vector of numbers the computer can work with directly. The same goes for a key, \(\mathbf{k}\). To compute how relevant \(\mathbf{k}\) is to \(\mathbf{q}\), we can compute the inner (dot) product of the two vectors - a standard way to measure similarity between vectors. The more similar two vectors are, the larger their dot product. We allow ourselves some liberty here and interpret this similarity as “relevance” of the key to the query. Because \(\mathbf{q}\) and \(\mathbf{k}\) are row vectors, their inner product is simply \(\mathbf{q}\mathbf{k}^T = z\), where \(z\) is a scalar representing relevance. (Note that I assume here that the query and key vectors are row vectors. This is just a convention, we could also use column vectors, but we would need to adjust the formulae accordingly. Column vectors are more common in the math and ML literature, but when programming we usually need to convert them to row vectors due to how vectors are efficiently stored in computer memory. Therefore, I thought we might as well use row vectors to start with.) How to apply this to embeddings? Recall from Figure 1 that we have several static embeddings of which we want to compute contextualized versions. We denoted the number of static input embeddings that are processed together as one sequence as \(seq_{\text{len}}\). Now imagine that for each static embedding \(\mathbf{x}_i\), a magical process creates three distinct vectors: a query \(\mathbf{q}_i\), a key \(\mathbf{k}_i\), and a value \(\mathbf{v}_i\) as visualized in Figure 5. (More on that in Section 8) Figure 5: Some operation creates three vectors from each static input embedding. We can interpret this as a database with two entries: database={\(\mathbf{k}_1\):\(\mathbf{v}_1\), \(\mathbf{k}_2\):\(\mathbf{v}_2\)}. We also have two queries, \(\mathbf{q}_1\) and \(\mathbf{q}_2\). Now, as explained, to compute how relevant the database entries are with respect to each query we compute the dot products between a query and each key. Thus, we use the dot product as similarity function. Concretely, to compute how relevant each key is for \(\mathbf{q}_1\) in our example, we compute the dot products \(\mathbf{q}_1 \mathbf{k}_1^T\) and \(\mathbf{q}_1 \mathbf{k}_2^T\). The results are two weights which we can collect in a vector \(\mathbf{z}_1\). Furthermore, if we stack all key row vectors into a matrix \(\mathbf{K}\), we can compute the relevance scores for \(\mathbf{q}_1\) against all keys at once using a vector-matrix product: \[ \mathbf{z}_1 = \mathbf{q}_1 \mathbf{K}^T \] And here is a visual for a better mental picture: Figure 6: The relevance of each database entry indexed by its key \(\mathbf{k}_1\), and \(\mathbf{k}_2\) with respect to \(\mathbf{q}_1\) is the dot product between the query and the according key, i.e. \(\mathbf{q}_1 \mathbf{k}_1\), and \(\mathbf{q}_1 \mathbf{k}_2\). We can conveniently compute this in one go by writing the computation as a vector-matrix product. The result is a vector \(\mathbf{z}\) which contains all the relevance scores for the specific query. These unnormalized entries in the vector \(\mathbf{z}_1\) are often referred to as scores. We have established that each individual score is a dot product of two vectors and the vectors are of dimension \(d_{\text{model}}\). This means that for each score we had to add \(d_{\text{model}}\) products together. The larger \(d_{\text{model}}\) the more numbers we add together the larger the score. But we don’t want very large scores as they can negatively affect the model training, thus we need to keep their values in check. This can be achieved by normalizing each score by dividing it by the square root of the vector dimension \(d_{\text{model}}\), thus, \(\mathbf{z}_1 / \sqrt{d_{\text{model}}}\). More on this in the optional chapter Why scaling. Also, recall from the preceding discussion that we want the attention weights for a query to sum to one. We can achieve this by passing the normalized scores vector through the softmax function. The result of this - finally - are the attention weights \(\boldsymbol{\alpha}_1\): \[ \boldsymbol{\alpha}_1 = \text{softmax}\left(\dfrac{\mathbf{z}_1}{\sqrt{d_{\text{model}}}}\right) \] Note that \(\boldsymbol{\alpha}_1\) is a vector that contains all the attention weights for \(\mathbf{q}_1\) with respect to all keys in the database. The sum of all entries is 1. Why scaling? As a reminder: we are looking at single-head scaled dot-product attention. Here is where the term scaled comes from. We must divide the vector \(\mathbf{z}_1\) by \(\sqrt{d_{\text{model}}}\), the square root of the query vector’s dimension. To see why we need it, consider what happens when the query and key vectors have many dimensions, thus \(d_{\text{model}}\) is a large number. The dot product \(\mathbf{q}_1 \cdot \mathbf{k}_i\) sums \(d_{\text{model}}\) individual products \(q_{1,j} k_{i,j}\). If we assume each component of \(\mathbf{q}_1\) and \(\mathbf{k}_i\) is independent, zero-mean, and unit-variance, then each individual product \(q_{1,j} k_{i,j}\) also has variance 1 - but the dot product itself, being a sum of \(d_{\text{model}}\) such independent terms, has variance \[ \text{Var}(\mathbf{q}_1 \cdot \mathbf{k}_i) = \sum_{j=1}^{d_{\text{model}}} \text{Var}(q_{1,j} k_{i,j}) = d_{\text{model}}. \] In practice, this means the scores in \(\mathbf{z}_1\) become more spread out as \(d_{\text{model}}\) grows — even when no key is genuinely more relevant than another, some scores will land noticeably higher than others purely by chance. This is a problem for softmax, since it’s an exponential function: even a modest gap between the largest score and the rest gets amplified dramatically. The result is that softmax becomes overconfident, pushing nearly all the weight onto a single key and squeezing the rest toward zero — a phenomenon called saturation. Saturated softmax outputs have tiny gradients almost everywhere, which slows down or stalls learning. (You can refer to this post if you need to brush up on the derivative of the softmax function.) Dividing every score by \(\sqrt{d_{\text{model}}}\) fixes this: since the variance of the scores grows with \(d_{\text{model}}\), dividing by \(\sqrt{d_{\text{model}}}\) brings it back down to a constant, regardless of dimensionality — keeping softmax’s input in a well-behaved range no matter how large \(d_{\text{model}}\) is. How to compute the attention weights for all queries using matrix notation? Thus far we have computed the attention weights for a single query. Note that with a tiny change in our computation we can compute the attention weights for all queries in one go. For that we simply collect all query vectors in a matrix \(\mathbf{Q}\) as shown the figure below. The output is the matrix \(\mathbf{Z}\) with each row i being the relevance scores for query i. Figure 7: Computing all attention weights at once using matrix notation. Again, to compute the attention weights we need to normalize the scores and pass each row through softmax (row-wise softmax). This can now be compactly written as: \[ \mathbf{A} = \text{softmax}\left(\frac{\mathbf{QK}^T }{ \sqrt{d_{\text{model}}}}\right) \] Thus now \(\mathbf{A}\) is the matrix with each row \(\boldsymbol{\alpha}_i\) being the attention weights for query i. Scaled dot-product attention (Single-Head) As mentioned in the introduction, we start with the simpler single-head version first. Once we master how it works, scaling up to multi-head attention will be easy. Now that we can compute the attention weights, we are ready to compute the contextual embeddings—which is our ultimate goal. Recall from our database example that a contextual embedding is the weighted sum of the values of each database entry, with the weights being the attention weights. In our setup, this means the contextual embedding for the static input embedding \(\mathbf{x}_1\) is the weighted sum of the value vectors \(\mathbf{v}_1\) and \(\mathbf{v}_2\), weighted by the entries in the vector \(\boldsymbol{\alpha}_1\). Thus, the contextual embedding \(\mathbf{y}_1\) is simply: \[ \mathbf{y}_1 = \alpha_{1,1} \mathbf{v}_1 + \alpha_{1,2} \mathbf{v}_2 \] where \(\alpha_{1,1}\) denotes the first component of \(\boldsymbol{\alpha}_1\) and \(\alpha_{1,2}\) the second. Again, we can collect the row vectors \(\mathbf{v}_i\) into a matrix \(\mathbf{V}\) and write the sum above as a vector-matrix product: \[ \mathbf{y}_1 = \boldsymbol{\alpha}_1 \mathbf{V}, \qquad \boldsymbol{\alpha}_1 \in \mathbb{R}^{1 \times seq_{\text{len}}}, \quad \mathbf{V} \in \mathbb{R}^{seq_{\text{len}} \times d_{\text{model}}} \] The result \(\mathbf{y}_1 \in \mathbb{R}^{1 \times d_{\text{model}}}\) is a vector of the exact same dimension as the static input embedding it replaces—just now enriched with context. That’s our contextual embedding! Analogous to what we’ve done before, we can collect all weight vectors \(\boldsymbol{\alpha}_i\) into the matrix \(\mathbf{A}\). The product \(\mathbf{AV}\) then gives all contextual embeddings \(\mathbf{Y}\) in one go. This is visualized below for two database entries: Figure 8: Computing all contextual embeddings at once using matrix notation, shown for two database entries. Each weight vector consists of individual weights, here denoted as \(\alpha_{1,1}\), \(\alpha_{1,2}\), etc. The contextual embedding of static input vector \(i\) is the sum of value vectors weighted by the corresponding attention weight. Putting it all together gives us the exact scaled dot-product attention formula from Vaswani et al.: \[ \text{Attention}(\mathbf{Q}, \mathbf{K}, \mathbf{V}) = \mathbf{AV} = \text{softmax}\left(\frac{\mathbf{Q}\mathbf{K}^T}{\sqrt{d_{\text{model}}}}\right)\mathbf{V} \] That’s it! Take a moment to pat yourself on the back—we’ve just covered the core mechanics of attention and as such the heart of the transformer architecture. We aren’t quite done yet, though. Two remaining questions stand in our way: Where do these vectors \(\mathbf{q}_i\), \(\mathbf{k}_i\), and \(\mathbf{v}_i\) actually come from? What is multi-head attention? Neither of these is difficult now that you’ve mastered single-head attention. We will answer the first question immediately and the second in the promised follow-up article. Where do the vectors come from? Earlier I said that there is a process that creates \(\mathbf{q}_i\), \(\mathbf{k}_i\), and \(\mathbf{v}_i\) for each static input vector \(\mathbf{x}_i\). In Figure 5 it was called a “magical operation”. Ok, I’ll be honest - it’s not magic. It’s simply a linear projection and here is how it’s done: We initialize three weight matrices \(\mathbf{W}_Q\), \(\mathbf{W}_K\), and \(\mathbf{W}_V\) with small random values. We then multiply each static input embedding \(\mathbf{x}_i\) with these matrices to obtain the query, key, and value vectors: \[\mathbf{q}_i = \mathbf{x}_i \cdot \mathbf{W}_Q\] \[\mathbf{k}_i = \mathbf{x}_i \cdot \mathbf{W}_K\] \[\mathbf{v}_i = \mathbf{x}_i \cdot \mathbf{W}_V\] During training the model learns the values inside these matrices so that they produce contextual embeddings that minimize the loss. This is the core principle we introduced in Section 3.2. As before, we can use matrix notation to compute all query, key and value vectors in a single pass. By stacking all static input embeddings \(\mathbf{x}_i\) into a sequence matrix \(\mathbf{X}\) we compute \[\mathbf{Q} = \mathbf{X} \cdot \mathbf{W}_Q\] \[\mathbf{K} = \mathbf{X} \cdot \mathbf{W}_K\] \[\mathbf{V} = \mathbf{X} \cdot \mathbf{W}_V\] Each row in the resulting matrices represents the query, key, or value vector for the corresponding input embedding. Figure 9 visualizes the operation \(\mathbf{XW}_Q = \mathbf{Q}\). The matrices \(\mathbf{K}\) and \(\mathbf{V}\) are computed identically using \(\mathbf{W}_K\) and \(\mathbf{W}_V\) . To perform this multiplication, \(\mathbf{W}_Q\) must have \(d_{\text{model}}\) rows, but we are free to choose the number of columns, which we denote as \(d_{\text{head}}\). The resulting query, key, and value vectors will have dimension \(d_{\text{head}}\). In the simplest setup, we set \(d_{\text{head}} = d_{\text{model}}\). This is the single-head attention setup we have been exploring so far. We’ll look at other choices in the follow-up article about multi-head attention. Figure 9: Visualizing \(\mathbf{Q} = \mathbf{X}\mathbf{W}_Q\) for two input embeddings: Each row in \(\mathbf{X}\) is a static input embedding \(\mathbf{x}_i\) of dimension \(d_{\text{model}}\). A sequence of \(seq_{\text{len}}\) embeddings forms the matrix \(\mathbf{X}\) of shape \(seq_{\text{len}} \times d_{\text{model}}\). Multiplying \(\mathbf{X}\) by \(\mathbf{W}_Q\) yields matrix \(\mathbf{Q}\) of shape \(seq_{\text{len}} \times d_{\text{head}}\), where row \(i\) contains the query vector \(\mathbf{q}_i\) for static input embedding \(\mathbf{x}_i\). Abstract Queries vs. Human Questions Something I personally found confusing when first reading about transformers was the frequent mention of “questions that tokens ask about other tokens.” Without seeing the math behind the analogy, I couldn’t comprehend what these “questions” were or where they came from. Here, we have already established that these “questions” are query vectors (\(\mathbf{q}_i\))—vectors of numbers, not natural language sentences. They are produced by multiplying a static input embedding \(\mathbf{x}_i\) by \(\mathbf{W}_Q\) (as shown in Section 8). But why is it valid to interpret a query \(\mathbf{q}_i\) as a “question,” other than that it fits nicely into the database analogy? To see why this claim holds up, let’s look at how a query vector is created. Figure 10 schematically shows a 3-dimensional static embedding \(\mathbf{x}_i\) (colored rectangles) multiplied by \(\mathbf{W}_Q\). The components of \(\mathbf{W}_Q\) are represented as black squares whose sizes correspond to their values: larger squares represent larger weights. To compute \(\mathbf{x}_i \mathbf{W}_Q\), each component of \(\mathbf{x}_i\) is weighted by the corresponding component in \(\mathbf{W}_Q\). The resulting query vector \(\mathbf{q}_i\) (bottom row) consists of the sum of these weighted entries. In other words, the components of \(\mathbf{q}_i\) are the dot products \(\mathbf{x}_i\mathbf{w}_{q,1}\) and \(\mathbf{x}_i \mathbf{w}_{q,2}\). (With \(\mathbf{w}_{q,1}\), and \(\mathbf{w}_{q,2}\) denoting the first and second column vector in \(\mathbf{W}\).) Figure 10: Showing how \(\mathbf{W}_Q\) acts as a filter on \(\mathbf{x}_i\). See text for further information. Notice how the first component of \(\mathbf{q}_i\) is dominated by the green entry (the third component of \(\mathbf{x}_i\)), while other components are dampened nearly to zero. We can clearly see how \(\mathbf{W}_Q\) acts as a filter: it can isolate and amplify specific features of \(\mathbf{x}_i\) while silencing others. So, how does a filtered vector become a “question”? The answer lies in what happens next when we compute the dot product between a query and a key (\(\mathbf{q}_i \cdot \mathbf{k}_j\)). A dot product outputs a high score only when two vectors share large numbers in the same positions. Because \(\mathbf{W}_Q\) filtered \(\mathbf{x}_i\) to emphasize a specific trait, the resulting vector \(\mathbf{q}_i\) will only match strongly with keys (\(\mathbf{k}_j\)) that possess that exact same trait. In other words: Producing a query vector that isolates a specific feature and then computing the dot product with all keys (as in \(\mathbf{q}_i\mathbf{K}^T\)) is mathematically equivalent to broadcasting: “Who else in this sequence has this feature?” This is why calling \(\mathbf{q}_i\) a “question” is justified. How to code single-head attention? For clarity, here is an overview of the steps involved in scaled dot-product attention: Figure 11: Single-head scaled dot-product attention at a glance. We can implement this directly in Python using NumPy: import numpy as np # Set seed for reproducibility np.random.seed(42) # --------------------------------------------------------------- # Create matrix X with some inputs: 2 words (seq_len=2) with # embedding dimension d_model=3 # --------------------------------------------------------------- X_words = np.array([[1.0, 2.0, 3.0], [3.0, 3.0, 2.0]]) seq_len, d_model = X_words.shape # --------------------------------------------------------------- # Create the projection matrices whose weights will be learned. # Matrices shape: (d_model, d_model) # --------------------------------------------------------------- W_Q = np.random.randn(d_model, d_model) W_K = np.random.randn(d_model, d_model) W_V = np.random.randn(d_model, d_model) # --------------------------------------------------------------- # 1) Create 3 vectors from each static input embedding: # I.e. Linear Projections: Q = X @ W_Q, K = X @ W_K, V = X @ W_V # Matrix shape: # (seq_len, d_model) @ (d_model, d_model) -> (seq_len, d_model) # --------------------------------------------------------------- Q = X_words @ W_Q K = X_words @ W_K V = X_words @ W_V print("Queries (Q):\n", Q, "\n") print("Keys (K):\n", K, "\n") print("Values (V):\n", V, "\n") # --------------------------------------------------------------- # 2) Apply each query to each key to obtain scores # I.e. Unscaled Dot-Products & Variance Scaling # Matrix shape: # (seq_len, d_model) @ (d_model, seq_len) -> (seq_len, seq_len) # --------------------------------------------------------------- scores_raw = Q @ K.T # --------------------------------------------------------------- # 3) Scale and Normalize # Convert raw similarity scores into row-wise probability # distribution # --------------------------------------------------------------- scores_scaled = scores_raw / np.sqrt(d_model) def softmax(x): exp_x = np.exp(x - np.max(x, axis=-1, keepdims=True)) # Numerical stability trick return exp_x / exp_x.sum(axis=-1, keepdims=True) A = softmax(scores_scaled) print("Attention Weights Matrix (A):\n", A, "\n") # --------------------------------------------------------------- # 4) Compute weighted sum # I.e. Final Output Matrix (Y = A @ V) # Matrix shape: # (seq_len, seq_len) @ (seq_len, d_model) -> (seq_len, d_model) # --------------------------------------------------------------- Y = A @ V print("Contextualized Output Embeddings (Y):\n", Y) Queries (Q): [[ 8.28041231 1.69573314 -1.22900853] [ 9.21765766 0.41761643 0.30170597]] Keys (K): [[-0.660378 -7.32847154 -2.97282342] [ 1.22899189 -9.15575605 -5.94344809]] Values (V): [[-2.99272485 -0.94447952 -4.83682834] [-4.49016658 -3.81248131 -2.17928541]] Attention Weights Matrix (A): [[8.68363148e-05 9.99913164e-01] [1.12005616e-04 9.99887994e-01]] Contextualized Output Embeddings (Y): [[-4.49003654 -3.81223226 -2.17951618] [-4.48999886 -3.81216008 -2.17958307]] In practice, you won’t need to write the attention steps manually. Frameworks like PyTorch provide optimized implementations of the mechanism. Note that PyTorch expects inputs to include a batch dimension. Here, we wrap our single input sequence in a batch of size 1: import torch import torch.nn as nn # Set seed for reproducibility torch.manual_seed(42) # --------------------------------------------------------------- # Create matrix X with some inputs: 2 words (seq_len=2) with # embedding dimension d_model=3 # --------------------------------------------------------------- X_words = torch.tensor([[1.0, 2.0, 3.0], [3.0, 3.0, 2.0]]) seq_len, d_model = X_words.shape # --------------------------------------------------------------- # PyTorch layers expect batched inputs: # (batch_size, sequence_length, d_model) # --------------------------------------------------------------- X_batch = X_words.unsqueeze(0) # Shape becomes [1, 2, 3] # --------------------------------------------------------------- # Instantiate PyTorch's MultiheadAttention with num_heads=1 for # single-head batch_first=True expects # shape: (batch_size, sequence_length, d_model) # --------------------------------------------------------------- attn_layer = nn.MultiheadAttention(embed_dim=d_model, num_heads=1, batch_first=True) # --------------------------------------------------------------- # Compute Self-Attention (passing X in for Query, Key, and Value) # Now torch does all the required work and simply # spits out the contextualized embeddings as Y. # Returns: (contextualized_outputs, attention_weights) # --------------------------------------------------------------- Y_batch, A_batch = attn_layer(query=X_batch, key=X_batch, value=X_batch) # --------------------------------------------------------------- # Remove the batch dimension for inspection # --------------------------------------------------------------- Y = Y_batch.squeeze(0) A = A_batch.squeeze(0) print("Attention Weights Matrix (A):\n", A.detach().numpy(), "\n") print("Contextualized Output Embeddings (Y):\n", Y.detach().numpy()) Attention Weights Matrix (A): [[0.8597921 0.14020787] [0.91717255 0.08282746]] Contextualized Output Embeddings (Y): [[0.703761 0.64951515 0.80751586] [0.7178543 0.6376342 0.79714787]] Important Note on Numerical Outputs: If you compare the output matrix \(Y\) from the NumPy implementation to the PyTorch results, you will notice the numbers differ, for two reasons. First, even though we set np.random.seed(42) and torch.manual_seed(42), NumPy and PyTorch use different underlying algorithms for pseudo-random number generation, so the initialized projection weights (\(W_Q, W_K, W_V\)) start from different raw values. Second, the nn.MultiheadAttention module doesn’t stop at \(\mathbf{AV}\): it passes the result through one more learned linear layer (out_proj) before returning it. This extra projection is a real part of the transformer architecture—we simply haven’t introduced it yet, since it only becomes relevant once we combine multiple attention passes, which is exactly the subject of the follow-up article. For now, just be aware that the PyTorch output isn’t our formula with different weights; it’s our formula plus one additional learned step on top. Summary As promised in the introduction, we have tackled both the how and the why of single-head scaled dot-product attention. We have seen how attention transforms static input embeddings into contextualized representations (summarized in Figure 11). We now also understand why each step takes place: Query, key, and value vectors are created via learned linear projections (\(\mathbf{W}_Q, \mathbf{W}_K, \mathbf{W}_V\)). Their purpose is to construct weighted averages of the input vectors. Their names come from the database analogy. Queries act as feature filters—asking “questions” about other tokens—which is realized mathematically through the score matrix \(\mathbf{QK}^T\). Scaling by \(\sqrt{d_{\text{model}}}\) prevents numerical instability and softmax saturation by keeping score variance constant regardless of embedding size. Softmax converts these scaled scores into row-wise probability distributions (attention weights) between 0 and 1 that sum to 1. Multiplying by \(\mathbf{V}\) forms the weighted sum of the value vectors attached to each token. Thus, for Single-Head Scaled Dot-Product Attention, we have: \[\text{Attention}(\mathbf{Q}, \mathbf{K}, \mathbf{V}) = \text{softmax}\left(\frac{\mathbf{QK}^T}{\sqrt{d_{\text{model}}}}\right)\mathbf{V}\] What’s Next? While a single attention pass is powerful, it has a major shortcoming: it produces only one query vector per static input embedding. Thus, each token can ask exactly one question. In natural language, a single question is rarely enough to capture every contextual nuance. A token might need to ask one question about grammatical structure and another about semantic meaning. In the next article, we will see how running multiple attention passes in parallel—a setup known as Multi-Head Attention—solves this by allowing the model to project inputs into several feature spaces at once. References [1] A. Vaswani et al., “Attention is all you need,” in Advances in neural information processing systems, 2017, pp. 5998–6008. [2] D. Bahdanau, K. Cho, and Y. Bengio, “Neural machine translation by jointly learning to align and translate,” Sep. 2014, doi: 10.48550/ARXIV.1409.0473. [3] M. Baroni, G. Dinu, and G. Kruszewski, “Don’t count, predict! A systematic comparison of context-counting vs. Context-predicting semantic vectors,” in Proceedings of the 52nd annual meeting of the association for computational linguistics (volume 1: Long papers), K. Toutanova and H. Wu, Eds., Baltimore, Maryland: Association for Computational Linguistics, Jun. 2014, pp. 238–247. doi: 10.3115/v1/P14-1023. [4] M. Apidianaki, “From word types to tokens and back: A survey of approaches to word meaning representation and interpretation,” Computational Linguistics, vol. 49, no. 2, pp. 465–523, Jun. 2022, doi: 10.1162/coli_a_00474. [5] T. Mikolov, K. Chen, G. Corrado, and J. Dean, “Efficient estimation of word representations in vector space.” 2013. Available: https://arxiv.org/abs/1301.3781 [6] J. Pennington, R. Socher, and C. Manning, “GloVe: Global vectors for word representation,” in Proceedings of the 2014 conference on empirical methods in natural language processing (EMNLP), A. Moschitti, B. Pang, and W. Daelemans, Eds., Doha, Qatar: Association for Computational Linguistics, Oct. 2014, pp. 1532–1543. doi: 10.3115/v1/D14-1162. [7] Y.-H. H. Tsai, S. Bai, M. Yamada, L.-P. Morency, and R. Salakhutdinov, “Transformer dissection: An unified understanding for transformer’s attention via the lens of kernel,” in Proceedings of the 2019 conference on empirical methods in natural language processing and the 9th international joint conference on natural language processing (EMNLP-IJCNLP), Hong Kong, China: Association for Computational Linguistics, 2019, pp. 4344–4353. [8] E. Garcia-Portugues, “Nonparametric statistics.” https://egarpor.github.io/NP-UC3M/kre-i.html. Share this: Share on Facebook (Opens in new window) Facebook Share on X (Opens in new window) X Like this:Like Loading… Discover more from Master the Math Subscribe to get the latest posts sent to your email. Type your email… Subscribe Previous Understanding the Softmax Function Derivative Related Posts What is a Logit? Machine Learning / By Markelic The Sigmoid and its Derivative Leave a Comment / Machine Learning / By Markelic Leave a ReplyCancel reply
