Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

BERT Architecture › Transformer Encoder Layers

Michael BrenndoerferUpdated July 16, 202512 min read

Language AI HandbookAll chapters

Part 24 of 67BERT and Variants

Chapter 169BERT Architecture › Transformer Encoder Layers

Reading section 4 of 12 · Transformer Encoder Layers · 12 minSection 4 of 12 · 12 min

Section 4 of 12 · 12 min read

BERT Architecture › Transformer Encoder LayersOne section · Sections in this chapter4 / 12 ↓

Sections in this chapter12

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

The core of BERT consists of stacked transformer encoder blocks. Each block applies multi-head self-attention followed by a position-wise feed-forward network, with residual connections and layer normalization around each sub-layer. This section covers both components in depth, including the subtle design choices that affect training stability and final performance.

Understanding each layer requires appreciating what information flows through it. The input to each layer is a matrix of shape (sequence length, 768), where each row is a 768-dimensional vector representing one token's current representation. The layer transforms this matrix to produce a new matrix of the same shape: same number of tokens, same dimensionality, but with updated representations that incorporate information from across the sequence via attention. The key design constraint is that this shape-preserving property must hold at every layer so that layers can be stacked arbitrarily.

Layer StructureLink Copied

Each transformer layer follows a consistent pattern:

Multi-head self-attention with residual connection
Layer normalization
Feed-forward network with residual connection
Layer normalization

BERT uses "post-norm" architecture, where layer normalization follows each sub-layer rather than preceding it. This differs from the "pre-norm" variant used in some later models like GPT-2.

The post-norm versus pre-norm distinction matters for training dynamics. In post-norm (used by BERT), the residual path and the transformed path are summed first, then normalized. This means the scale of the residual and the transformed signal can drift relative to each other during training, which sometimes leads to training instability in very deep networks. Pre-norm places the layer normalization before the transformation. This ensures each sub-layer sees normalized inputs, which tends to stabilize gradient flow in deep networks. The reason BERT used post-norm was simply that it followed the original Vaswani et al. transformer design without modification. Later work showed that pre-norm allows training of much deeper models without careful learning rate warmup, but for 12 or 24 layers, post-norm works well with appropriate hyperparameter tuning.

In[12]:

Code

class BertSelfAttention(nn.Module):
    """Multi-head self-attention for BERT."""

    def __init__(
        self, hidden_size: int = 768, num_heads: int = 12, dropout: float = 0.1
    ):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = hidden_size // num_heads

        # Q, K, V projections
        self.query = nn.Linear(hidden_size, hidden_size)
        self.key = nn.Linear(hidden_size, hidden_size)
        self.value = nn.Linear(hidden_size, hidden_size)

        # Output projection
        self.output = nn.Linear(hidden_size, hidden_size)
        self.dropout = nn.Dropout(dropout)

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        batch_size, seq_len, _ = hidden_states.shape

        # Project to Q, K, V
        q = self.query(hidden_states)
        k = self.key(hidden_states)
        v = self.value(hidden_states)

        # Reshape for multi-head attention: (batch, seq, heads, head_dim)
        q = q.view(
            batch_size, seq_len, self.num_heads, self.head_dim
        ).transpose(1, 2)
        k = k.view(
            batch_size, seq_len, self.num_heads, self.head_dim
        ).transpose(1, 2)
        v = v.view(
            batch_size, seq_len, self.num_heads, self.head_dim
        ).transpose(1, 2)

        # Compute attention scores
        scores = torch.matmul(q, k.transpose(-2, -1)) / (self.head_dim**0.5)

        # Apply attention mask if provided
        if attention_mask is not None:
            scores = scores + attention_mask

        # Softmax and dropout
        attention_probs = F.softmax(scores, dim=-1)
        attention_probs = self.dropout(attention_probs)

        # Apply attention to values
        context = torch.matmul(attention_probs, v)

        # Reshape back: (batch, seq, hidden)
        context = (
            context.transpose(1, 2).contiguous().view(batch_size, seq_len, -1)
        )

        # Output projection
        output = self.output(context)

        return output, attention_probs

The attention computation follows the standard scaled dot-product formula. Given the QKV matrices, attention computes a weighted combination of values where the weights depend on query-key similarity. Before looking at the formula, consider the intuition: for each token position, we want to produce an updated representation that incorporates information from other relevant positions. Relevance is determined by how similar the current token's "query" vector is to each other token's "key" vector. The "value" vector at each position contains the actual content to be aggregated. Think of the attention mechanism as a soft dictionary lookup: the query looks up relevant entries in the dictionary (keys), and the returned values are a weighted combination of the associated values.

