I Built a Small LLM From Scratch — Here’s What I Finally Understood About Transformers

Have you ever wondered what actually happens inside an LLM when you type:

"Hello, how are" Enter fullscreen mode Exit fullscreen mode and it predicts:

"you" Enter fullscreen mode Exit fullscreen mode We often hear words like Transformer, Attention, Tokens, RoPE, GQA, RMSNorm, and SwiGLU.

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.

By the end, you should understand the basic flow:

Text ↓ Tokens ↓ Embeddings ↓ Transformer ↓ Attention ↓ Feed Forward Network ↓ Prediction ↓ Next Token Enter fullscreen mode Exit fullscreen mode Let's break it down step by step.

A language model learns to predict what comes next.

Input: "The cat" Prediction: "sat" Enter fullscreen mode Exit fullscreen mode Then:

"The cat sat" Prediction: "on" Enter fullscreen mode Exit fullscreen mode Then:

"The cat sat on" Prediction: "the" Enter fullscreen mode Exit fullscreen mode 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.

Before the model can learn anything, we have a problem.

A neural network doesn't directly understand:

Hello Enter fullscreen mode Exit fullscreen mode It needs numbers.

H → 20 e → 43 l → 50 l → 50 o → 53 Enter fullscreen mode Exit fullscreen mode So:

"Hello" Enter fullscreen mode Exit fullscreen mode might become:

[20, 43, 50, 50, 53] Enter fullscreen mode Exit fullscreen mode This is called tokenization.

In this particular project, we use a simple character-level tokenizer.

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 ) } Enter fullscreen mode Exit fullscreen mode Then encoding becomes:

def encode ( text ): return [ char_to_idx [ c ] for c in text ] Enter fullscreen mode Exit fullscreen mode And decoding:

def decode ( ids ): return "". join ( idx_to_char [ i ] for i in ids ) Enter fullscreen mode Exit fullscreen mode So we now have:

Human language ↓ Characters ↓ Numbers Enter fullscreen mode Exit fullscreen mode That's the first step.

"To be or not to be" Enter fullscreen mode Exit fullscreen mode We can create training examples by shifting the sequence by one position.

Input: To be o Target: o be or Enter fullscreen mode Exit fullscreen mode Why?

Because each position teaches the model what comes next.

Input seen Target T o To [space] To [space] b To [space] b e Enter fullscreen mode Exit fullscreen mode So the model repeatedly gets:

Previous characters → Next character Enter fullscreen mode Exit fullscreen mode This is called next-token prediction.

The training data in our implementation creates these shifted input/target sequences automatically.

[20, 43, 50, 50, 53] Enter fullscreen mode Exit fullscreen mode 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.

Token ID ↓ Embedding table ↓ Vector Enter fullscreen mode Exit fullscreen mode For our model, each token becomes a vector of size 256.

50 Enter fullscreen mode Exit fullscreen mode we might have:

[0.12, -0.43, 0.87,...] Enter fullscreen mode Exit fullscreen mode with 256 values.

Now the neural network has a useful numerical representation to work with.

A Transformer processes the token representations and allows tokens to interact with one another.

"The dog chased the ball because it was moving." Enter fullscreen mode Exit fullscreen mode What does "it" refer to?

A model needs to understand relationships between different words/tokens.

Which other tokens should I pay attention to when processing this token?

Q = Query K = Key V = Value Enter fullscreen mode Exit fullscreen mode A simple way to think about them:

The model calculates how strongly a Query should attend to different Keys.

Q × Kᵀ ↓ attention scores ↓ softmax ↓ attention weights ↓ weights × V ↓ attention output Enter fullscreen mode Exit fullscreen mode In code, the core idea looks like:

scores = ( q @ k. transpose ( - 2, - 1 )) * scale weights = F. softmax ( scores, dim =- 1 ) out = weights @ v Enter fullscreen mode Exit fullscreen mode But there is an important problem.

During language generation, the model shouldn't be allowed to see future tokens.

I love machine learning Enter fullscreen mode Exit fullscreen mode When predicting "machine", the model shouldn't already know "learning".

