xbe.at

NLP and transformers / Transformer architecture

Why Transformers Don't Have Vanishing Gradients

11 min read, 14 sections

Contents
  1. Understanding Transformers vs RNNs
  2. The Vanishing Gradient Problem in RNNs
  3. Transformer Architecture Overview
  4. Self-Attention Mechanism
  5. Parallel Processing in Transformers
  6. Positional Encoding
  7. Residual Connections and Layer Normalization
  8. Multi-Head Attention
  9. Gradient Flow in Transformers
  10. Comparison of Gradient Flow: RNN vs Transformer
  11. Real-life Example: Machine Translation
  12. Real-life Example: Text Summarization
  13. Conclusion
  14. Additional Resources

Understanding Transformers vs RNNs

Transformers and Recurrent Neural Networks (RNNs) are both powerful architectures for sequence processing tasks. However, Transformers have gained popularity due to their ability to handle long-range dependencies more effectively. One key advantage of Transformers is their resistance to the vanishing gradient problem, which often plagues RNNs. Let's explore why this is the case.

import numpy as np
import matplotlib.pyplot as plt

def plot_gradient_flow(rnn_gradients, transformer_gradients):
    plt.figure(figsize=(10, 6))
    plt.plot(rnn_gradients, label='RNN')
    plt.plot(transformer_gradients, label='Transformer')
    plt.title('Gradient Flow in RNN vs Transformer')
    plt.xlabel('Time Steps')
    plt.ylabel('Gradient Magnitude')
    plt.legend()
    plt.show()

# Simulate gradient flow
time_steps = 100
rnn_gradients = np.exp(-np.arange(time_steps) * 0.1)
transformer_gradients = np.ones(time_steps)

plot_gradient_flow(rnn_gradients, transformer_gradients)

The Vanishing Gradient Problem in RNNs

RNNs process sequences step by step, maintaining a hidden state that carries information from previous time steps. However, as the sequence length increases, the influence of early inputs on later outputs diminishes exponentially. This phenomenon is known as the vanishing gradient problem.

import torch
import torch.nn as nn

class SimpleRNN(nn.Module):
    def __init__(self, input_size, hidden_size):
        super(SimpleRNN, self).__init__()
        self.rnn = nn.RNN(input_size, hidden_size, batch_first=True)
    
    def forward(self, x):
        output, _ = self.rnn(x)
        return output

# Create a simple RNN
input_size = 10
hidden_size = 20
seq_length = 100
batch_size = 1

rnn = SimpleRNN(input_size, hidden_size)
input_tensor = torch.randn(batch_size, seq_length, input_size)

# Forward pass
output = rnn(input_tensor)

# Calculate gradients
loss = output.sum()
loss.backward()

# Print gradient norms for different time steps
for i in range(0, seq_length, 10):
    print(f"Time step {i}: Gradient norm = {rnn.rnn.weight_ih_l0.grad[0, i].norm().item():.6f}")

Transformer Architecture Overview

Transformers, introduced in the "Attention Is All You Need" paper, take a different approach to sequence processing. They rely on self-attention mechanisms to capture relationships between all positions in a sequence, regardless of their distance.

import torch
import torch.nn as nn

class TransformerBlock(nn.Module):
    def __init__(self, d_model, nhead):
        super(TransformerBlock, self).__init__()
        self.attention = nn.MultiheadAttention(d_model, nhead)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.feed_forward = nn.Sequential(
            nn.Linear(d_model, d_model * 4),
            nn.ReLU(),
            nn.Linear(d_model * 4, d_model)
        )
    
    def forward(self, x):
        attn_output, _ = self.attention(x, x, x)
        x = self.norm1(x + attn_output)
        ff_output = self.feed_forward(x)
        x = self.norm2(x + ff_output)
        return x

# Create a Transformer block
d_model = 512
nhead = 8
transformer = TransformerBlock(d_model, nhead)

# Example input
seq_length = 100
batch_size = 1
x = torch.randn(seq_length, batch_size, d_model)

# Forward pass
output = transformer(x)
print(f"Input shape: {x.shape}")
print(f"Output shape: {output.shape}")

