Post

LLMs from Scratch Part 2: Attention is All You Need

Posted Updated
By Jayden
12 min read

Part 1 ended with a question: what if the decoder could look at any part of the input directly, instead of routing everything through one compressed vector? This post follows the answer from the first attention mechanism to the full Transformer.

Table of Contents

0. What If the Model Could Look Back?

The Seq2Seq bottleneck was a compression problem. The encoder had to squeeze an input of any length into one fixed-size vector, and the decoder had to rebuild meaning from that vector alone. The information it needed could be anywhere in the input (paragraph 3, then again in paragraph 47), but the decoder had no way to go back and look.

The fix sounds obvious in hindsight: give the decoder direct access to every encoder hidden state, not just the last one. At each generation step, let it look across the whole input and decide which parts matter most right now.

Encoder hidden states Revenue beat Q3 estimates Seq2Seq: one compressed vector context c Attention: weighted access to all states 5% 82% 10% 3% decoder step generating "a dépassé" dépassé

Figure 1. The core difference. Seq2Seq routes all context through a single compressed vector. Attention lets the decoder reach back directly to every encoder state, with learned weights that change at each generation step.

Bahdanau et al. (2015) introduced this idea of giving the decoder dynamic, weighted access to the full encoder output. It improved translation quality a lot, especially on long sentences. It also contained an idea that would later spread well beyond Seq2Seq: the attention weights are computed on the fly as a learned function of the current context.

1. The Attention Mechanism

The attention mechanism formalizes this idea using three components: a Query, a set of Keys, and a set of Values.

The intuition comes from information retrieval. The Query is what you’re looking for. Each Key describes what a piece of information contains. The Value is the content itself. Attention scores how well the Query matches each Key, turns those scores into weights, and returns a weighted sum of the Values.

Query Q "what do I need now?" dot product Key₁ "Revenue" Key₂ "beat" Key₃ "estimates" scores 0.12 0.81 0.07 softmax Value₁ [revenue vector] Value₂ [beat vector] ← 81% Value₃ [estimates vector]

Figure 2. Attention as retrieval (illustrative numbers). The Query is matched against each Key to produce a score; softmax turns scores into weights; the output is a weighted sum of the Values. In this example, “beat” receives 81% of the weight.

In Seq2Seq with attention, the Query comes from the decoder’s current hidden state, and the Keys and Values both come from the encoder’s hidden states. Bahdanau et al. scored each pair with a small feed-forward network (additive attention); Luong et al. (2015) proposed simpler dot-product scores. The Transformer (Vaswani et al., 2017) uses the scaled dot-product form, which in matrix form, with $n$ input tokens and $m$ output steps, is

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

where $Q \in \mathbb{R}^{m \times d_k}$, $K \in \mathbb{R}^{n \times d_k}$, $V \in \mathbb{R}^{n \times d_v}$.

$QK^T$ computes the dot product between every Query and every Key at once, giving an $m \times n$ matrix of raw similarity scores. Softmax normalizes each row into a probability distribution over the $n$ input positions, and multiplying by $V$ returns a weighted sum of value vectors, one per output step.

Why divide by $\sqrt{d_k}$? If the components of $q$ and $k$ have unit variance, their dot product has variance $d_k$. With large $d_k$ the scores get large, softmax saturates, almost all the weight lands on one position, and the gradients become tiny. Dividing by $\sqrt{d_k}$ brings the variance back to about 1 regardless of dimension.

2. Self-Attention

The attention above connects two different sequences: decoder queries attending to encoder keys and values. The 2017 Transformer built its whole architecture on running the same computation inside a single sequence, with every position attending to every other position. This is self-attention. The idea itself had appeared earlier (for example Cheng et al., 2016; Parikh et al., 2016; Lin et al., 2017), but the Transformer was the first model to rely on it entirely, without recurrence.

In self-attention, the Query, Key, and Value matrices are all derived from the same input $X$ through learned linear projections

\[Q = XW^Q, \quad K = XW^K, \quad V = XW^V\]

where $W^Q, W^K, W^V \in \mathbb{R}^{d_{model} \times d_k}$ are learned weight matrices. The same formula applies

\[\text{Self-Attention}(X) = \text{softmax}\!\left(\frac{XW^Q (XW^K)^T}{\sqrt{d_k}}\right)\!XW^V\]

The question is no longer “which encoder states matter for my current decoder step?” but “which other positions in this sequence matter for understanding each position?” That makes self-attention a general tool for relationships within a sequence.

Sentence: "The stock didn't recover because it was too volatile." The stock didn't recover because it was too volatile 88%: "it" refers to "stock" 9% Self-attention resolves the pronoun without any hand-crafted coreference rules

