Working of Decoders in Transformers

A decoder in deep learning, especially in Transformer architectures, is the part of the model responsible for generating output sequences from encoded representations. In sequence-to-sequence tasks like machine translation, text summarization, or image captioning, the decoder takes the output from the encoder and converts it into a target language or format. It does this step-by-step, attending to both the encoded input and the already generated outputs.

Decoders in Transformers

• Autoregressive generation: Predicts one token at a time, using previously generated tokens.
• Masked self-attention: Prevents information leakage from future tokens using Masked self-attention.
• Encoder-decoder attention: Aligns output tokens with relevant parts of the input. Also uses parallel pre-processing.
• Stacked architecture: Typically has multiple identical layers (e.g., 6 in original Transformer).
• Positional encoding: Adds order information to input embeddings.
• Flexible output: Can be used for both classification and generation tasks.

Role of Decoders

• The encoder transforms the input sequence into a vector representation.
• The decoder takes this representation and produces the output sequence, attending to both: Itself, Encoder's output.

Working Principle

1. Input Embeddings are passed into the decoder with positional encodings.
2. Masked Self-Attention Layer ensures the model can’t “see” future tokens.
3. Encoder-Decoder Attention allows the decoder to focus on relevant input tokens.
4. Feedforward Layers refine representations.
5. A linear layer maps the final output to the vocabulary space.
6. Softmax provides a probability distribution over tokens for generation.

Components of Transformer Decoder

Each decoder layer contains:

1. Masked Multi-Head Self-Attention: Computes attention on previously generated tokens. Uses a causal mask to prevent future information leakage.
2. Multi-Head Encoder-Decoder Attention: Attends to encoder outputs.
3. Feedforward Network: Applies two linear transformations with a ReLU in between.
4. Layer Normalization and Residual Connections: Stabilize training and speed up convergence.
5. Positional Encoding: Adds token position information.

Example Use case

In English-to-French translation, the encoder processes the English sentence, and the decoder generates the French sentence one word at a time, using previously generated words and attention to the encoded sentence.

Mathematical Representation

• Masked Self-Attention:

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

• Encoder-Decoder Attention:

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

• Feedforward Network:

\text{FFN}(x) = \text{ReLU}(xW_1 + b_1)W_2 + b_2

Each decoder layer can be defined as:

x = \text{LayerNorm}(x + \text{MaskedSelfAttention}(x)) \\

x = \text{LayerNorm}(x + \text{CrossAttention}(x, E)) \\

x = \text{LayerNorm}(x + \text{FFN}(x))

Where,

• X: Input to the decoder
• E: Encoder output

Transformer Decoder Implementation

1. Imports

PyTorch and Math libraries are imported for model building and numerical operations.

import torch
import torch.nn as nn
import torch.nn.functional as F
import math

2. PositionalEncoding class

• Adds sinusoidal positional information to token embeddings.
• Helps the model understand token positions since transformers lack recurrence.
• Values are added to embeddings before input to the attention layers.

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=5000):
        super().__init__()
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len).unsqueeze(1).float()
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(1)  
        self.register_buffer('pe', pe)

    def forward(self, x):
        return x + self.pe[:x.size(0)]

3. TransformerDecoderLayer class

Defines one decoder layer containing:

• Masked multi-head self-attention to attend to previous tokens.
• Multi-head encoder-decoder attention to focus on encoder output.
• Feedforward network for non-linear transformation.
• Layer normalization and dropout for training stability.

class TransformerDecoderLayer(nn.Module):
    def __init__(self, d_model, nhead, dim_ff, dropout=0.1):
        super().__init__()
        self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
        self.multihead_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
        self.linear1 = nn.Linear(d_model, dim_ff)
        self.dropout = nn.Dropout(dropout)
        self.linear2 = nn.Linear(dim_ff, d_model)

        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.norm3 = nn.LayerNorm(d_model)
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)
        self.dropout3 = nn.Dropout(dropout)

    def forward(self, tgt, memory, tgt_mask=None, memory_mask=None):
        tgt2 = self.self_attn(tgt, tgt, tgt, attn_mask=tgt_mask)[0]
        tgt = self.norm1(tgt + self.dropout1(tgt2))

        tgt2 = self.multihead_attn(tgt, memory, memory, attn_mask=memory_mask)[0]
        tgt = self.norm2(tgt + self.dropout2(tgt2))

        tgt2 = self.linear2(F.relu(self.linear1(tgt)))
        tgt = self.norm3(tgt + self.dropout3(tgt2))
        return tgt

4. TransformerDecoder class

• Builds the complete decoder by stacking multiple decoder layers.
• Converts token indices to embeddings.
• Adds positional encodings.
• Applies a sequence of decoder layers.
• Uses a final linear layer to map outputs to vocabulary logits.

class TransformerDecoder(nn.Module):
    def __init__(self, num_layers, d_model, nhead, dim_ff, vocab_size, dropout=0.1, max_len=5000):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.pos_encoder = PositionalEncoding(d_model, max_len)
        self.layers = nn.ModuleList([
            TransformerDecoderLayer(d_model, nhead, dim_ff, dropout)
            for _ in range(num_layers)
        ])
        self.output_layer = nn.Linear(d_model, vocab_size)

    def forward(self, tgt, memory, tgt_mask=None, memory_mask=None):
        tgt = self.embedding(tgt)
        tgt = self.pos_encoder(tgt)
        for layer in self.layers:
            tgt = layer(tgt, memory, tgt_mask, memory_mask)
        return self.output_layer(tgt)

5. Hyperparameter setup

The hyperparameter setup includes embedding dimension size, attention heads, feedforward layer hidden size, decoder layers, output tokens, input shape for dummy test.

d_model = 128
nhead = 4
dim_ff = 512
num_layers = 2
vocab_size = 5000
seq_len = 10
batch_size = 4

6. Model instantiation

An instance of the TransformerDecoder is created using defined hyperparameters.

decoder = TransformerDecoder(num_layers, d_model, nhead, dim_ff, vocab_size)

Sample input:

• tgt: random integers simulating target token indices.
• memory: random tensor simulating encoder output.

tgt = torch.randint(0, vocab_size, (seq_len, batch_size))  
memory = torch.rand(seq_len, batch_size, d_model)

Forward pass:

• Inputs are passed through the decoder to get output logits.
• Output shape is (seq_len, batch_size, vocab_size), suitable for classification of each token position over the vocabulary.

out = decoder(tgt, memory)
print(out.shape)  
print(out)

Output

You can download the source code from here.

Applications

• Machine Translation
• Text Summarization
• Speech-to-Text systems
• Code generation by LLMs