Capítulo 1: Mecanismo de Atenção

Introdução

O mecanismo de atenção é a inovação central que torna os Transformers possíveis. Ele permite que os modelos foquem dinamicamente nas partes relevantes da entrada ao produzir cada saída, superando as limitações de vetores de contexto de tamanho fixo.

A Intuição por Trás da Atenção

Uma Analogia Humana

Pense em como você lê uma frase complexa:

"O animal não atravessou a rua porque ele estava muito cansado."

Para entender a que "ele" se refere, você mentalmente presta atenção em "o animal" e não em "a rua". É exatamente isso que a atenção faz computacionalmente.

A Ideia Central

A atenção calcula uma soma ponderada de valores, onde os pesos são determinados por quão relevante cada valor é para uma determinada query:

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

Query, Key, Value: As Três Projeções

Entendendo Q, K, V

Componente	Papel	Analogia
Query (Q)	O que estamos procurando	Uma busca
Key (K)	O que cada posição oferece	Entradas de um índice
Value (V)	O conteúdo real	Os documentos

Definição Matemática

Dada uma entrada $X \in \mathbb{R}^{n \times d}$ (sequência de n tokens, cada um com d dimensões):

$$Q = XW_Q, \quad K = XW_K, \quad V = XW_V$$

Onde:

• $W_Q \in \mathbb{R}^{d \times d_k}$
• $W_K \in \mathbb{R}^{d \times d_k}$
• $W_V \in \mathbb{R}^{d \times d_v}$

Implementação

import torch
import torch.nn as nn
import math

class AttentionProjections(nn.Module):
    def __init__(self, d_model, d_k, d_v):
        super().__init__()
        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_v, bias=False)
    
    def forward(self, x):
        """
        x: (batch, seq_len, d_model)
        Returns Q, K, V projections
        """
        Q = self.W_Q(x)  # (batch, seq_len, d_k)
        K = self.W_K(x)  # (batch, seq_len, d_k)
        V = self.W_V(x)  # (batch, seq_len, d_v)
        return Q, K, V

Atenção por Produto Escalar Escalado

Cálculo Passo a Passo

Passo 1: Calcule os Scores de Atenção

$$\text{scores} = QK^T$$

Isso nos dá uma matriz onde a entrada $(i, j)$ representa o quanto a posição $i$ deve prestar atenção à posição $j$.

# Q: (batch, seq_len, d_k)
# K: (batch, seq_len, d_k)
scores = torch.matmul(Q, K.transpose(-2, -1))
# scores: (batch, seq_len, seq_len)

Passo 2: Escale os Scores

$$\text{scaled\_scores} = \frac{QK^T}{\sqrt{d_k}}$$

Why Scale?

Sem o escalonamento, para grandes $d_k$, os produtos escalares podem crescer muito em magnitude, levando o softmax para regiões com gradientes extremamente pequenos. Escalar por $\sqrt{d_k}$ mantém a variância estável.

d_k = Q.size(-1)
scaled_scores = scores / math.sqrt(d_k)

Passo 3: Aplique o Softmax

$$\text{attention\_weights} = \text{softmax}(\text{scaled\_scores})$$

attention_weights = torch.softmax(scaled_scores, dim=-1)
# Each row sums to 1

Passo 4: Soma Ponderada dos Valores

$$\text{output} = \text{attention\_weights} \cdot V$$

output = torch.matmul(attention_weights, V)
# output: (batch, seq_len, d_v)

Implementação Completa

def scaled_dot_product_attention(Q, K, V, mask=None):
    """
    Compute scaled dot-product attention.
    
    Args:
        Q: Query tensor (batch, seq_len, d_k)
        K: Key tensor (batch, seq_len, d_k)
        V: Value tensor (batch, seq_len, d_v)
        mask: Optional mask tensor
    
    Returns:
        output: Attention output (batch, seq_len, d_v)
        attention_weights: Attention weights (batch, seq_len, seq_len)
    """
    d_k = Q.size(-1)
    
    # Step 1 & 2: Compute scaled scores
    scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k)
    
    # Optional: Apply mask (for causal attention)
    if mask is not None:
        scores = scores.masked_fill(mask == 0, float('-inf'))
    
    # Step 3: Softmax
    attention_weights = torch.softmax(scores, dim=-1)
    
    # Step 4: Weighted sum
    output = torch.matmul(attention_weights, V)
    
    return output, attention_weights

Visualizando a Atenção

