Как учитывается маскирование отступов в головах внимания трансформера?

В чисто образовательных целях моя цель – реализовать базовую архитектуру Transformer с нуля. До сих пор я сосредоточился на энкодере для классификационных задач и предположил, что все образцы в пакете имеют одинаковую длину. Это означает, что мне не нужно было беспокоиться о маскировании.

Однако теперь я хочу поддерживать маскирование. Мне нравится думать, что я понимаю роль, например, целевой маски, чтобы порядок не мог “заглядывать в будущее”. Я генерирую эту маску следующим образом:

source_batch = torch.LongTensor([
    [1, 2, 3, 0, 0, 0],
    [1, 2, 3, 4, 5, 6],
    [1, 2, 3, 4, 5, 0]
])

batch_size, seq_len = source_batch.shape

def generate_tgt_mask(size):
    return torch.triu(torch.ones(seq_len, seq_len) * float('-inf'), diagonal=1)

print(generate_tgt_mask(seq_len))

что дает:

tensor([[0., -inf, -inf, -inf, -inf, -inf],
        [0.,   0., -inf, -inf, -inf, -inf],
        [0.,   0.,   0., -inf, -inf, -inf],
        [0.,   0.,   0.,   0., -inf, -inf],
        [0.,   0.,   0.,   0.,   0., -inf],
        [0.,   0.,   0.,   0.,   0.,   0.]])

что должно быть ожидаемым результатом, когда я проверяю документацию PyTorch. Эта маска имеет форму (L,L), где L – это длина последовательности исходной или целевой последовательности. Снова это соответствует документации.

Я использую эту маску в своей реализации Scaled Dot Product Attention следующим образом, что должно быть в согласии с многими другими реализациями, которые я видел:

class Attention(nn.Module):
    ### Реализует Scaled Dot Product Attention

    def __init__(self):
        super().__init__()

    def forward(self, Q, K, V, mask=None, dropout=None):
        # Все формы: (batch_size, seq_len, hidden_size)

        # Выполняем Q*K^T (* - это скалярное произведение здесь)
        # Нам нужно использовать torch.matmul, так как мы работаем с пакетами!
        out = torch.matmul(Q, K.transpose(1, 2)) # => форма: (B, L, L)

        # Делим на масштабирующий фактор
        out = out / (Q.shape[-1] ** 0.5)

        # Необязательно: src_mask/tgt_mask (форма: (L, L); значения маски представлены как -inf)
        if mask is not None:
            out += mask.unsqueeze(0) # Broadcast, так как это одна и та же маска для всех образцов в пакете

        # Применяем softmax
        out = f.softmax(out, dim=-1)

        # Необязательно: Dropout
        if dropout is not None:
            out = nn.Dropout(out, dropout)

        # Умножаем на значения V
        out = torch.matmul(out, V)

        return out

До сих пор все хорошо… по крайней мере, я так думаю. Однако моя проблема сейчас заключается в маске для обработки паддинга (например, src_key_padding_mask). Из разных уроков по использованию nn.Transformer эта маска может быть сгенерирована следующим образом:

pad_token_index = 0

src_key_padding_mask = (source_batch != pad_token_index)

print(src_key_padding_mask)

что дает:

tensor([[ True,  True,  True, False, False, False],
        [ True,  True,  True,  True,  True,  True],
        [ True,  True,  True,  True,  True, False]])

с формой (N,L), что снова соответствует документации.

Что мне сейчас не хватает: Как мне нужно включить эту матрицу в свою реализацию Attention?

Интуитивно, я предположил бы, что маскирующая матрица будет содержать -inf для каждой позиции, связанной с паддингом. Например, глядя на первую последовательность в моем примере пакета выше, я бы предположил, что маскирующая матрица будет выглядеть так:

tensor([[0.,   0.,   0.,   -inf, -inf, -inf],
        [0.,   0.,   0.,   -inf, -inf, -inf],
        [0.,   0.,   0.,   -inf, -inf, -inf],
        [-inf, -inf, -inf, -inf, -inf, -inf],
        [-inf, -inf, -inf, -inf, -inf, -inf],
        [-inf, -inf, -inf, -inf, -inf, -inf]])

И действительно, некоторые – но не все – примеры кода, которые реализуют архитектуру Transformer с нуля, создают маскирующую матрицу для паддинга таким образом. Применение этой матрицы к оценкам, очевидно, также устанавливает оценки в 0, то есть последние 3 строки все 0.

Однако, после применения softmax, последние 3 строки теперь все содержат значение 1/6. Например, для source_batch выше я получаю

tensor([[[0.1989, 0.4297, 0.3714, 0.0000, 0.0000, 0.0000],
         [0.4334, 0.2225, 0.3440, 0.0000, 0.0000, 0.0000],
         [0.2880, 0.2284, 0.4836, 0.0000, 0.0000, 0.0000],
         [0.1667, 0.1667, 0.1667, 0.1667, 0.1667, 0.1667],
         [0.1667, 0.1667, 0.1667, 0.1667, 0.1667, 0.1667],
         [0.1667, 0.1667, 0.1667, 0.1667, 0.1667, 0.1667]],
       ...
       (другие 2 образца пакета не показаны)

Что я упускаю здесь? Я уверен, что это что-то тривиальное, но я просто не могу это сейчас увидеть.

Вот как работает реализация fairseq.

Сначала она превращает матрицу BH x T x S в B x T x S, где B – это размер пакета, H – количество голов, T – длина цели, а S – длина источника. Вы не используете многоголовое внимание, поэтому можете пропустить это. (Также, когда выполняется самовнимание T == S, что вы называете L). Как последний шаг, чтобы разбить данные обратно на BH x T x S после применения маски.

Их key_padding_mask является тем, что вы называете src_key_padding_mask, и имеет размер B x L. Следующий шаг – превратить это в B x 1 x 1 x L. Снова, так как вы не используете многоголовое внимание, вам нужно использовать unsqueeze(1) только один раз, чтобы получить B x 1 x L.

Затем они используют masked_fill. Таким образом, это становится однострочным выражением:

out = out.masked_fill( src_key_padding_mask.unsqueeze(1), float("-inf") )

Это идет после вашего текущего кода mask и перед вызовом softmax.

Кстати, я написал выше, ссылаясь на свою локальную копию, над которой я поработал и упростил ее. Смотря на версию на github (ссылка выше), я вижу, что они используют view(), чтобы добавить 5-е измерение, а затем используют unsqueeze() трижды. Я не могу сказать вам, почему; что-то связано с тем, когда kv_bsz отличается от bsz. Это только в блоке not is_tpu, так что просто какая-то умная оптимизация?