BERT: an encoder that reads both directions

By Hiva Mohammadzadeh · Stanford MSCS, UC Berkeley EECS

What makes BERT different

BERT stands for Bidirectional Encoder Representations from Transformers, and the word that matters most in that name is bidirectional. Before BERT, most language models read text left-to-right (or right-to-left). BERT reads the entire sequence at once. When the model processes a token, it has access to every token to its left and every token to its right, so the context is complete.

BERT is an encoder-only model. It does not have the decoder stack you see in GPT-style architectures. There is no autoregressive generation; the entire input is processed in parallel. This design makes BERT effective for understanding tasks: classification, named entity recognition, question answering, and anything else where you need to comprehend the full input before producing an output.

The base configuration I work with uses 128 tokens as the sequence length, an embedding dimension of 768, 12 attention heads, and a feed-forward filter size of 3072. The model consists of 12 stacked transformer encoder blocks, each containing two sub-layers: a multi-head self-attention mechanism and a position-wise fully connected feed-forward network.

Training: masked language modeling and next-sentence prediction

BERT's training happens in two distinct phases.

Pre-training

The model is pre-trained on large, diverse text corpora with an unsupervised learning approach. Two objectives drive this phase:

• Masked Language Modeling (MLM): Some percentage of the input tokens are replaced with a special [MASK] token, and the model is trained to predict the original token. This forces the model to learn deep bidirectional representations; it cannot just memorize a left-to-right pattern because the masked token could be anywhere.
• Next Sentence Prediction (NSP): The model receives pairs of sentences and must predict whether the second sentence actually follows the first in the original text, or whether it is a random sentence. This teaches the model to understand relationships between sentences, which matters for downstream tasks like question answering and natural language inference.

Fine-tuning

After pre-training, you fine-tune BERT on your specific downstream task by adding a task-specific output layer. The pre-trained weights give the model a strong starting point, and the fine-tuning phase adapts those representations to your particular problem. BERT is practical because the expensive pre-training happens once, and fine-tuning is comparatively cheap.

The architecture

BERT's architecture has two key ingredients: the transformer encoder stack and the task-specific output layer on top.

Each encoder block follows the same pattern: self-attention, then a residual connection and layer normalization, then a feed-forward network, then another residual connection and layer normalization. This pattern repeats 12 times. The output of the final block feeds into whatever task-specific head you have added for fine-tuning.

Inside the transformer encoder

Each transformer encoder block has two main components: the attention mechanism and the feed-forward network. The supporting parts — residual connections, layer normalization, and the activation function — matter just as much for making the model trainable.

• Self-attention lets the model attend to different parts of the input sequence. Every token can look at every other token and decide how much to weight each one.
• The feed-forward network (FFN) processes the output of the normalization layer to fit it better to the next attention layer. It applies a nonlinear transformation independently to each position.
• Residual connections wrap both sub-layers to avoid the vanishing gradient problem. The input to each sub-layer is added to its output before normalization.
• Layer normalization improves convergence speed by normalizing activations across the feature dimension.
• GELU is the nonlinearity used in the feed-forward layers. It performs better than ReLU and other activation functions in transformer architectures, and the smoothness of its gradient helps during training.

Dimensions of the weight matrices

The dimensions cause a lot of confusion, so I lay them out precisely.

Input embedding

The input embedding matrix has shape:

For parallel attention heads, this gets reshaped to:

Self-attention weight matrices (W_q, W_k, W_v, W_o)

Each projection weight matrix is:

Projected Q, K, V matrices

After multiplying the input embeddings by the weight matrices:

Feed-forward network weights

The FFN expands the representation from 768 to 3072 (a 4x expansion), applies GELU, then projects it back down to 768. This expand-then-compress pattern is standard in transformer architectures.

The attention block

The full computation flow through the multi-head self-attention mechanism uses the notation E_d = 768, L = 128, h = 12, d = E_d/h = 64:

The step-by-step data flow:

1. Input: X with shape (L, E_d) = (128, 768)
2. Project: Multiply by W_q(E_d, d, h), W_k(E_d, d, h), W_v(E_d, d, h) to get Q(L, d, h), K(L, d, h), V(L, d, h)
3. Transpose K: K becomes K^T(h, d, L) for the dot product
4. Scaled dot-product: Q · K^T produces attention scores with shape (L, L, h)
5. Scale: Divide by √d = √64 = 8
6. Softmax: Apply softmax to get normalized attention weights, shape (L, L, h)
7. Apply to values: Multiply scores by V to get z_h(L, d, h)
8. Concatenate heads: Merge the h dimension back to get z(L, E_d)
9. Output projection: Multiply by W_o(E_d, E_d) to get out(L, E_d)
10. Residual connection: Add input X to get y(L, E_d)
11. Layer normalization: Normalize to get the final output(L, E_d)