Matriz de Atenção

Para uma frase como "O gato sentou":

         O    gato  sentou
O     [ 0.7   0.2   0.1 ]
gato  [ 0.3   0.5   0.2 ]
sentou[ 0.2   0.4   0.4 ]

Cada linha mostra o quanto aquela posição presta atenção a todas as outras posições.

Código para Visualizar

import matplotlib.pyplot as plt
import seaborn as sns

def visualize_attention(attention_weights, tokens):
    """
    Visualize attention weights as a heatmap.
    """
    plt.figure(figsize=(10, 8))
    sns.heatmap(
        attention_weights.detach().numpy(),
        xticklabels=tokens,
        yticklabels=tokens,
        cmap='Blues',
        annot=True,
        fmt='.2f'
    )
    plt.xlabel('Keys')
    plt.ylabel('Queries')
    plt.title('Attention Weights')
    plt.show()

Atenção Causal (Mascarada)

Para modelos de linguagem autoregressivos, cada posição deve prestar atenção apenas às posições anteriores.

A Máscara Causal

def create_causal_mask(seq_len):
    """
    Create a causal mask for autoregressive attention.
    Position i can only attend to positions <= i.
    """
    mask = torch.tril(torch.ones(seq_len, seq_len))
    return mask

# Example for seq_len=4:
# [[1, 0, 0, 0],
#  [1, 1, 0, 0],
#  [1, 1, 1, 0],
#  [1, 1, 1, 1]]

Por Que Usar Máscara Causal?

Durante o treinamento, queremos que o modelo preveja o próximo token apenas com base nos tokens anteriores, não nos futuros (que ele não terá no momento da inferência).

def causal_attention(Q, K, V):
    """
    Attention with causal masking.
    """
    seq_len = Q.size(1)
    mask = create_causal_mask(seq_len).to(Q.device)
    return scaled_dot_product_attention(Q, K, V, mask)

Complexidade da Atenção

Complexidade de Tempo

$$O(n^2 \cdot d)$$

Onde:

• $n$ = comprimento da sequência
• $d$ = dimensão do modelo

Essa complexidade quadrática no comprimento da sequência é o motivo pelo qual modelos de contexto longo são desafiadores.

Complexidade de Memória

$$O(n^2)$$

Para a própria matriz de atenção (armazenando todos os scores pareados).

Técnicas de Otimização

Diversas abordagens tratam a complexidade quadrática da atenção:

1. Flash Attention: Implementação eficiente em memória
2. Linear Attention: Aproxima a atenção em O(n)
3. Sparse Attention: Atenção apenas para um subconjunto de posições
4. Sliding Window: Atenção local com tokens globais

Juntando Tudo

class Attention(nn.Module):
    """
    Complete attention module.
    """
    def __init__(self, d_model, d_k=None, d_v=None, causal=False):
        super().__init__()
        d_k = d_k or d_model
        d_v = d_v 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_v, bias=False)
        self.W_O = nn.Linear(d_v, d_model, bias=False)
        
        self.causal = causal
        self.d_k = d_k
    
    def forward(self, x, mask=None):
        batch_size, seq_len, _ = x.size()
        
        # Project to Q, K, V
        Q = self.W_Q(x)
        K = self.W_K(x)
        V = self.W_V(x)
        
        # Compute attention
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)
        
        # Apply causal mask if needed
        if self.causal:
            causal_mask = torch.tril(torch.ones(seq_len, seq_len, device=x.device))
            scores = scores.masked_fill(causal_mask == 0, float('-inf'))
        
        # Apply additional mask if provided
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))
        
        # Softmax and weighted sum
        attn_weights = torch.softmax(scores, dim=-1)
        output = torch.matmul(attn_weights, V)
        
        # Output projection
        output = self.W_O(output)
        
        return output, attn_weights

Principais Pontos

Summary

1. Atenção calcula somas ponderadas com base na relevância
2. Q, K, V são três projeções da entrada com papéis distintos
3. Escalonamento por $\sqrt{d_k}$ estabiliza o treinamento
4. Máscara causal impede olhar para tokens futuros
5. Complexidade é $O(n^2)$ no comprimento da sequência

Próximos Passos

No próximo capítulo, vamos explorar como a atenção é usada dentro de uma única sequência (self-attention) e como múltiplos mecanismos de atenção são combinados (multi-head attention).

---

← Visão Geral do Módulo

Capítulo 2: Self & Multi-Head Attention →