How Transformers work and why Attention is the key

If you work in modern AI or NLP, you cannot avoid the Transformer. Introduced in the landmark 2017 paper “Attention Is All You Need” by Vaswani et al. (building on earlier attention concepts by Bahdanau et al. in 2015), this architecture fundamentally changed how machines understand human language.

To appreciate why the Transformer is so revolutionary, we need to look at what came before it, how it processes text, and what happens under the hood during a full forward pass.

Prerequisites: this post assumes you're already familiar with basic neural network concepts, including forward passes, embeddings, and backpropagation.

Why did we need the Transformer in the first place?

Before 2017, the standard approach for sequence modeling relied on Recurrent Neural Networks (RNNs) and Long Short-Term Memory (LSTMs). But these architectures had two massive bottlenecks:

• Context loss, or the long-term dependency problem: RNNs process text sequentially, one token after another. So by the time an RNN reaches the end of a long paragraph, it has often forgotten the context of the first few words. While LSTMs introduced “forget gates” to retain important information, earlier details still get diluted across long paragraphs, making long-range context easy to lose.
• Sequential processing problem: since RNNs process text word-by-word, they were extremely slow to train. Because every step depends directly on the one before it, you can't train sequences in parallel across GPUs, leaving massive compute power sitting idle!

The Transformer architecture solved both of these by tossing out recurrence entirely and processing everything at once.

Tokens and embeddings

Before we dive into how a Transformer processes text, we have to translate the words into math.

First, we break the raw text down into tokens. A token can be an entire word, or just a part of one, like a prefix or a suffix. Each unique token in our vocabulary is assigned an integer ID. The same word always gets assigned to the same ID.

A neural network can't do complex calculus on simple integer IDs, though. So we pass these IDs through an Embedding Layer, which converts each token into a high-dimensional vector (in other words, a dense list of numbers) that captures its semantic meaning. In the original paper, each word was embedded into a 512-dimensional vector ($d_{\text{model}} = 512$).

What is “Attention”?

In Machine Learning, Attention is almost synonymous with the word in everyday English. In human language, the meaning of a word changes depending on the words around it. Think about the word “model.” It could refer to a fashion model, a mathematical model, or even a machine learning model! The only way to know what a word means is by looking at the surrounding context.

Similarly, “Attention” is the mechanism that lets a model look at the entire sentence and calculate exactly how much importance, or influence, every word should have on every other word.

$$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$$

Each component of this formula is explained in the sections below.

Calculating Query, Key and Value

To compute this Attention score, the Transformer creates three new vectors for each word by multiplying the original embedding matrix by three tunable weight matrices ($W_q$, $W_k$, and $W_v$), which later get adjusted during backpropagation.

Let's understand this with an example:

The tired engineer drank the cold coffee

If the model is currently processing the word “coffee”, then:

• Query (Q): denotes what the current word is looking for, to better understand itself. For instance, the Query for “coffee” would say: “I am a beverage and a noun. I'm looking for any adjectives that describe my state, or verbs that tell me what is happening to me.”

• Key (K): denotes what other words in the sentence have to offer about themselves. The Key for “cold” would say: “I'm an adjective for temperature, meant to modify a physical object.”

• Value (V): denotes the underlying meaning that gets passed along if there is a match. The Value for “cold” is the mathematical representation of low temperature.

So when the Query of “coffee” checks the Keys of every other word in the sentence, it finds a great match with the Key of “cold.” Because of this high score, the model takes the Value of “cold” and combines it with the mathematical representation of “coffee.”

By the end of this step, the word “coffee” has been updated to represent cold coffee!

The math behind the architecture

We've seen the formula for calculating Attention. Now let's connect every piece.

$$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$$

1. Dot product ($QK^T$): in vector algebra, a dot product signifies similarity. Here it denotes how similar each word is to every other word in the sentence. We multiply the Query matrix by the transpose of the Key matrix (transpose matches the dimensions so the matrices are eligible for the matmul operation).

   If you map this out on a grid for the sentence “Your cat is a great cat,” you'd see high scores where “your” meets “your,” and where “cat” meets “cat.”

   	YOUR	CAT	IS	A	GREAT	CAT	Σ
   YOUR	0.612	0.081	0.062	0.051	0.114	0.080	1
   CAT	0.042	0.584	0.051	0.023	0.092	0.208	1
   IS	0.052	0.081	0.640	0.041	0.122	0.064	1
   A	0.061	0.042	0.053	0.710	0.081	0.053	1
   GREAT	0.031	0.245	0.042	0.030	0.521	0.131	1
   CAT	0.038	0.212	0.041	0.025	0.114	0.570	1

