Constructing the Transformer Decoder

Constructing the Transformer Decoder | CodeSignal Learn

Introduction

Welcome to the fifth and final lesson in Deconstructing the Transformer Architecture! Congratulations on reaching this milestone: you've journeyed through multi-head attention mechanisms, mastered feed-forward networks and normalization techniques, understood positional encodings, and successfully built a complete Transformer encoder. Now, we tackle the final piece of the puzzle: the Transformer Decoder Layer . While the encoder processes input sequences to create rich representations, the decoder has a more complex responsibility. It must generate output sequences one token at a time while carefully managing what information it can access. This requires two distinct attention mechanisms: masked self-attention, which prevents the decoder from "peeking ahead" at future tokens, and cross-attention, which allows the decoder to focus on relevant parts of the encoder's output. By the end of this lesson, you'll understand how decoders balance autoregressive generation with encoder context, completing your understanding of the full Transformer architecture.

Understanding the Decoder's Unique Role

The decoder layer operates under fundamentally different constraints than the encoder, reflecting its role in sequence generation rather than sequence understanding. While an encoder can attend to all positions simultaneously because the entire input sequence is available, a decoder must generate outputs step by step, only accessing tokens it has already produced. This autoregressive nature requires a careful balance: the decoder needs enough context to make informed predictions while being prevented from accessing future information that would make training trivial and inference impossible. This challenge manifests in the decoder's architecture through three distinct sub-layers instead of the encoder's two. The first sub-layer implements masked self-attention, where each position can only attend to earlier positions in the target sequence. The second sub-layer introduces cross-attention, allowing decoder positions to attend to all encoder outputs, enabling the model to focus on relevant source information when generating each target token. The final sub-layer uses the same position-wise feed-forward network as the encoder, processing the combined self and cross-attention information to produce the final representations. The interplay between these mechanisms is what makes sequence-to-sequence tasks possible: translation, summarization, and dialogue generation all rely on this careful orchestration of attention patterns. Understanding this architecture helps explain why Transformer-based models excel at tasks requiring both comprehension of input context and coherent generation of output sequences.

Join the 1M+ learners on CodeSignal

Be a part of our community of 1M+ users who develop and demonstrate their skills on CodeSignal

Building the Decoder Layer Foundation

Implementing the Three-Stage Forward Pass

i

Creating the Decoder Stack

Validating the Implementation