Self-Attention Mechanism

The key to Transformers' success in avoiding vanishing gradients lies in their self-attention mechanism. This mechanism allows each position in the sequence to attend to all other positions directly, creating short paths for gradient flow.

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

class SelfAttention(nn.Module):
    def __init__(self, embed_size, heads):
        super(SelfAttention, self).__init__()
        self.embed_size = embed_size
        self.heads = heads
        self.head_dim = embed_size // heads

        self.values = nn.Linear(self.head_dim, self.head_dim, bias=False)
        self.keys = nn.Linear(self.head_dim, self.head_dim, bias=False)
        self.queries = nn.Linear(self.head_dim, self.head_dim, bias=False)
        self.fc_out = nn.Linear(heads * self.head_dim, embed_size)

    def forward(self, values, keys, query):
        N = query.shape[0]
        value_len, key_len, query_len = values.shape[1], keys.shape[1], query.shape[1]

        # Split the embedding into self.heads pieces
        values = values.reshape(N, value_len, self.heads, self.head_dim)
        keys = keys.reshape(N, key_len, self.heads, self.head_dim)
        query = query.reshape(N, query_len, self.heads, self.head_dim)

        values = self.values(values)
        keys = self.keys(keys)
        queries = self.queries(query)

        # Scaled dot-product attention
        energy = torch.einsum("nqhd,nkhd->nhqk", [queries, keys])
        attention = F.softmax(energy / (self.embed_size ** (1/2)), dim=3)
        out = torch.einsum("nhql,nlhd->nqhd", [attention, values]).reshape(
            N, query_len, self.heads * self.head_dim
        )

        out = self.fc_out(out)
        return out

# Example usage
embed_size = 256
heads = 8
self_attention = SelfAttention(embed_size, heads)

# Sample input
x = torch.randn(32, 10, embed_size)  # (batch_size, seq_len, embed_size)
output = self_attention(x, x, x)

print(f"Input shape: {x.shape}")
print(f"Output shape: {output.shape}")

Parallel Processing in Transformers

Unlike RNNs, which process sequences sequentially, Transformers can process entire sequences in parallel. This parallel processing allows for efficient computation and helps maintain gradient flow across long sequences.

import torch
import torch.nn as nn

class ParallelTransformer(nn.Module):
    def __init__(self, d_model, nhead, num_layers):
        super(ParallelTransformer, self).__init__()
        self.transformer_encoder = nn.TransformerEncoder(
            nn.TransformerEncoderLayer(d_model, nhead),
            num_layers
        )
    
    def forward(self, src):
        return self.transformer_encoder(src)

# Create a parallel Transformer
d_model = 512
nhead = 8
num_layers = 6
transformer = ParallelTransformer(d_model, nhead, num_layers)

# Example input
seq_length = 100
batch_size = 32
src = torch.randn(seq_length, batch_size, d_model)

# Measure processing time
import time
start_time = time.time()
output = transformer(src)
end_time = time.time()

print(f"Input shape: {src.shape}")
print(f"Output shape: {output.shape}")
print(f"Processing time: {end_time - start_time:.4f} seconds")

Positional Encoding

To incorporate sequence order information, Transformers use positional encodings. These encodings are added to the input embeddings, allowing the model to distinguish between different positions in the sequence without introducing sequential dependencies.

import torch
import numpy as np
import matplotlib.pyplot as plt

def positional_encoding(seq_len, d_model):
    position = np.arange(seq_len)[:, np.newaxis]
    div_term = np.exp(np.arange(0, d_model, 2) * -(np.log(10000.0) / d_model))
    pos_encoding = np.zeros((seq_len, d_model))
    pos_encoding[:, 0::2] = np.sin(position * div_term)
    pos_encoding[:, 1::2] = np.cos(position * div_term)
    return torch.FloatTensor(pos_encoding)

# Generate positional encodings
seq_len = 100
d_model = 512
pos_encoding = positional_encoding(seq_len, d_model)

# Visualize the positional encodings
plt.figure(figsize=(10, 6))
plt.imshow(pos_encoding.numpy(), cmap='viridis', aspect='auto')
plt.colorbar()
plt.title('Positional Encoding')
plt.xlabel('Embedding Dimension')
plt.ylabel('Sequence Position')
plt.show()