The feed-forward network

The FFN is a position-wise operation; it processes each token independently through the same two-layer network.

The data flow, where E_d = 768, L = 128, and F_l = 3072:

1. Input: y with shape (L, E_d) = (128, 768), the attention block output
2. First linear transformation: W(E_d, F_l) projects to shape (L, F_l) = (128, 3072)
3. GELU activation: Applied element-wise, shape unchanged at (L, F_l)
4. Second linear transformation: W(F_l, E_d) projects back to shape (L, E_d) = (128, 768)
5. Residual connection: Add the original input y
6. Layer normalization: Final output shape (L, E_d) = (128, 768)

The pattern has three steps: linear projection, nonlinearity, and linear projection. The 4x expansion to 3072 gives the network enough capacity to learn useful intermediate representations before it compresses back down.

The complete encoder block

The complete BERT encoder block chains attention and FFN with their respective residual connections and normalizations:

Input X(128, 768) flows through self-attention, gets added back (residual), normalized, then flows through the FFN, gets added back again (residual), and normalized one more time. The output is (128, 768), the same shape as the input. Stack this 12 times and you have BERT Base.

Counting FLOPs

The computational cost of BERT matters for deployment, hardware planning, and optimization. I assume every dot product requires 1 multiplication and 1 addition (2 FLOPs per multiply-accumulate).

Attention block: weight matrix multiplications

Multiplying the input embedding matrix by the weight matrices for Q, K, and V:

Attention block: multi-head self-attention

The dot products within the attention mechanism across 12 heads:

Feed-forward network

Two dense layers with the 768 to 3072 expansion and back:

Total per transformer block

For the full 12-block BERT Base, multiply by 12: roughly 22.3 billion FLOPs per forward pass. The feed-forward network dominates the compute at about 65% of the total, so it is the first place to optimize.

From RNNs to transformers to BERT

Before BERT, the NLP world relied on Recurrent Neural Networks (RNNs). RNNs had an encoder that took the input sequence and the previous hidden state to output the next hidden state, and a decoder that generated words one at a time. They worked, but they had serious problems: sequential processing made them slow to train and impossible to parallelize, they required fixed-order input processing, and they struggled with long-range dependencies.

The 2017 paper "Attention Is All You Need" replaced RNNs with self-attention, which lets the model attend to different parts of the input sequence regardless of position. This made long-term dependencies easier to model and enabled parallel processing during training.

BERT took this encoder architecture and showed that bidirectional pre-training, reading text in both directions at once, produced better representations than left-to-right or right-to-left approaches alone.

Beyond BERT: RoBERTa

RoBERTa proposed several improvements to BERT's pre-training:

• Dynamic masking: Instead of using the same static mask, RoBERTa randomly masks different tokens at each training step, so the model cannot rely on one masking pattern
• Larger-scale pre-training: Trained on 160GB of text (10x BERT's 16GB)
• Longer training: Maximum sequence length of 512 tokens, trained for 100 epochs
• Removed NSP: Dropped the Next Sentence Prediction objective entirely
• Larger batch sizes and dynamic learning rates

RoBERTa outperformed BERT on most benchmark NLP tasks, showing that BERT was significantly undertrained and that the training recipe matters as much as the architecture.

Key takeaways

• BERT is encoder-only and bidirectional. It reads the full sequence at once, which is why it excels at understanding tasks rather than generation tasks.
• Two-phase training (pre-train with MLM/NSP, then fine-tune) is what makes BERT practical. The expensive pre-training is done once; fine-tuning is cheap.
• Each transformer block has two sub-layers: multi-head self-attention and a position-wise feed-forward network, both wrapped in residual connections and layer normalization.
• The dimensions are systematic: 768 embedding size, 12 heads with 64 dimensions each, and a 3072-dimensional FFN hidden layer (4x expansion).
• GELU is the activation function of choice, with smoother gradients and better convergence than ReLU.
• The FFN dominates the FLOP count at roughly 1.2 billion FLOPs per block, compared to about 654 million for the full attention mechanism.
• A single transformer block costs about 1.86 billion FLOPs, the basis for reasoning about BERT's computational budget.

Based on my CS199 Supervised Independent Study at UC Berkeley and the presentations I created in 2023.