Let's validate our complete implementation with comprehensive tests that demonstrate both individual decoder functionality and full encoder-decoder interaction. First, we test the decoder layer in isolation: Pythondef test_decoder_layer(): """Test single Transformer Decoder Layer""" print("Testing Transformer Decoder Layer...") # Configuration batch_size = 2 src_seq_len = 8 tgt_seq_len = 6 d_model = 64 num_heads = 8 d_ff = 256 # Create sample inputs torch.manual_seed(42) decoder_input = torch.randn(batch_size, tgt_seq_len, d_model) encoder_output = torch.randn(batch_size, src_seq_len, d_model) # Create causal mask for decoder self-attention causal_mask = create_causal_mask(tgt_seq_len) # Create decoder layer decoder_layer = TransformerDecoderLayer(d_model, num_heads, d_ff) # Forward pass output, (self_attn_weights, cross_attn_weights) = decoder_layer( decoder_input, encoder_output, causal_mask) print(f"Decoder input shape: {decoder_input.shape}") print(f"Encoder output shape: {encoder_output.shape}") print(f"Decoder output shape: {output.shape}") print(f"Self-attention weights shape: {self_attn_weights.shape}") print(f"Cross-attention weights shape: {cross_attn_weights.shape}") # Verify dimensions assert output.shape == decoder_input.shape assert self_attn_weights.shape == (batch_size, num_heads, tgt_seq_len, tgt_seq_len) assert cross_attn_weights.shape == (batch_size, num_heads, tgt_seq_len, src_seq_len) return decoder_layerdef test_decoder_layer(): """Test single Transformer Decoder Layer""" print("Testing Transformer Decoder Layer...") # Configuration batch_size = 2 src_seq_len = 8 tgt_seq_len = 6 d_model = 64 num_heads = 8 d_ff = 256 # Create sample inputs torch.manual_seed(42) decoder_input = torch.randn(batch_size, tgt_seq_len, d_model) encoder_output = torch.randn(batch_size, src_seq_len, d_model) # Create causal mask for decoder self-attention causal_mask = create_causal_mask(tgt_seq_len) # Create decoder layer decoder_layer = TransformerDecoderLayer(d_model, num_heads, d_ff) # Forward pass output, (self_attn_weights, cross_attn_weights) = decoder_layer( decoder_input, encoder_output, causal_mask) print(f"Decoder input shape: {decoder_input.shape}") print(f"Encoder output shape: {encoder_output.shape}") print(f"Decoder output shape: {output.shape}") print(f"Self-attention weights shape: {self_attn_weights.shape}") print(f"Cross-attention weights shape: {cross_attn_weights.shape}") # Verify dimensions assert output.shape == decoder_input.shape assert self_attn_weights.shape == (batch_size, num_heads, tgt_seq_len, tgt_seq_len) assert cross_attn_weights.shape == (batch_size, num_heads, tgt_seq_len, src_seq_len) return decoder_layer Running test_decoder_layer produces this output: textTesting Transformer Decoder Layer... Decoder input shape: torch.Size([2, 6, 64]) Encoder output shape: torch.Size([2, 8, 64]) Decoder output shape: torch.Size([2, 6, 64]) Self-attention weights shape: torch.Size([2, 8, 6, 6]) Cross-attention weights shape: torch.Size([2, 8, 6, 8])Testing Transformer Decoder Layer... Decoder input shape: torch.Size([2, 6, 64]) Encoder output shape: torch.Size([2, 8, 64]) Decoder output shape: torch.Size([2, 6, 64]) Self-attention weights shape: torch.Size([2, 8, 6, 6]) Cross-attention weights shape: torch.Size([2, 8, 6, 8]) These results validate that our decoder correctly preserves target dimensions while processing encoder context. The self-attention weights show the expected 6 × 6 pattern for target positions, while cross-attention weights demonstrate proper source-target interaction with shape 6 × 8. Now, let's test the complete encoder-decoder pipeline: Pythondef test_encoder_decoder_interaction(): """Test complete encoder-decoder interaction""" print("Testing Encoder-Decoder Interaction...") # Configuration vocab_size = 1000 num_layers = 2 d_model = 64 num_heads = 8 d_ff = 256 src_seq_len = 10 tgt_seq_len = 8 batch_size = 2 # Create sample sequences torch.manual_seed(42) src_tokens = torch.randint(0, vocab_size, (batch_size, src_seq_len)) tgt_tokens = torch.randint(0, vocab_size, (batch_size, tgt_seq_len)) # Create components token_embedding = TokenEmbedding(vocab_size, d_model) positional_encoding = PositionalEncoding(d_model, dropout=0.1) encoder = TransformerEncoder(num_layers, d_model, num_heads, d_ff) decoder = TransformerDecoder(num_layers, d_model, num_heads, d_ff) # Process source sequence through encoder src_embeddings = positional_encoding(token_embedding(src_tokens)) encoder_output, _ = encoder(src_embeddings) # Process target sequence through decoder tgt_embeddings = positional_encoding(token_embedding(tgt_tokens)) causal_mask = create_causal_mask(tgt_seq_len) decoder_output, _ = decoder(tgt_embeddings, encoder_output, causal_mask) print(f"Source tokens shape: {src_tokens.shape}") print(f"Target tokens shape: {tgt_tokens.shape}") print(f"Encoder output shape: {encoder_output.shape}") print(f"Decoder output shape: {decoder_output.shape}") return decoder_outputdef test_encoder_decoder_interaction(): """Test complete encoder-decoder interaction""" print("Testing Encoder-Decoder Interaction...") # Configuration vocab_size = 1000 num_layers = 2 d_model = 64 num_heads = 8 d_ff = 256 src_seq_len = 10 tgt_seq_len = 8 batch_size = 2 # Create sample sequences torch.manual_seed(42) src_tokens = torch.randint(0, vocab_size, (batch_size, src_seq_len)) tgt_tokens = torch.randint(0, vocab_size, (batch_size, tgt_seq_len)) # Create components token_embedding = TokenEmbedding(vocab_size, d_model) positional_encoding = PositionalEncoding(d_model, dropout=0.1) encoder = TransformerEncoder(num_layers, d_model, num_heads, d_ff) decoder = TransformerDecoder(num_layers, d_model, num_heads, d_ff) # Process source sequence through encoder src_embeddings = positional_encoding(token_embedding(src_tokens)) encoder_output, _ = encoder(src_embeddings) # Process target sequence through decoder tgt_embeddings = positional_encoding(token_embedding(tgt_tokens)) causal_mask = create_causal_mask(tgt_seq_len) decoder_output, _ = decoder(tgt_embeddings, encoder_output, causal_mask) print(f"Source tokens shape: {src_tokens.shape}") print(f"Target tokens shape: {tgt_tokens.shape}") print(f"Encoder output shape: {encoder_output.shape}") print(f"Decoder output shape: {decoder_output.shape}") return decoder_output Running test_encoder_decoder_interaction produces this output: textTesting Encoder-Decoder Interaction... Source tokens shape: torch.Size([2, 10]) Target tokens shape: torch.Size([2, 8]) Encoder output shape: torch.Size([2, 10, 64]) Decoder output shape: torch.Size([2, 8, 64])Testing Encoder-Decoder Interaction... Source tokens shape: torch.Size([2, 10]) Target tokens shape: torch.Size([2, 8]) Encoder output shape: torch.Size([2, 10, 64]) Decoder output shape: torch.Size([2, 8, 64]) This demonstrates successful end-to-end processing through our complete Transformer implementation, validating that raw token sequences are properly transformed through embeddings, encoders, and decoders while respecting all architectural constraints.

