Mi a különbség a key padding mask és az attention mask között a PyTorch-ban?

A PyTorch-ban, amikor transzformátorokkal és figyelemmel (attention) dolgozik, gyakran használnak maszkokat a különböző hosszúságú szekvenciák helyes feldolgozásához és a nem kívánt pozíciókra való figyelem megakadályozásához.

• Key Padding Mask: ez egy maszk, amely megmutatja, hogy mely pozíciók a szekvenciában padding (kitöltés), így a modell ezeket nem veszi figyelembe a figyelem számításakor. Általában egy logikai tensor, ahol a True padding pozíciót jelez. Ez alkalmazható a kulcsokra (keys) és értékekre (values) a figyelem mechanizmusában, hogy figyelmen kívül hagyja a padding tokeneket.

• Attention Mask: egy általánosabb maszk, amely szabályozza, hogy mely pozíciók befolyásolhatják a jelenlegi pozíciót a figyelem számításakor. Például, automatikus szöveg generálási feladatoknál egy olyan maszkot használnak, amely megakadályozza, hogy előre nézzen (jövőbeli tokenek), így elkerülve az információ szivárgását. Ez a maszk lehet egy háromszögmátrix, ahol bizonyos kapcsolatok tiltottak.

Példa:

import torch
from torch.nn import MultiheadAttention

# Tegyük fel, van egy szekvenciánk paddinggel
key_padding_mask = torch.tensor([[False, False, True, True]])  # True paddinghez

# Figyelem maskája, hogy megakadályozza a jövőbeli tokenek figyelembevételét
attn_mask = torch.triu(torch.ones(4, 4), diagonal=1).bool()  # Felső háromszög maszk

mha = MultiheadAttention(embed_dim=8, num_heads=2)
query = torch.rand(4, 1, 8)  # (szekvencia hossz, batch, embed dim)
key = value = query

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

Így a key padding mask a padding tokenek figyelmen kívül hagyására szolgál, míg az attention mask az figyelem szerkezetének irányítására (pl. a jövőbeli pozíciókra való figyelem megakadályozására).