T0 T1 T2 T3 T0 ✓ ✗ ✗ ✗ T1 ✓ ✓ ✗ ✗ T2 ✓ ✓ ✓ ✗ T3 ✓ ✓ ✓ ✓ Enter fullscreen mode Exit fullscreen mode ✓ means the token can attend there.

itself + everything before it Enter fullscreen mode Exit fullscreen mode but never the future.

That's what makes next-token prediction possible.

The implementation applies a triangular causal mask before softmax.

Dog bites man Enter fullscreen mode Exit fullscreen mode and:

Man bites dog Enter fullscreen mode Exit fullscreen mode The same words are present, but their positions are different.

A Transformer therefore needs positional information.

Instead of simply adding a position vector, RoPE rotates parts of the Query and Key representations.

Position 3 → Position 5 distance = 2 Position 7 → Position 9 distance = 2 Enter fullscreen mode Exit fullscreen mode 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.

Now we have another engineering problem.

Attention can become expensive, especially when models become large.

This is where Grouped Query Attention (GQA) can help.

Q0 → K0, V0 Q1 → K1, V1 Q2 → K2, V2 Q3 → K3, V3... Enter fullscreen mode Exit fullscreen mode With GQA, multiple Query heads share Key and Value heads.

8 Query heads 2 KV heads Enter fullscreen mode Exit fullscreen mode Conceptually:

Q0 ─┐ Q1 ─┤ Q2 ─┤──→ K0, V0 Q3 ─┘ Q4 ─┐ Q5 ─┤ Q6 ─┤──→ K1, V1 Q7 ─┘ Enter fullscreen mode Exit fullscreen mode So four Query heads share the same KV head.

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.

Modern LLM architecture isn't only about making models powerful. It's also about making them efficient.

Attention isn't the only part of a Transformer block.

After attention, we have a Feed-Forward Network (FFN).

Input ↓ Linear ↓ ReLU ↓ Linear ↓ Output Enter fullscreen mode Exit fullscreen mode SwiGLU uses a gating mechanism.

┌──→ SiLU ──┐ Input ────────┤ × ──→ Output └──→ Up ─────┘ Enter fullscreen mode Exit fullscreen mode The code is essentially:

gate = F. silu ( self. w_gate ( x )) up = self. w_up ( x ) output = self. w_down ( gate * up ) Enter fullscreen mode Exit fullscreen mode 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.

You'll notice another component appearing before attention and the FFN:

RMSNorm Enter fullscreen mode Exit fullscreen mode 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.

Input ↓ Calculate RMS ↓ Normalize ↓ Learnable scale ↓ Output Enter fullscreen mode Exit fullscreen mode A simplified implementation:

rms = torch. sqrt ( x. pow ( 2 ). mean ( dim =- 1, keepdim = True ) + eps ) output = ( x / rms ) * weight Enter fullscreen mode Exit fullscreen mode RMSNorm is simpler than LayerNorm because it doesn't center the values by subtracting the mean.

Input │ ├── RMSNorm ↓ Attention │ ↓ + ←──────── Input │ ↓ x' │ ├── RMSNorm ↓ SwiGLU │ ↓ + ←──────── x' │ ↓ Output Enter fullscreen mode Exit fullscreen mode In code:

x = x + self. attention ( self. attn_norm ( x ), rope_cos, rope_sin ) x = x + self. ffn ( self. ffn_norm ( x ) ) Enter fullscreen mode Exit fullscreen mode 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.

Token IDs ↓ Token Embedding ↓ ┌─────────────────────┐ │ Transformer │ │ Block │ │ │ │ RMSNorm │ │ ↓ │ │ GQA + RoPE │ │ ↓ │ │ Residual │ │ ↓ │ │ RMSNorm │ │ ↓ │ │ SwiGLU │ │ ↓ │ │ Residual │ └─────────────────────┘ ↓ × 4 Blocks ↓ Final RMSNorm ↓ Language Head ↓ Logits ↓ Probabilities ↓ Next Character Enter fullscreen mode Exit fullscreen mode Our implementation uses four Transformer blocks.