\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V

where:

$Q$ : the query matrix with shape (sequence length, head dimension), representing what each position is "looking for"
$K$ : the key matrix with shape (sequence length, head dimension), representing what each position "offers" for matching
$V$ : the value matrix with shape (sequence length, head dimension), containing the information to aggregate
$K^T$ : the transpose of $K$ , enabling the matrix multiplication $QK^T$ that produces similarity scores
$d_k$ : the head dimension (64 in BERT), used for scaling
$\sqrt{d_k}$ : the scaling factor that prevents dot products from growing too large
$\text{softmax}(\cdot)$ : normalizes scores to a probability distribution over positions

The scaling by $\sqrt{d_k}$ keeps training stable. Without it, the dot products $QK^T$ grow in magnitude with the dimension, pushing the softmax into regions where gradients vanish. With 64-dimensional heads, dot products could easily reach values of 8-10, making the softmax output nearly one-hot and preventing gradient flow.

Out[13]:

Visualization

Histogram of unscaled attention scores showing wide spread from -15 to 15. — Without scaling, attention scores have high variance, causing softmax to produce near-one-hot distributions.

Histogram of scaled attention scores showing concentrated distribution from -2 to 2. — With scaling by the square root of dimension, scores have controlled variance, producing softer attention distributions.

The histograms above demonstrate this effect. Raw dot products of 64-dimensional vectors have variance around 64, producing scores that span a wide range. After dividing by $\sqrt{64} = 8$ , the variance drops to approximately 1, keeping scores in a range where softmax produces meaningful probability distributions rather than near-deterministic outputs.

Feed-Forward NetworkLink Copied

Each layer's feed-forward network expands the representation to 4x the hidden dimension, applies a non-linearity, then projects back. The feed-forward network operates independently on each token position. Unlike the attention mechanism, which mixes information across positions, the feed-forward network applies the same transformation to each position separately. This makes it position-wise in the precise sense: the token at position 5 gets exactly the same function applied to it as the token at position 23, with no interaction between them.

The reason for the 4x expansion is not fully understood theoretically, but empirically it works well and the pattern has been reproduced across nearly every transformer architecture since. The expansion from 768 to 3,072 dimensions (or 1,024 to 4,096 in BERT-Large) creates a larger intermediate space where the model can express complex non-linear transformations before projecting back down. Some interpretability research suggests that feed-forward networks function as key-value memories: the first linear layer's weights act as keys that activate for specific input patterns, and the second linear layer's weights act as values that add specific information to the representation when those patterns appear. Under this view, the feed-forward network is where BERT stores factual knowledge it acquired during pretraining.

In[14]:

Code

class BertFeedForward(nn.Module):
    """Position-wise feed-forward network for BERT."""

    def __init__(
        self,
        hidden_size: int = 768,
        intermediate_size: int = 3072,
        dropout: float = 0.1,
    ):
        super().__init__()
        self.dense1 = nn.Linear(hidden_size, intermediate_size)
        self.dense2 = nn.Linear(intermediate_size, hidden_size)
        self.dropout = nn.Dropout(dropout)

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        hidden = self.dense1(hidden_states)
        hidden = F.gelu(hidden)  # BERT uses GELU activation
        hidden = self.dense2(hidden)
        hidden = self.dropout(hidden)
        return hidden

BERT uses GELU (Gaussian Error Linear Unit) activation rather than ReLU. GELU provides a smooth approximation to the gating mechanism, allowing small negative values to pass through while still providing non-linearity. The function can be understood as multiplying each input by the probability that a standard normal random variable would be less than that input:

\text{GELU}(x) = x \cdot \Phi(x) = x \cdot \frac{1}{2}\left[1 + \text{erf}\left(\frac{x}{\sqrt{2}}\right)\right]

where:

$x$ : the input value to the activation function
$\Phi(x)$ : the cumulative distribution function (CDF) of the standard normal distribution, giving the probability that a standard normal random variable is less than $x$
$\text{erf}(\cdot)$ : the error function, a mathematical function related to the normal distribution's CDF
$\sqrt{2}$ : a scaling constant that converts from the standard error function to the normal CDF

Unlike ReLU, which abruptly zeroes out all negative inputs, GELU provides a smooth transition. For large positive $x$ , $\Phi(x) \approx 1$ , so GELU $(x) \approx x$ . For large negative $x$ , $\Phi(x) \approx 0$ , so GELU $(x) \approx 0$ . The smooth transition around zero means small negative values can still contribute, which empirically improves training dynamics in transformer models.

