Transformer Decoder Stack

A transformer decoder stack is a sequence of N identical decoder layers that transforms an embedded target sequence into contextual representations suitable for predicting the next token at every position simultaneously during training, and one token at a time during inference. Each decoder layer contains up to three sublayers: a masked (causal) self-attention sublayer, an optional cross-attention (also called encoder–decoder attention) sublayer, and a position-wise feed-forward network sublayer. Every sublayer is wrapped in a residual connection followed by layer normalization, so that the output of a sublayer is `LayerNorm(x + Sublayer(x))` (post-norm, as in the original transformer) or `x + Sublayer(LayerNorm(x))` (pre-norm, common in later models such as GPT-2 and beyond). The defining property of the decoder is the causal mask applied inside self-attention: when computing attention scores for target position i, all positions j > i are set to negative infinity before the softmax, so their attention weight becomes zero. This enforces the autoregressive property — the representation at position i depends only on tokens at positions 1 through i — which is exactly what is needed to train the model with teacher forcing on whole sequences in parallel while guaranteeing that at generation time the model never uses information it would not have during real left-to-right decoding. When the decoder is paired with an encoder (as in machine translation, summarization, or any conditional sequence-to-sequence task), a second attention sublayer, cross-attention, sits between the masked self-attention and the feed-forward sublayer. Here the queries come from the decoder's own (causally-masked) hidden states, but the keys and values come from the encoder's final output representations. This lets every decoder position "look at" the entire source sequence — with no masking, since the whole source is available — while still respecting causality among target tokens. Architectures that drop the encoder entirely and remove cross-attention, keeping only masked self-attention and the feed-forward sublayer, are called decoder-only models; this is the design used by GPT and virtually all modern large language models.

Step 1 — Input representation. The target sequence of tokens is mapped to embeddings and combined with positional encodings (sinusoidal or learned), since self-attention itself has no notion of order.

Step 2 — Masked self-attention sublayer. Each decoder layer first computes multi-head self-attention over the target sequence, with an upper-triangular mask that prevents attending to future positions. Formally, for queries Q, keys K, values V of dimension $d_k$ per head,

where $M_{ij} = -\infty$ for $j > i$ and $0$ otherwise. The additive mask is applied before the softmax so that masked entries become exactly zero after exponentiation.

Step 3 — Residual + normalization. The self-attention output is added back to its input and normalized: `x = LayerNorm(x + SelfAttn(x))`.

Step 4 — Cross-attention sublayer (if an encoder is present). Multi-head attention is computed again, but queries come from the current decoder state while keys and values come from the encoder's output `H_enc`. No causal mask is needed here, because the entire source sequence is legitimately visible at every decoding step.

Step 5 — Residual + normalization wraps the cross-attention sublayer exactly as in step 3.

Step 6 — Feed-forward sublayer. A position-wise two-layer MLP with a nonlinearity (ReLU or GELU) is applied identically and independently to each position, followed by another residual-and-normalization wrap.

Step 7 — Stack repetition. Steps 2–6 constitute one decoder layer; this block is repeated N times (N = 6 in the original transformer, tens to over a hundred in modern large models), with each layer's output feeding the next layer's input.

Step 8 — Output projection. The final layer's output is projected through a linear layer (often tied to the input embedding matrix) and a softmax to produce a probability distribution over the vocabulary for the next token, at every position. At inference time, generation is inherently sequential: the model produces one token, appends it to the input, and repeats — this is why efficient decoders cache previously computed keys and values (KV-caching) so each new token is generated in O(n) additional work rather than recomputing attention over the whole prefix from scratch. ``` function DECODER_LAYER(x, H_enc, use_cross_attention):

# x: [seq_len, d_model] target hidden states so far

# H_enc: [src_len, d_model] encoder output, or None for decoder-only

# 1. Masked self-attention sublayer a = MULTI_HEAD_ATTENTION(query=x, key=x, value=x, mask=CAUSAL_MASK(len(x))) x = LAYER_NORM(x + a)

# 2. Optional cross-attention sublayer if use_cross_attention and H_enc is not None: c = MULTI_HEAD_ATTENTION(query=x, key=H_enc, value=H_enc, mask=None) x = LAYER_NORM(x + c)

