Capítulo 2: Self-Attention e Multi-Head Attention

Introdução

Aproveitando nosso entendimento do mecanismo de atenção, este capítulo explora como a atenção é aplicada quando uma sequência presta atenção a si mesma (self-attention) e por que o uso de múltiplas cabeças de atenção em paralelo melhora drasticamente a capacidade do modelo.

Self-Attention: O Núcleo dos Transformers

O que é Self-Attention?

Em self-attention, uma sequência presta atenção a si mesma. Cada posição na entrada pode prestar atenção a todas as outras posições na mesma entrada.

Input:  "The cat sat on the mat"

Self-Attention: Cada palavra presta atenção a TODAS as palavras
├─ "The" presta atenção a ["The", "cat", "sat", "on", "the", "mat"]
├─ "cat" presta atenção a ["The", "cat", "sat", "on", "the", "mat"]
├─ "sat" presta atenção a ["The", "cat", "sat", "on", "the", "mat"]
└─ ... e assim por diante

Por que Self-Attention?

Self-attention permite:

1. Dependências de longo alcance: Conexões diretas entre quaisquer duas posições
2. Processamento paralelo: Todas as posições são computadas simultaneamente
3. Contexto dinâmico: Pesos diferentes para diferentes entradas

Formulação Matemática

Para self-attention, Q, K e V vêm todos da mesma entrada:

$$\text{SelfAttention}(X) = \text{softmax}\left(\frac{XW_Q(XW_K)^T}{\sqrt{d_k}}\right)XW_V$$

Implementação

import torch
import torch.nn as nn
import math

class SelfAttention(nn.Module):
    """
    Self-attention layer where Q, K, V all come from the same input.
    """
    def __init__(self, d_model, d_k=None):
        super().__init__()
        d_k = d_k or d_model
        
        self.W_Q = nn.Linear(d_model, d_k, bias=False)
        self.W_K = nn.Linear(d_model, d_k, bias=False)
        self.W_V = nn.Linear(d_model, d_k, bias=False)
        self.d_k = d_k
    
    def forward(self, x, mask=None):
        """
        x: Input tensor (batch, seq_len, d_model)
        mask: Optional attention mask
        """
        # All projections from the same input
        Q = self.W_Q(x)
        K = self.W_K(x)
        V = self.W_V(x)
        
        # Compute attention scores
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)
        
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))
        
        attention_weights = torch.softmax(scores, dim=-1)
        output = torch.matmul(attention_weights, V)
        
        return output, attention_weights

Comparando Tipos de Atenção

Self-Attention vs. Cross-Attention

Tipo	Fonte de Q	Fonte de K, V	Caso de Uso
Self-Attention	Entrada X	Entrada X	Encoder, Decoder
Cross-Attention	Decoder	Encoder	Modelos Encoder-Decoder

Exemplo de Cross-Attention

class CrossAttention(nn.Module):
    """
    Cross-attention where Q comes from one source, K and V from another.
    Used in encoder-decoder architectures.
    """
    def __init__(self, d_model, d_k=None):
        super().__init__()
        d_k = d_k or d_model
        
        self.W_Q = nn.Linear(d_model, d_k, bias=False)
        self.W_K = nn.Linear(d_model, d_k, bias=False)
        self.W_V = nn.Linear(d_model, d_k, bias=False)
        self.d_k = d_k
    
    def forward(self, query_input, key_value_input, mask=None):
        """
        query_input: Source for Q (e.g., decoder state)
        key_value_input: Source for K, V (e.g., encoder output)
        """
        Q = self.W_Q(query_input)
        K = self.W_K(key_value_input)
        V = self.W_V(key_value_input)
        
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)
        
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))
        
        attention_weights = torch.softmax(scores, dim=-1)
        output = torch.matmul(attention_weights, V)
        
        return output, attention_weights

Multi-Head Attention

A Motivação

Uma única cabeça de atenção só consegue capturar um tipo de relação. Multi-head attention executa múltiplas operações de atenção em paralelo, permitindo que o modelo preste atenção simultaneamente a informações de diferentes subespaços de representação.

Intuição

Cabeças diferentes podem aprender padrões diferentes:

• Cabeça 1: Relações sintáticas (sujeito-verbo)
• Cabeça 2: Relações posicionais (palavras próximas)
• Cabeça 3: Relações semânticas (sinônimos)
• Cabeça 4: Correferência (pronomes para substantivos)

Definição Matemática

$$\text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, ..., \text{head}_h)W^O$$

Onde cada cabeça é:

$$\text{head}_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V)$$

Dimensões dos Parâmetros

Para $h$ cabeças com dimensão do modelo $d_{model}$:

• Cada cabeça tem dimensão $d_k = d_v = d_{model} / h$
• O total de parâmetros permanece aproximadamente o mesmo que na atenção de uma única cabeça

# Exemplo: d_model = 512, h = 8
d_model = 512
num_heads = 8
d_k = d_model // num_heads  # 64 por cabeça

Implementação Completa

