Czym jest maska padding klucza a maska uwagi w PyTorch?

W PyTorch, podczas pracy z transformatorami i mechanizmem uwagi (attention) często używa się masek do poprawnego przetwarzania sekwencji o różnej długości i zapobiegania uwadze na niepożądane pozycje.

• Key Padding Mask: to maska wskazująca, które pozycje w sekwencji są paddingiem (wypełnieniem), aby model ich nie brał pod uwagę przy obliczaniu uwagi. Zazwyczaj jest to tensor logiczny, gdzie True oznacza pozycję paddingu. Jest stosowana do kluczy (keys) i wartości (values) w mechanizmie uwagi, aby ignorować tokeny paddingu.

• Attention Mask: bardziej ogólna maska, która kontroluje, które pozycje mogą wpływać na bieżącą pozycję podczas obliczania uwagi. Na przykład, w zadaniach automatycznego generowania tekstu, używa się maski, która zabrania patrzenia do przodu (tokens przyszłe), aby zapobiec wyciekom informacji. Taka maska może być macierzą trójkątną, gdzie pewne połączenia są zabronione.

Przykład:

import torch
from torch.nn import MultiheadAttention

# Załóżmy, że mamy sekwencję z paddingiem
key_padding_mask = torch.tensor([[False, False, True, True]])  # True dla paddingu

# Maska uwagi zapobiegająca patrzeniu do przodu
attn_mask = torch.triu(torch.ones(4, 4), diagonal=1).bool()  # Macierz trójkątna górna

mha = MultiheadAttention(embed_dim=8, num_heads=2)
query = torch.rand(4, 1, 8)  # (długość sekwencji, partia, wymiar osadzenia)
key = value = query

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

W ten sposób, key padding mask służy do ignorowania tokenów padding, a attention mask do kontrolowania struktury uwagi (np. zapobiegania uwadze na pozycje przyszłe).