Figure 3. Self-attention resolving a pronoun in a financial sentence (illustrative numbers). When processing “it”, the model puts most of its weight on “stock”. In a trained model, patterns like this are learned from data rather than from explicit coreference rules.

Take a sentence like “The central bank’s rate hike surprised the market, and it pushed bond yields sharply higher.” To resolve “it” to the rate hike rather than the market, the model has to look at the structure of the whole sentence. Self-attention can do this because every position can attend to every other position directly, with nothing forced through a sequential bottleneck.

Self-attention and the cross-attention used in Seq2Seq run the same computation. The only difference is where the inputs come from. In self-attention, $Q$, $K$, and $V$ all come from the same sequence. In cross-attention, $Q$ comes from one sequence (the decoder) and $K$, $V$ from another (the encoder).

3. Multi-Head Attention

A single self-attention computation gives the model one set of relationships: one learned similarity function, one set of weight matrices. A sentence like “Revenue beat estimates, but guidance disappointed” has several overlapping structures at once: syntax (subject and verb), sentiment (positive vs. negative), and numbers (earnings figures vs. guidance). One head would have to cover all of them.

The Transformer’s answer is to run several attention computations in parallel, each with its own projection matrices

\[\text{head}_i = \text{Attention}(QW_i^Q,\, KW_i^K,\, VW_i^V)\]

and then concatenate the results and project back to the model dimension

\[\text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, \ldots, \text{head}_h)\, W^O\]
Input X Head 1 syntactic relations Head 2 sentiment polarity Head 3 numeric references Concat h₁ ‖ h₂ ‖ h₃ 3 × d_k dims Linear W^O → d_model Output rich multi-perspective representation

Figure 4. Multi-head attention runs h parallel attention computations, each with independent projection matrices. The outputs are concatenated and projected back to d_model.

Each head operates in a lower-dimensional subspace of size $d_k = d_{model}/h$, so the total computation remains comparable to a single full-dimensional attention. With $h = 8$ heads and $d_{model} = 512$, each head works in a 64-dimensional space.

Multi-head attention is not the same computation repeated $h$ times. Each head has its own $W_i^Q$, $W_i^K$, $W_i^V$, so each can learn to look for something different. Studies of trained Transformers (e.g. Clark et al., 2019; Voita et al., 2019) found heads that track specific syntactic relations, rare words, or positions, without any explicit supervision. They also found that many heads can be pruned with little loss, so not every head ends up specialized.

4. Positional Encoding

Self-attention as described so far has no sense of order. The score between positions $i$ and $j$ depends only on the vectors at those positions, not on which comes first. Shuffle the input tokens and the outputs are shuffled in the same way (self-attention is permutation-equivariant), so “Revenue beat expectations” and “Expectations beat revenue” look like the same bag of tokens.

This is the cost of processing everything in parallel. The LSTM got position for free from the order in which tokens were fed in. The Transformer processes all tokens at once, so position has to be added back explicitly.

The Transformer paper’s solution was to construct a fixed function that maps each position to a unique vector and add it to the token embedding before the first attention layer

\[PE_{(pos,\, 2i)} = \sin\!\left(\frac{pos}{10000^{2i/d_{model}}}\right), \qquad PE_{(pos,\, 2i+1)} = \cos\!\left(\frac{pos}{10000^{2i/d_{model}}}\right)\]
−1 −0.5 0 0.5 1 position 0 20 40 60 80 100 dim i=0 (fast, sin) dim i=50 (medium, sin) dim i=100 (slow, cos)

Figure 5. Positional encoding curves for three embedding dimensions. Low-index dimensions oscillate rapidly (fine-grained position), high-index dimensions change slowly (coarse position). Every position maps to a unique combination across all dimensions.

Each dimension oscillates at a different frequency. Low-index dimensions complete many cycles across the sequence, while high-index dimensions barely move. Each position gets its own combination of values across all $d_{model}$ dimensions. The relative offset is also encoded: the vector at position $pos + k$ is a linear function (a rotation) of the vector at $pos$, which makes it easy for the model to learn to reason about distance.

Why sin and cos rather than learned position embeddings, as BERT later used? Vaswani et al. tried both and got nearly identical results. They chose sinusoids because the function is defined for any $pos$, so it might extrapolate to sequences longer than those seen in training, while learned embeddings stop at the maximum training length.

5. The Full Transformer

With self-attention, multi-head attention, and positional encoding in place, the Transformer assembles these into two stacked modules: an encoder and a decoder. Both share the same internal structure, with one key difference.

Each encoder layer contains two sub-layers: multi-head self-attention followed by a position-wise feed-forward network

\[\text{FFN}(x) = \max(0,\; xW_1 + b_1)\, W_2 + b_2\]

Each sub-layer is wrapped with a residual connection and layer normalization

\[\text{output} = \text{LayerNorm}(x + \text{Sublayer}(x))\]

