Building a Transformer Decoder Block

Hello! Welcome back to our deep dive into the building blocks of Audio AI.

In our last session, we constructed a complete Transformer encoder block, the workhorse behind models that analyze and understand sequences. We saw how self-attention, feed-forward networks, residual connections, and layer normalization come together to form a powerful, stackable unit.

Today, we will build its counterpart: the Transformer decoder block. While the encoder's job is to create a rich representation of an input sequence, the decoder's job is to generate a new sequence, one element at a time, while paying attention to that encoded representation. This is the core mechanism that enables sequence-to-sequence tasks like speech-to-text or machine translation.

Your learning outcome for this lesson is to: Assemble a Transformer decoder block, including masked self-attention, cross-attention, and a feed-forward network.

We'll dissect the unique three-part structure of the decoder block, focusing on the new and crucial concept of cross-attention, which allows the decoder to "look at" the encoder's output.

1. The Anatomy of a Transformer Decoder Block

At a high level, the decoder block looks very similar to the encoder block, but with a key addition. An encoder block has two sub-layers (self-attention and FFN). A decoder block has three:

1. Masked Multi-Head Self-Attention: This allows each position in the decoder's input sequence to attend to previous positions in the same sequence.
2. Multi-Head Cross-Attention: This is the new component. It allows each position in the decoder to attend to all positions in the encoder's output. This is how the decoder conditions its generation on the source sequence.
3. Position-wise Feed-Forward Network: Just like in the encoder, this layer provides further processing for each position independently.

Each of these three sub-layers is wrapped in the same "Add & Norm" component (a residual connection and layer normalization) that we saw in the encoder.

The data flows through a single decoder block like this:

1. Input passes through Masked Self-Attention.
2. Result is passed through an Add & Norm layer.
3. The new result passes through Cross-Attention, where it also receives the encoder's output.
4. Result is passed through another Add & Norm layer.
5. The new result passes through the Feed-Forward Network.
6. Result is passed through a final Add & Norm layer.

Let's explore the two attention mechanisms in detail.

2. Sub-layer 1: Masked Self-Attention

The first sub-layer is a multi-head self-attention mechanism, just like in the encoder. However, there's a critical difference: it's masked.

The purpose of the decoder is to generate a sequence auto-regressively, meaning the prediction for the current token can only depend on the previously generated tokens. To enforce this during training (where the entire target sequence is fed in at once), we apply a causal mask or look-ahead mask.

This mask prevents each position from attending to subsequent positions. For example, when processing the 3rd token of the output sequence, it can attend to tokens 1, 2, and 3, but not 4, 5, etc.

3. Sub-layer 2: The Crucial Cross-Attention

This is where the decoder gets its context from the input sequence. Cross-attention is structurally identical to self-attention, but its inputs for Query, Key, and Value are different:

• Query (Q): The query vectors come from the output of the previous sub-layer (the masked self-attention layer). This represents what the decoder is "looking for" at its current state.
• Key (K) and Value (V): The key and value vectors come from the final output of the encoder stack. They represent the complete, context-rich information from the source sequence.

Think of it like this: the decoder's query Q asks a question ("What part of the source audio is most relevant for generating the next word?"), and the encoder's K and V provide the information to answer that question.

Important: No causal mask is applied in cross-attention. The decoder should be able to attend to any and all parts of the source sequence at every generation step.

Now, let's see the detailed mechanism.

4. Assembling the Decoder Block in PyTorch

Now that we understand all the components—masked self-attention, cross-attention, and the familiar feed-forward network—we can assemble them into a single nn.Module.

The implementation will look very similar to the EncoderBlock from our previous lesson, but the forward pass will be more complex, orchestrating the flow of data between the three sub-layers and the encoder's output (memory).

The QuarkML blog has an exceptionally clean and well-documented implementation of a DecoderLayer. We will use this as our primary reference for building our module.

Here is a slightly simplified PyTorch implementation based on the principles from the article. This code uses a Post-LN (Layer Normalization after the sub-layer) structure, as seen in the original Transformer paper and the QuarkML article.

import torch
import torch.nn as nn




# Assume MultiHeadAttention and FeedForward classes are defined as in previous lessons.
# Let's use stubs for clarity.
class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super().__init__()



        # In a real implementation, this would contain the logic from previous lessons
        self.proj = nn.Linear(d_model, d_model)
    def forward(self, query, key, value, mask=None):



        # A real implementation would use mask
        return self.proj(query)

class FeedForward(nn.Module):
    def __init__(self, d_model, d_ff):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.ReLU(),
            nn.Linear(d_ff, d_model)
        )
    def forward(self, x):
        return self.net(x)




