從零搭建小型大語言模型，梳理 Transformer 的運作流程
I Built a Small LLM From Scratch — Here’s What I Finally Understood About Transformers
作者以 PyTorch 從零搭建小型大語言模型，逐步解説文字如何轉成 token，再經 Transformer 處理以預測下一個 token。示例模型以莎士比亞文本訓練，使用字元級 tokenizer、4 個 Transformer blocks、8 個 Query heads、2 個 KV heads，embedding size 為 256。
How a Tiny LLM Works: Building a Language Model From Scratch in PyTorch
Have you ever wondered what actually happens inside an LLM when you type:
"Hello, how are"
and it predicts:
"you"
We often hear words like Transformer, Attention, Tokens, RoPE, GQA, RMSNorm, and SwiGLU.
But what do these things actually do?
I recently worked through the implementation of a small language model from scratch, and I found that the easiest way to understand an LLM is to follow one piece of text through the entire model.
So in this post, we'll do exactly that.
By the end, you should understand the basic flow:
Text
↓
Tokens
↓
Embeddings
↓
Transformer
↓
Attention
↓
Feed Forward Network
↓
Prediction
↓
Next Token
Let's break it down step by step.
1. First: What Is a Language Model?
A language model learns to predict what comes next.
For example:
Input:
"The cat"
Prediction:
"sat"
Then:
"The cat sat"
Prediction:
"on"
Then:
"The cat sat on"
Prediction:
"the"
So at its simplest, a language model is doing:
Given the previous tokens, predict the next token.
That's the central idea behind autoregressive language models.
Our small model will learn this task using Shakespeare text.
2. Computers Don't Understand Text
Before the model can learn anything, we have a problem.
A neural network doesn't directly understand:
Hello
It needs numbers.
So we convert characters into IDs.
For example:
H → 20
e → 43
l → 50
l → 50
o → 53
So:
"Hello"
might become:
[20, 43, 50, 50, 53]
This is called tokenization.
In this particular project, we use a simple character-level tokenizer.
The basic idea is:
chars = sorted(set(text))
char_to_idx = {
c: i for i, c in enumerate(chars)
}
idx_to_char = {
i: c for i, c in enumerate(chars)
}
Then encoding becomes:
def encode(text):
return [char_to_idx[c] for c in text]
And decoding:
def decode(ids):
return "".join(idx_to_char[i] for i in ids)
So we now have:
Human language
↓
Characters
↓
Numbers
That's the first step.
3. How Does the Model Know What to Predict?
Suppose our text is:
"To be or not to be"
We can create training examples by shifting the sequence by one position.
Input: To be o
Target: o be or
Why?
Because each position teaches the model what comes next.
For example:
Input seen Target
T o
To [space]
To [space] b
To [space] b e
So the model repeatedly gets:
Previous characters → Next character
This is called next-token prediction.
The training data in our implementation creates these shifted input/target sequences automatically.
4. Numbers Go Into an Embedding Layer
We now have IDs such as:
[20, 43, 50, 50, 53]
But these IDs are just labels.
The number 50 isn't inherently more meaningful than 20.
So we convert each token ID into a vector.
This is called a token embedding.
Think of it like:
Token ID
↓
Embedding table
↓
Vector
For our model, each token becomes a vector of size 256.
So instead of:
50
we might have:
[0.12, -0.43, 0.87, ...]
with 256 values.
Now the neural network has a useful numerical representation to work with.
5. But How Does the Model Understand Relationships Between Tokens?
This is where the Transformer comes in.
A Transformer processes the token representations and allows tokens to interact with one another.
The most important component is:
Attention
Let's use a simple example:
"The dog chased the ball because it was moving."
What does "it" refer to?
A model needs to understand relationships between different words/tokens.
Attention helps the model decide:
Which other tokens should I pay attention to when processing this token?
6. Query, Key and Value
Attention uses three representations:
Q = Query
K = Key
V = Value
A simple way to think about them:
- Query: What information am I looking for?
- Key: What information do I contain?
- Value: What information should I provide?
The model calculates how strongly a Query should attend to different Keys.
The basic attention calculation is:
Q × Kᵀ
↓
attention scores
↓
softmax
↓
attention weights
↓
weights × V
↓
attention output
In code, the core idea looks like:
scores = (q @ k.transpose(-2, -1)) * scale
weights = F.softmax(scores, dim=-1)
out = weights @ v
But there is an important problem.
During language generation, the model shouldn't be allowed to see future tokens.
7. Causal Attention: Don't Look Into the Future
Suppose the sequence is:
I love machine learning
When predicting "machine", the model shouldn't already know "learning".
So we use a causal mask.
It looks conceptually like this:
T0 T1 T2 T3
T0 ✓ ✗ ✗ ✗
T1 ✓ ✓ ✗ ✗
T2 ✓ ✓ ✓ ✗
T3 ✓ ✓ ✓ ✓
✓ means the token can attend there.
✗ means it cannot.
So each token can see:
itself + everything before it
but never the future.
That's what makes next-token prediction possible.
The implementation applies a triangular causal mask before softmax.
8. Why Does the Model Need Positional Information?
Imagine these two sentences:
Dog bites man
and:
Man bites dog
The same words are present, but their positions are different.
A Transformer therefore needs positional information.
One modern approach is:
RoPE — Rotary Positional Embeddings
Instead of simply adding a position vector, RoPE rotates parts of the Query and Key representations.
The important intuition is:
Position 3 → Position 5
distance = 2
Position 7 → Position 9
distance = 2
The model can capture relationships based on relative position.
RoPE applies rotations to pairs of dimensions in the Query and Key vectors.
You don't need to memorize the rotation formula to understand the main idea:
RoPE gives attention information about token positions and relative distances.
9. Why Do We Need GQA?
Now we have another engineering problem.
Attention can become expensive, especially when models become large.
This is where Grouped Query Attention (GQA) can help.
Normally, we might have:
Q0 → K0, V0
Q1 → K1, V1
Q2 → K2, V2
Q3 → K3, V3
...
With GQA, multiple Query heads share Key and Value heads.
Our small model uses:
8 Query heads
2 KV heads
Conceptually:
Q0 ─┐
Q1 ─┤
Q2 ─┤──→ K0, V0
Q3 ─┘
Q4 ─┐
Q5 ─┤
Q6 ─┤──→ K1, V1
Q7 ─┘
So four Query heads share the same KV head.
Why is that useful?
Because fewer Key/Value heads means a smaller KV cache and less memory usage during inference.
The implementation expands the two KV heads to match the eight Query heads after applying RoPE.
The important lesson:
Modern LLM architecture isn't only about making models powerful. It's also about making them efficient.
10. What Happens After Attention?
Attention isn't the only part of a Transformer block.
After attention, we have a Feed-Forward Network (FFN).
In our model, we use:
SwiGLU
A simple older FFN might look like:
Input
↓
Linear
↓
ReLU
↓
Linear
↓
Output
SwiGLU uses a gating mechanism.
Conceptually:
┌──→ SiLU ──┐
Input ────────┤ × ──→ Output
└──→ Up ─────┘
The code is essentially:
gate = F.silu(self.w_gate(x))
up = self.w_up(x)
output = self.w_down(gate * up)
The multiplication acts like a gate.
It allows the network to learn which information should pass through more strongly.
The implementation uses two parallel projections followed by element-wise multiplication and a final projection back to the model dimension.
11. Why RMSNorm?
You'll notice another component appearing before attention and the FFN:
RMSNorm
Neural networks can become difficult to train when activations become poorly scaled.
Normalization helps keep these values under control.
RMSNorm works by calculating the root mean square of the values and scaling them.
Conceptually:
Input
↓
Calculate RMS
↓
Normalize
↓
Learnable scale
↓
Output
A simplified implementation:
rms = torch.sqrt(
x.pow(2).mean(dim=-1, keepdim=True) + eps
)
output = (x / rms) * weight
RMSNorm is simpler than LayerNorm because it doesn't center the values by subtracting the mean.
12. Residual Connections: Why Do We Add the Input Back?
Our Transformer block looks like:
Input
│
├── RMSNorm
↓
Attention
│
↓
+ ←──────── Input
│
↓
x'
│
├── RMSNorm
↓
SwiGLU
│
↓
+ ←──────── x'
│
↓
Output
In code:
x = x + self.attention(
self.attn_norm(x),
rope_cos,
rope_sin
)
x = x + self.ffn(
self.ffn_norm(x)
)
The x + ... part is the residual connection.
Instead of forcing every layer to completely transform the representation, the layer can learn a useful change and add it to what was already there.
This pattern helps make deeper networks easier to train.
13. Now We Can Build the Whole Transformer
We can combine everything:
Token IDs
↓
Token Embedding
↓
┌─────────────────────┐
│ Transformer │
│ Block │
│ │
│ RMSNorm │
│ ↓ │
│ GQA + RoPE │
│ ↓ │
│ Residual │
│ ↓ │
│ RMSNorm │
│ ↓ │
│ SwiGLU │
│ ↓ │
│ Residual │
└─────────────────────┘
↓
× 4 Blocks
↓
Final RMSNorm
↓
Language Head
↓
Logits
↓
Probabilities
↓
Next Character
Our implementation uses four Transformer blocks.
The model configuration includes:
Vocabulary: 65 characters
Embedding size: 256
Transformer blocks: 4
Query heads: 8
KV heads: 2
Context length: 256
The complete model connects token embeddings, Transformer blocks, final normalization and the language-model head.
14. How Does the Model Actually Learn?
Now we have a model.
But initially, it knows nothing.
So we train it.
The basic training loop is:
1. Get training data
↓
2. Forward pass
↓
3. Calculate loss
↓
4. Calculate gradients
↓
5. Update weights
↓
6. Repeat
In code:
xb, yb = get_batch(...)
logits, loss = model(xb, yb)
optimizer.zero_grad()
loss.backward()
optimizer.step()
Our training setup uses:
Batch size: 64
Learning rate: 3e-4
Optimizer: AdamW
Training steps: 3000
Gradient clipping is also used:
torch.nn.utils.clip_grad_norm_(
model.parameters(),
max_norm=1.0
)
The purpose is to prevent very large gradients from destabilizing training.
15. What Is the Loss?
The model outputs scores called logits.
For our character vocabulary, the model produces a score for each possible character.
For example:
a → 1.2
b → -0.8
c → 2.4
d → 0.1
...
We convert these scores into probabilities.
Then we compare the predicted probability distribution with the correct next character.
This is where cross-entropy loss is used.
Simply put:
If the model gives a high probability to the correct answer, the loss is smaller.
If it gives a low probability to the correct answer, the loss is larger.
The model then uses the loss to calculate gradients and update its weights.
16. Training vs Validation
A model can memorize its training data.
So we don't only look at training loss.
We also evaluate on validation data.
Think of it like studying for an exam:
Training data
= Questions you practiced
Validation data
= New questions you haven't practiced
If training loss keeps decreasing but validation loss stops improving or gets worse, that can indicate overfitting.
During evaluation, we disable dropout and don't calculate gradients:
@torch.no_grad()
def estimate_loss():
model.eval()
...
model.train()
The implementation evaluates both training and validation splits to monitor the loss.
17. How Does Text Generation Work?
After training, we can finally generate text.
Suppose we give the model:
ROMEO:
The model predicts the next character.
Maybe:
ROMEO: W
Then we feed the new sequence back into the model:
ROMEO: W
It predicts another character:
ROMEO: Wh
Then:
ROMEO: What
And so on.
So generation is:
Prompt
↓
Predict next token
↓
Add token to prompt
↓
Predict again
↓
Add token
↓
Repeat
This is called autoregressive generation.
The implementation repeatedly takes the last-position logits, converts them to probabilities, samples the next token, and appends it to the sequence.
18. What Is Temperature?
Suppose the model predicts:
A → 70%
B → 20%
C → 8%
D → 2%
Should we always choose A?
Not necessarily.
Temperature controls how sharp or random the probability distribution becomes.
Low temperature
More predictable
Higher temperature
More varied
For example:
0.3 → conservative
0.8 → balanced
1.5 → more creative
So temperature doesn't change what the model learned.
It changes how we sample from what the model learned.
19. The Whole Process in One Picture
If all of this feels like a lot, remember just this:
TEXT
↓
Tokenization
↓
Token IDs
↓
Embeddings
↓
┌───────────────────┐
│ Transformer Block │
│ │
│ RMSNorm │
│ ↓ │
│ Attention │
│ ↓ │
│ Residual │
│ ↓ │
│ RMSNorm │
│ ↓ │
│ SwiGLU │
│ ↓ │
│ Residual │
└───────────────────┘
↓
More Blocks
↓
Final Norm
↓
Logits
↓
Probabilities
↓
Predict Next Token
↓
Add Token to Input
↓
Repeat
That's the basic idea behind an autoregressive Transformer language model.
20. What I Think Is Most Important to Remember
If you're learning Transformers for the first time, don't try to memorize every equation immediately.
Start with these questions:
What is tokenization?
Converts text into numbers.
What are embeddings?
Turns token IDs into useful vectors.
What is attention?
Allows tokens to use information from other tokens.
What is causal masking?
Prevents the model from looking at future tokens.
What is RoPE?
Provides positional information through rotations of Query and Key representations.
What is GQA?
Allows multiple Query heads to share Key/Value heads, reducing KV-cache size.
What is SwiGLU?
A gated feed-forward network used inside the Transformer block.
What is RMSNorm?
Normalizes activations using their root mean square.
What is a residual connection?
Adds the previous representation back to the transformed representation.
What is training?
Adjusting the model's weights to reduce prediction error.
What is generation?
Repeatedly predicting and adding the next token.
If you understand these ten ideas, you already have a much stronger mental model of a Transformer.
21. The Most Useful Lesson I Got From Building It
Before working through this implementation, I knew many of these terms individually.
But knowing:
RoPE
GQA
RMSNorm
SwiGLU
Attention
isn't enough.
The important part is understanding how they work together.
A language model isn't one magical algorithm.
It's a collection of components, each solving a particular problem:
Tokenization
→ How do we represent text?
Embeddings
→ How do we represent tokens as vectors?
Attention
→ How do tokens interact?
Causal masking
→ How do we prevent cheating?
RoPE
→ How does the model know positions?
GQA
→ How can attention become more memory efficient?
SwiGLU
→ How does the network transform information?
RMSNorm
→ How do we keep activations stable?
Residual connections
→ How do we make deep networks easier to train?
Training
→ How does the model learn?
Generation
→ How does the model produce text?
Once you look at it this way, the architecture becomes much easier to understand.
Conclusion
The easiest way I found to understand an LLM was not to start with:
“Let's learn all the Transformer equations.”
Instead, I started with a much simpler question:
How does
"Hello"eventually become the next predicted token?
Then I followed the data:
Hello
↓
Token IDs
↓
Embeddings
↓
Attention
↓
RoPE
↓
GQA
↓
SwiGLU
↓
Transformer blocks
↓
Logits
↓
Probability
↓
Next token
That made the architecture much less mysterious.
I'm still learning, but building a small model forced me to stop treating LLMs as a black box and start thinking about data flow, tensor shapes, architecture, training, and inference.
And that's probably the most useful advice I can give to someone learning AI:
Don't just learn what a component is. Learn what problem it solves and follow the data through it.
If you're learning Transformers too, I hope this makes the topic a little easier to understand.
Next step: take one of these components — attention, RoPE, GQA, or SwiGLU — and implement it yourself. That's where the concepts really start to stick.
來源：DEV Community · dev.to