Understanding Transformer From the Ground Up: How Language Models Predict the Next Word

The Essence of Language Models: Next-Token Prediction

The core technology behind large language models (LLMs) is the Transformer — the "T" in ChatGPT. Understanding how Transformer works is the first step toward exploring what's really happening inside a language model. This article walks through the complete pipeline of how an input sentence is transformed into a probability distribution over the next token.

At its core, a language model is doing next-token prediction — given some text, predict what word or token is most likely to come next. It's essentially a function with a massive number of parameters, all learned automatically from enormous amounts of training data.

Historically, early approaches used N-gram models — predicting the Nth word based on the preceding N-1 words (Bigram looks at one word back, Trigram looks at two). These models are simple and efficient, but suffer from the data sparsity problem: many word combinations never appear in training data, leading to zero-probability estimates. The deep learning era introduced RNNs (Recurrent Neural Networks), which use a "hidden state" to theoretically capture arbitrarily long context. In practice, however, vanishing gradients limit effective memory to short distances. More critically, RNNs must process input sequentially, making it impossible to fully leverage the parallel computing power of modern GPUs — and that's precisely why Transformer eventually replaced them.

The Transformer pipeline can be broken down into key steps: Tokenization → Input Layer (understanding each token) → Attention (understanding context) → Feed-forward Network (integrating information) → Output Layer (outputting a probability distribution). Attention + Feed-forward together form a Transformer Block, and a full model stacks many such blocks — representing repeated rounds of "thinking" by the model.

Step 1: Tokenization — Splitting Text into Tokens

The fundamental unit a language model actually processes is the token. Both inputs and outputs are in terms of tokens. The very first thing the model does with a sentence is split it into a sequence of tokens.

Take "Introduction of Generative AI" as an example: after tokenization, "Introduction" might be split into "Int" and "roduction", "Generative" into "Gener" and "ative", while "of" and "AI" each become a single token.

To perform this splitting, you need a token vocabulary prepared in advance. Different language models use different vocabularies: many Chinese models treat each Chinese character as a single token, while ChatGPT is designed for multiple languages and may require several tokens to represent a single Chinese character.

The vocabulary can be defined manually or derived automatically using algorithms. The most representative approach is BPE (Byte Pair Encoding) — originally a data compression algorithm, introduced to NLP in 2016. The process works as follows: start by splitting all text into individual characters, count the frequency of all adjacent character pairs, merge the most frequent pair into a new token, and repeat until the vocabulary reaches its target size. For example, if "e" and "s" frequently appear adjacent, they merge into "es"; if "es" and "t" are frequently adjacent, they merge into "est". This bottom-up greedy strategy automatically discovers meaningful morphemes and word roots — "Probabilistic" gets split into two tokens, and the year "1980" into "198" and "0". OpenAI's tiktoken and Google's SentencePiece are both BPE-based. No trainable parameters are involved in this step.

Step 2: Embedding — Turning Tokens into Vectors

After tokenization, the next step is to understand the meaning of each token by representing it as a vector — a process called Embedding.

The idea of embeddings traces back to Bengio et al.'s 2003 neural language model, but the concept gained widespread attention with Google's Word2Vec in 2013. Word2Vec uses self-supervised tasks to automatically map semantically similar words to nearby locations in vector space, giving rise to the famous vector arithmetic: King − Man + Woman ≈ Queen. If tokens remained as symbols, there would be no way to capture that Run and Jump are related, or that Cat and Apple are not. Once tokens are represented as vectors, semantically similar tokens appear closer together in vector space — visualizing these vectors shows all animals clustered together, all verbs clustered together, and so on.

In practice, the Transformer maintains a lookup table mapping tokens to embeddings. The embedding process is essentially a table lookup. Every vector in this table is a model parameter, learned from training data — no manual specification required. The embedding dimension (i.e., vector length) is a key hyperparameter: GPT-3 uses 12,288 dimensions, and larger models generally use higher-dimensional embeddings.

Adding Positional Information

Semantic meaning alone isn't enough — the same token in different positions within a sentence can carry entirely different meanings. This is why Positional Embedding is introduced: each position gets a unique vector, which is added to the corresponding token's embedding.

The implementation of positional embeddings has evolved over time. The original Transformer paper used fixed sinusoidal encoding, with the advantage of generalizing to longer sequences not seen during training. BERT and the GPT series switched to learnable positional embeddings, treating positions as ordinary parameters trained alongside the model. More recently, mainstream open-source models (LLaMA, Qwen, ChatGLM, etc.) have widely adopted RoPE (Rotary Position Embedding), which encodes positional information as rotation angles applied to vectors. This causes attention scores to naturally reflect the relative distance between tokens, yielding better performance on long-text generalization.

Step 3: Attention — Understanding Context

The embeddings produced so far have a critical flaw — they ignore context. Whether "bank" means a financial institution or a riverbank, or whether "apple" refers to a fruit or the tech company, the token embedding is identical. The module that solves this problem is the celebrated Attention mechanism.

After processing through Attention, the model incorporates the full sentence context so that the same token receives a different vector representation depending on its surroundings. These are called Contextualized Embeddings.

A common misconception is worth clearing up: the 2017 landmark paper Attention is All You Need did not invent Attention. Attention had long been used in language models, but was always thought to require RNNs. The paper's real contribution was proving that RNNs could be eliminated entirely — Attention alone was sufficient for excellent language modeling. This was shocking to the research community at the time, because RNNs had long been seen as indispensable.

