The Geometry of BERT

BERT is a stack of Transformer encoder blocks. The architecture is well-known, but the reasoning behind each component — the projections, the scaling, the softmax, the multi-head split — is often presented as received wisdom rather than derived from first principles. This article walks through the encoder’s math, one piece at a time, with concrete numbers.

The full path from a sentence to BERT’s output goes through four stages:

1. Tokenization and embedding — The sentence is split into tokens using a fixed 30,000-token vocabulary. Special tokens are added: [CLS] at the start (whose output will be used for classification) and [SEP] between sentences (used when BERT processes sentence pairs). Each token is mapped to a learned 768-dimensional embedding vector.
2. Positional encoding — A learned position vector (position 1, position 2, …, up to 512) is added to each token embedding so the model knows where each token sits in the sequence.
3. Transformer blocks — The sequence of vectors passes through 12 identical blocks. Each block does two things in sequence: multi-head self-attention (tokens exchange information) and a feed-forward network (each token processed independently). Every step is wrapped in a residual connection and layer normalization.
4. Classification head — The output vector for the [CLS] token is extracted and passed through a single linear layer to produce class probabilities.

Only the attention mechanism is non-trivial; everything else is a standard deep learning building block. This article explains the attention mechanism from first principles, then summarizes the rest.

For reference, BERT-base uses these dimensions throughout:

Words as Points in a Vector Space

The starting assumption is that words live in a linear space. Each word is mapped to a vector $\mathbf{x} \in \mathbb{R}^d$ where $d$ is the embedding dimension (768 in BERT-base). Two words with similar meanings have vectors that are close under dot product:

This is the standard model. “King” minus “man” plus “woman” lands near “queen” because the vector differences encode semantic relationships as directions. The linear structure buys us the ability to add and compare meanings, which the Transformer relies on.

The problem is that a word’s meaning depends on its neighbors. “Bank” near “river” and “bank” near “money” should occupy different positions in the space, but they start at the same embedding. The role of the Transformer encoder is to move each token’s vector to a position that reflects its context.

Why Position Matters

Before the encoder can use context, there is a more basic problem. The attention mechanism has no sense of order. If you swap two words, the dot products between their vectors remain the same — attention sees a bag of words. “The dog bit the man” and “The man bit the dog” contain the same words, so the attention scores are identical — the model cannot tell who bit whom.

The fix is to add a position-dependent signal to each embedding before the encoder sees it. BERT uses learned position embeddings: a separate vector for each position (position 1, position 2, …, up to 512), added to the token embedding. After this addition, the same word “bank” at position 3 and position 7 has a different input vector because the position component differs. The attention mechanism can then learn to use position as part of its relevance metric — and since the positions are learned from data, the model decides what positional relationships matter.

Attention as a Weighted Average by Relevance

The core operation is: for each token, compute a weighted average of all tokens’ vectors, where the weight of token $j$ relative to token $i$ is the relevance of $j$ to $i$. If “river” is highly relevant to “bank” in a given sentence, its vector should pull “bank” toward it.

The relevance score needs to come from the vectors themselves — there is no external source of information about which words relate. The natural choice in a vector space is the dot product:

But using the raw embeddings has a problem: a token can only measure similarity to itself. Token $i$ asking about $\mathbf{x}_j$ and token $j$ asking about $\mathbf{x}_i$ produce the same result — the relationship is symmetric. Language needs asymmetric relevance: “river” is highly informative about “bank” in a riverside context, but “bank” is much less informative about “river.” A token’s role as asker and its role as answerer need different representations.

The fix is to project each token into three separate spaces before the interaction:

The Query $\mathbf{q}_i$ represents “what token $i$ is looking for,” the Key $\mathbf{k}_j$ represents “what token $j$ contains,” and the Value $\mathbf{v}_j$ represents “what token $j$ contributes when attended to.” The score becomes:

Because $W_Q$ and $W_K$ are different learned matrices, token $i$ can define relevance differently from token $j$. The three-projection design is the minimal change that buys asymmetric relevance while keeping Value separate from the scoring mechanism — Value controls the information passed, not the weight itself. The projections $W_Q$, $W_K$, and $W_V$ are all learned, so the model discovers what kinds of relevance matter for the task.

The raw scores $s_{ij}$ are unbounded. To use them as weights, they need to be converted to a probability distribution over the sequence — that is what softmax does:

