§3.5The Classification Head & The Decoder

We’ve built the entire encoder. Two things remain to round out the architecture: how to turn the encoder’s per-token outputs into a single prediction (a head), and what makes the decoder different.

Adding a classification head

Transformer models split cleanly into a task-independent body body The task-independent transformer trunk that produces hidden states, to which different task-specific heads are attached. defined in ch. 3 — open in glossary (everything we built in §3.1–3.4) and a task-specific head. The encoder gives us one hidden state per token, but a text classifier needs one prediction. The convention is to take the first token — the [CLS] cls token A special token prepended to the input whose final hidden state is conventionally used to represent the whole sequence for classification. defined in ch. 3 — open in glossary token — and run it through a dropout and a linear layer:

class TransformerForSequenceClassification(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.encoder = TransformerEncoder(config)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)
        self.classifier = nn.Linear(config.hidden_size, config.num_labels)

    def forward(self, x):
        x = self.encoder(x)[:, 0, :]   # take the [CLS] token's hidden state
        x = self.dropout(x)
        x = self.classifier(x)
        return x

config.num_labels = 3
encoder_classifier = TransformerForSequenceClassification(config)
encoder_classifier(inputs.input_ids).size()   # torch.Size([1, 3])

For each example we now get unnormalized logits over the three classes — exactly the shape a classifier needs. Swap the head, keep the body, and the same transformer serves masked-language-modeling, token classification, or the question-answering reader of Chapter 7.

The decoder

The decoder looks like the encoder but has two attention sublayers instead of one. Click through the block:

Masked self-attention — no peeking ahead

The decoder generates one token at a time, so at position $i$ it must only see positions $\le i$. If it could see the future, it would “cheat” during training by copying the answer. Masked self-attention masked self-attention Decoder self-attention with a causal mask so each position attends only to itself and earlier positions, preventing it from seeing future tokens. defined in ch. 3 — open in glossary enforces this with a causal mask causal mask A lower-triangular mask that sets attention scores to future tokens to −∞, so softmax gives them zero weight. defined in ch. 3 — open in glossary : a lower-triangular matrix of ones.

seq_len = inputs.input_ids.size(-1)
mask = torch.tril(torch.ones(seq_len, seq_len)).unsqueeze(0)
mask[0]
# tensor([[1., 0., 0., 0., 0.],
#         [1., 1., 0., 0., 0.],
#         [1., 1., 1., 0., 0.],
#         [1., 1., 1., 1., 0.],
#         [1., 1., 1., 1., 1.]])

Wherever the mask is zero (the future), we set the score to $-\infty$ before the softmax. Since $e^{-\infty} = 0$, those positions get exactly zero weight:

scores.masked_fill(mask == 0, -float("inf"))
# tensor([[[26.8082,    -inf,    -inf,    -inf,    -inf],
#          [-0.6981, 26.9043,    -inf,    -inf,    -inf],
#          [-2.3190,  1.2928, 27.8710,    -inf,    -inf],
#          [-0.5897,  0.3497, -0.3807, 27.5488,    -inf],
#          [ 0.5275,  2.0493, -0.4869,  1.6100, 29.0893]]])

Here is the masked attention on our running sentence. The upper triangle is greyed out — every token attends only backward:

Folding masking into the §3.2 attention function is a one-line change:

def scaled_dot_product_attention(query, key, value, mask=None):
    dim_k = query.size(-1)
    scores = torch.bmm(query, key.transpose(1, 2)) / sqrt(dim_k)
    if mask is not None:
        scores = scores.masked_fill(mask == 0, float("-inf"))
    weights = F.softmax(scores, dim=-1)
    return weights.bmm(value)

Encoder–decoder (cross) attention

The second sublayer is encoder–decoder attention encoder-decoder attention Cross-attention in the decoder where the queries come from the decoder while the keys and values come from the encoder's output. defined in ch. 3 — open in glossary . It runs multi-head attention where the queries come from the decoder but the keys and values come from the encoder’s output. This is how information crosses from the source sequence into the generated one.

That completes every moving part of the Transformer. Next we step back and survey the family of models that grew from this architecture.