← Back to articles
AI-generated image

A look inside the Transformer: computing attention by hand

The Transformer is what computes next-token probabilities inside an LLM. See its overall structure, and compute its core, attention, in a small piece of code.

Written by AI · 10 min · For readers of “How LLMs work”
TransformerAttentionLLM
What this article covers

Check the structure of the part that produces next-token probabilities (the Transformer) with figures and code. The core is attention, which decides how much to look at each token.

What we look at in this article

The previous article explained that an LLM computes “the probability of each next token” and picks from it. This article looks inside the part that computes those probabilities: the Transformer. The Transformer is a structure proposed in the 2017 paper “Attention Is All You Need.” It was built around attention, without recurrence or convolutions. Many of today’s LLMs are based on this structure. However, this article does not cover the detailed internal structure of any individual model, Claude included. What it describes is the general structure shown in the paper.

The figures and numbers are simplifications for understanding. Real models are orders of magnitude larger in layers and vector size, and the specifics vary by model. The paper proposed a configuration for translation; what is widely used for LLMs is a form close to its output side (decoder-style).

The big picture: flowing from bottom to top

The figure below shows the flow of a typical decoder-style Transformer. The colored dots suggest data flowing from bottom to top (they are not shown if you have reduced motion enabled).

Overview of a Transformer (decoder-style)Next-token probabilitiesSoftmaxtemperature acts hereLineara score for every vocabulary entry× N layersAdd & NormFeed-forwardprocesses each token separatelyAdd & NormMasked self-attentiondecides which tokens to look atskipskipEmbedding + position infoturns tokens into vectorsthemodelpredicts
Figure 1. The flow of a typical decoder-style Transformer. Data flows from the input at the bottom to the probabilities at the top. The framed part is stacked many times.
  • Turn tokens into vectors (lists of numbers): the embedding. Position information is added here too
  • Attention: each token decides how much to look at each earlier token, and gathers information
  • Feed-forward: process each token’s vector separately
  • Add & Norm and the skip (residual) connections: add the original input back so deep stacks are easier to train
  • Stack this set many times (× N layers in the figure)
  • Finally, produce a score for every vocabulary entry (Linear) and turn it into probabilities with softmax. The temperature from the previous article acts on the scores before this softmax

Tokens become vectors

Inside the model, tokens are handled as lists of numbers (vectors). Training adjusts them so that tokens used in similar ways get nearby vectors. Attention alone can’t tell the order of tokens, so position information is added too. There are several ways to represent position, and they vary by model.

Attention: deciding which tokens to look at

The purpose of attention is to compute “which earlier tokens matter for deciding the meaning of the current token.” From each token it makes three vectors. Q (query) is “what am I looking for,” K (key) is “what do I have,” and V (value) is “what I hand over.” The calculation has four steps.

  • Take the dot product of Q and K to score how well tokens match. Divide by d\sqrt{d} so values don’t grow too large
  • Hide future tokens (the mask), so that training to predict the next token doesn’t peek at the answer
  • Apply softmax to each row to make weights that sum to 1
  • Add up V using the weights to get the token’s new vector

As a formula, this is:

softmax ⁣(QK⊤d)V\mathrm{softmax}\!\left(\frac{QK^{\top}}{\sqrt{d}}\right)V

Here dd is the vector dimension.

Checking it in code

I implemented attention at a size you can follow by hand (5 tokens, 4 dimensions, 1 head). The weights are untrained random numbers, so the pattern itself means nothing. What matters is the shape of the calculation. I tried several random seeds and picked one where the differences are easy to see (a different seed gives a different pattern). You can switch between the JavaScript and Python versions of the code with the tabs. The random numbers come from the same formula, so the output is identical (we ran them on Node.js 26.1.0 and Python 3.11.5).

attention.py (the same calculation in Python)
# Compute self-attention for a single head by hand (the same calculation as attention.mjs, in Python).
import math

tokens = ['the', 'model', 'predicts', 'the', 'next']
d = 4                                    # vector size (real models are far larger)


def rng(seed):
    """A random number generator using the same formula as the JavaScript version: the same seed gives the same sequence."""
    state = seed & 0xFFFFFFFF

    def random():
        nonlocal state
        state = (state + 0x6D2B79F5) & 0xFFFFFFFF
        t = ((state ^ (state >> 15)) * (1 | state)) & 0xFFFFFFFF
        t = ((t + (((t ^ (t >> 7)) * (61 | t)) & 0xFFFFFFFF)) & 0xFFFFFFFF) ^ t
        return ((t ^ (t >> 14)) & 0xFFFFFFFF) / 4294967296

    return random