How Attention Works: The QKV Framework

The core logic of Attention is: take a sequence of vectors as input, output a sequence of vectors of equal length, with each output vector enriched with contextual information.

Understanding the full implementation requires the QKV (Query-Key-Value) framework. Each token's embedding is passed through three different linear transformation matrices to produce three vectors: Query, Key, and Value. This framework uses the analogy of a database query: the Query is what you're looking for, the Key is the index, and the Value is the actual information retrieved. The process has two steps:

Step 1: Compute relevance. Take the current token's Query vector and compute the dot product with every token's Key vector. Divide by the square root of the dimension (to scale and prevent vanishing gradients), then apply Softmax normalization to obtain Attention Weights.

Step 2: Weighted sum. Use the Attention Weights to compute a weighted sum of all tokens' Value vectors, producing the output vector for this position. In formula form: Attention(Q, K, V) = Softmax(QK^T / √d_k) · V. In short: first determine what's relevant, then aggregate that relevant information.

Causal Attention and Multi-Head Attention

With 5 vectors, you'd theoretically compute 5×5 = 25 Attention Weights, forming the Attention Matrix. In practice, language models only compute attention between each token and the tokens to its left (preceding it) — this is called Causal Attention.

Additionally, "relevance" can be understood from multiple dimensions: Dog and Cat are related because they're both animals; Dog and Bark are related because of a subject-action relationship. A single attention module can't capture all these dimensions simultaneously. This is why Multi-Head Attention was introduced — multiple sets of attention modules (typically at least 16), each with different parameters, produce their own set of attention weights and weighted sums, allowing the model to understand token relationships from multiple perspectives simultaneously.

Feed-Forward Network and Stacked Blocks

After Attention produces multiple output vectors, a Feed-Forward Network (FFN) integrates them into a single vector. Research has shown that FFN parameters typically account for more than two-thirds of the entire model, and FFNs are thought to act as "knowledge stores" — factual knowledge such as "Paris is the capital of France" is encoded in the FFN weight matrices, while the Attention layers handle information retrieval and routing. A standard FFN consists of two linear transformations with a nonlinear activation function in between, and the hidden layer dimension is usually 4× the embedding dimension.

The MoE (Mixture of Experts) architecture, which has gained prominence in recent years, revolutionizes the FFN: instead of a single FFN, it uses multiple parallel "expert" FFNs, activating only a small subset of them for each token. This dramatically increases total parameter count (theoretically storing more knowledge) while keeping actual computation constant. Mixtral, DeepSeek, and reportedly GPT-4 all adopt this approach.

Attention + Feed-Forward together form one Transformer Block (commonly referred to as a Layer in the literature). The input embeddings pass through multiple blocks in sequence, being progressively refined — each layer producing representations with richer contextual integration. The vectors at the end of the final block are passed to the Output Layer (containing a Linear Transform and Softmax), which produces a probability distribution over the next token.

Why Only Look Left? Why Is Long Context Hard?

The reason language models only compute left-side attention comes from their next-token prediction logic: the model generates W1, appends W1 to the input to generate W2, and so on until an end token is reached. Because the input is the same each time, previously computed attention results don't need to be recomputed — they can be reused directly. In engineering terms, this corresponds to the critical optimization of KV Cache: when generating each new token, the model only needs to compute the new token's Query; the Key and Value vectors for all previous tokens are already cached in GPU memory and read directly, dramatically reducing computation.

Would computing full attention over both sides yield better results? Experiments have given a clear answer: some models did compute full attention, but the results were no better — and ChatGPT, using only left-side attention, performs remarkably well. The current consensus is that left-side attention is sufficient.

This also reveals the core bottleneck of long-context processing. Attention requires computing pairwise relationships between all tokens, so the number of computations scales with the square of the text length. With 100K tokens, you need 100K² attention computations. Computational cost grows quadratically with length, and KV Cache memory usage also grows linearly with sequence length — for a large model handling 128K context, KV Cache alone can reach tens of gigabytes. To address this, techniques like GQA (Grouped Query Attention) and MQA (Multi-Query Attention) have been developed, allowing multiple queries to share a single set of KV pairs, dramatically reducing memory usage. This is the most challenging technical obstacle behind the race to extend context window lengths.

Frontiers: Beyond Transformer

Optimizing Attention remains an active research area. Key directions include:

• Accelerating Attention computation: various engineering and algorithmic optimizations targeting quadratic complexity;
• Infinite-length Attention: multiple research efforts exploring entirely new approaches to break the length barrier;
• Train Short, Test Long: training on short sequences while achieving effective long-context inference;
• Alternative architectures: Transformer may not be the final answer. The Mamba series, based on State Space Models (SSMs), introduces a key innovation: a "selective mechanism" that allows model parameters to dynamically adjust based on the current input, deciding what information is worth retaining and what can be forgotten. This gives the model content-aware capabilities similar to Attention, but with inference complexity that is linear in sequence length (O(n)) rather than Transformer's O(n²) — a significant advantage on very long sequences, representing something of a "RNN renaissance." Jamba combines Mamba and Transformer to capture the strengths of both.

---

Once you understand how a Transformer works — how a sentence is processed step by step through Tokenization, Embedding, Attention, and Feed-Forward layers to ultimately yield a probability distribution over the next token — you've truly opened the door to understanding what's happening inside a language model.