Encoder-Decoder Data Flow Breakdown

Introduction

In our previous lessons, we've systematically constructed the fundamental building blocks of the Transformer:

• Input embeddings and positional encodings to represent tokens and their order.
• Multi-head self-attention to create context-aware representations.
• Position-wise feed-forward networks (FFNs) to add representational depth.
• Residual connections and Layer Normalization ("Add & Norm") to ensure stable and effective training of deep models.

Today, we will assemble these pieces. Our learning outcome is to trace the data flow through a full encoder and decoder block, including all sub-layers and residual connections. By the end of this lesson, you will have a complete, step-by-step mental model of how information is processed by the core computational units of the original Transformer architecture. This is the final step in cementing your foundational understanding before we move on to more advanced concepts.

We will primarily trace the flow using the Pre-LayerNorm (Pre-LN) architecture, which you'll recall from our last lesson is the modern standard for training stability. While some diagrams (including the original paper) depict the Post-LN variant, understanding the Pre-LN flow is more relevant for modern implementations.

The Transformer Architecture: A Bird's-Eye View

Let's start with the canonical diagram from the "Attention Is All You Need" paper. This gives us the complete map for our data flow exploration.

Our focus today is on a single "gray box" from the encoder stack and a single "gray box" from the decoder stack.

Part 1: Data Flow Through the Encoder Block

The encoder's job is to take an input sequence and produce a rich, contextual representation of it. Each encoder block processes a sequence of vectors and outputs a sequence of vectors of the exact same shape.

An encoder block consists of two main sub-layers:

1. Multi-Head Self-Attention
2. Position-wise Feed-Forward Network

Let's trace the path of a tensor X with shape (batch_size, seq_len, d_model) through a single encoder block using the Pre-LN architecture.

1. Input: X enters the block. This is either the sum of token embeddings and positional encodings (for the first block) or the output of the previous encoder block.
2. First 'Add & Norm' (Pre-LN for Attention):
   • Normalization: X is first passed through a LayerNorm layer: X_norm1 = ln1(X).
   • Self-Attention: The normalized tensor X_norm1 is used to compute Queries, Keys, and Values for the multi-head attention mechanism: attn_output = self_attention(query=X_norm1, key=X_norm1, value=X_norm1).
   • Residual Connection: The output of the attention layer is added back to the original input of the block: X_residual1 = X + attn_output.
3. Second 'Add & Norm' (Pre-LN for FFN):
   • Normalization: The result from the first residual connection is normalized: X_norm2 = ln2(X_residual1).
   • Feed-Forward Network: The normalized tensor X_norm2 is passed through the FFN: ffn_output = ffn(X_norm2).
   • Residual Connection: The FFN's output is added back to its input: output = X_residual1 + ffn_output.
4. Output: The tensor output, which still has the shape (batch_size, seq_len, d_model), is the final output of the block. It then serves as the input to the next encoder block.

The following resource provides a clear implementation of this flow.

For a more dynamic visualization of this process, the following video walks through each component of the encoder.

Part 2: Data Flow Through the Decoder Block

The decoder's role is more complex. It must generate an output sequence token by token, while also paying attention to the final representation produced by the encoder. This is why it has three sub-layers:

1. Masked Multi-Head Self-Attention
2. Encoder-Decoder Cross-Attention
3. Position-wise Feed-Forward Network

The decoder block takes two primary inputs:

• Y: The output from the previous decoder layer, shape (batch_size, target_seq_len, d_model).
• encoder_output: The final output from the entire encoder stack, shape (batch_size, **source_seq_len**, d_model).

Here is the data flow:

1. Input: Y enters the block.
2. First 'Add & Norm' (Masked Self-Attention):
   • Normalization: Y is normalized: Y_norm1 = ln1(Y).
   • Masked Self-Attention: The block performs self-attention on Y_norm1. Crucially, a look-ahead mask is applied to prevent any token from attending to future tokens in the output sequence. masked_attn_output = masked_self_attention(query=Y_norm1, key=Y_norm1, value=Y_norm1).
   • Residual Connection: Y_residual1 = Y + masked_attn_output.
