Decoder Layer, Cross-Attention, and Teacher Forcing

2682 words

13 min read

Visual companion

Python

Type and operator map

Python Week 1: the first filter for runtime behavior

Open gallery Related surface

View

Revision summary

What this note is really saying

Short form

# Decoder Layer, Cross-Attention, and Teacher Forcing ## 🎯 Learning Objectives - Explain the three sublayers in the Transformer decoder - Distinguish between masked self-attention and cross-attention - Understand teacher forcing and its role in training - Implement causal masking for autoregressive generation - Com...

Decoder Layer, Cross-Attention, and Teacher Forcing

🎯 Learning Objectives

Explain the three sublayers in the Transformer decoder
Distinguish between masked self-attention and cross-attention
Understand teacher forcing and its role in training
Implement causal masking for autoregressive generation
Compare training-time vs inference-time decoding

📋 Prerequisites

Transformer architecture (Week 1)
Self-attention and QKV computation
Multi-head attention

1. 📖 Core Content

1.1 Decoder Architecture

The Transformer decoder has three sublayers (not two like the encoder): (Diagram)

Sublayer	Purpose	Q Source	K/V Source	Mask?
Masked Self-Attention	Process generated tokens	Decoder	Decoder	Yes (causal)
Cross-Attention	Look at input	Decoder	Encoder	No
Feed-Forward	Non-linear transform	-	-	-

1.2 Masked Self-Attention

Intuition

When generating text, the model should only see words it has already generated — not future words. If you're generating "The cat sat," when predicting "sat," you should only know "The" and "cat." Masked self-attention enforces this by setting attention scores for future positions to

-\infty

before softmax.

Implementation

\text{MaskedAttention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}} + M\right)V

Where

M

is a causal mask: 0 & \text{if } i \geq j \text{ (allowed)} \\ -\infty & \text{if } i < j \text{ (masked)} \end{cases}

#### Worked Example: Masked Attention for 4 Tokens Let's trace through masked attention with 4 tokens: **Step 1**: Compute raw scores (unmasked):

S = \begin{bmatrix} 2.0 & 0.5 & 1.0 & 0.1 \\ 0.3 & 1.5 & 0.8 & 2.0 \\ 1.2 & 0.7 & 1.8 & 0.4 \\ 0.9 & 1.1 & 0.6 & 1.4 \end{bmatrix}

**Step 2**: Add causal mask:

M = \begin{bmatrix} 0 & -\infty & -\infty & -\infty \\ 0 & 0 & -\infty & -\infty \\ 0 & 0 & 0 & -\infty \\ 0 & 0 & 0 & 0 \end{bmatrix}

S + M = \begin{bmatrix} 2.0 & -\infty & -\infty & -\infty \\ 0.3 & 1.5 & -\infty & -\infty \\ 1.2 & 0.7 & 1.8 & -\infty \\ 0.9 & 1.1 & 0.6 & 1.4 \end{bmatrix}

