Transformer Architecture Complete Guide: From Tokenization to Training — A Hands-On Walkthrough

A hands-on breakdown of Transformer's four core modules, from tokenization to building a working text classifier with PyTorch.
This article builds toward a working text classification model by systematically explaining Transformer's four core modules: the embedding layer (token vectorization), attention mechanism (context association), feed-forward network (knowledge storage), and residual connections with normalization (training stability). It covers the QKV three-role model and causal masking in depth, explains why positional encoding is necessary, and clarifies the independent-parameter principle in stacked layers. The training section covers next-token prediction as a self-supervised objective, padding strategies, and the key trick of using only the last token's output for classification. The final model — just 110K parameters — reaches 99% training accuracy in 15 epochs.
Why You Should Deeply Understand the Transformer Architecture
The large model landscape is booming. From the GPT series to various hundred-billion-parameter models, the names differ but the underlying architecture is remarkably consistent — Transformer. As the industry consensus goes: "Modern large models may differ in details, but over 95% of their core structure is built on the Transformer architecture."
Mastering the Transformer is therefore a prerequisite for large model development. This article systematically breaks down Transformer's four core modules, pairs the explanation with PyTorch implementations, and walks you through the complete pipeline — from tokenization to model training — ultimately building a working text classification model.



The Four Core Modules of Transformer
Before diving into technical details, let's establish a high-level architectural overview. The Transformer is composed of four core modules that work in concert:
- Feed-Forward Network (FFN): The knowledge storage layer. Facts the model learns — such as "Beijing is the capital of China" — are stored as parameters here.
- Attention Mechanism: The context association layer. Enables information exchange between tokens, capturing semantic dependencies in text.
- Embedding Layer: The token vectorization module. Converts discrete text symbols into continuous vector representations the model can process.
- Residual Connections & Normalization (Add & Norm): The training stability layer. Accelerates convergence and prevents vanishing or exploding gradients.
Understanding the role of each module will make every implementation detail that follows much clearer.
Tokenization and Embeddings: From Text to Vectors
Neural networks cannot process raw text directly — the first step is tokenization. Each character maps to a unique ID in the vocabulary, and the vocabulary size equals the total number of tokens. For example, a vocabulary of 10,000 Chinese characters has a vocab size of 10,000.
How Token IDs Map to Word Vectors
Each token ID maps to a fixed-dimensional vector (an array of numbers). These vectors start as random values and are progressively refined during training through backpropagation, enabling them to accurately represent the semantic features of their corresponding tokens. PyTorch's Embedding layer handles this mapping: given a sequence of token IDs, it outputs a matrix of corresponding word vectors.
Take the sentence "中国的首都是北京" ("Beijing is the capital of China") as an example. The seven characters correspond to seven token IDs, which the embedding layer transforms into seven word vectors — forming the model's initial input.
Why Positional Encoding Is Necessary
Word vectors alone carry no positional information. The same word in different positions often serves different semantic roles, which is why Positional Encoding is needed. A maximum sequence length is defined (e.g., 50), and a positional vector is generated for each position — either via random initialization or sinusoidal functions.
Final input = word vector + positional vector. The two are added element-wise, fusing both content and positional information.
It's worth noting that positional encoding has many implementation variants. The tokenizer and positional encoding scheme used in production systems may differ from those in the original paper — this is a natural part of technical evolution.
Deep Dive into Attention: The QKV Three-Role Model
The attention mechanism is Transformer's core innovation and a staple of technical interviews. The key challenge is understanding how a single vector can simultaneously play three different roles.
Understanding QKV Through an Analogy
Think of it like an employee wearing multiple hats in a company:
- Query (Q) — The Interviewer: Represents the current token actively querying the relevance of other tokens.
- Key (K) — The Interviewee: The identifying information that other tokens query and match against.
- Value (V) — The Executor: The actual content being transmitted.
The same token vector is projected through three independent weight matrices (Wq, Wk, Wv) to produce Q, K, and V vectors. Although they originate from the same input, the different projections give them distinct functional roles.
The Three-Step Attention Calculation
Using the token "北" in "中国的首都是[北]京" as an example:
- Relevance Scoring: The current token's Q vector is dot-producted with the K vectors of all preceding tokens (including itself) to produce raw relevance scores. After Softmax normalization, attention weights w0 through w4 are obtained.
- Information Aggregation: The attention weights are used to compute a weighted sum of the V vectors of each token, producing a new vector that incorporates context from previous tokens.
- Output Projection: A linear transformation is applied to the weighted sum to extract key features as the final output of this attention layer.
Causal Masking
Due to the autoregressive nature of text generation, each token can only attend to its predecessors — it cannot "see" future tokens. This is implemented in code using an upper-triangular mask matrix: torch.triu generates a matrix with 1s in the upper triangle and 0s in the lower triangle, ensuring the i-th token can only attend to positions ≤ i. This design reflects the natural temporal order of language.
Feed-Forward Networks and Stacked Layers: Storing and Transferring Knowledge
Once the attention mechanism integrates contextual information, the Feed-Forward Network (FFN) transforms that information into the model's internalized knowledge. The "intelligence" of large models is stored as parameters in the FFN weight matrices.
The Stacked Architecture
A set of vectors is processed by the attention layer to produce new vectors, then passed through the FFN to yield the next-stage representations. Matching input and output dimensions is the critical design choice — it allows Transformer blocks to be stacked arbitrarily.
A single complete layer follows this flow: Attention → Add & Norm → FFN → Add & Norm. Repeating this structure N times produces an N-layer Transformer encoder or decoder.
Key point: each layer has entirely independent parameters. The computation logic is the same, but the weight matrices differ. Shallower layers capture local patterns and basic features; deeper layers extract abstract semantics and complex relationships. The hundreds of billions of parameters in large models largely come from stacking dozens of these layers.
Training Strategy and Text Classification in Practice
The pretraining objective for large models is Next Token Prediction. Given "中", predict "国"; given "中国", predict the next character. This self-supervised learning approach allows models to learn language patterns from vast amounts of unlabeled text.
The Nature of Classification Tasks
Predicting the next token is fundamentally a multi-class classification problem, with the number of classes equal to the vocabulary size (e.g., a 10,000-token vocab produces a 10,000-dimensional output). For simplicity in this tutorial, we implement a five-class text classification task (e.g., categories like everyday knowledge and geography).
Key technique: for classification, only the output vector of the last token in the sequence is used. Because the last token has aggregated all prior information through attention, it carries a semantic representation of the entire sentence. The outputs of earlier tokens are not needed for final classification.
Handling Padding
In practice, sentences vary in length, but GPUs require uniform-length batches. The solution is padding: set a fixed length (e.g., 50), pad short sentences with zeros, and truncate long ones. These padding positions carry no real semantic meaning — they exist purely to satisfy the requirements of the compute framework. In code, you need to locate the position of the "true last token" for classification and ignore the padding positions.
Verifying Training Performance
Experiment results show that after 10 training epochs, training accuracy reaches 88%; after 15 epochs, it reaches 99%. On a small test set, 100% accuracy was achieved (limited by dataset size). The model has approximately 110,000 parameters in total — about 16,000 in the attention module and 33,000 in the FFN. The proportion of parameters in the embedding layer tends to decrease relative to model scale as model size grows.
Key Takeaways
The Transformer's working mechanism can be summarized as: the attention layer uses the QKV three-role mechanism to capture contextual associations and compress semantic information into vector representations; the feed-forward network then stores and retrieves the knowledge encoded in those vectors. This architecture is the foundation of today's mainstream large models, and the various innovations in subsequent models are all incremental refinements built on top of it. Deeply understanding the Transformer is a critical step toward mastering modern deep learning.
Related articles

Supply Chain Hardware Implants: The Most Dangerous Security Threat You're Overlooking
A deep dive into supply chain hardware implant attacks: how they work, historical cases, and defense strategies. Learn why hardware backdoors are nearly undetectable and how to build a zero-trust defense.

Apple M6 and M5 Ultra Chips Unveiled: What the Major AI Performance Boost Really Means
Apple launches M6 and M5 Ultra chips with dramatically enhanced Neural Engine and on-device AI performance. A deep dive into architecture upgrades, unified memory, and real-world impact.

Fine-Tuning LLMs to Mimic Real Human Chat Styles: A Guide to Building Emotion-Aware Datasets
How to fine-tune an LLM to mimic real human chat styles? This guide covers emotion labeling, context-aware datasets, LoRA fine-tuning, and iterative optimization.