# --- This is our main class for this lesson ---
class TransformerDecoderBlock(nn.Module):
    def __init__(self, d_model: int, num_heads: int, d_ff: int, dropout: float = 0.1):
        """
        Args:
            d_model: Dimensionality of the input (embedding dimension).
            num_heads: Number of attention heads.
            d_ff: Hidden dimension of the FFN.
            dropout: Dropout probability.
        """
        super().__init__()
        



        # Sub-layer 1: Masked Self-Attention
        self.self_attn = MultiHeadAttention(d_model, num_heads)
        self.norm1 = nn.LayerNorm(d_model)
        self.dropout1 = nn.Dropout(dropout)
        



        # Sub-layer 2: Cross-Attention
        self.cross_attn = MultiHeadAttention(d_model, num_heads)
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout2 = nn.Dropout(dropout)
        



        # Sub-layer 3: Feed-Forward Network
        self.feed_forward = FeedForward(d_model, d_ff)
        self.norm3 = nn.LayerNorm(d_model)
        self.dropout3 = nn.Dropout(dropout)

    def forward(self, 
                x: torch.Tensor, 
                memory: torch.Tensor, 
                tgt_mask: torch.Tensor = None, 
                memory_mask: torch.Tensor = None) -> torch.Tensor:
        """
        Args:
            x: Target sequence tensor from previous decoder layer.
               Shape: (batch_size, tgt_seq_len, d_model)
            memory: Final output tensor from the encoder stack.
                    Shape: (batch_size, src_seq_len, d_model)
            tgt_mask: The causal mask for the target sequence self-attention.
            memory_mask: The padding mask for the source sequence in cross-attention.
        """



        # 1. Masked Self-Attention sub-layer
        attn_output = self.self_attn(query=x, key=x, value=x, mask=tgt_mask)
        x = x + self.dropout1(attn_output)
        x = self.norm1(x)
        



        # 2. Cross-Attention sub-layer
        # The query is from the decoder, the key and value are from the encoder's memory
        cross_attn_output = self.cross_attn(query=x, key=memory, value=memory, mask=memory_mask)
        x = x + self.dropout2(cross_attn_output)
        x = self.norm2(x)
        



        # 3. Feed-Forward sub-layer
        ff_output = self.feed_forward(x)
        x = x + self.dropout3(ff_output)
        x = self.norm3(x)
        
        return x


```grasp
{
  "type": "exercise",
  "id": "bb052512-1723-45dc-9ec3-709e4324bb34"
}

--- Example Usage ---

d_model = 512
num_heads = 8
d_ff = 2048
batch_size = 4
tgt_seq_len = 80
src_seq_len = 100

Create a decoder block

decoder_block = TransformerDecoderBlock(d_model, num_heads, d_ff)

Create dummy inputs

decoder_input = torch.randn(batch_size, tgt_seq_len, d_model)
encoder_output = torch.randn(batch_size, src_seq_len, d_model)

Create a dummy causal mask for the decoder's self-attention

tgt_mask = torch.triu(torch.ones(tgt_seq_len, tgt_seq_len), diagonal=1).bool()

Forward pass

output = decoder_block(decoder_input, encoder_output, tgt_mask=tgt_mask)

print(f"Decoder input shape: {decoder_input.shape}")
print(f"Encoder output shape: {encoder_output.shape}")
print(f"Final output shape: {output.shape}")

Expected output:

Decoder input shape: torch.Size([4, 80, 512])

Encoder output shape: torch.Size([4, 100, 512])

Final output shape: torch.Size([4, 80, 512])

Notice that the output of the decoder block has the same shape as its input `x`. This shape preservation is what allows us to stack these blocks deeply to form a full decoder.


```grasp
{
  "type": "exercise",
  "id": "7033bbad-112e-4f9e-ae46-0dfbf7f700a9"
}

Conclusion

You have now successfully assembled a Transformer decoder block! This is a significant step, as you now understand the mechanics of both halves of the original Transformer architecture.

Key Takeaways:

• A Transformer decoder block has three sub-layers: masked self-attention, cross-attention, and a feed-forward network.
• Masked self-attention looks at previous tokens in the target sequence, preserving the auto-regressive property required for generation.
• Cross-attention is the bridge between the encoder and decoder. Its queries come from the decoder, while its keys and values come from the encoder, allowing it to "read" the source sequence.
• The Feed-Forward Network provides additional computation, just as it does in the encoder.
• The entire block is held together by residual connections and layer normalization (Add & Norm) after each sub-layer.

Preview of the Next Lesson:
With both the EncoderBlock and DecoderBlock in our toolkit, we are finally ready to put everything together. In the next lesson, we will construct a full encoder-decoder Transformer model in PyTorch for sequence-to-sequence tasks. We will see how to stack our blocks, add the final output layer, and manage the flow of data through the entire model from source sequence input to target sequence prediction.