Transformer Decoder Block Architecture

Hello! Welcome to the next lesson on the Transformer architecture.

In our previous session, we constructed the Transformer encoder block, which is responsible for creating rich, contextual representations of an input sequence. We saw how multi-head self-attention, feed-forward networks, residual connections, and layer normalization work together to achieve this.

Today, we shift our focus to the other half of the model: the decoder. The decoder's job is to take the encoder's output and generate a new sequence, token by token. This is where tasks like language translation happen. This lesson will focus on the learning outcome: Describe the architecture of a single Transformer decoder block, including masked self-attention and cross-attention.

We'll find that the decoder block shares many components with the encoder, but with two crucial modifications that give it its generative power.

A High-Level View of the Decoder

Let's start by looking at where the decoder fits into the full Transformer architecture.

Now, let's zoom in on a single decoder block to see its internal structure.

As you can see, the decoder block is slightly more complex than the encoder block. It consists of three primary sub-layers:

1. Masked Multi-Head Self-Attention: This layer allows the decoder to look at the tokens it has already generated in the output sequence.
2. Encoder-Decoder Cross-Attention: This layer allows the decoder to look at the encoded representation of the input sequence.
3. Position-wise Feed-Forward Network: This is identical in structure to the one in the encoder.

Let's break down each of these components, starting with the first and most unique attention mechanism.

1. Masked Multi-Head Self-Attention

The first sub-layer is a self-attention mechanism, but with a critical twist: masking. The purpose of the decoder is to generate a sequence one token at a time. This is known as being autoregressive. For example, when translating "I like pizza" to "Me gusta la pizza," to predict the word "gusta," the model should only know about "Me" and the start-of-sequence token, not the future words "la" and "pizza."

During training, we can process the entire target sequence (" Me gusta la pizza ") at once for efficiency. However, we must prevent the model from "cheating" by looking ahead at future tokens. This is where masking comes in.

To understand why this is so important and how we can still train the decoder in parallel, please watch the following video segment.

To summarize the key idea from the video:

• The Problem: We want fast, parallel training, but the autoregressive nature of the decoder means each prediction depends on the previous ones. Feeding the whole target sequence at once during training would cause data leakage, where the model learns to simply copy the next word instead of predicting it.
• The Solution: We apply a look-ahead mask to the attention scores inside the decoder's first self-attention layer. This mask sets the attention scores for all future tokens to negative infinity ($-\infty$).
• The Effect: When the softmax function is applied, the scores corresponding to future tokens become zero. This forces each position in the decoder to only attend to itself and previous positions, perfectly simulating the one-by-one generation process while still allowing for parallel computation across the entire sequence.

This masked self-attention mechanism ensures the decoder properly learns the conditional dependencies of the target language.

2. Encoder-Decoder Cross-Attention

The second sub-layer is where the decoder finally interacts with the encoder. After the decoder has processed its own partial output sequence using masked self-attention, it needs to consult the input sentence to decide what to generate next. This is the job of cross-attention.

This layer is a standard multi-head attention mechanism, but its inputs are "crossed":

• The Query (Q) vector comes from the output of the previous decoder layer (the masked self-attention layer). It represents the "question" the decoder is asking, like "Based on what I've generated so far, what should I look for in the source sentence?"
• The Key (K) and Value (V) vectors both come from the output of the final encoder block. They represent the full context of the input sentence.

This allows the decoder to weigh the importance of each word in the input sequence when generating the next word in the output sequence.

The following video segment explains this crucial interaction.

The key insight here is that unlike older RNN-based encoder-decoder models that compressed the entire input into a single "context vector", the Transformer decoder can look at the entire encoded input sequence at every single step. Through cross-attention, it dynamically decides which input words are most relevant for the current prediction.

For a concise, text-based summary of the decoder block's structure, including this cross-attention layer, you can review the following resource.

3. Position-wise Feed-Forward Network

The third and final sub-layer is the position-wise feed-forward network (FFN). This component is identical in structure and function to the one we studied in the encoder block. Its purpose is to perform further non-linear processing on the output from the cross-attention layer.

As a brief reminder:

• It consists of two linear layers with a ReLU activation in between.
• It is applied independently to each position in the sequence.
• It is followed by the standard "Add & Norm" step.

We don't need to go into depth here, as its role is the same as in the encoder: to process the information gathered by the attention layers.

Conclusion

In this lesson, we have dissected the Transformer decoder block. While it shares familiar components with the encoder, its unique design is what enables autoregressive generation.

Key Takeaways:

• A single decoder block has three sub-layers: Masked Self-Attention, Cross-Attention, and a Feed-Forward Network. Each is followed by a residual connection and layer normalization.
• Masked Self-Attention is the decoder's internal "memory." It attends to previously generated tokens in its own sequence, using a look-ahead mask to prevent it from seeing future tokens during training.
• Cross-Attention is the bridge to the encoder. It allows the decoder to consult the entire encoded input sequence at each step, with queries coming from the decoder and keys/values coming from the encoder.
• The Feed-Forward Network provides additional processing capacity, just as it does in the encoder.

Preview of the Next Lesson

We now understand the internal workings of both the encoder and decoder blocks. In the next lesson, we will zoom out and discuss how these blocks are stacked and how the full encoder and decoder stacks interact to form the complete Transformer model, ready for tasks like machine translation.