The softmax is not an arbitrary choice. It emerges from the principle of maximum entropy: given a set of expected scores, the least-committal probability distribution that satisfies them is the Gibbs distribution $\alpha_i \propto \exp(s_i / \tau)$, where $\tau$ is the temperature. Softmax is the special case at $\tau = 1$ (and we will see shortly that $\sqrt{d_k}$ plays the role of temperature). The derivation is a short Lagrange multiplier exercise: maximize $-\sum_i p_i \log p_i$ subject to $\sum_i p_i s_i = \text{constant}$ and $\sum_i p_i = 1$, and the optimal $p_i$ is proportional to $\exp(s_i / \tau)$.

The output for token $i$ is then the weighted average of the Value vectors $\mathbf{v}_j = W_V \mathbf{x}_j$:

Value is the information that token $j$ contributes when attended to, shaped by its own learned projection $W_V$.

Writing the whole operation in matrix form for a sequence of $n$ tokens:

The scaling factor $\sqrt{d_k}$ is the subject of the next section.

Why the Scaling Factor?

Suppose the components of $\mathbf{q}$ and $\mathbf{k}$ are independent random variables with mean 0 and variance 1 (as they will be after layer normalization and random initialization). Their dot product is:

Each term has mean 0 and variance 1, so the sum has mean 0 and variance $d_k$. The standard deviation is $\sqrt{d_k}$. This holds exactly at initialization; after training, $\mathbf{q}$ and $\mathbf{k}$ develop correlations, but the scaling remains necessary because the argument captures the core problem — the variance grows with $d_k$ regardless.

At $d_k = 64$ (BERT-base), this means raw dot products typically range from roughly $-3 \times 8 = -24$ to $+24$. When these scores enter the softmax, the exponentiation amplifies the largest one dramatically. For a sequence of moderate length, the attention distribution becomes nearly one-hot — all weight goes to one token, and every other gradient vanishes.

Consider three tokens with raw scores $[5, 20, 5]$. Without scaling:

The second token gets all the attention. The first and third contribute nothing to the output and receive no gradient signal.

Now divide by $\sqrt{d_k}$ (effectively $\sqrt{20} \approx 4.47$, giving scores $[1.12, 4.47, 1.12]$):

Every token now contributes to both the output and the gradient. The scaling factor $\sqrt{d_k}$ is a temperature parameter: it controls how sharply the attention distribution focuses. Large values of $\sqrt{d_k}$ produce diffuse attention; small values produce focused attention. The choice $\sqrt{d_k}$ is exactly the standard deviation of the raw scores, restoring unit variance and preventing premature collapse.

Q/K/V as a Learned Metric

The score $\mathbf{q}_i \cdot \mathbf{k}_j$ can be rewritten as:

where $M = W_Q^T W_K$ is a $d \times d$ matrix. This is a bilinear form — a learned quadratic measure of similarity on the embedding space. The model does not use a fixed metric like cosine distance; it learns the metric that best captures the notion of “relevance” for the task. The same structure will repeat across multiple heads, each learning a different metric in parallel.

A Worked Example: 3 Tokens, 4 Dimensions

We can trace every number through a tiny encoder to see exactly what happens. Consider three tokens with 4-dimensional embeddings:

Each row is a token. The embeddings have been chosen to make the patterns visible — tokens 1 and 2 share some structure, token 3 is distinct.

For simplicity, let the projection matrices be the identity ($W_Q = W_K = W_V = I$), so $\mathbf{q}_i = \mathbf{k}_i = \mathbf{v}_i = \mathbf{x}_i$. This removes the learned metric and lets us see the geometry of the embeddings directly.

Step 1 — compute the raw similarity scores $S = X X^T$:

Token 1 has similarity 2 with itself, 0 with token 2, and 1 with token 3. Token 2 has similarity 0 with token 1 — they are orthogonal in this 4-dimensional space.

Step 2 — scale by $\sqrt{d_k} = \sqrt{4} = 2$:

Step 3 — apply softmax row-wise. The first row $[1, 0, 0.5]$ exponentiates to $[2.72, 1, 1.65]$, summing to 5.37:

Second row $[0, 1, 0.5]$ gives the same result:

Third row $[0.5, 0.5, 1]$ exponentiates to $[1.65, 1.65, 2.72]$, sum 6.02:

Step 4 — compute the output $Z = A X$:

Each output token has moved in the embedding space. Token 1, originally $(1,0,1,0)$, is now at $(0.82, 0.50, 0.51, 0.19)$ — pulled slightly toward tokens 2 and 3 by the attention weights. Token 3, originally $(1,1,0,0)$, is now at $(0.72, 0.72, 0.27, 0.27)$ — its own weight is only 0.45, so it moves more than the others.

This is contextualization in miniature. Before the encoder, every occurrence of a given word has the same embedding. After the encoder, each occurrence has been pulled toward the vectors of the tokens it attends to. “Bank” near “river” ends up in a different region of the embedding space than “bank” near “money” — not because the model knows what “river” or “money” means semantically, but because the attention weights redistribute each token’s position toward its specific context.

Now compare with the unscaled attention. Without dividing by $\sqrt{d_k}$, row 1 of the score matrix is $[2, 0, 1]$, which exponentiates to $[7.39, 1, 2.72]$, sum 11.11:

The distribution is more concentrated on token 1 itself — the gradient signal to other tokens is weaker. With scaling, the distribution is $[0.51, 0.19, 0.31]$ — significantly more spread, giving each token meaningful gradient contribution. This is the difference between a model that learns from all tokens and one that mostly ignores context.

Why Multi-Head?

A single attention distribution captures one notion of relevance. But relevance is multi-faceted: a word might be related to another syntactically (subject-verb), semantically (synonymy), or positionally (nearby words in the sentence).

Multi-head attention runs the same mechanism $h$ times in parallel, each with its own learned projections $W_Q^h, W_K^h, W_V^h$ and therefore its own bilinear form $M_h$. The outputs of all heads are concatenated and projected back to the original dimension:

Each head learns a different metric on the embedding space. One head might learn high weights for adjacent words (capturing local syntax), another for semantically similar words across long distances (capturing coreference). The concatenation ensures the model can simultaneously move each token along multiple relevance dimensions.

In BERT-base, $h = 12$ heads with $d_k = 64$ per head, giving a total attention dimension of $12 \times 64 = 768$, matching the embedding dimension. Each head operates on a 64-dimensional slice of the embedding space — a different learned geometry.

With the encoder architecture in place, the question is how to train it. BERT learns representations from unlabeled text by masking tokens and predicting them from context.

Cross-Entropy and What BERT Learns

BERT is trained on the Masked Language Model objective: randomly mask 15% of tokens and predict them from context. The 15% rate is a heuristic balance — too few masks starves the training signal per sequence, too many makes the sentence unintelligible and the model learns to ignore context rather than use it. The loss is cross-entropy between the predicted distribution over the vocabulary and the true token:

Cross-entropy can be decomposed:

where $H(p)$ is the entropy of the true distribution (always 0 for a single correct token) and $KL(p || q)$ is the Kullback-Leibler divergence from the prediction to the truth. Minimizing cross-entropy is equivalent to minimizing the KL divergence — bringing the model’s predicted distribution closer to the correct one.

For a concrete example, suppose the true token is “river” and the model predicts $[0.2, 0.5, 0.2, 0.1]$ over a 4-token vocabulary for simplicity. The cross-entropy is $-\log(0.5) = 0.69$. After training, the model predicts $[0.02, 0.94, 0.02, 0.02]$ and the loss drops to $-\log(0.94) = 0.06$.

Cross-entropy is used rather than mean squared error because the output is a probability distribution over the vocabulary, and cross-entropy is the natural distance between distributions. MSE would penalize all incorrect tokens equally — the model would waste gradient on pushing down the 29,999 wrong tokens instead of focusing on raising the correct one. Cross-entropy concentrates the loss on the correct token’s probability: a 50% confidence costs 0.69, while 10% costs 2.30, so the gradient drives the model to be increasingly confident about the right answer. Combined with softmax, the gradient takes the clean form $\partial \mathcal{L} / \partial z_i = p_i - y_i$ — the difference between prediction and truth — producing a useful signal from every example and never saturating.

The scale of the task matters. The Transformer’s vocabulary is 30,000 tokens, so each prediction requires narrowing down from $\log_2(30000) \approx 15$ bits of uncertainty. With 15% masking on a 512-token sequence, BERT gets roughly $512 \times 0.15 \approx 77$ predictions per sequence. That is $77 \times 15 = 1155$ bits of training signal per sequence — which is why it works, but also why it is sample-inefficient compared to alternatives like ELECTRA (which gets a prediction signal from every token).