**Step 3**: Apply softmax row-wise: Row 1: $e^{2.0}/(e^{2.0}+0+0+0) = 7.389/7.389 = 1.0$ Row 2: $e^{0.3}/(e^{0.3}+e^{1.5}) = 1.350/(1.350+4.482) = 0.231, 0.769$ Row 3: $e^{1.2}/(e^{1.2}+e^{0.7}+e^{1.8}) = 3.320/(3.320+2.014+6.050) = 0.292, 0.177, 0.531$ Row 4: All four values unmasked → normal softmax A = \\begin{bmatrix} 1.0 & 0 & 0 & 0 \\ 0.231 & 0.769 & 0 & 0 \\ 0.292 & 0.177 & 0.531 & 0 \\ 0.273 & 0.249 & 0.203 & 0.275 \\end{bmatrix} Token 0 only sees itself. Token 1 sees tokens 0 and 1. Token 2 sees tokens 0-2. Token 3 sees all tokens. \`\`\`python # runnable import numpy as np def create_causal_mask(seq_len): """Create upper-triangular causal mask""" mask = np.triu(np.ones((seq_len, seq_len)) * -1e9, k=1) return mask def masked_attention(Q, K, V, mask=None): """Scaled dot-product attention with optional mask""" d_k = Q.shape\[-1\] scores = np.dot(Q, K.T) / np.sqrt(d_k) if mask is not None: scores = scores + mask # Softmax exp_scores = np.exp(scores - np.max(scores, axis=-1, keepdims=True)) attn_weights = exp_scores / np.sum(exp_scores, axis=-1, keepdims=True) output = np.dot(attn_weights, V) return output, attn_weights # Example with 4 tokens, d_k=4 np.random.seed(42) Q = np.random.randn(4, 4) K = np.random.randn(4, 4) V = np.random.randn(4, 4) mask = create_causal_mask(4) output, attn = masked_attention(Q, K, V, mask) np.set_printoptions(precision=3, suppress=True) print("Attention weights with causal mask:") print(np.round(attn, 3)) print("\\nRow sums:", np.round(attn.sum(axis=-1), 6)) # Should all be 1.0 \`\`\` ### 1.3 Cross-Attention #### Intuition The decoder needs to know about the input sequence to generate appropriate output. In machine translation, the English decoder needs to look at the French encoder's representations. Cross-attention enables this: the decoder queries the encoder's key-value pairs. #### Formal Definition In cross-attention: - **Q** comes from the decoder (previous sublayer output) - **K** and **V** come from the encoder (final encoder layer output) \\text{CrossAttn}(Q_{dec}, K_{enc}, V_{enc}) = \\text{softmax}\\left(\\frac{Q_{dec}K_{enc}^T}{\\sqrt{d_k}}\\right)V_{enc} \#### Key Properties 1. **No masking needed**: The decoder can attend to all encoder positions (the entire input is known) 2. **Cross-attention output size**: Same as decoder sequence length, not encoder length 3. **Each decoder token can attend to all encoder tokens**: Enables flexible alignment #### Worked Example Encoder output (2 tokens, d_model=4): H_{enc} = \\begin{bmatrix} 1 & 0 & 2 & 1 \\ 3 & 1 & 0 & 2 \\end{bmatrix}

Decoder input so far (3 tokens, d_model=4):

H_{dec} = \\begin{bmatrix} 0 & 2 & 1 & 0 \\ 1 & 1 & 2 & 1 \\ 2 & 0 & 1 & 3 \\end{bmatrix} Cross-attention: 1. Project H_dec through W^Q → Q_dec (3, d_k) 2. Project H_enc through W^K → K_enc (2, d_k) 3. Project H_enc through W^V → V_enc (2, d_v) 4. Compute $Q_{dec}K_{enc}^T$ → (3, 2) matrix 5. Scale, softmax → attention weights (3, 2) 6. Weighted sum → output (3, d_v) Each decoder token gets a weighted combination of encoder tokens. ### 1.4 Teacher Forcing #### The Problem During training, the decoder is an **autoregressive model** — it generates one token at a time, using previously generated tokens as input for the next step. But if the model makes a wrong prediction at step 1, all subsequent steps are trained on incorrect inputs, creating a cascade of errors. #### The Solution: Teacher Forcing Teacher forcing feeds the **ground truth** (correct) token as input to the next decoder step, regardless of what the model predicted. \`\`\`mermaid graph LR subgraph "Without Teacher Forcing (Inference)" direction LR A1\[BOS\] --> D1\[Decoder Step 1\] D1 --> P1\[Pred: "cat"\] P1 --> D2\[Decoder Step 2\] D2 --> P2\[Pred: "dog"\] P2 --> D3\[Decoder Step 3\] end subgraph "With Teacher Forcing (Training)" direction LR B1\[BOS\] --> E1\[Decoder Step 1\] E1 --> T1\[GT: "the" → input next\] T1 --> E2\[Decoder Step 2\] E2 --> T2\[GT: "cat" → input next\] T2 --> E3\[Decoder Step 3\] end \`\`\` **Why teacher forcing works:** 1. Faster convergence: the model learns from correct context 2. Stable training: avoids error accumulation 3. Parallel training: all positions can be computed simultaneously (shifted right) **The loss**: Cross-entropy between predicted tokens and ground truth tokens, computed at each position. #### Teacher Forcing in Practice \`\`\`python # runnable import numpy as np def teacher_forcing_training_step(model, encoder_output, target_tokens): """ Simulate a teacher forcing training step Args: model: Transformer model encoder_output: Encoder representations target_tokens: Ground truth output tokens (including BOS) Returns: loss: Cross-entropy loss for this step """ seq_len = len(target_tokens) # Prepare decoder input (shifted right: remove last token, prepend BOS) decoder_input = target_tokens\[:-1\] # All except last # Forward pass through decoder (all positions in parallel!) decoder_output = model.decode(decoder_input, encoder_output) # Compare with target (all except BOS) predictions = decoder_output # (seq_len-1, vocab_size) targets = target_tokens\[1:\] # (seq_len-1,) # Cross-entropy loss loss = 0 for t in range(seq_len - 1): pred_probs = np.exp(predictions\[t\]) / np.sum(np.exp(predictions\[t\])) loss += -np.log(pred_probs\[targets\[t\]\] + 1e-10) return loss / (seq_len - 1) print("In teacher forcing, the entire target sequence is processed in parallel.") print("The causal mask ensures position t can't see position t+1.") \`\`\` ### 1.5 Training vs Inference | Aspect | Training (Teacher Forcing) | Inference (Autoregressive) | |--------|---------------------------|---------------------------| | Input | Ground truth tokens | Predicted tokens | | Speed | All positions parallel | Sequential (one at a time) | | Accuracy | Uses correct context | May use wrong context | | Loss | Computed per position | Evaluated on final output | | Mask | Same causal mask | Same causal mask | ### 1.6 Exposure Bias Teacher forcing creates a mismatch between training and inference called **exposure bias**: during training, the model always sees correct inputs; during inference, it sees its own (potentially incorrect) predictions. This means errors compound during generation. **Mitigations:** 1. **Scheduled sampling**: Gradually replace ground truth with model predictions during training 2. **Beam search**: Maintain multiple hypotheses at inference time 3. **Noise injection**: Add noise to training inputs to make the model robust to imperfect context ### 1.7 Why This Matters Understanding the decoder is crucial because: - **All modern LLMs (GPT, LLaMA, Claude) are decoder-only**: No encoder, just masked self-attention - **Inference efficiency**: The decoder's sequential nature is why LLMs generate tokens one at a time - **KV-cache optimization**: Cross-attention and self-attention can be cached for faster inference - **Prompt engineering**: Understanding decoder behavior explains why context window matters --- ## 4. 📐 Key Formulas / Concepts | Concept | Formula | Notes | |---------|---------|-------| | Masked attention | $\\text{softmax}((QK^T + M)/\\sqrt{d_k})V$ | M is causal mask | | Causal mask | $M_{ij} = -\\infty$ if $i < j$ | Prevent future token access | | Cross-attention | $\\text{softmax}(Q_{dec}K_{enc}^T/\\sqrt{d_k})V_{enc}$ | Q from decoder, K,V from encoder | | Decoder input | Shifted output + positional encoding | During training: ground truth | | Training loss | $\\sum_t -\\log P(y_t | y_{<t}, X)$ | Cross-entropy per position | | Inference | $y_t \\sim P(y | y_{<t}, X)$ | Sample/argmax per position | --- ## 5. ⚠️ Common Pitfalls ### Pitfall 1: Thinking cross-attention is the same as self-attention **The mistake**: Assuming all attention mechanisms in the decoder are self-attention. **Correction**: The decoder has two different attention types: 1. Masked self-attention (first sublayer): Q, K, V all from decoder. Bidirectional within the decoder but causally masked. 2. Cross-attention (second sublayer): Q from decoder, K, V from encoder. Not masked. ### Pitfall 2: Forgetting the "shifted right" in decoder input **The mistake**: Feeding the exact target sequence as decoder input without shifting. **Correction**: During teacher forcing, the decoder input is the target sequence shifted right by one position: - Target: \[BOS, the, cat, sat, EOS\] - Decoder input: \[BOS, the, cat, sat\] (remove last) - Target for loss: \[the, cat, sat, EOS\] (remove BOS) Position t predicts position t+1. ### Pitfall 3: Confusing teacher forcing with distillation **The mistake**: Thinking teacher forcing and knowledge distillation are the same. **Correction**: - **Teacher forcing**: Using ground truth as input during training (standard supervised learning) - **Knowledge distillation**: Using a larger model's predictions as training targets for a smaller model They're completely different concepts despite the word "teacher" appearing in both. --- ## 6. 📝 Practice Questions > **Q1: For a sequence of length 5 with causal masking, how many non-zero entries are in the attention matrix (before softmax)?** > > The causal mask allows i ≥ j (including diagonal). For n=5: > - Row 0: 1 allowed (position 0 attends to itself) > - Row 1: 2 allowed (positions 0, 1) > - Row 2: 3 allowed (positions 0, 1, 2) > - Row 3: 4 allowed > - Row 4: 5 allowed > > Total: 1 + 2 + 3 + 4 + 5 = 15 = n(n+1)/2 = 5(6)/2 = 15 non-zero entries. > > Total entries: 25. Fraction used: 15/25 = 60%. > **Q2: Why can't we use teacher forcing during inference?** > > During inference, we don't have ground truth tokens. The whole point of generation is to produce new tokens. We must use the model's own predictions as input for the next step. Teacher forcing requires knowing the correct output, which is only available during training with labeled data. > > Imagine translating "Hello" to French: during training we know the answer is "Bonjour," so we can feed it. During inference with a new sentence, we don't know the translation. > **Q3: In cross-attention, what happens if the decoder sequence is longer than the encoder sequence?** > > The attention matrix has shape (decoder_len, encoder_len). Each decoder token attends to all encoder tokens. The output has shape (decoder_len, d_v). This is fine — multiple decoder tokens can attend to the same encoder token. In translation, this happens when the output language has more words than the input. > > Example: English "I study" (2 tokens) → German "Ich lerne" (2 tokens, same length) vs "I love you" (3 tokens) → French "Je t'aime" (3 tokens or "Je t'aime" = 3 tokens). > **Q4: How would you modify the causal mask to implement prefix LM (bidirectional on prefix, causal on continuation)?** > > For a prefix LM: the first k tokens (prefix) can attend bidirectionally, and the remaining tokens are causally masked. > > Mask: > - First k rows: all 0 (bidirectional within prefix) > - Remaining rows: causal (0 for i ≥ j, -∞ for i < j) > >

M_{ij} = \begin{cases} > 0 & \text{if } j \leq k \text{ or } i \geq j \\ > -\infty & \text{otherwise} > \end{cases}

> > This is used in models like T5 and UniLM for tasks where the prefix is a prompt that should see all context. > **Q5: A decoder-only model (like GPT) has no cross-attention. How does it incorporate input information?** > > In decoder-only models, the input is simply concatenated with the output in a single sequence. The prompt/context is the "input" and the generated tokens are the "output." Causal masking ensures the prompt can attend to all prompt tokens (bidirectional within the prompt if no mask is applied to the prompt region), and the generated tokens can only attend to earlier tokens. > > Actually, in GPT-style models, the causal mask is applied across the entire sequence. So the prompt tokens are also causally masked (each prompt token only attends to previous prompt tokens). This is why adding more context to the beginning of a prompt works — it's always visible to later tokens. > **Q6: What is the gradient flow through a decoder layer with residual connections? Why does this help?** > > With residual connections, the decoder layer output is: >

x_{out} = x_{in} + \text{Sublayer}(x_{in})

> > During backpropagation: >

\frac{\partial L}{\partial x_{in}} = \frac{\partial L}{\partial x_{out}} \cdot \frac{\partial x_{out}}{\partial x_{in}} = \frac{\partial L}{\partial x_{out}} \cdot \left(1 + \frac{\partial \text{Sublayer}}{\partial x_{in}}\right)$$ > > The "1" term ensures gradients can flow directly through the residual path without vanishing, even if

\\frac{\\partial \\text{Sublayer}}{\\partial x_{in}}

is small. This is crucial for training deep (12+ layer) decoders. > Q7: In teacher forcing, the decoder processes all positions in parallel. How does this work with causal masking? > > The decoder computes attention for all positions simultaneously using matrix operations. The causal mask ensures that position i only attends to positions j ≤ i, even though all positions are processed at once. > > Conceptually: it's as if we run the decoder step-by-step, but the mask makes the parallel computation equivalent to the sequential one. The only difference is training efficiency — parallel is much faster. > Q8: Why does the decoder have 3 add-and-norm layers while the encoder has only 2? > > The decoder has an additional sublayer (cross-attention) compared to the encoder. Each sublayer is wrapped with a residual connection and layer normalization. So: > - Encoder: 2 sublayers (self-attention + FFN) → 2 add-and-norm > - Decoder: 3 sublayers (masked self-attention + cross-attention + FFN) → 3 add-and-norm > Q9: How would you implement scheduled sampling (mixing ground truth and predictions) in a decoder? > > At each training step with probability ε (which decays over time): > - Use model's predicted token as input (instead of ground truth) > - Otherwise use ground truth (standard teacher forcing) > > ```python > def scheduled_sampling_decoder_step(decoder, prev_token, encoder_output, epsilon): > """Mix ground truth and predictions""" > if np.random.random() < epsilon: > # Use model's prediction > logits = decoder(prev_token, encoder_output) > predicted_token = np.argmax(logits) > return predicted_token > else: > # Use ground truth (teacher forcing) > return ground_truth_token > ``` > > Schedule: ε starts high (e.g., 0.5) and decays to 0 during training. > Q10: In a 6-layer decoder, how many total attention mechanisms process each token during generation? > > Each decoder layer has 2 attention mechanisms (masked self-attention + cross-attention). With 6 layers: > Total = 6 × 2 = 12 attention mechanisms per token > > But during inference, previously computed attention values (keys and values) can be cached (KV-cache), so each new token only needs to compute attention with the new key-value pairs. This makes inference O(n) per token instead of O(n²). --- ## 7. 🔗 Cross-References - Next: [Layer Normalization & Residual Connections](/courses/bsda5004/notes/06-layer-norm-residual) - Advanced: KV-Cache Optimization (Week 10) - Video: BSDA5004 Week 2-3 transcripts Join Discord PreviousPositional EncodingNextLayer Norm & Residual

What this note is really saying

Decoder Layer, Cross-Attention, and Teacher Forcing

🎯 Learning Objectives

📋 Prerequisites

1. 📖 Core Content

1.1 Decoder Architecture

1.2 Masked Self-Attention

Intuition

Implementation

Intelligence Hub