# Print a sample of positional encoding values
print("Sample positional encoding values:")
print(pos_encoding[0, :10])  # First 10 values for the first position
print(pos_encoding[10, :10])  # First 10 values for the 11th position

Residual Connections and Layer Normalization

Transformers employ residual connections and layer normalization to facilitate gradient flow and stabilize training. These techniques help maintain the model's performance across deep architectures.

import torch
import torch.nn as nn

class TransformerBlock(nn.Module):
    def __init__(self, d_model, nhead):
        super(TransformerBlock, self).__init__()
        self.attention = nn.MultiheadAttention(d_model, nhead)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.feed_forward = nn.Sequential(
            nn.Linear(d_model, d_model * 4),
            nn.ReLU(),
            nn.Linear(d_model * 4, d_model)
        )
    
    def forward(self, x):
        # Self-attention with residual connection and layer normalization
        attn_output, _ = self.attention(x, x, x)
        x = self.norm1(x + attn_output)
        
        # Feed-forward with residual connection and layer normalization
        ff_output = self.feed_forward(x)
        x = self.norm2(x + ff_output)
        
        return x

# Create a Transformer block
d_model = 512
nhead = 8
transformer = TransformerBlock(d_model, nhead)

# Example input
seq_length = 100
batch_size = 1
x = torch.randn(seq_length, batch_size, d_model)

# Forward pass
output = transformer(x)

print(f"Input shape: {x.shape}")
print(f"Output shape: {output.shape}")
print(f"Layer norm 1 parameters: {transformer.norm1.weight.shape}")
print(f"Layer norm 2 parameters: {transformer.norm2.weight.shape}")

Multi-Head Attention

Multi-head attention allows Transformers to capture different types of relationships between sequence elements simultaneously. This mechanism enhances the model's ability to process complex dependencies efficiently.

import torch
import torch.nn as nn

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super(MultiHeadAttention, self).__init__()
        self.d_model = d_model
        self.num_heads = num_heads
        self.head_dim = d_model // num_heads
        
        self.query = nn.Linear(d_model, d_model)
        self.key = nn.Linear(d_model, d_model)
        self.value = nn.Linear(d_model, d_model)
        self.out = nn.Linear(d_model, d_model)
        
    def forward(self, query, key, value):
        N = query.shape[0]
        
        # Linear projections
        Q = self.query(query).view(N, -1, self.num_heads, self.head_dim).transpose(1, 2)
        K = self.key(key).view(N, -1, self.num_heads, self.head_dim).transpose(1, 2)
        V = self.value(value).view(N, -1, self.num_heads, self.head_dim).transpose(1, 2)
        
        # Scaled dot-product attention
        scores = torch.matmul(Q, K.transpose(-2, -1)) / torch.sqrt(torch.tensor(self.head_dim, dtype=torch.float32))
        attention = torch.softmax(scores, dim=-1)
        
        # Apply attention to values
        out = torch.matmul(attention, V).transpose(1, 2).reshape(N, -1, self.d_model)
        out = self.out(out)
        
        return out

# Create a Multi-Head Attention module
d_model = 512
num_heads = 8
mha = MultiHeadAttention(d_model, num_heads)

# Example input
seq_length = 100
batch_size = 32
x = torch.randn(batch_size, seq_length, d_model)

# Forward pass
output = mha(x, x, x)

print(f"Input shape: {x.shape}")
print(f"Output shape: {output.shape}")
print(f"Number of attention heads: {num_heads}")
print(f"Dimension per head: {d_model // num_heads}")

Gradient Flow in Transformers

The direct connections between all positions in the sequence, coupled with residual connections and layer normalization, allow gradients to flow more easily through the network. This design mitigates the vanishing gradient problem observed in RNNs.

import torch
import torch.nn as nn
import matplotlib.pyplot as plt