The masking objective forces the model to use both sides of context simultaneously. Consider the sentence “The man went to the [MASK] to deposit money.” A model that only looked left of the mask would see “The man went to the” — not enough to distinguish “bank” from “store” from “window.” One that only looked right would see “to deposit money” — which narrows it down considerably. The combination of both directions is what makes the representations rich. This bidirectionality is the defining difference between BERT and autoregressive models like GPT, which only see left context and must predict the next token without knowing what comes after.

The training objective alone does not make the Transformer work — it depends on the components that stabilize and transform representations within each block. Two remain to be explained.

The Other Components: Layer Normalization and the Feed-Forward Network

The full path through one Transformer block involves two operations wrapped in residual connections, with layer normalization applied after each addition.

Layer normalization takes a single token’s vector and rescales it to have mean 0 and variance 1, then applies a learned shift and scale:

where $\mu$ and $\sigma$ are the mean and standard deviation of the components of $\mathbf{x}$, and $\gamma, \beta$ are learned vectors of the same dimension. It is applied per token independently, not across the batch (unlike batch normalization). The purpose is to keep the distributions stable across layers, preventing the values from growing or shrinking uncontrollably as they pass through the stack.

The feed-forward network is two linear transformations with a non-linearity in between:

GELU (the Gaussian Error Linear Unit, $\text{GELU}(x) = x \cdot \Phi(x)$ where $\Phi$ is the standard normal CDF) is a smooth version of ReLU: it heavily suppresses negative values while preserving positive ones, with a soft transition between the two regimes. The non-linearity prevents the entire stack from collapsing into a single linear transformation.

In BERT-base, $W_1$ expands the dimension from 768 to 3072 (a factor of 4), and $W_2$ projects it back to 768. The FFN has no interaction between tokens — each token’s vector is transformed independently. This is the complement to attention: attention mixes information between tokens, and the FFN processes what attention has gathered within each token’s representation.

Both operations are wrapped in residual connections and normalized:

The residual connection $X + F(X)$ has a simple but critical effect on gradients. The derivative with respect to $X$ is:

The identity term $I$ ensures that gradients can flow backward through the network even if $\partial F / \partial X$ is small. In a 12-layer stack, without residual connections, early layers would receive vanishingly small gradients and would not learn. The residual bypass is what makes deep Transformer training tractable.

The Architecture in One Picture

The Transformer encoder is a stack of blocks. For each block:

1. Project the input sequence $X \in \mathbb{R}^{n \times d}$ into Query, Key, and Value spaces via learned linear transformations, one set per head.
2. Compute attention scores as a learned bilinear form $X (W_Q^h)^T W_K^h X^T$, one per head.
3. Scale by $\sqrt{d_k}$ to control the softmax temperature and prevent gradient collapse.
4. Apply softmax to produce the attention distribution, then take the weighted sum of Values.
5. Concatenate all heads and project back to dimension $d$.
6. Add the input (residual connection) and apply layer normalization.
7. Apply the feed-forward network: expand from $d$ to $4d$, apply GELU, project back to $d$.
8. Add the input from step 6 (residual connection) and apply layer normalization.
9. After the last block, extract the [CLS] token’s vector and pass it through a linear layer to produce class probabilities.

For step 9, the [CLS] token was prepended to the input as a placeholder — it has no meaningful token embedding of its own. After 12 blocks of contextualization, its vector has absorbed information from the entire sequence. A single linear layer maps it to class scores:

where $W \in \mathbb{R}^{\text{num\_classes} \times 768}$ and $\mathbf{b} \in \mathbb{R}^{\text{num\_classes}}$ are learned. Each component of the logit vector is a dot product between the [CLS] vector and a learned weight vector for that class, plus a bias. The softmax converts the logits to a probability distribution over the classes.

For a concrete example, suppose the [CLS] vector is $(0.5, -0.2, 0.1, 0.8)$ in a 4-dimensional embedding, and we have 3 classes with weights:

The logits are $W \mathbf{x}_{[\text{CLS}]} + \mathbf{b}$:

The components, worked out one at a time:

Applying softmax:

For a 2-class sentiment task, this would be the model’s confidence that the input is positive vs negative. The entire path from sentence to probability is now traced: token embeddings → attention → FFN → [CLS] vector → linear layer → softmax.

Each component solves a specific problem: projections enable asymmetric relevance, scaling enables stable gradients, multi-head enables multiple relevance types, residuals enable depth, layer normalization keeps distributions stable, and the feed-forward network processes what attention has gathered. Every number in the design follows from a constraint, not from convention.