Encoder-Only Transformer in C++ (LibTorch)

The Mental Model: A Transformer Reader

GPT is a decoder-only Transformer: it reads left-to-right and predicts the next token. An encoder-only Transformer is different. It is a reader. It sees the whole input at once, lets every token attend to every other token, and returns contextual representations that are useful for classification, retrieval, tagging, or ranking.

This is the family that includes BERT-style models. Instead of asking "what token comes next?", we usually ask "what does this sequence mean?" or "what label belongs to each token?"

What Changes Compared with GPT?

The implementation is still built from the same ingredients: embeddings, positional encoding, multi-head self-attention, feed-forward networks, residual connections, and LayerNorm. The big difference is the attention mask and the final head.

1. Bidirectional Attention Head

The single-head attention code from the GPT post used a lower-triangular mask. For an encoder, remove that causal mask. Optionally keep a padding mask so real tokens do not attend to padding tokens.

#include <torch/torch.h>
#include <cmath>
#include <memory>
#include <string>
#include <vector>

class BidirectionalAttentionHead : public torch::nn::Module {
    int n_embd, head_size;
    torch::nn::Linear K{nullptr}, Q{nullptr}, V{nullptr};

public:
    BidirectionalAttentionHead(int n_embd, int head_size)
        : n_embd(n_embd), head_size(head_size) {
        K = register_module("K", torch::nn::Linear(
            torch::nn::LinearOptions(n_embd, head_size).bias(false)));
        Q = register_module("Q", torch::nn::Linear(
            torch::nn::LinearOptions(n_embd, head_size).bias(false)));
        V = register_module("V", torch::nn::Linear(
            torch::nn::LinearOptions(n_embd, head_size).bias(false)));
    }

    torch::Tensor forward(torch::Tensor x,
                          torch::Tensor attention_mask = torch::Tensor()) {
        auto q = Q(x);  // (B, T, head_size)
        auto k = K(x);  // (B, T, head_size)
        auto v = V(x);  // (B, T, head_size)

        auto scores = torch::matmul(q, k.transpose(-2, -1));
        scores = scores / std::sqrt((double)head_size);  // (B, T, T)

        // Optional padding mask: 1 = keep token, 0 = ignore padding token.
        if (attention_mask.defined()) {
            auto keep = attention_mask.unsqueeze(1).to(torch::kBool); // (B, 1, T)
            scores = scores.masked_fill(keep.logical_not(), -1e9);
        }

        auto weights = torch::softmax(scores, -1);
        return torch::matmul(weights, v);  // (B, T, head_size)
    }
};

There is no torch::tril here. Every row of the attention matrix can place weight on every non-padding column. That is exactly what makes encoder representations bidirectional.

2. Multi-Head Attention and EncoderBlock

Multi-head attention repeats the head several times, concatenates the outputs, and projects back to n_embd. The encoder block then wraps this attention layer and the feed-forward layer with LayerNorm and residual connections.

class MultiHeadSelfAttention : public torch::nn::Module {
    int num_heads, head_size, n_embd;
    std::vector<std::shared_ptr<BidirectionalAttentionHead>> heads;
    torch::nn::Linear projection{nullptr};

public:
    MultiHeadSelfAttention(int num_heads, int head_size, int n_embd)
        : num_heads(num_heads), head_size(head_size), n_embd(n_embd) {
        for (int i = 0; i < num_heads; i++) {
            auto head = std::make_shared<BidirectionalAttentionHead>(n_embd, head_size);
            register_module("head_" + std::to_string(i), head);
            heads.push_back(head);
        }

        projection = register_module("projection",
            torch::nn::Linear(head_size * num_heads, n_embd));
    }

    torch::Tensor forward(torch::Tensor x,
                          torch::Tensor attention_mask = torch::Tensor()) {
        std::vector<torch::Tensor> head_outputs;
        for (int i = 0; i < num_heads; i++) {
            head_outputs.push_back(heads[i]->forward(x, attention_mask));
        }
        auto concat = torch::cat(head_outputs, -1);
        return projection(concat);
    }
};

