The Encoder-Only Transformer

The encoder-only transformer takes a sequence of tokens, converts them to embeddings, adds positional information, then passes them through a stack of identical layers. Each layer has two sub-components: multi-head self-attention and a position-wise feed-forward network, with layer normalization and residual connections wrapping each.

Input Representation

Each input token $x_i$ is mapped to a $d_{\text{model}}$-dimensional embedding. Positional encodings are added to inject sequence order, since self-attention is inherently permutation-invariant:

Explanation: Each position gets a unique fingerprint of interleaved sine and cosine waves at geometrically spaced frequencies, so the model can distinguish any two positions and even generalize to unseen sequence lengths. Here $pos$ is the token position (0-indexed), $i$ is the dimension index, and $d_{\text{model}}$ is the total embedding dimension (e.g. 768 in BERT-base).

Scaled Dot-Product Attention

The core computation packs the sequence into three matrices — queries $Q$, keys $K$, and values $V$ — each produced by learned linear projections of the input:

Explanation: Each token's query vector scores its compatibility with every key in the sequence, the scores are normalized into a probability distribution, and then that distribution is used to compute a weighted sum of value vectors.

• $Q \in \mathbb{R}^{n \times d_k}$: query matrix, one row per token
• $K \in \mathbb{R}^{n \times d_k}$: key matrix, one row per token
• $V \in \mathbb{R}^{n \times d_v}$: value matrix, one row per token
• $d_k$: key/query dimension — the $\sqrt{d_k}$ scaling prevents dot products from growing so large that softmax gradients vanish

Multi-Head Attention

Rather than running one attention function, the model runs $h$ attention heads in parallel, each with its own projections, then concatenates and re-projects:

Explanation: Each head learns to attend to a different type of relationship (syntax, coreference, semantics, etc.) and the output projection mixes them back into a single $d_{\text{model}}$-dimensional space. where $\text{head}_i = \text{Attention}(QW_i^Q,\, KW_i^K,\, VW_i^V)$ with $W_i^Q, W_i^K \in \mathbb{R}^{d_{\text{model}} \times d_k}$, $W_i^V \in \mathbb{R}^{d_{\text{model}} \times d_v}$, and $W^O \in \mathbb{R}^{hd_v \times d_{\text{model}}}$. Typically $d_k = d_v = d_{\text{model}}/h$.

Position-Wise Feed-Forward Network

After attention, each token representation is passed independently through a two-layer MLP with a GELU or ReLU nonlinearity:

Explanation: This per-token MLP (with a large hidden dimension, typically $4 \times d_{\text{model}}$) acts as the model's long-term memory and feature transformer, applying the same transformation to every position independently. $W_1 \in \mathbb{R}^{d_{\text{model}} \times d_{ff}}$ expands the representation; $W_2 \in \mathbb{R}^{d_{ff} \times d_{\text{model}}}$ projects it back. BERT-base uses $d_{ff} = 3072$ with $d_{\text{model}} = 768$.

Full Layer with Residuals

Each encoder layer wraps both sub-components with Pre-LN or Post-LN residual connections:

Explanation: The residual stream $x$ is never destroyed — attention and FFN layers only add corrections to it, which is why gradient flow remains healthy across 12–24 layers and why probing the residual stream at intermediate layers reveals progressively richer linguistic structure.

Masked Language Modeling Objective

Encoder-only models are typically pre-trained with Masked Language Modeling (MLM): 15% of input tokens are replaced by `[MASK]`, and the model must predict the original token from bidirectional context. Unlike causal language modeling, the loss is computed only over masked positions: ```python import torch import torch.nn as nn from transformers import BertTokenizer, BertForMaskedLM tokenizer = BertTokenizer.from_pretrained("bert-base-uncased") model = BertForMaskedLM.from_pretrained("bert-base-uncased")

# Mask a token manually text = "The [MASK] sat on the mat." inputs = tokenizer(text, return_tensors="pt") with torch.no_grad(): outputs = model(**inputs) logits = outputs.logits # shape: (1, seq_len, vocab_size)

# Get prediction at the [MASK] position mask_idx = (inputs.input_ids == tokenizer.mask_token_id).nonzero(as_tuple=True)[1] predicted_token_id = logits[0, mask_idx].argmax(-1) print(tokenizer.decode(predicted_token_id)) # → "cat" ```

A Minimal Encoder-Only Block in PyTorch

```python import torch import torch.nn as nn import torch.nn.functional as F import math class EncoderBlock(nn.Module): def __init__(self, d_model=768, n_heads=12, d_ff=3072, dropout=0.1): super().__init__() assert d_model % n_heads == 0 self.n_heads = n_heads self.d_k = d_model // n_heads

# Attention projections (Q, K, V packed together for efficiency) self.qkv_proj = nn.Linear(d_model, 3 * d_model) self.out_proj = nn.Linear(d_model, d_model)

# Feed-forward self.ff1 = nn.Linear(d_model, d_ff) self.ff2 = nn.Linear(d_ff, d_model) self.norm1 = nn.LayerNorm(d_model) self.norm2 = nn.LayerNorm(d_model) self.drop = nn.Dropout(dropout) def forward(self, x, attention_mask=None): B, T, C = x.shape

# --- Multi-head self-attention --- qkv = self.qkv_proj(x).reshape(B, T, 3, self.n_heads, self.d_k) qkv = qkv.permute(2, 0, 3, 1, 4) # (3, B, heads, T, d_k) q, k, v = qkv[0], qkv[1], qkv[2] scale = math.sqrt(self.d_k) attn = (q @ k.transpose(-2, -1)) / scale # (B, heads, T, T) if attention_mask is not None:

# attention_mask: (B, T), 1=keep, 0=pad mask = attention_mask[:, None, None, :] # broadcast attn = attn.masked_fill(mask == 0, float('-inf')) attn = F.softmax(attn, dim=-1) attn = self.drop(attn) out = (attn @ v).transpose(1, 2).reshape(B, T, C) out = self.out_proj(out)

# Residual + LayerNorm x = self.norm1(x + self.drop(out))

# --- Feed-forward --- ff = self.ff2(self.drop(F.gelu(self.ff1(x)))) x = self.norm2(x + self.drop(ff)) return x ```