Building a Transformer from Scratch: Core Architecture and PyTorch Implementation

Transformer: The Foundation of Modern Large Language Models

The Transformer architecture is the backbone of today's large language models. Whether it's GPT, Claude, or other models with hundreds of billions of parameters, 95% of their structure is built on the Transformer. Understanding this architecture is an essential step toward mastering LLM development.

This article walks you through every component of the Transformer — from the underlying principles to a working implementation — using complete PyTorch code.

The Four Core Components of Transformer

The Transformer is built from four key modules, each with a distinct role:

Feed-Forward Network (FFN): The Knowledge Store

The FFN is responsible for storing knowledge. When a large model knows facts like "Beijing is the capital of China," that knowledge lives here. It transforms contextual information into implicit knowledge, stored as vectors in the network's parameters. In MoE (Mixture of Experts) architectures, optimizations around knowledge storage are primarily focused on this layer.

Attention Mechanism: The Context Associator

Attention handles contextual reasoning. When processing a sentence, the model needs to understand the context established by earlier tokens — Attention is the mechanism that makes this possible. It allows each token to "see" and relate to all preceding tokens.

Embedding: Text Vectorization

The Embedding layer converts text into vectors. Since models can't process characters or letters directly, this layer maps each token to a numerical vector that serves as the model's actual input.

Add & Norm: The Training Stabilizer

Layer normalization improves training stability and accelerates convergence. By adding residual connections and normalization after each sub-module, it ensures that even deep networks can be trained effectively.

From Text to Vectors: Tokenization and Positional Encoding

The Tokenization Pipeline

Every character or letter maps to a unique ID. For example, the character "小" might map to ID 75 — this mapping is defined by the vocabulary. The vocabulary size determines how many distinct tokens the model can handle.

In PyTorch, this is implemented via nn.Embedding:

embedding = nn.Embedding(vocab_size, d_model)

Here, vocab_size is the vocabulary size and d_model is the vector dimensionality. Initially, each ID maps to a random vector that gets refined through training.

Why Positional Encoding Matters

The same word can carry different meanings depending on its position in a sentence. So beyond word vectors, each position also gets its own vector. If the maximum sequence length is 50 tokens, there are 50 positional vectors.

The final input to the model = token embedding + positional embedding

For a 7-character sentence, the model generates 7 token vectors and 7 positional vectors (positions 0 through 6), which are summed to produce the actual input.

A Deep Dive into the Attention Mechanism

Three Roles: The Essence of QKV

The core insight of Attention is that each token simultaneously plays three roles, achieved through three separate linear transformations:

• Q (Query): The active searcher — looks for information relevant to itself
• K (Key): The passive match target — gets searched and matched by others
• V (Value): The information provider — supplies the actual semantic content

Think of it like an employee's three roles at a company: interviewer (selecting candidates), job applicant (being evaluated), and regular employee (doing the work).

Step-by-Step Computation

Using a token mid-sentence as an example:

1. Compute similarity: Take the current token's Q vector and compute dot products with the K vectors of all preceding tokens (including itself) to get attention scores
2. Normalize weights: Apply softmax to convert scores into a probability distribution, yielding weights w0 through w4
3. Weighted sum: Use these weights to compute a weighted sum of the V vectors, producing a new vector enriched with contextual information
4. Linear projection: Apply a linear transformation to extract more useful representations

Key code:

# Compute attention weights
scores = torch.matmul(Q, K.transpose(-2, -1))
weights = F.softmax(scores, dim=-1)
# Weighted sum
output = torch.matmul(weights, V)

The Mask Mechanism: Preventing Information Leakage

During training, each token should only see what came before it — not future tokens. This is enforced with an upper-triangular mask matrix:

mask = torch.triu(torch.ones(seq_len, seq_len), diagonal=1)

The first token sees only itself, the second token sees the first two, and the last token can see all previous context. This mirrors the autoregressive nature of language generation.

Stacking Layers: The Decoder Architecture

A single Decoder layer follows this computation flow:

1. Input vectors → Self-Attention → Add & Norm
2. → Feed-Forward Network (FFN) → Add & Norm
3. Output vectors

Because both input and output are sequences of vectors, these layers can be stacked arbitrarily deep. Models like GPT-3 stack dozens of these layers to reach hundreds of billions of parameters — but the core structure is just this pattern repeated.

Code implementation:

for layer in self.decoder_layers:
    x = layer(x, mask)  # Each layer has independent parameters

Note: although the computation method is the same across layers, each layer has its own independent parameters. Shallow layers capture direct features; deeper layers learn complex semantics.

Training Objective: Next Token Prediction

How Large Models Are Trained

The core task is predicting the next token. Given one token, predict the next; given two tokens, predict the third. Every position is predicting what comes after it.

This training approach gives the model the ability to "continue" any sequence. Accurately predicting what comes next means the model genuinely understands the current context — similar to how a skilled trader who can predict tomorrow's market movement must have deep insight into historical patterns.

Text Classification: A Practical Example

This article demonstrates a simpler application: text classification. Assume there are 5 categories — the model outputs a 5-dimensional vector representing the probability of each class, and the class with the highest value is the prediction.

Key insight: only the output vector of the last token is needed for classification, since it has attended to all preceding context. The outputs from earlier tokens are discarded.

Full Code Walkthrough

Model Definition

class TransformerClassifier(nn.Module):
    def __init__(self, vocab_size, d_model, n_layers, n_classes):
        super().__init__()
        self.token_embedding = nn.Embedding(vocab_size, d_model)
        self.position_embedding = nn.Embedding(max_len, d_model)
        self.decoder_layers = nn.ModuleList([
            DecoderLayer(d_model) for _ in range(n_layers)
        ])
        self.classifier = nn.Linear(d_model, n_classes)

Training Tip: Handling Padding

Sentences vary in length, but GPUs require uniform input sizes. The solution:

• Short sentences: pad with zeros at the end
• Long sentences: truncate

In the code, the index of the last non-zero token is found dynamically, and that position's output vector is used for classification.

Training Results

Training for 15 epochs on the demo dataset achieves 99% accuracy. While the dataset is small, it validates the architecture effectively.

Parameter breakdown:

• Total parameters: ~110,000
• Attention: ~16,000 (for capturing context)
• FFN: ~33,000 (for storing knowledge)
• Embedding: the remainder (token and positional encodings)

As the number of layers increases, the proportion of parameters in Attention and FFN grows significantly, while the Embedding share shrinks.

Key Takeaways

If you're asked about the Attention mechanism in an interview, don't just draw a diagram — be able to articulate what each module actually does:

1. Attention parameters: capture context by building associations between tokens
2. FFN parameters: store knowledge by transforming contextual information into predictable outputs
3. QKV separation: a single vector playing three roles through three distinct transformations is the essence of Attention
4. Layer stacking: repeated Decoder layers progressively extract higher-level semantic features

From Transformer to Large Language Models

Modern LLMs (GPT, Claude, LLaMA, etc.) are all built on this core architecture. The differences come down to implementation details:

• Attention variants (Multi-Head, Grouped-Query, etc.)
• Normalization approaches (LayerNorm, RMSNorm, etc.)
• Positional encoding schemes (RoPE, ALiBi, etc.)
• Activation functions (GeLU, SwiGLU, etc.)

But the foundation stays the same. Master the architecture covered in this article and you'll have the key to understanding any large language model. Keep this as a reference, work through the code hands-on, and make sure you can not only follow the reasoning — but write it yourself.