The decoder's forward pass implements a carefully orchestrated three-stage process that transforms target sequences while incorporating encoder context:

Python

def forward(self, x, encoder_output, self_attention_mask=None, cross_attention_mask=None):
    """Forward pass through decoder layer"""
    # 1. Masked Multi-Head Self-Attention + Add & Norm
    self_attn_output, self_attn_weights = self.self_attention(x, x, x, self_attention_mask)
    x = self.add_norm1(x, self_attn_output)
    
    # 2. Encoder-Decoder Multi-Head Attention + Add & Norm
    cross_attn_output, cross_attn_weights = self.cross_attention(
        x, encoder_output, encoder_output, cross_attention_mask)
    x = self.add_norm2(x, cross_attn_output)
    
    # 3. Position-wise Feed-Forward + Add & Norm
    ff_output = self.feed_forward(x)
    x = self.add_norm3(x, ff_output)
    
    return x, (self_attn_weights, cross_attn_weights)

The first stage implements masked self-attention, where the self_attention_mask enforces the autoregressive property by preventing each position from attending to future positions. This mask is typically a lower triangular matrix where position $i$ can only attend to positions $j \leq i$ . The mechanism ensures that during training, even though the entire target sequence is available, the model learns to generate each token based only on previous context.

The second stage introduces cross-attention, the decoder's most distinctive feature. Here, queries come from the decoder state x (what we want to generate), while keys and values come from encoder_output (what we're conditioning on). This allows each decoder position to selectively focus on relevant parts of the source sequence. For example, when translating "The cat sat" to "Le chat s'assit," the decoder generating "chat" can focus strongly on "cat" in the encoder output.

The final stage processes these enriched representations through the feed-forward network, applying non-linear transformations to the combined self-context and source-context information. By returning both attention weight matrices, we enable analysis of the model's decision-making process.