BERT -- Bidirectional Encoder¶

BERT (Bidirectional Encoder Representations from Transformers) fundamentally changed NLP when it was introduced by Devlin et al. in 2018¹. Unlike the decoder-only, left-to-right models that dominate text generation, BERT is an encoder-only model that attends bidirectionally -- every token can attend to every other token in the sequence, including future positions. This design makes BERT unsuitable for autoregressive generation but exceptionally powerful for understanding tasks: classification, named entity recognition, question answering, and sentence embeddings.

1. Architecture Overview¶

BERT uses the encoder half of the original Transformer architecture (Vaswani et al., 2017)² with post-norm residual connections (LayerNorm applied after the residual addition, unlike the pre-norm design in LLaMA).

2. Key Innovations¶

2.1 Bidirectional Attention¶

In causal (autoregressive) models, the attention mask prevents token \( i \) from attending to any token \( j > i \). BERT removes this constraint:

2.2 Masked Language Modeling (MLM)¶

Instead of predicting the next token (autoregressive), BERT is trained to predict masked tokens:

2.3 Next Sentence Prediction (NSP)¶

BERT also learns to predict whether two segments are consecutive in the original text:

2.4 Three-Part Input Embedding¶

BERT's input is the sum of three embedding types:

where:

• \( E_{\text{token}} \) maps tokens to dense vectors
• \( E_{\text{segment}} \) distinguishes sentence A from sentence B
• \( E_{\text{position}} \) provides learned absolute position embeddings

3. Architecture Diagram¶

flowchart TD
    INPUT["[CLS] token_1 ... token_n [SEP] sent_b_1 ... [SEP]"]

    subgraph EMBEDDING["Input Embedding"]
        TOK["Token Embedding"]
        SEG["Segment Embedding\n(sentence A=0, B=1)"]
        POS["Position Embedding\n(learned, absolute)"]
        SUM["Sum"]
        TOK --> SUM
        SEG --> SUM
        POS --> SUM
    end

    INPUT --> EMBEDDING
    SUM --> BLOCKS

    subgraph BLOCKS["L x TransformerBlock (Post-Norm)"]
        MHA["Multi-Head Attention\n(bidirectional, no causal mask)"]
        ADD1["Residual Add"]
        LN1["LayerNorm"]
        FFN["Feed-Forward (GELU)"]
        ADD2["Residual Add"]
        LN2["LayerNorm"]

        MHA --> ADD1 --> LN1
        LN1 --> FFN --> ADD2 --> LN2
    end

    BLOCKS --> POOLER

    subgraph POOLER["Output Heads"]
        CLS["[CLS] -> Pooler -> Classification"]
        MLM_HEAD["Token logits -> MLM Prediction"]
    end

4. Configuration Parameters¶

5. Mathematical Formulation¶

5.1 Post-Norm Transformer Block¶

BERT uses post-norm (original Transformer style), which differs from the pre-norm design used in LLaMA and most modern decoders:

[ h' = \text{LayerNorm}(x + \text{MHA}(x)) ] [ h'' = \text{LayerNorm}(h' + \text{FFN}(h')) ]

5.2 Feed-Forward Network¶

BERT was the first major model to use GELU activation instead of ReLU in the feed-forward layers.

5.3 Pooler¶

For classification tasks, BERT uses the hidden state of the [CLS] token, passed through a pooler layer:

6. Zig Implementation¶

6.1 BERTConfig¶

pub const BERTConfig = struct {
    n_layers: u32,
    d_model: u32,
    n_heads: u32,
    d_ff: u32,
    vocab_size: u32 = 30522,
    max_seq_len: u32 = 512,
    type_vocab_size: u32 = 2,     // segment types (sentence A, B)
    norm_eps: f32 = 1e-12,
    activation: ActivationType = .gelu,
    use_post_norm: bool = true,   // BERT uses post-norm

    pub fn headDim(self: BERTConfig) u32 {
        return self.d_model / self.n_heads;
    }
};

6.2 BERT Embedding¶

pub const BERTEmbedding = struct {
    token_embedding: Tensor(f32),      // [vocab_size, d_model]
    segment_embedding: Tensor(f32),    // [type_vocab_size, d_model]
    position_embedding: Tensor(f32),   // [max_seq_len, d_model]
    layernorm: LayerNorm,

    pub fn forward(
        self: *BERTEmbedding,
        token_ids: []const u32,
        segment_ids: []const u32,
        position_ids: []const u32,
    ) !Tensor(f32) {
        var output = try self.token_embedding.lookup(token_ids);

        // Sum all three embeddings
        for (0..token_ids.len) |i| {
            const seg = self.segment_embedding.row(segment_ids[i]);
            const pos = self.position_embedding.row(position_ids[i]);
            for (0..self.d_model) |j| {
                output.data[i * self.d_model + j] += seg[j] + pos[j];
            }
        }

        return self.layernorm.forward(output);
    }
};

6.3 Bidirectional Attention¶

pub const BidirectionalAttention = struct {
    wq: Linear,
    wk: Linear,
    wv: Linear,
    wo: Linear,
    n_heads: u32,

    pub fn forward(self: *BidirectionalAttention, x: Tensor(f32)) !Tensor(f32) {
        const q = self.wq.forward(x);
        const k = self.wk.forward(x);
        const v = self.wv.forward(x);

        // Reshape to [seq, n_heads, d_head] and transpose
        // Compute attention scores WITHOUT causal mask
        const scores = matmul(q, transpose(k));
        const scale = 1.0 / @sqrt(@as(f32, @floatFromInt(self.headDim())));
        const scaled = elementScale(scores, scale);

        // No causal mask -- all positions attend to all positions
        const attn_weights = softmax(scaled, .last_dim);
        const context = matmul(attn_weights, v);

        return self.wo.forward(context);
    }
};

6.4 Post-Norm Block¶

pub const BERTBlock = struct {
    attention: BidirectionalAttention,
    ffn: FeedForward,
    ln1: LayerNorm,
    ln2: LayerNorm,

    pub fn forward(self: *BERTBlock, x: Tensor(f32)) !Tensor(f32) {
        // Post-norm: LayerNorm(x + sublayer(x))
        const attn_out = try self.attention.forward(x);
        const normed1 = self.ln1.forward(residualAdd(x, attn_out));

        const ffn_out = try self.ffn.forward(normed1);
        return self.ln2.forward(residualAdd(normed1, ffn_out));
    }
};

7. Variants¶

8. Educational Value¶

9. References¶