Kas yra key padding mask ir attention mask PyTorch?

„PyTorch“ dažnai naudojant transformerius ir dėmesio mechanizmą (attention), dažnai taikomos kaukės, siekiant teisingai apdoroti skirtingo ilgio sekas ir išvengti dėmesio nepageidaujamoms pozicijoms.

• Key Padding Mask — kaukė, kuri nurodo, kurios sekos pozicijos yra užpildymo (padding) žymės, kad modelis jų neatsižvelgtų skaičiuodamas dėmesį. Paprastai tai yra loginis tensoras, kur True reiškia užpildymo poziciją. Ji taikoma raktams (keys) ir reikšmėms (values) dėmesio mechanizme, kad būtų ignoruojami užpildymo tokenai.

• Attention Mask — labiau bendro pobūdžio kaukė, kuri kontroliuoja, kurios pozicijos gali turėti įtakos dabartinei pozicijai skaičiuojant dėmesį. Pavyzdžiui, teksto automatinės generacijos užduotyse naudojama kaukė, draudžianti žiūrėti į ateities tokenus, siekiant išvengti informacijos nutekėjimo. Tokia kaukė gali būti trikampė matrica, kurioje tam tikros jungtys yra uždraustos.

Pavyzdys:

import torch
from torch.nn import MultiheadAttention

# Tarkime, turime seką su užpildymu
key_padding_mask = torch.tensor([[False, False, True, True]])  # True reiškia užpildymą

# Dėmesio kaukė, siekiant išvengti dėmesio ateities tokenams
attn_mask = torch.triu(torch.ones(4, 4), diagonal=1).bool()  # Viršutinė trikampė kaukė

mha = MultiheadAttention(embed_dim=8, num_heads=2)
query = torch.rand(4, 1, 8)  # (sekos ilgis, partija, įterpimo dim)
key = value = query

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

Taigi, key padding mask naudojama ignoruoti užpildymo tokenus, o attention mask — kontroliuoti dėmesio struktūrą (pavyzdžiui, užkirsti kelią dėmesiui į ateities pozicijas).