Do I need to understand attention mechanisms to build LLM applications?

Not to write prompts or call an API — but understanding the pipeline (tokens -> embeddings -> transformer blocks -> next-token probabilities -> sampling) explains why context windows, temperature, and token costs behave the way they do, which does matter for building reliable applications. For the attention math itself, see the dedicated Transformers in PyTorch lesson.

What is a transformer block?

One repeated unit of the model: a layer-norm, a self-attention sublayer, another layer-norm, and a feed-forward (MLP) sublayer, each wrapped in a residual connection (the sublayer's output is added back to its input, not used to replace it). A model stacks many of these blocks — distilgpt2 uses 6, larger models use dozens to over a hundred.

What does self-attention actually compute?

For every token, a weighted average of every other token's representation, where the weights ('attention scores') are learned based on how relevant each token is to the current one. This is what lets 'it' in 'the trophy didn't fit in the suitcase because it was too big' resolve to 'trophy' rather than 'suitcase' — the model attends more to the correct antecedent.

Why does the model predict 'the' or 'home' instead of the obviously correct word?

A small model like distilgpt2 (82M parameters, 6 transformer blocks) has far less capacity and training data than a production model, so its next-token probabilities are noisier and less factually reliable. The mechanism shown here — tokens through embeddings through transformer blocks to a probability distribution — is identical to how a large production model works; only the scale and training data differ.

What is the difference between the embedding layer and a transformer block?

The embedding layer is a lookup table: one fixed vector per vocabulary entry, with no awareness of surrounding tokens. A transformer block is where tokens exchange information with each other via self-attention — the representation of 'bank' only becomes context-aware ('river bank' vs 'savings bank') after passing through at least one block.

How an LLM Actually Works | SkillByExample

A model never sees “The capital of France is” as a sentence. It sees a list of integers, turns each into a vector, runs those vectors through the same stack of transformations dozens of times, and ends up with one probability per vocabulary entry for “what comes next.” Tracing that path through a real, small, locally-runnable model makes each step concrete.

Text
 |
 v
Tokenizer -> token IDs
 |
 v
Embedding lookup -> one vector per token
 |
 v
Transformer block x N  (attention + feed-forward, repeated)
 |
 v
Final hidden state -> logits (one score per vocabulary entry)
 |
 v
Softmax -> next-token probabilities
 |
 v
Sampling -> next token -> append -> repeat

From Tokens to Embeddings

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

tok = AutoTokenizer.from_pretrained("distilgpt2")
model = AutoModelForCausalLM.from_pretrained("distilgpt2")
model.eval()

text = "The capital of France is"
ids = tok(text, return_tensors="pt")
print(f"1. Text:      {text!r}")
print(f"2. Token IDs: {ids['input_ids'].tolist()[0]}")
print(f"3. Tokens:    {[tok.decode([t]) for t in ids['input_ids'][0]]}")

embed_layer = model.transformer.wte
embeddings = embed_layer(ids["input_ids"])
print(f"4. Embedding shape: {tuple(embeddings.shape)}  (batch, sequence, hidden_dim)")

$ python trace_embeddings.py
1. Text:      'The capital of France is'
2. Token IDs: [464, 3139, 286, 4881, 318]
3. Tokens:    ['The', ' capital', ' of', ' France', ' is']
4. Embedding shape: (1, 5, 768)  (batch, sequence, hidden_dim)

Five tokens became five 768-dimensional vectors, one lookup per ID — no context between them yet. 768 is a hyperparameter of this specific model (its “hidden size”), not something derived from the input.

Inside a Transformer Block

from transformers import AutoModelForCausalLM, AutoConfig

config = AutoConfig.from_pretrained("distilgpt2")
print(f"Transformer blocks (layers): {config.n_layer}")
print(f"Attention heads per block:   {config.n_head}")
print(f"Hidden (embedding) size:     {config.n_embd}")
print(f"Vocabulary size:             {config.vocab_size}")
print(f"Max context length:          {config.n_positions}")

model = AutoModelForCausalLM.from_pretrained("distilgpt2")
total_params = sum(p.numel() for p in model.parameters())
print(f"Total parameters:            {total_params:,}")

block0 = model.transformer.h[0]
print()
print("First transformer block's sublayers:")
for name, module in block0.named_children():
    n_params = sum(p.numel() for p in module.parameters())
    print(f"  {name:12} {type(module).__name__:20} {n_params:,} params")

$ python trace_block.py
Transformer blocks (layers): 6
Attention heads per block:   12
Hidden (embedding) size:     768
Vocabulary size:             50257
Max context length:          1024
Total parameters:            81,912,576

First transformer block's sublayers:
  ln_1         LayerNorm            1,536 params
  attn         GPT2Attention        2,362,368 params
  ln_2         LayerNorm            1,536 params
  mlp          GPT2MLP              4,722,432 params

That’s the whole recipe for one block: layer-norm, self-attention (12 heads splitting the 768-dimensional space into 12 parallel 64-dimensional attention computations), another layer-norm, then a feed-forward network — and both the attention and feed-forward outputs get added back to the block’s input (a residual connection) rather than replacing it. Stack this block 6 times and every token’s representation has been updated by information from every other token, 6 times over. For the actual attention-score math — the weighted-average computation inside attn — see Self-Attention From Scratch, which implements it directly in PyTorch.

From Hidden States to a Next Token

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

tok = AutoTokenizer.from_pretrained("distilgpt2")
model = AutoModelForCausalLM.from_pretrained("distilgpt2")
model.eval()

ids = tok("The capital of France is", return_tensors="pt")

with torch.no_grad():
    out = model(**ids)
logits = out.logits[0, -1]
print(f"5. Logits shape: {tuple(logits.shape)}  (one score per vocabulary entry)")

probs = torch.softmax(logits, dim=-1)
top5 = torch.topk(probs, 5)
print("6. Top 5 next-token probabilities:")
for prob, idx in zip(top5.values, top5.indices):
    print(f"   {tok.decode([idx]):>12} -> {prob.item():.4f}")

$ python trace_logits.py
5. Logits shape: (50257,)  (one score per vocabulary entry)
6. Top 5 next-token probabilities:
            the -> 0.1177
           home -> 0.0823
            now -> 0.0499
              a -> 0.0493
          being -> 0.0356

After 6 transformer blocks, the last token’s 768-dimensional vector gets projected to 50,257 raw scores (logits) — one per vocabulary entry — and softmax turns those into a probability distribution. Notice “Paris” isn’t in the top 5: distilgpt2 is a genuinely small, weakly trained model (82M parameters vs. tens or hundreds of billions in production models), so its probabilities are noisy. The pipeline that produced this distribution — tokenize, embed, run through stacked transformer blocks, project to vocabulary logits, softmax — is exactly the same pipeline a production model runs; only the scale changes. What happens to this probability distribution next (how a specific token gets chosen from it) is covered in Inference and Decoding.

Practice

1. Print the shape of every transformer block's output to confirm the sequence length never changes.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

tok = AutoTokenizer.from_pretrained("distilgpt2")
model = AutoModelForCausalLM.from_pretrained("distilgpt2")
model.eval()

ids = tok("The capital of France is", return_tensors="pt")
with torch.no_grad():
    out = model(**ids, output_hidden_states=True)

for i, h in enumerate(out.hidden_states):
    print(f"After block {i}: {tuple(h.shape)}")

After block 0: (1, 5, 768)
After block 1: (1, 5, 768)
After block 2: (1, 5, 768)
After block 3: (1, 5, 768)
After block 4: (1, 5, 768)
After block 5: (1, 5, 768)
After block 6: (1, 5, 768)

7 hidden-state snapshots for 6 blocks: the first is the raw embeddings (before any block runs), the rest are each block’s output. The shape never changes — every block refines the same 5 768-dimensional vectors, it never adds or removes tokens.

2. Compare the top-5 next-token predictions for two different prompts.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

tok = AutoTokenizer.from_pretrained("distilgpt2")
model = AutoModelForCausalLM.from_pretrained("distilgpt2")
model.eval()

for prompt in ["My favorite programming language is", "The weather today is"]:
    ids = tok(prompt, return_tensors="pt")
    with torch.no_grad():
        logits = model(**ids).logits[0, -1]
    top3 = torch.topk(torch.softmax(logits, dim=-1), 3)
    words = [tok.decode([i]).strip() for i in top3.indices]
    print(f"{prompt!r} -> {words}")

'My favorite programming language is' -> ['Ruby', 'Python', 'the']
'The weather today is' -> ['a', 'the', 'very']

Different prompts genuinely produce different probability distributions — the model is conditioning on the input, not returning a fixed answer. “Programming language” actually surfaces “Ruby” and “Python” near the top, while “weather” falls back to bland function words — distilgpt2 has picked up a real topical association in one case and not the other, which is exactly the kind of unevenness a small, weakly trained model produces.

3. Confirm the embedding layer is a pure lookup table (same token, same vector, regardless of position).

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

tok = AutoTokenizer.from_pretrained("distilgpt2")
model = AutoModelForCausalLM.from_pretrained("distilgpt2")

ids = tok("I like cats and I like cats", return_tensors="pt")
print(ids["input_ids"])
print([tok.decode([t]) for t in ids["input_ids"][0]])
embeddings = model.transformer.wte(ids["input_ids"])[0]
# "cats" is token index 2 and 6 in this sentence
same_token_same_vector = torch.equal(embeddings[2], embeddings[6])
print(f"'cats' occurrences have identical embeddings: {same_token_same_vector}")

tensor([[   40,   588, 11875,   290,   314,   588, 11875]])
['I', ' like', ' cats', ' and', ' I', ' like', ' cats']
'cats' occurrences have identical embeddings: True

” cats” appears twice (positions 2 and 6, both with the same leading-space token ID) and gets the exact same embedding vector both times — the embedding layer alone has no notion of position or context. Positional information is injected separately (GPT-2 adds a learned position embedding on top), and contextual meaning only emerges after the transformer blocks run.

From Tokens to Embeddings

Inside a Transformer Block

From Hidden States to a Next Token

Practice

Frequently Asked Questions