class EncoderBlock : public torch::nn::Module {
    std::shared_ptr<MultiHeadSelfAttention> self_attn;
    std::shared_ptr<FeedForward> ff;
    torch::nn::LayerNorm ln1{nullptr}, ln2{nullptr};

public:
    EncoderBlock(int num_heads, int head_size, int n_embd, int d_ff) {
        self_attn = std::make_shared<MultiHeadSelfAttention>(
            num_heads, head_size, n_embd);
        ff = std::make_shared<FeedForward>(n_embd, d_ff);

        register_module("self_attn", self_attn);
        register_module("ff", ff);
        ln1 = register_module("ln1",
            torch::nn::LayerNorm(torch::nn::LayerNormOptions({n_embd})));
        ln2 = register_module("ln2",
            torch::nn::LayerNorm(torch::nn::LayerNormOptions({n_embd})));
    }

    torch::Tensor forward(torch::Tensor x,
                          torch::Tensor attention_mask = torch::Tensor()) {
        x = x + self_attn->forward(ln1(x), attention_mask);
        x = x + ff->forward(ln2(x));
        return x;  // (B, T, n_embd)
    }
};

3. Full Encoder-Only Model

The complete model is now a short pipeline: token embedding, positional encoding, a stack of encoder blocks, final LayerNorm, then a task-specific head. The example below uses a [CLS] token at position 0 and a classification head.

class EncoderOnlyTransformer : public torch::nn::Module {
    int vocab_size, n_embd, num_heads, head_size, d_ff, block_size;
    int num_layers, num_classes;

    std::shared_ptr<PositionalEncoding> pe;
    std::vector<std::shared_ptr<EncoderBlock>> blocks;

    torch::nn::Embedding token_emb{nullptr};
    torch::nn::LayerNorm final_ln{nullptr};
    torch::nn::Linear classifier{nullptr};

public:
    EncoderOnlyTransformer(int vocab_size, int n_embd, int num_heads,
                           int head_size, int d_ff, int block_size,
                           int num_layers, int num_classes)
        : vocab_size(vocab_size), n_embd(n_embd), num_heads(num_heads),
          head_size(head_size), d_ff(d_ff), block_size(block_size),
          num_layers(num_layers), num_classes(num_classes) {
        pe = std::make_shared<PositionalEncoding>(n_embd);

        token_emb = register_module("token_emb",
            torch::nn::Embedding(vocab_size, n_embd));
        final_ln = register_module("final_ln",
            torch::nn::LayerNorm(torch::nn::LayerNormOptions({n_embd})));
        classifier = register_module("classifier",
            torch::nn::Linear(n_embd, num_classes));

        for (int layer = 0; layer < num_layers; layer++) {
            auto block = std::make_shared<EncoderBlock>(
                num_heads, head_size, n_embd, d_ff);
            register_module("encoder_block_" + std::to_string(layer), block);
            blocks.push_back(block);
        }
    }

    torch::Tensor encode(torch::Tensor token_ids,
                         torch::Tensor attention_mask = torch::Tensor()) {
        auto x = token_emb(token_ids);       // (B, T, n_embd)
        x = x + pe->forward(x);             // add positional signal

        for (int layer = 0; layer < num_layers; layer++) {
            x = blocks[layer]->forward(x, attention_mask);
        }
        return final_ln(x);                 // (B, T, n_embd)
    }

    torch::Tensor forward(torch::Tensor token_ids,
                          torch::Tensor attention_mask = torch::Tensor()) {
        auto hidden = encode(token_ids, attention_mask);
        auto cls = hidden.select(1, 0);      // (B, n_embd), position 0
        return classifier(cls);             // (B, num_classes)
    }
};

Interactive Animation: Encoder Forward Pass

Step through the model and watch how the sequence stays shape-compatible until the final task head pools the [CLS] vector.

Shape Trace

For a concrete classifier with B=4, T=32, n_embd=128, num_layers=3, and num_classes=5:

Training the Classifier

For sequence classification, the loss is ordinary cross-entropy over the class logits:

auto logits = model.forward(token_ids, attention_mask);  // (B, num_classes)
auto loss = torch::nn::functional::cross_entropy(logits, labels);

optimizer.zero_grad();
loss.backward();
optimizer.step();

For token classification, return per-token logits with shape (B, T, num_tags). Then flatten batch and time before cross-entropy, and ignore padding positions with the same padding mask you used in attention.

Common Pitfalls

What Comes Next

Encoder-only models are perfect when the full input is known up front. The next full architecture is the encoder-decoder Transformer: an encoder reads the source sequence once, and a decoder generates the target sequence while cross-attending to the encoder memory.