2. Scaling factor: this is a crucial step. As the vector dimension ($d_k$) grows, the dot products grow massively too. Pushed through softmax, large numbers get shoved toward the extremes (basically 0 or 1), where gradients flatten out to almost zero. Dividing by $\sqrt{d_k}$ keeps those values in check. Without it, the gradients vanish (very commonly known as the “vanishing gradient” problem) and the model straight-up stops learning!

3. Softmax function: the dot product results in large positive and negative numbers, so we wrap them in a softmax function to normalize the scores into probabilities between 0 and 1.

4. Multiplying the Value matrix: we then multiply these normalized scores by the Value matrix (V). Words with high similarity scores pass their “Value” forward, while words with a lesser score or near 0 are rightfully ignored.

Multi-Head Attention

Instead of calculating attention just once, the Transformer splits the token representation across multiple “heads” operating in parallel (8 heads in the original paper). Think of each head as looking at the sentence through a different lens; one might focus on grammatical structure, another on pronoun references (figuring out what “it” refers to), another on emotional tone.

Once all 8 heads finish computing their scores, we concatenate them back into one giant vector and pass it through a final projection matrix ($W_O$).

$$\text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, \text{head}_2, \dots, \text{head}_h)W^O$$

$$\text{where} \quad \text{head}_i = \text{Attention}\left(Q W_i^Q, K W_i^K, V W_i^V\right)$$

where:

• $W_i^Q \in \mathbb{R}^{d_{\text{model}} \times d_k}$, $W_i^K \in \mathbb{R}^{d_{\text{model}} \times d_k}$, and $W_i^V \in \mathbb{R}^{d_{\text{model}} \times d_v}$ are learnable projection matrices for head $i$.
• $W^O \in \mathbb{R}^{h d_v \times d_{\text{model}}}$ is the final output projection matrix.

In the 2017 paper, with $d_{\text{model}} = 512$ and $h = 8$ heads, each head's dimension was scaled down to $d_k = d_v = d_{\text{model}} / h = 64$.

Once all 8 heads finish, we concatenate them side-by-side into a matrix of size $h \times d_v = 512$ and pass them through the final projection matrix $W^O$, combining the independent insights from all 8 heads back into a single $d_{\text{model}}$-dimensional representation without increasing computational cost.

Positional Encoding

Unlike RNNs and LSTMs, Transformers process every token in parallel, so they have no built-in concept of sequence. If we don't add a sense of position, the model would treat “The dog bit the man” the same way as “The man bit the dog,” which is a massive flaw.

Positional Encoding fixes this by adding positional information into the embeddings using sine and cosine waves:

$$PE_{(pos, 2i)} = \sin\left(\frac{pos}{10000^{\frac{2i}{d_{\text{model}}}}}\right) \qquad \text{for even indices}$$

$$PE_{(pos, 2i+1)} = \cos\left(\frac{pos}{10000^{\frac{2i}{d_{\text{model}}}}}\right) \qquad \text{for odd indices}$$

where:

• $pos$: token position in the input sequence ($pos = 0, 1, 2, \dots$)
• $i$: dimension index along the vector ($0 \le i < d_{\text{model}}/2$)
• $d_{\text{model}}$: total dimension size of the token embedding (e.g. 512)

But why sine and cosine waves?

The authors chose sine and cosine waves for two specific mathematical reasons:

1. Because of trigonometric identities like $\sin(\alpha + \beta) = \sin\alpha \cos\beta + \cos\alpha \sin\beta$, the encoding for position $(pos+k)$ can be represented as a linear function of position $pos$, making it easy for self-attention to learn relative distances between words regardless of where they appear in the sentence.

2. Since these functions are continuous and periodic across different frequencies, the model can theoretically extrapolate to sequence lengths longer than those it saw during training.

An important property of these positional encodings is that they depend entirely on index positions and not on the actual words! They are also static and are computed once and reused across every single sentence, batch and epoch.

Consider these two very different sentences:

• “Cat drinks milk”
• “I love chocolates”

Assuming 0-indexed positions, “drinks” and “love” both sit at $pos=1$. Even though they have completely different meanings and embedding vectors, they still get the exact same positional encoding vector! When we add these vectors together (Embedding + Positional Encoding) for each sentence, the model gets both the semantic meaning of the word and its exact position in the sentence, at once!

Residual networks and Normalization

Inside the Transformer blocks, two features keep training stable:

• Residual connections: instead of forcing a sublayer to compute an entirely new output from scratch, a residual connection adds the original input directly back to the output:

  $$\text{output} = x + \text{sublayer}(x)$$

  The network only needs to calculate the “residual” part. This prevents the vanishing gradient problem and ensures original signals aren't lost in deep networks, thus preventing signal degradation.

To put it simply, a residual connection is the model saying: add the new context on top of what's already there.

