7 O Mecanismo de Atenção (Self-Attention)

O mecanismo de Self-Attention (Auto-Atenção) é o componente fundamental que distingue a arquitetura Transformer das redes recorrentes (RNNs) e convolucionais (CNNs) anteriores. Enquanto as RNNs processam sequências passo a passo (o que dificulta a paralelização e a retenção de dependências de longo prazo), o Self-Attention permite que o modelo relacione cada palavra na sequência de entrada com todas as outras palavras simultaneamente.

Este capítulo disseca a Scaled Dot-Product Attention, a fórmula matemática que rege essa operação, e implementa o mascaramento necessário para modelos de linguagem generativos (como a família GPT).

7.1 1. Conceitos Fundamentais: Query, Key e Value

Para calcular a atenção, projetamos a entrada (embeddings) em três vetores distintos para cada token. Esta abordagem é inspirada em sistemas de recuperação de informação:

Query ($Q$ - Consulta): O que o token atual está procurando? (Representa o foco atual).
Key ($K$ - Chave): O que o token atual oferece? (Serve como um identificador para correspondência).
Value ($V$ - Valor): Qual é o conteúdo real da informação? (A informação que será agregada se a chave corresponder à consulta).

Em um mecanismo de Self-Attention, $Q$, $K$ e $V$ originam-se da mesma fonte (a saída da camada anterior), multiplicados por matrizes de pesos aprendíveis ($W^Q, W^K, W^V$).

7.2 2. A Fórmula Matemática

A atenção é calculada como uma soma ponderada dos valores ($V$), onde o peso atribuído a cada valor é determinado pela compatibilidade entre a consulta ($Q$) e a chave correspondente ($K$).

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

7.2.1 Passo a Passo do Cálculo:

MatMul ($QK^T$): Calcula o produto escalar entre a Query e a Key transposta. Isso resulta em uma Matriz de Scores (ou afinidade). Se o produto escalar é alto, os vetores estão alinhados, indicando alta relevância entre os tokens.
Scale ($\frac{1}{\sqrt{d_k}}$): Dividimos os scores pela raiz quadrada da dimensão das chaves ($d_k$).
- Por que escalar? Sem isso, para valores altos de $d_k$, o produto escalar cresce em magnitude, empurrando a função Softmax para regiões onde os gradientes são extremamente pequenos (vanishing gradients), dificultando o treinamento.
Mask (Opcional): Aplica-se uma máscara (substituindo valores por $-\infty$) para impedir que o modelo “veja” certos tokens (ex: tokens futuros ou padding).
Softmax: Converte os scores em probabilidades (pesos de atenção) que somam 1.
**MatMul ($ V$):** Multiplica-se os pesos pelos Valores ($V$). Isso filtra as informações irrelevantes e preserva as relevantes.

7.3 3. Fluxo de Dados Visual

O diagrama abaixo ilustra o fluxo de tensores dentro do bloco de atenção escalar.

graph TD
    subgraph Inputs
    Q[Query (Q)]
    K[Key (K)]
    V[Value (V)]
    end

    Q --> MatMul1[MatMul (Q x K^T)]
    K --> MatMul1
    
    MatMul1 --> Scale[Scale (1 / sqrt(d_k))]
    
    Scale --> MaskNode{Aplicar Máscara?}
    Mask[Máscara Causal / Padding] -.-> MaskNode
    
    MaskNode -- Sim --> MaskFill[Fill com -inf]
    MaskNode -- Não --> Softmax
    
    MaskFill --> Softmax[Softmax]
    
    Softmax --> MatMul2[MatMul (Weights x V)]
    V --> MatMul2
    
    MatMul2 --> Output[Output Contextualizado]

    style MatMul1 fill:#f9f,stroke:#333,stroke-width:2px
    style Softmax fill:#ff9,stroke:#333,stroke-width:2px
    style MatMul2 fill:#f9f,stroke:#333,stroke-width:2px

7.4 4. Mascaramento (Masking)

Existem dois tipos críticos de máscaras no Transformer:

Padding Mask: Ignora tokens de preenchimento (pad tokens) usados para igualar o tamanho das sentenças em um batch.
Look-ahead (Causal) Mask: Essencial para decodificadores (como no GPT). Ao prever o token na posição $t$, o modelo não pode ter acesso aos tokens nas posições $t+1, t+2, \dots$.