# 3. Position-wise feed-forward sublayer f = FEED_FORWARD(x) # Linear -> GELU/ReLU -> Linear, applied per position x = LAYER_NORM(x + f) return x

function DECODER_STACK(target_tokens, H_enc, num_layers, use_cross_attention): x = EMBED(target_tokens) + POSITIONAL_ENCODING(len(target_tokens)) for layer_index in 1..num_layers: x = DECODER_LAYER(x, H_enc, use_cross_attention) logits = LINEAR_PROJECTION(x) # [seq_len, vocab_size] return SOFTMAX(logits) # next-token distribution at each position

function CAUSAL_MASK(n): M = zero matrix of shape [n, n] for i in 1..n: for j in (i+1)..n: M[i][j] = -infinity return M

function GENERATE(H_enc, max_len, use_cross_attention): tokens = [START_TOKEN] kv_cache = {} # per-layer cached keys/values for generated prefix for step in 1..max_len: x = DECODER_STACK_STEP(tokens, H_enc, kv_cache, use_cross_attention) next_token = ARGMAX_OR_SAMPLE(x[-1]) # distribution at final position tokens.append(next_token) if next_token == END_TOKEN: break return tokens ``` The pseudocode's key correctness-critical detail is `CAUSAL_MASK`: it is strictly upper-triangular (excluding the diagonal), so position i attends to itself and all earlier positions, never later ones. Cross-attention deliberately passes `mask=None`, because masking the source would incorrectly hide legitimate context.

Time complexity per layer. Let n be the target sequence length, m the source sequence length (when cross-attention is present), and d the model dimension.

• Masked self-attention: computing $QK^\top$ costs $O(n^2 d)$, since it is an $n \times d$ matrix times a $d \times n$ matrix; the softmax and weighted sum with V cost the same order. The causal mask itself is applied in $O(n^2)$ but does not change the asymptotic cost, which is dominated by the matrix multiplications.
• Cross-attention: queries are $n \times d$, keys/values are $m \times d$, so the cost is $O(n m d)$.
• Feed-forward sublayer: applied independently per position with hidden width $d_{ff}$ (typically $4d$), costing $O(n \cdot d \cdot d_{ff}) = O(n d^2)$ up to a constant factor.

Summing across N layers, total training-time (full-sequence, teacher-forced) complexity is

Space complexity is dominated by the attention score matrices, $O(n^2)$ per head-batch for self-attention (or $O(nm)$ for cross-attention) plus $O(n d)$ for activations, per layer; with N layers and H heads this is $O(N H n^2)$ in the worst case if all intermediate score matrices are retained (as needed for backpropagation), though techniques like memory-efficient/FlashAttention-style kernels reduce the materialized memory to roughly $O(n)$ per query block while preserving the same FLOP count.

Why parallel training is correct despite sequential generation. During training with teacher forcing, the entire target sequence is fed in at once and processed in parallel — this is only valid because the causal mask guarantees the representation and loss at position i are mathematically identical to what would be computed if the model had generated tokens $1, \dots, i-1$ one at a time and only just now produced token i. Formally, define $h_i^{(\ell)}$ as the hidden state at position i after layer $\ell$. The causal mask enforces

with no dependence on $h_j^{(\ell-1)}$ for $j > i$, by induction on layer depth $\ell$: the base case ($\ell=0$, the embeddings) trivially satisfies this, and each subsequent layer's masked self-attention plus position-wise feed-forward preserves the property, since attention weights to future positions are exactly zero and the feed-forward sublayer never mixes information across positions at all. This invariant is precisely what licenses computing an entire sequence's loss in one forward pass and matches token-by-token autoregressive generation.

Why residual connections and layer normalization make deep stacking trainable. Each sublayer computes $x' = \text{LayerNorm}(x + \text{Sublayer}(x))$. The additive residual term means the gradient of the loss with respect to $x$ contains an identity term:

