Что такое key padding mask vs attention mask в PyTorch?

В PyTorch при работе с трансформерами и вниманием (attention) часто используются маски для корректной обработки последовательностей разной длины и предотвращения внимания к нежелательным позициям.

• Key Padding Mask — маска, которая указывает, какие позиции в последовательности являются паддингом (заполнителями), чтобы модель не учитывала их при вычислении внимания. Обычно это булевый тензор, где True означает позицию паддинга. Она применяется к ключам (keys) и значениям (values) в механизме внимания, чтобы игнорировать паддинговые токены.

• Attention Mask — более общая маска, которая контролирует, какие позиции могут влиять на текущую позицию при вычислении внимания. Например, в задачах автогенерации текста используется маска, запрещающая смотреть вперед (future tokens), чтобы предотвратить утечку информации. Такая маска может быть треугольной матрицей, где запрещены определённые связи.

Пример:

import torch
from torch.nn import MultiheadAttention

# Допустим, у нас есть последовательность с паддингом
key_padding_mask = torch.tensor([[False, False, True, True]])  # True для паддинга

# Attention mask для предотвращения внимания к будущим токенам
attn_mask = torch.triu(torch.ones(4, 4), diagonal=1).bool()  # Верхнетреугольная маска

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

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

Таким образом, key padding mask используется для игнорирования паддинговых токенов, а attention mask — для контроля структуры внимания (например, предотвращения внимания на будущие позиции).