Transformer decoder architecture for text generation

What's actually being asked

Design one decoder block of a Transformer decoder used for autoregressive generation (like GPT), explaining every sub-layer in order: causal masked self-attention, optional cross-attention, residual+LayerNorm, the feed-forward network, the output projection, and how it all becomes next-token prediction.

Answer

Overview

A Transformer decoder generates a sequence one token at a time, each new token conditioned only on the tokens already generated. It is a stack of N identical decoder blocks sitting on top of an input embedding layer, followed by one shared output projection. GPT-style models (decoder-only) drop the cross-attention sub-layer; the original Transformer (used for translation) keeps it, because there it also needs to look at the encoder's output.

1. Input representation

Each input token id is looked up in an embedding table and added to a positional encoding (sinusoidal, per the canonical formula PE(pos,2i) = sin(pos/10000^(2i/d_model)), PE(pos,2i+1) = cos(pos/10000^(2i/d_model))), since attention itself has no notion of order — without this the model couldn't tell "dog bites man" from "man bites dog".

2. Causal masked self-attention

Self-attention is computed the usual way: Q = XW_Q, K = XW_K, V = XW_V, Attention(Q,K,V) = softmax(QK^T/√d_k) V. The √d_k scaling keeps the dot products from blowing up and saturating softmax into near-zero-gradient regions.

The causal mask is what makes this autoregressive: before the softmax, every entry of QK^T where key-position j > query-position i is set to -∞ (large negative number in practice). After softmax those positions become exactly 0, so token i can only attend to tokens 0..i, never to future tokens. This is non-negotiable for generation — without it the model would "see the answer" during training and collapse at inference, where future tokens genuinely don't exist yet.

This runs as multi-head attention: head_i = Attention(QW_Q^i, KW_K^i, VW_V^i), MultiHead(Q,K,V) = Concat(head_1,...,head_h) W_O — multiple heads so different heads can specialize (one tracks syntax, another long-range coreference, etc.).

3. Cross-attention (optional)

Only present in encoder-decoder Transformers (translation, summarization). Same formula as self-attention, but now Q comes from the decoder's own (masked) hidden states while K and V come from the encoder's output. Matlab, decoder har step pe encoder ke output mein "lookup" karta hai — "which part of the source sentence is relevant for the word I'm generating right now." A decoder-only model like GPT has no encoder at all, so it skips this sub-layer entirely and only ever has masked self-attention.

4. Residual connections + layer normalization

Every sub-layer (masked self-attn, cross-attn, FFN) is wrapped as x_out = LayerNorm(x + Sublayer(x)). The residual (x +) lets gradients flow directly through the stack even if a sub-layer's output is initially near zero, which is what makes training deep stacks (dozens of blocks) actually converge. LayerNorm normalizes activations across the feature dimension, stabilizing the scale at every layer. GPT-family models actually use pre-LN (x + Sublayer(LayerNorm(x))) rather than the original post-LN — worth a one-line mention if asked, since it's a known training-stability improvement, but the mechanism (residual + normalization, wrapped around every sub-layer) is what matters conceptually.

5. Feed-forward network (FFN)

A position-wise two-layer MLP applied identically (same weights) to every position: FFN(x) = max(0, xW_1 + b_1) W_2 + b_2 (ReLU in the original paper; GPT-style models commonly use GELU instead, same shape of idea). It expands to a larger hidden dimension and projects back, giving the model extra non-linear capacity beyond what attention alone provides.

6. Output projection + next-token prediction

After the final decoder block, the hidden state at the last position is projected to vocabulary size: logits = h W_vocab + b (often W_vocab is tied to the input embedding matrix to save parameters), then P(token) = softmax(logits).

• Training: teacher forcing — feed the real previous tokens, compute cross-entropy loss between the predicted distribution and the actual next token, summed/averaged over the sequence.
• Inference: there's no "real" next token yet, so the loop is autoregressive — sample (or take argmax of) a token from P, append it to the input, feed the whole thing back in, repeat until an end-of-sequence token or a length limit.

Worked example — the causal mask, concretely

For a 4-token sequence, the raw attention-score matrix QK^T is 4×4. The causal mask zeroes out (sets to -∞ pre-softmax) everything above the diagonal:

allowed =  1 0 0 0
           1 1 0 0
           1 1 1 0
           1 1 1 1

Row 3 (0-indexed) can attend to all 4 positions (it's the last generated token so far); row 0 can only attend to itself. This lower-triangular pattern is the causal mask — every decoder, whatever its size, uses exactly this structure.

Exam tip

If the question says "GPT" specifically, say clearly there is no cross-attention — only masked self-attention + FFN, repeated N times. Examiners often test whether you can tell decoder-only (GPT: generation) from encoder-decoder (T5/original Transformer: translation, summarization) decoders apart — the differentiator is the presence of the cross-attention sub-layer, nothing else changes structurally.