• Layer Normalization (LN): after the residual addition, values are normalized. If a layer outputs [1.5, 1.0, 4.5, 7.0], normalization finds the mean and standard deviation, subtracts the mean, and divides by the standard deviation. This scales the values around 0, keeping the data flowing consistently through the network without being too large.

Pre-LN vs Post-LN: back in 2017, the authors put Layer Norm after the residual addition (Post-LN). It worked fine for 6 layers, but once it scaled up to deeper models, training turned into a nightmare! Modern models almost always flip this to Pre-LN, normalizing before attention or the feed-forward layer: $\text{output} = x + \text{SubLayer}(\text{LayerNorm}(x))$. This keeps the residual path unblocked, making deep networks much easier to train!

The Feed-Forward Network

Attention gets most of the praise naturally, but the position-wise Feed-Forward Network (FFN) sitting right behind it handles a massive chunk of the actual heavy lifting.

Under the hood it's two linear layers with a non-linear activation (like ReLU or GELU) in between. The first layer expands the vector dimension from $d_{\text{model}}$ up to $4 \times d_{\text{model}}$ (512 → 2048 in the original paper), and the second compresses it back down.

This expansion-compression bottleneck is exactly why FFN layers hold approximately two-thirds of a Transformer's total parameter count, and why research suggests most of a model's actual factual memory is stored directly inside these FFN weights!

You can think of self-attention like sitting in a meeting and gathering notes from everyone else in the room. The FFN, on the other hand, is you going back to your desk, sitting quietly, and processing what you just heard!

The full forward pass

Putting it all together: when a prompt goes into a standard Encoder-Decoder Transformer, here's exactly what happens, chronologically.

1. Input Processing: raw text is tokenized and converted to input embeddings. Positional encodings are added.
2. Encoder stack: the vectors enter the Multi-Head Attention block, letting every word contextualize itself against the rest of the sentence. The output goes through an Add & Norm step, then a Feed-Forward Network, then another Add & Norm step. This repeats for several layers (N=6 in the paper).
3. Decoder stack: the target sequence goes through its own embedding and positional encoding, then hits a Masked Multi-Head Attention layer. To stop the model from cheating and looking at future words, future token values are masked with $-\infty$. Because $e^{-\infty} = 0$, the softmax function zeroes out their probabilities!

   	YOUR	CAT	IS	A	GREAT	CAT
   YOUR	0.84	−∞	−∞	−∞	−∞	−∞
   CAT	0.31	1.45	−∞	−∞	−∞	−∞
   IS	0.12	0.98	0.54	−∞	−∞	−∞
   A	0.05	0.22	0.18	0.76	−∞	−∞
   GREAT	0.10	1.82	0.25	0.14	1.12	−∞
   CAT	0.28	2.10	0.15	0.09	1.64	0.85

4. Cross-Attention: the decoder runs a second attention block where its own vectors act as the Query ($Q$), while pulling the Key ($K$) and Value ($V$) directly from the final output of the Encoder block. This is where the actual translation work happens.

   Consider translating “I love chocolates” into French (“J'adore les chocolats”):

   • Queries (Q) come from the Decoder: after generating “J'adore” (“I love”), the French decoder prepares to predict the direct object. Its Query vector at this step asks: “What item in the source sentence is being loved?”
   • Keys (K) and Values (V) come from the Encoder: the English encoder processed “I love chocolates” and created contextual representations for each word.
   • Match: the decoder's Query for “object being loved” computes a high dot-product score against the English Key for “chocolates.” It uses that attention score to get the semantic vector (Value) of “chocolates,” allowing the decoder to correctly output “les chocolats” next.

   Without Cross-Attention, the decoder would generate words in isolation. It lets the decoder examine the original sentence at every step to retrieve the exact context it needs.

5. Output layer, logits, and the generation loop: after the final FFN, the hidden state vector (size $d_{\text{model}}$) goes through a Linear Layer (the LM Head) that projects it up to the size of the vocabulary ($V$). This produces logits (raw, unnormalized confidence scores) for every word in the dictionary. A softmax function then converts these logits into a probability distribution. Finally, the model picks the next token (via greedy selection or sampling techniques like temperature and top-$p$), appends that new token back to the input sequence, and runs the entire forward pass again to predict the next word.

What about GPT and Claude?

The architecture walked through here is the original 2017 Encoder-Decoder design, built primarily for translation tasks. GPT, LLaMA, and Claude, however, are decoder-only models! They threw away the encoder block entirely and just stack masked self-attention layers to predict the next tokens.

The autoregressive loop is where most inference cost lives, and KV caching is the first optimization worth digging into. From there, understanding the $O(N^2)$ memory wall with FlashAttention is a must for keeping systems fast. The architecture keeps evolving, but every new paper that comes out still relies on the same core intuitions walked through here.