Vocabulary: 65 characters Embedding size: 256 Transformer blocks: 4 Query heads: 8 KV heads: 2 Context length: 256 Enter fullscreen mode Exit fullscreen mode The complete model connects token embeddings, Transformer blocks, final normalization and the language-model head.

1. Get training data ↓ 2. Forward pass ↓ 3. Calculate loss ↓ 4. Calculate gradients ↓ 5. Update weights ↓ 6. Repeat Enter fullscreen mode Exit fullscreen mode In code:

xb, yb = get_batch logits, loss = model ( xb, yb ) optimizer. zero_grad () loss. backward () optimizer. step () Enter fullscreen mode Exit fullscreen mode Our training setup uses:

Batch size: 64 Learning rate: 3e-4 Optimizer: AdamW Training steps: 3000 Enter fullscreen mode Exit fullscreen mode Gradient clipping is also used:

torch. nn. utils. clip_grad_norm_ ( model. parameters (), max_norm = 1.0 ) Enter fullscreen mode Exit fullscreen mode The purpose is to prevent very large gradients from destabilizing training.

For our character vocabulary, the model produces a score for each possible character.

a → 1.2 b → -0.8 c → 2.4 d → 0.1... Enter fullscreen mode Exit fullscreen mode 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.

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.

Training data = Questions you practiced Validation data = New questions you haven't practiced Enter fullscreen mode Exit fullscreen mode 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 () Enter fullscreen mode Exit fullscreen mode The implementation evaluates both training and validation splits to monitor the loss.

After training, we can finally generate text.

ROMEO: Enter fullscreen mode Exit fullscreen mode The model predicts the next character.

ROMEO: W Enter fullscreen mode Exit fullscreen mode Then we feed the new sequence back into the model:

ROMEO: W Enter fullscreen mode Exit fullscreen mode It predicts another character:

ROMEO: Wh Enter fullscreen mode Exit fullscreen mode Then:

ROMEO: What Enter fullscreen mode Exit fullscreen mode And so on.

Prompt ↓ Predict next token ↓ Add token to prompt ↓ Predict again ↓ Add token ↓ Repeat Enter fullscreen mode Exit fullscreen mode 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.

A → 70% B → 20% C → 8% D → 2% Enter fullscreen mode Exit fullscreen mode Should we always choose A?

Temperature controls how sharp or random the probability distribution becomes.

More predictable Enter fullscreen mode Exit fullscreen mode Higher temperature More varied Enter fullscreen mode Exit fullscreen mode For example:

0.3 → conservative 0.8 → balanced 1.5 → more creative Enter fullscreen mode Exit fullscreen mode So temperature doesn't change what the model learned.

It changes how we sample from what the model learned.

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 Enter fullscreen mode Exit fullscreen mode That's the basic idea behind an autoregressive Transformer language model.

If you're learning Transformers for the first time, don't try to memorize every equation immediately.

Allows tokens to use information from other tokens.

Prevents the model from looking at future tokens.

Provides positional information through rotations of Query and Key representations.

Allows multiple Query heads to share Key/Value heads, reducing KV-cache size.

A gated feed-forward network used inside the Transformer block.

Normalizes activations using their root mean square.

Adds the previous representation back to the transformed representation.

Adjusting the model's weights to reduce prediction error.

Repeatedly predicting and adding the next token.

If you understand these ten ideas, you already have a much stronger mental model of a Transformer.

Before working through this implementation, I knew many of these terms individually.

RoPE GQA RMSNorm SwiGLU Attention Enter fullscreen mode Exit fullscreen mode 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? Enter fullscreen mode Exit fullscreen mode Once you look at it this way, the architecture becomes much easier to understand.

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?

Hello ↓ Token IDs ↓ Embeddings ↓ Attention ↓ RoPE ↓ GQA ↓ SwiGLU ↓ Transformer blocks ↓ Logits ↓ Probability ↓ Next token Enter fullscreen mode Exit fullscreen mode 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.