A máscara causal é uma matriz triangular superior (excluindo a diagonal) preenchida com $-\infty$. Quando somada aos scores antes do Softmax: \[ \text{softmax}(x + (-\infty)) = 0 \] Isso garante que a atenção para tokens futuros seja zero.

7.5 5. Implementação em PyTorch

Abaixo, implementamos uma classe robusta para a Scaled Dot-Product Attention, capaz de lidar com máscaras causais.

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

class ScaledDotProductAttention(nn.Module):
    """
    Calcula a atenção escalonada por produto escalar.
    Fórmula: Softmax(QK^T / sqrt(d_k)) * V
    """
    
    def __init__(self, dropout_rate=0.1):
        super(ScaledDotProductAttention, self).__init__()
        self.dropout = nn.Dropout(dropout_rate)

    def forward(self, query, key, value, mask=None):
        """
        Args:
            query: Tensor de forma (Batch, Heads, Seq_Len_Q, Head_Dim)
            key:   Tensor de forma (Batch, Heads, Seq_Len_K, Head_Dim)
            value: Tensor de forma (Batch, Heads, Seq_Len_K, Head_Dim)
            mask:  Tensor de máscara (Batch, 1, Seq_Len_Q, Seq_Len_K) ou similar.
                   Valores True indicam posições a serem mascaradas (ou 1s/0s dependendo da lógica).
        
        Returns:
            context: Tensor de saída ponderado
            attention_weights: Pesos de atenção (para visualização)
        """
        
        # 1. Obter dimensões
        d_k = query.size(-1)
        
        # 2. Calcular Scores (Q * K^T)
        # Transpomos as duas últimas dimensões de K para permitir a multiplicação
        scores = torch.matmul(query, key.transpose(-2, -1))
        
        # 3. Escalonamento (Scaling)
        scores = scores / math.sqrt(d_k)
        
        # 4. Aplicação da Máscara (Se fornecida)
        if mask is not None:
            # Substitui valores onde a máscara é 0 (ou True, dependendo da implementação) por -inf
            # Aqui assumimos que mask=0 significa "esconder"
            scores = scores.masked_fill(mask == 0, float('-inf'))
        
        # 5. Softmax para obter probabilidades
        # Aplicado na última dimensão (dimensão da sequência alvo)
        attention_weights = F.softmax(scores, dim=-1)
        
        # 6. Dropout (Opcional, mas comum para regularização)
        attention_weights = self.dropout(attention_weights)
        
        # 7. Multiplicação pelos Valores (Weights * V)
        context = torch.matmul(attention_weights, value)
        
        return context, attention_weights

# --- Exemplo de Uso com Máscara Causal ---

def create_causal_mask(seq_len):
    """Cria uma máscara triangular inferior para impedir 'olhar para o futuro'."""
    # tril retorna a parte triangular inferior da matriz
    mask = torch.tril(torch.ones(seq_len, seq_len))
    return mask

# Configuração de teste
batch_size = 1
heads = 1
seq_len = 4
d_k = 8

# Dados fictícios
q = torch.randn(batch_size, heads, seq_len, d_k)
k = torch.randn(batch_size, heads, seq_len, d_k)
v = torch.randn(batch_size, heads, seq_len, d_k)

# Criar máscara causal (1s visíveis, 0s ocultos)
causal_mask = create_causal_mask(seq_len) 
# Expande dimensões para broadcasting: (1, 1, seq_len, seq_len)
causal_mask = causal_mask.unsqueeze(0).unsqueeze(0)

# Instanciar e rodar
attention_layer = ScaledDotProductAttention()
output, weights = attention_layer(q, k, v, mask=causal_mask)

print("Pesos de Atenção (com máscara causal):")
print(weights[0][0]) 
# Note que a parte superior direita será 0.

7.5.1 Análise da Saída do Código

Se executarmos o código acima, a matriz de pesos de atenção (weights) terá uma estrutura triangular inferior. Por exemplo, para uma sequência de tamanho 4:

O token 1 só atende ao token 1.
O token 2 atende aos tokens 1 e 2.
O token 3 atende aos tokens 1, 2 e 3.
As posições onde $j > i$ (futuro) terão peso $0$ (após o softmax de $-\infty$).