random = rng(2)


def matrix(rows, cols):
    return [[random() * 2 - 1 for _ in range(cols)] for _ in range(rows)]


def mul(A, B):
    return [[sum(a * b for a, b in zip(row, col)) for col in zip(*B)] for row in A]


def transpose(A):
    return [list(col) for col in zip(*A)]


def softmax(row):
    top = max(row)
    exps = [math.exp(v - top) for v in row]
    total = sum(exps)
    return [v / total for v in exps]


X = matrix(len(tokens), d)               # a vector for each token (the embedding)
Wq, Wk, Wv = matrix(d, d), matrix(d, d), matrix(d, d)
Q, K, V = mul(X, Wq), mul(X, Wk), mul(X, Wv)

# 1. Dot each token's Q with every token's K, then divide by √d
scores = [[v / math.sqrt(d) for v in row] for row in mul(Q, transpose(K))]
# 2. Hide future tokens (set the upper-right part to -∞)
masked = [[v if j <= i else -math.inf for j, v in enumerate(row)] for i, row in enumerate(scores)]
# 3. Apply softmax to each row to get weights that sum to 1
weights = [softmax(row) for row in masked]
# 4. Add up V using the weights to get each token's new vector
output = mul(weights, V)

print('weights:')
for token, row in zip(tokens, weights):
    print(token.ljust(9), ' '.join(f'{w:.2f}' for w in row))
print('row sums:', ' '.join(f'{sum(row):.2f}' for row in weights))
print('output of "next":', ' '.join(f'{v:.2f}' for v in output[4]))

Reading the result

Running it gives the following (we ran it on Node.js 26.1.0; the seed is fixed, so every run is the same).

Output
weights:
the       1.00 0.00 0.00 0.00 0.00
model     0.54 0.46 0.00 0.00 0.00
predicts  0.22 0.10 0.68 0.00 0.00
the       0.21 0.27 0.29 0.23 0.00
next      0.15 0.08 0.52 0.02 0.23
row sums: 1.00 1.00 1.00 1.00 1.00
output of "next": 0.29 0.92 0.90 -0.29
Attention weights (5 tokens × 5 tokens)↓ LookingBeing looked at →thethemodelmodelpredictspredictsthethenextnext1.00××××0.540.46×××0.220.100.68××0.210.270.290.23×0.150.080.520.020.2301weightFuture tokens (hidden)
Figure 2. The weights as actually computed (rows look, columns are looked at). The upper right is future tokens, which are hidden, and every row sums to 1. The outlined row (next) is examined in the next figure.

The figure is a table of these weights. Rows are the token doing the looking, and columns are the token being looked at. Three things can be read from it.

  • The upper right is all 0, because future tokens are hidden
  • Each row sums to 1.00, because softmax produces the weights
  • The first row is only itself (1.00), because the first token has nothing else to look at

The weights are still random, so the weight gathering on predicts is a coincidence. In training, the matrices that make Q, K and V (WQW_Q, WKW_K, WVW_V) are adjusted so the model predicts the next token well.

Looking closely at one row: “next”

Let’s follow the last row (next). next can look at five tokens, itself included. After the dot-product scores go through softmax, predicts gets a weight of 0.52, next itself 0.23, and the first the 0.15. Adding up V with those weights gives next’s new vector (0.29, 0.92, 0.90, -0.29).

How much “next” looks at each earlier tokenHow much “next” looks at each tokenscoreweightthe0.420.15model-0.200.08predicts1.640.52the-1.600.02next0.850.23Add up V using the weightsNew vector for “next”0.290.920.90-0.29
Figure 3. The calculation for the last row (next). The dot-product scores go through softmax to become weights, and adding up V with those weights gives the new vector for next. The numbers are from the code above.

After attention

After attention gathers information from the surroundings, feed-forward processes each token’s vector and the skip connection adds the original input back. This set is stacked many times. Real models split attention into several “heads” computed in parallel so that each can look at a different relationship; this code is one such head. After the last layer, a score for every vocabulary entry is produced from the last token’s vector. Applying softmax to those scores gives the “next-token probabilities” from the previous article. Temperature acts on the scores before that softmax.

With this, the previous article’s description, “choose the next token by probability,” is now connected to the calculation that produces those probabilities.

Official references

Vaswani et al. — Attention Is All You Need (arXiv:1706.03762) ↗Claude Docs — Glossary ↗
PAIBOT / DEVELOPMENT JOURNALBack to articles →

Find what you need

Searches titles, tags and body text esc to close