Kas ir key padding mask un attention mask PyTorch?

PyTorch bieži izmanto transformatorus un uzmanības mehānismu (attention), izmantojot maskas, lai pareizi apstrādātu dažādu garumu secības un novērstu uzmanības pievēršanu nevēlamām pozīcijām.

• Key Padding Mask — maska, kas norāda, kuras pozīcijas secībā ir aizpildījums (padding), lai modelis tās neņemtu vērā aprēķinot uzmanību. Parasti tas ir loģisks tensor, kur True nozīmē aizpildījuma pozīciju. To piemēro atslēgām (keys) un vērtībām (values) uzmanības mehānismā, lai ignorētu aizpildījuma tokenus.

• Attention Mask — plašāka maska, kas kontrolē, kuras pozīcijas var ietekmēt pašreizējo pozīciju aprēķinot uzmanību. Piemēram, teksta automātiskās ģenerēšanas uzdevumos tiek izmantota maska, kas aizliedz skatīties uz nākotnes tokeniem, lai novērstu informācijas noplūdi. Šāda maska var būt trijstūra matrica, kurā ir aizliegtas noteiktas saites.

Piemērs:

import torch
from torch.nn import MultiheadAttention

# Pieņemsim, mums ir secība ar aizpildījumu
key_padding_mask = torch.tensor([[False, False, True, True]])  # True nozīmē aizpildījumu

# Uzmanības maska, lai novērstu uzmanību uz nākotnes tokeniem
attn_mask = torch.triu(torch.ones(4, 4), diagonal=1).bool()  # Augšējā trijstūra maska

mha = MultiheadAttention(embed_dim=8, num_heads=2)
query = torch.rand(4, 1, 8)  # (secības garums, partija, embed dim)
key = value = query

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

Tādējādi, key padding mask tiek izmantota, lai ignorētu aizpildījuma tokenus, bet attention mask — kontrolētu uzmanības struktūru (piemēram, novēršot uzmanību uz nākotnes pozīcijām).