class SimpleTransformer(nn.Module):
    def __init__(self, d_model, nhead, num_layers):
        super(SimpleTransformer, self).__init__()
        self.transformer = nn.TransformerEncoder(
            nn.TransformerEncoderLayer(d_model, nhead),
            num_layers
        )
    
    def forward(self, x):
        return self.transformer(x)

# Create a simple Transformer
d_model, nhead, num_layers = 512, 8, 6
transformer = SimpleTransformer(d_model, nhead, num_layers)

# Example input
seq_length, batch_size = 100, 1
x = torch.randn(seq_length, batch_size, d_model, requires_grad=True)

# Forward pass and backward pass
output = transformer(x)
loss = output.sum()
loss.backward()

# Collect and plot gradient norms
grad_norms = [param.grad.norm().item() for param in transformer.parameters() if param.grad is not None]
plt.figure(figsize=(10, 6))
plt.bar(range(len(grad_norms)), grad_norms)
plt.title('Gradient Norms in Transformer Layers')
plt.xlabel('Layer')
plt.ylabel('Gradient Norm')
plt.show()

print(f"Number of layers with gradients: {len(grad_norms)}")
print(f"Mean gradient norm: {sum(grad_norms) / len(grad_norms):.4f}")

Comparison of Gradient Flow: RNN vs Transformer

To visualize the difference in gradient flow between RNNs and Transformers, we can compare the gradient magnitudes across time steps for both architectures.

import torch
import torch.nn as nn
import matplotlib.pyplot as plt

class SimpleRNN(nn.Module):
    def __init__(self, input_size, hidden_size):
        super(SimpleRNN, self).__init__()
        self.rnn = nn.RNN(input_size, hidden_size, batch_first=True)
    
    def forward(self, x):
        return self.rnn(x)[0]

class SimpleTransformer(nn.Module):
    def __init__(self, d_model, nhead):
        super(SimpleTransformer, self).__init__()
        self.transformer = nn.TransformerEncoder(
            nn.TransformerEncoderLayer(d_model, nhead),
            num_layers=1
        )
    
    def forward(self, x):
        return self.transformer(x)

# Parameters
input_size = hidden_size = d_model = 64
nhead = 4
seq_length = 100
batch_size = 1

# Create models
rnn = SimpleRNN(input_size, hidden_size)
transformer = SimpleTransformer(d_model, nhead)

# Input tensor
x = torch.randn(batch_size, seq_length, input_size, requires_grad=True)

# RNN forward and backward pass
rnn_output = rnn(x)
rnn_loss = rnn_output.sum()
rnn_loss.backward()

# Transformer forward and backward pass
transformer_output = transformer(x.transpose(0, 1))
transformer_loss = transformer_output.sum()
transformer_loss.backward()

# Collect gradients
rnn_grads = x.grad.abs().mean(dim=(0, 2)).detach().numpy()
transformer_grads = x.grad.abs().mean(dim=(0, 2)).detach().numpy()

# Plot gradients
plt.figure(figsize=(10, 6))
plt.plot(rnn_grads, label='RNN')
plt.plot(transformer_grads, label='Transformer')
plt.title('Gradient Magnitudes: RNN vs Transformer')
plt.xlabel('Time Step')
plt.ylabel('Average Gradient Magnitude')
plt.legend()
plt.show()

Real-life Example: Machine Translation

Machine translation is one of the most common applications of Transformer models. Let's look at a simplified example of how a Transformer-based translation model might be structured.

import torch
import torch.nn as nn

class TranslationTransformer(nn.Module):
    def __init__(self, src_vocab_size, tgt_vocab_size, d_model, nhead, num_encoder_layers, num_decoder_layers):
        super(TranslationTransformer, self).__init__()
        self.src_embedding = nn.Embedding(src_vocab_size, d_model)
        self.tgt_embedding = nn.Embedding(tgt_vocab_size, d_model)
        self.transformer = nn.Transformer(d_model, nhead, num_encoder_layers, num_decoder_layers)
        self.fc_out = nn.Linear(d_model, tgt_vocab_size)
    
    def forward(self, src, tgt):
        src_embed = self.src_embedding(src)
        tgt_embed = self.tgt_embedding(tgt)
        out = self.transformer(src_embed, tgt_embed)
        return self.fc_out(out)