7.6 Conclusão do Capítulo

O mecanismo de Scaled Dot-Product Attention é eficiente computacionalmente, pois se baseia em operações de matriz altamente otimizadas. A introdução do fator de escala $\sqrt{d_k}$ garante estabilidade numérica, e o uso de máscaras permite que a mesma arquitetura seja utilizada tanto para codificação (entender todo o contexto) quanto para decodificação (gerar texto sequencialmente). No próximo capítulo, veremos como expandir este conceito para Multi-Head Attention, permitindo que o modelo foque em diferentes subespaços de representação simultaneamente.

7.7 Auditoria de formas e invariantes

Para $Q,K\in\mathbb{R}^{B\times H\times T\times d_k}$, o produto $QK^\top$ tem forma $(B,H,T,T)$. A softmax deve ser aplicada na última dimensão, pois cada consulta distribui peso sobre as chaves. Multiplicar por $V\in\mathbb{R}^{B\times H\times T\times d_v}$ devolve $(B,H,T,d_v)$.

Três testes simples capturam muitos bugs: cada linha dos pesos deve somar aproximadamente 1; posições futuras devem receber zero após a máscara causal; e o gradiente deve alcançar as projeções de $Q$, $K$ e $V$. Máscaras invertidas e softmax no eixo errado frequentemente produzem código executável, porém semanticamente incorreto.

Máscara antes da softmax

Substitua escores proibidos por um valor muito negativo antes da softmax. Zerar pesos depois dela quebra a normalização, a menos que a linha seja renormalizada.

7.8 Autoatenção, atenção cruzada e custo

Na autoatenção, consultas, chaves e valores são obtidos da mesma sequência. Na atenção cruzada, as consultas vêm de uma sequência, enquanto chaves e valores vêm de outra. Esse segundo caso aparece em arquiteturas codificador-decodificador e em modelos multimodais, nos quais representações de imagem ou áudio orientam a geração textual.

A matriz de escores possui $T\times T$ elementos por cabeça. Portanto, a atenção densa tem custo quadrático em memória em relação ao comprimento $T$. Atenção local, esparsa, por blocos e aproximações lineares procuram reduzir esse gargalo, geralmente trocando alcance global, exatidão ou simplicidade por eficiência.

Exemplo numérico mínimo

Para $q=[1,0]$, $k_1=[1,0]$ e $k_2=[0,1]$, os escores são $1$ e $0$. A softmax dá maior peso ao primeiro valor, mas não transforma a operação em uma escolha rígida: a saída continua sendo uma combinação ponderada dos dois vetores de valor.

7.9 Atenção eficiente sem mudar a fórmula

Uma implementação ingênua materializa toda a matriz $T\times T$ na memória global da GPU. FlashAttention reorganiza o cálculo em blocos pequenos que cabem em memória rápida, calcula estatísticas da softmax incrementalmente e evita armazenar grandes intermediários. O resultado matemático permanece sendo atenção exata; o ganho vem do padrão de acesso à memória.

Essa distinção é importante: complexidade assintótica e tempo real não são sinônimos. Uma operação com o mesmo número de multiplicações pode ser muito mais rápida ao reduzir transferências entre níveis de memória. Por isso, analisar apenas FLOPs não explica todo o desempenho.

7.9.1 Janela deslizante e alcance global

Na atenção de janela, cada token consulta somente uma vizinhança de largura $w$. O custo passa de aproximadamente $T^2$ para $Tw$, mas uma dependência distante precisa atravessar várias camadas ou usar tokens globais. A escolha de $w$ define um compromisso entre memória, latência e velocidade de propagação do contexto.

Modelos híbridos podem alternar camadas locais e globais. Outra possibilidade é usar atenção linear, que reorganiza ou aproxima o cálculo para evitar a matriz quadrática. Essas variantes não são substitutas universais: tarefas de recuperação distante devem fazer parte da avaliação.

Teste de contexto longo

Insira uma informação única no começo de um documento e faça perguntas em distâncias crescentes. Meça acerto e latência; apenas declarar o tamanho máximo da janela não revela quanto contexto o modelo utiliza efetivamente.