The residual connection keeps the original signal alongside the transformation and gives gradients a direct path back, the same idea that makes ResNets trainable at depth. Layer normalization rescales each token’s vector across its feature dimension, which keeps activations in a stable range. One detail: in this original Post-LN arrangement, LayerNorm sits after the addition, so the residual path isn’t a perfectly clean identity. Most later models moved LayerNorm inside the residual branch (Pre-LN), which tends to make deep Transformers easier to train.

Encoder (×N) Input Embedding + Pos. Enc. Add & Norm Multi-Head Self-Attention Add & Norm Feed-Forward FFN(x) = max(0, xW₁+b₁)W₂+b₂ Encoder Output Decoder (×N) Output Embedding + Pos. Enc. Masked Multi-Head Self-Attention Add & Norm Cross-Attention Q ← decoder, K/V ← encoder Add & Norm → Linear → Softmax → output

Figure 6. The full Transformer architecture. The encoder (left) builds a rich representation of the input through self-attention and FFN layers, each wrapped with residual connections. The decoder (right) adds a masked self-attention layer and a cross-attention layer that reads from the encoder output (orange dashed arrow).

The decoder has one more component: a cross-attention layer between the masked self-attention and the feed-forward network. Its Queries come from the decoder’s previous layer, and its Keys and Values come from the encoder’s output. This is how the decoder reads the input: at every step, each decoder position can attend to any part of the encoded input.

The decoder’s self-attention is masked. When generating position $t$, the model can only attend to positions it has already generated, not to future ones. Without the mask, it could predict position $t$ by peeking at position $t+1$, which would make training trivial and inference impossible.

The original Transformer used $N = 6$ encoder layers and $N = 6$ decoder layers, $d_{model} = 512$, $h = 8$ heads, and $d_{ff} = 2048$ for the feed-forward layer. Later models changed many details (decoder-only stacks, Pre-LN, different position encodings), but this is the template they started from.

Beyond better scores, the Transformer changed throughput. Self-attention is a matrix operation, so all positions in a sequence can be processed at once. Training on long sequences, painfully slow with LSTMs because of the step-by-step dependency, became a parallel matrix multiply. That is a big reason models with billions of parameters became practical: GPUs could finally be used efficiently for sequence modeling, and the large language models that followed were built on this.

References

  • Bahdanau, D., Cho, K., & Bengio, Y. (2015). Neural Machine Translation by Jointly Learning to Align and Translate. ICLR 2015. arxiv.org/abs/1409.0473
  • Luong, M.-T., Pham, H., & Manning, C. D. (2015). Effective Approaches to Attention-based Neural Machine Translation. EMNLP 2015. arxiv.org/abs/1508.04025
  • Cheng, J., Dong, L., & Lapata, M. (2016). Long Short-Term Memory-Networks for Machine Reading. EMNLP 2016. arxiv.org/abs/1601.06733
  • Parikh, A., Täckström, O., Das, D., & Uszkoreit, J. (2016). A Decomposable Attention Model for Natural Language Inference. EMNLP 2016. arxiv.org/abs/1606.01933
  • Lin, Z., Feng, M., dos Santos, C. N., et al. (2017). A Structured Self-attentive Sentence Embedding. ICLR 2017. arxiv.org/abs/1703.03130
  • Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, Ł., & Polosukhin, I. (2017). Attention Is All You Need. NeurIPS 2017. arxiv.org/abs/1706.03762
  • He, K., Zhang, X., Ren, S., & Sun, J. (2016). Deep Residual Learning for Image Recognition. CVPR 2016. arxiv.org/abs/1512.03385
  • Ba, J. L., Kiros, J. R., & Hinton, G. E. (2016). Layer Normalization. arxiv.org/abs/1607.06450
  • Xiong, R., Yang, Y., He, D., et al. (2020). On Layer Normalization in the Transformer Architecture. ICML 2020. arxiv.org/abs/2002.04745
  • Clark, K., Khandelwal, U., Levy, O., & Manning, C. D. (2019). What Does BERT Look At? An Analysis of BERT’s Attention. BlackboxNLP 2019. arxiv.org/abs/1906.04341
  • Voita, E., Talbot, D., Moiseev, F., Sennrich, R., & Titov, I. (2019). Analyzing Multi-Head Self-Attention: Specialized Heads Do the Heavy Lifting, the Rest Can Be Pruned. ACL 2019. arxiv.org/abs/1905.09418
  • Alammar, J. (2018). The Illustrated Transformer. jalammar.github.io
  • Goodfellow, I., Bengio, Y., & Courville, A. (2016). Deep Learning. MIT Press. deeplearningbook.org
AI, LLM
attention transformer self-attention multi-head-attention positional-encoding scaled-dot-product nlp
This post is licensed under CC BY 4.0 by the author.