# Example usage
src_vocab_size = 10000
tgt_vocab_size = 12000
d_model = 512
nhead = 8
num_encoder_layers = 6
num_decoder_layers = 6

model = TranslationTransformer(src_vocab_size, tgt_vocab_size, d_model, nhead, num_encoder_layers, num_decoder_layers)

# Dummy input (batch_size=2, seq_len=10)
src = torch.randint(0, src_vocab_size, (10, 2))
tgt = torch.randint(0, tgt_vocab_size, (12, 2))

output = model(src, tgt)
print(f"Output shape: {output.shape}")

Real-life Example: Text Summarization

Text summarization is another task where Transformers excel. Here's a simplified implementation of a Transformer-based summarization model.

import torch
import torch.nn as nn

class SummarizationTransformer(nn.Module):
    def __init__(self, vocab_size, d_model, nhead, num_encoder_layers, num_decoder_layers):
        super(SummarizationTransformer, self).__init__()
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.transformer = nn.Transformer(d_model, nhead, num_encoder_layers, num_decoder_layers)
        self.fc_out = nn.Linear(d_model, vocab_size)
    
    def forward(self, src, tgt):
        src_embed = self.embedding(src)
        tgt_embed = self.embedding(tgt)
        out = self.transformer(src_embed, tgt_embed)
        return self.fc_out(out)

# Example usage
vocab_size = 50000
d_model = 512
nhead = 8
num_encoder_layers = 6
num_decoder_layers = 6

model = SummarizationTransformer(vocab_size, d_model, nhead, num_encoder_layers, num_decoder_layers)

# Dummy input (batch_size=2, src_seq_len=100, tgt_seq_len=20)
src = torch.randint(0, vocab_size, (100, 2))
tgt = torch.randint(0, vocab_size, (20, 2))

output = model(src, tgt)
print(f"Output shape: {output.shape}")

Conclusion

Transformers have successfully addressed the vanishing gradient problem that plagued RNNs, enabling more effective processing of long sequences. Key features contributing to this success include:

  1. Self-attention mechanism
  2. Parallel processing of sequences
  3. Residual connections and layer normalization
  4. Positional encodings

These architectural innovations have led to significant improvements in various natural language processing tasks and have inspired numerous variations and extensions of the original Transformer model.

Additional Resources

For those interested in diving deeper into Transformer architecture and its variations, the following resources are recommended:

  1. "Attention Is All You Need" (Original Transformer paper): Vaswani, A., et al. (2017). arXiv:1706.03762 [cs.CL] https://arxiv.org/abs/1706.03762
  2. "BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding": Devlin, J., et al. (2018). arXiv:1810.04805 [cs.CL] https://arxiv.org/abs/1810.04805
  3. "Language Models are Few-Shot Learners" (GPT-3 paper): Brown, T., et al. (2020). arXiv:2005.14165 [cs.CL] https://arxiv.org/abs/2005.14165

These papers provide in-depth explanations of Transformer-based models and their applications in various natural language processing tasks.

More in Transformer architecture

  1. Transformers The Paper That Changed Machine Learning7 min

    Transformers have revolutionized natural language processing and machine learning. Introduced in the paper "Attention Is All You Need" by Vaswani et al....

  2. Transformers The Dominant AI Architecture for Sequential Tasks8 min

    Recurrent Neural Networks represent a fundamental architecture in deep learning designed to process sequential data by maintaining an internal state (memory)...

  3. Transformers Revolutionizing Natural Language Processing16 min

    The first step in implementing a Transformer is converting text into numerical representations. We'll create a basic tokenizer and embedder that transforms...

  4. Transformers for Regression! Expressive Enough!11 min

    Transformers, originally designed for natural language processing tasks, have shown remarkable success in various domains. This presentation explores their...

  5. Transformers Architecture Explained7 min

    A Transformer is a neural network architecture designed for processing sequential data, particularly in natural language processing tasks. It relies solely...

  6. Transformers and the NLP Revolution10 min

    The Transformer architecture, introduced in 2017, revolutionized Natural Language Processing (NLP). It addressed limitations of previous sequential models...