(up to the LayerNorm Jacobian, which is well-conditioned because normalization keeps activation statistics stable). This identity term guarantees gradients can flow unattenuated through arbitrarily many layers even if a particular sublayer's Jacobian is small, directly counteracting the vanishing-gradient problem that plagues naively stacked deep networks. Layer normalization additionally re-centers and re-scales each position's representation, keeping activation magnitudes stable regardless of depth, which stabilizes optimization and permits the large learning rates used in practice.

Tradeoffs. The quadratic $O(n^2)$ self-attention cost is the main scalability bottleneck for long sequences; cross-attention adds a further $O(nm)$ term that grows with source length. KV-caching removes redundant recomputation during generation but trades memory for speed, since keys and values for every previously generated token must be stored. Decoder-only models forgo the encoder and cross-attention entirely, saving computation and parameters, at the cost of needing to encode source and target in a single shared causal sequence (as GPT-style models do by concatenating prompt and continuation).

Consider a tiny English-to-French translation setup with target-side vocabulary indices and a source sentence "the cat sat" already encoded by an encoder into $H_{enc} \in \mathbb{R}^{3 \times d}$. We trace generation of the French translation "le chat s'est assis" one token at a time, focusing on what masked self-attention and cross-attention can see at each step.

Two details are worth checking numerically. First, at step 3 the causal mask for a 3-token target prefix is

so after adding $M$ to the raw attention scores $QK^\top / \sqrt{d_k}$ and applying softmax row-wise, row 3 (the query for "chat", the most recent position) has nonzero weight only on columns 1–3, and rows 1–2 never attend to column 3 at all — confirming "le" and ``'s representations are computed exactly as they would be if "chat" had not been generated yet, satisfying the invariant from the correctness argument above. Second, cross-attention at every step uses the same $H_{enc}$ (3 rows, one per source token) as keys and values, regardless of how long the target prefix has grown; there is no masking here because "the cat sat" is fully available at every decoding step. If the decoder's query at step 3 (representing "chat", the direct object) places high attention weight on the encoder row for "cat," this is exactly the alignment behavior — content words attending to their translational counterparts — that makes cross-attention interpretable as soft alignment, analogous to classical statistical machine translation alignments but learned end-to-end. During training, all five target positions are processed in a single forward pass (teacher forcing with the ground-truth prefix at each position), and the causal mask guarantees this parallel computation produces identical per-position hidden states to the step-by-step trace above — this equivalence is what makes transformer training so much faster than the strictly sequential training required by recurrent architectures.

The transformer decoder stack emerged from a lineage of work trying to remove the sequential bottleneck of recurrent sequence-to-sequence models while retaining their ability to condition generation on an entire source sequence. Early neural machine translation systems used recurrent encoders and decoders connected by an attention mechanism that let the decoder "look back" at encoder hidden states — this attention idea is the direct ancestor of cross-attention. The 2017 transformer paper's key insight was that self-attention alone, properly masked, could replace the recurrent decoder entirely, enabling full parallelization across sequence positions during training. Soon after, researchers observed that the decoder half of the transformer, used alone without any encoder or cross-attention, was already a powerful general-purpose sequence model — this observation launched the decoder-only lineage that includes GPT and its successors, which now dominate the landscape of generative language models.

The encoder–decoder variant of the decoder stack, with cross-attention intact, underlies neural machine translation systems, abstractive summarization models, speech-to-text systems that decode transcripts conditioned on audio encodings, and image-captioning systems that decode text conditioned on visual encodings. The decoder-only variant, with cross-attention removed, is the architecture of GPT-family models, and by extension the technology behind conversational assistants, code-generation tools, and most modern chat-based large language models, where the "encoder" role is absorbed by simply placing the prompt earlier in the same causally-masked sequence. KV-caching, a direct consequence of the causal-mask invariant proven above, is the standard optimization that makes real-time interactive generation from these models computationally feasible, and it is implemented in essentially every production inference engine for large language models today. Understanding the decoder stack is also the natural bridge to the Full Transformer Architecture, which combines this stack with a bidirectional encoder, and to GPT & Autoregressive Language Modeling, which studies the decoder-only specialization in depth.

Ask the tutor: How does removing cross-attention and concatenating prompt and continuation into one sequence let a decoder-only model like GPT still condition generation on context?