class MultiHeadAttention(nn.Module):
    """
    Multi-head attention mechanism.
    """
    def __init__(self, d_model, num_heads, dropout=0.1):
        super().__init__()
        assert d_model % num_heads == 0, "d_model must be divisible by num_heads"
        
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_heads
        
        # Linear projections for Q, K, V (all heads combined)
        self.W_Q = nn.Linear(d_model, d_model, bias=False)
        self.W_K = nn.Linear(d_model, d_model, bias=False)
        self.W_V = nn.Linear(d_model, d_model, bias=False)
        
        # Output projection
        self.W_O = nn.Linear(d_model, d_model, bias=False)
        
        self.dropout = nn.Dropout(dropout)
    
    def split_heads(self, x, batch_size):
        """
        Split the last dimension into (num_heads, d_k).
        Transpose to shape (batch, num_heads, seq_len, d_k)
        """
        x = x.view(batch_size, -1, self.num_heads, self.d_k)
        return x.transpose(1, 2)
    
    def forward(self, query, key, value, mask=None):
        batch_size = query.size(0)
        
        # Linear projections
        Q = self.W_Q(query)
        K = self.W_K(key)
        V = self.W_V(value)
        
        # Split into heads
        Q = self.split_heads(Q, batch_size)  # (batch, heads, seq, d_k)
        K = self.split_heads(K, batch_size)
        V = self.split_heads(V, batch_size)
        
        # Scaled dot-product attention for all heads in parallel
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)
        
        if mask is not None:
            # Expand mask for heads dimension
            if mask.dim() == 3:
                mask = mask.unsqueeze(1)
            scores = scores.masked_fill(mask == 0, float('-inf'))
        
        attention_weights = torch.softmax(scores, dim=-1)
        attention_weights = self.dropout(attention_weights)
        
        # Apply attention to values
        context = torch.matmul(attention_weights, V)
        
        # Concatenate heads
        context = context.transpose(1, 2).contiguous()
        context = context.view(batch_size, -1, self.d_model)
        
        # Final linear projection
        output = self.W_O(context)
        
        return output, attention_weights

Usando Multi-Head Attention

# Exemplo de uso
d_model = 512
num_heads = 8
seq_len = 100
batch_size = 32

mha = MultiHeadAttention(d_model, num_heads)

# Entrada (igual para self-attention)
x = torch.randn(batch_size, seq_len, d_model)

# Self-attention
output, attn_weights = mha(x, x, x)
print(f"Output shape: {output.shape}")  # (32, 100, 512)
print(f"Attention shape: {attn_weights.shape}")  # (32, 8, 100, 100)

Benefícios da Paralelização

Por que Transformers treinam rápido

Diferente dos RNNs, todas as posições em um Transformer podem ser computadas simultaneamente:

# RNN: Sequencial (lento)
for t in range(seq_len):
    h_t = rnn_cell(x[t], h_{t-1})

# Transformer: Paralelo (rápido)
output = multi_head_attention(x, x, x)  # Todas as posições de uma vez

Eficiência em GPU

Multi-head attention é especialmente eficiente em GPU:

# Todas as cabeças são computadas em uma única multiplicação de matrizes batelada
# Shape: (batch, heads, seq, d_k) @ (batch, heads, d_k, seq)
# → (batch, heads, seq, seq)
scores = torch.matmul(Q, K.transpose(-2, -1))  # Paralelo em todas as cabeças!

Visualizando Múltiplas Cabeças

def visualize_multi_head_attention(attention_weights, tokens, num_heads_to_show=4):
    """
    Visualize attention patterns from multiple heads.
    attention_weights: (batch, heads, seq, seq)
    """
    import matplotlib.pyplot as plt
    
    fig, axes = plt.subplots(1, num_heads_to_show, figsize=(4*num_heads_to_show, 4))
    
    for head_idx in range(num_heads_to_show):
        ax = axes[head_idx]
        attn = attention_weights[0, head_idx].detach().numpy()  # First batch item
        
        im = ax.imshow(attn, cmap='Blues')
        ax.set_xticks(range(len(tokens)))
        ax.set_yticks(range(len(tokens)))
        ax.set_xticklabels(tokens, rotation=45, ha='right')
        ax.set_yticklabels(tokens)
        ax.set_title(f'Head {head_idx + 1}')
    
    plt.tight_layout()
    plt.show()

Padrões de Atenção na Prática

Pesquisas mostram que cabeças diferentes aprendem padrões especializados:

Tipo de Padrão	O que captura	Exemplo
Diagonal	Mesma posição / identidade	Token presta atenção a si mesmo
Token anterior	Relações sequenciais	Padrões de modelos de linguagem
Primeiro token	Agregação global	Token CLS no BERT
Sintático	Estrutura gramatical	Concordância sujeito-verbo
Semântico	Relações de significado	Sinônimos, antônimos

Principais Pontos

Summary

1. Self-attention permite que cada posição preste atenção a todas as posições na mesma sequência
2. Cross-attention possibilita prestar atenção a uma sequência diferente (encoder-decoder)
3. Multi-head attention executa múltiplas operações de atenção em paralelo
4. Cabeças diferentes aprendem diferentes tipos de relações
5. Paralelização torna os Transformers eficientes em GPUs

Próximos Passos

No próximo capítulo, vamos explorar como os Transformers lidam com informações de posição através do positional encoding e como os tokens são representados por embeddings.

---

← Capítulo 1: Mecanismo de Atenção

Capítulo 3: Positional Encoding →