3. Second 'Add & Norm' (Encoder-Decoder Cross-Attention):
   • Normalization: Y_residual1 is normalized: Y_norm2 = ln2(Y_residual1).
   • Cross-Attention: This is the key step where the decoder consults the encoder's output.
     • The Query comes from the decoder's current state: Q is derived from Y_norm2.
     • The Key and Value come from the encoder's final output: K and V are derived from encoder_output.
     • The attention is calculated: cross_attn_output = cross_attention(query=Y_norm2, key=encoder_output, value=encoder_output).
   • Residual Connection: Y_residual2 = Y_residual1 + cross_attn_output.
4. Third 'Add & Norm' (FFN):
   • Normalization: Y_residual2 is normalized: Y_norm3 = ln3(Y_residual2).
   • Feed-Forward Network: The result is passed through the FFN: ffn_output = ffn(Y_norm3).
   • Residual Connection: output = Y_residual2 + ffn_output.
5. Output: The tensor output, shape (batch_size, target_seq_len, d_model), is passed to the next decoder block.

Again, let's look at a reference implementation and a visual walkthrough.

The video from Umar Jamil also continues into the decoder, showing how it connects with the encoder.

Part 3: The Full Picture - Training vs. Inference

Tracing the data through a single block is one thing; understanding how the full model operates during training versus inference is essential for any practitioner.

Training: A Parallel Process with "Teacher Forcing"

During training, we have both the source sequence (e.g., an English sentence) and the full target sequence (e.g., its French translation). This allows for a highly parallelized forward pass:

1. Encoder: The entire source sequence is fed through the encoder stack in one go. The final output is a memory tensor (batch, source_len, d_model).
2. Decoder: The entire ground-truth target sequence (shifted right by one position and prepended with a [START] token) is fed to the decoder at once. The causal mask in the first attention sub-layer ensures that predictions for position i only depend on outputs up to position i-1.
3. Loss Calculation: The decoder's output is compared against the actual target sequence, a loss is computed (e.g., cross-entropy), and gradients are backpropagated to update all model weights.

This process is called teacher forcing because the decoder is always fed the correct previous tokens from the ground-truth data, rather than its own (potentially incorrect) predictions.

Inference: An Autoregressive, Step-by-Step Loop

During inference, we only have the source sequence. We must generate the target sequence one token at a time.

1. Encoder: The source sequence is encoded once. This encoder_output memory is reused for every step of the decoding process.
2. Decoder (Time Step t=1): The decoder input is just the [START] token. It processes this, performs cross-attention with the encoder_output, and generates a probability distribution over the entire vocabulary. We select the next token (e.g., via argmax or sampling). Let's say it's "Le".
3. Decoder (Time Step t=2): The decoder input is now [START], "Le". It processes this sequence, again using the same encoder_output, and generates the next token, e.g., "chat".
4. ...and so on. This loop continues, appending the newly generated token to the decoder's input at each step, until an [END] token is produced or a maximum length is reached.

This video provides the clearest explanation of this critical difference.

Conclusion

Congratulations on completing this deep dive into the Transformer's data flow. You have now connected all the individual components into a coherent, end-to-end process.

Key Takeaways:

• Encoder Block Flow: A two-step process of Self-Attention followed by a Feed-Forward Network, with each step wrapped in an Add & Norm layer. Its purpose is to build a rich representation of the input.
• Decoder Block Flow: A three-step process: Masked Self-Attention (to understand what's been generated so far), Encoder-Decoder Cross-Attention (to consult the input), and a Feed-Forward Network. Each step is also wrapped in Add & Norm.
• Cross-Attention is Key: This is the mechanism by which the decoder integrates information from the source sequence, using its own state as a Query to probe the Key and Value representations from the encoder.
• Training vs. Inference: Training is a parallel, one-shot process using teacher forcing. Inference is a sequential, autoregressive loop where the model feeds its own output back as input for the next step.

What's Next?

With a complete architectural understanding, we are now equipped to analyze the model from a practical engineering perspective. In the next lesson, we will move on to a very concrete task: calculating the number of parameters in a transformer model given its configuration. This skill is essential for estimating model size, memory requirements, and computational cost.