O que é a máscara de padding de chave vs máscara de atenção no PyTorch?

No PyTorch, ao trabalhar com transformadores e atenção, frequentemente usam-se máscaras para o processamento correto de sequências de diferentes comprimentos e para evitar atenção a posições indesejadas.

• Máscara de preenchimento de chave (Key Padding Mask): é uma máscara que indica quais posições na sequência são preenchimento (padding), para que o modelo não as considere ao calcular a atenção. Normalmente, é um tensor booleano onde True indica uma posição de preenchimento. É aplicada às chaves (keys) e valores (values) no mecanismo de atenção para ignorar os tokens de preenchimento.

• Máscara de atenção (Attention Mask): uma máscara mais geral que controla quais posições podem influenciar a posição atual ao calcular a atenção. Por exemplo, em tarefas de geração automática de texto, usa-se uma máscara que impede olhar para frente (tokens futuros) para evitar vazamento de informação. Essa máscara pode ser uma matriz triangular, onde certas ligações são proibidas.

Exemplo:

import torch
from torch.nn import MultiheadAttention

# Suponha que temos uma sequência com preenchimento
key_padding_mask = torch.tensor([[False, False, True, True]])  # True para preenchimento

# Máscara de atenção para impedir atenção a tokens futuros
attn_mask = torch.triu(torch.ones(4, 4), diagonal=1).bool()  # Máscara triangular superior

mha = MultiheadAttention(embed_dim=8, num_heads=2)
query = torch.rand(4, 1, 8)  # (comprimento da sequência, lote, dimensão do embed)
key = value = query

output, attn_weights = mha(query, key, value, attn_mask=attn_mask, key_padding_mask=key_padding_mask)

Assim, a máscara de preenchimento de chave é usada para ignorar tokens de preenchimento, e a máscara de atenção para controlar a estrutura da atenção (por exemplo, impedir atenção às posições futuras).