Notice that GELU is not monotonically increasing for all inputs. Around $x = -0.17$ , GELU reaches a small local minimum before rising. This means the activation function provides a slight regularizing effect: values near zero are "dampened" more than ReLU would dampen them, while large negative values are suppressed similarly. The practical consequence is that GELU-based models tend to converge to better solutions on NLP tasks than ReLU-based models, though the theoretical explanation for this remains an active research area. GPT-2, GPT-3, and many other major language models have adopted GELU or its close relatives following BERT's example.

Out[15]:

Visualization

Line plot comparing GELU and ReLU activation functions from -4 to 4. — GELU provides a smooth transition compared to ReLU's hard cutoff at zero. The shaded region shows where GELU allows small negative values to pass through, potentially preserving useful gradient information.

Complete Encoder LayerLink Copied

Combining attention and feed-forward with residual connections and layer normalization:

In[16]:

Code

class BertLayer(nn.Module):
    """Single BERT encoder layer."""

    def __init__(
        self,
        hidden_size: int = 768,
        num_heads: int = 12,
        intermediate_size: int = 3072,
        dropout: float = 0.1,
    ):
        super().__init__()
        self.attention = BertSelfAttention(hidden_size, num_heads, dropout)
        self.feed_forward = BertFeedForward(
            hidden_size, intermediate_size, dropout
        )
        self.attention_norm = nn.LayerNorm(hidden_size, eps=1e-12)
        self.output_norm = nn.LayerNorm(hidden_size, eps=1e-12)
        self.dropout = nn.Dropout(dropout)

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        # Self-attention with residual
        attention_output, attention_probs = self.attention(
            hidden_states, attention_mask
        )
        attention_output = self.dropout(attention_output)
        hidden_states = self.attention_norm(hidden_states + attention_output)

        # Feed-forward with residual
        ff_output = self.feed_forward(hidden_states)
        hidden_states = self.output_norm(hidden_states + ff_output)

        return hidden_states, attention_probs

Out[17]:

Console

Input shape: torch.Size([2, 10, 768])
Output shape: torch.Size([2, 10, 768])
Attention probabilities shape: torch.Size([2, 12, 10, 10])

The layer maintains the same tensor shape between input and output, which is essential for stacking layers and enabling residual connections. The attention probabilities tensor indexes the batch, attention head, query position, and key position, confirming that each of the 12 heads computes its own attention pattern over the sequence.

Residual connections enable training of deep networks. They allow gradients to flow directly backward through the network, mitigating the vanishing gradient problem. Without residuals, training a 12 or 24-layer network would be extremely difficult. To understand why, consider what happens to a gradient as it backpropagates through a long chain of transformations. Each transformation multiplies the gradient by a Jacobian matrix. If these Jacobian matrices consistently have values less than 1, the gradient shrinks exponentially with depth, becoming too small to update early layers effectively. Residual connections create a "shortcut" in the gradient computation: the gradient of the loss with respect to early layers includes a direct path through the identity function (the residual) that does not pass through the potentially problematic Jacobian matrices of the transformer sub-layers. This is why you can train a 24-layer BERT-Large successfully while a 24-layer network without residuals would require extreme care with initialization and learning rates.

Worked Example: Residual Connection Gradient Flow

Consider a single residual block: $\mathbf{h}' = \text{LayerNorm}(\mathbf{h} + f(\mathbf{h}))$ , where $f$ is the attention or feed-forward sub-layer. The gradient of the loss $\mathcal{L}$ with respect to $\mathbf{h}$ involves:

\frac{\partial \mathcal{L}}{\partial \mathbf{h}} = \frac{\partial \mathcal{L}}{\partial \mathbf{h}'} \cdot \frac{\partial \text{LayerNorm}}{\partial (\mathbf{h} + f(\mathbf{h}))} \cdot \left(I + \frac{\partial f(\mathbf{h})}{\partial \mathbf{h}}\right)

The term $I + \frac{\partial f(\mathbf{h})}{\partial \mathbf{h}}$ contains the identity matrix $I$ because of the residual connection. Even if the Jacobian $\frac{\partial f(\mathbf{h})}{\partial \mathbf{h}}$ is small or poorly conditioned, the $I$ term ensures the gradient is at least as large as if $f$ were the identity function. Across 12 or 24 such blocks, this prevents gradient vanishing.