Многоголовый механизм внимания

class torch.nn.MultiheadAttention(embed_dim, num_heads, dropout=0.0, bias=True, add_bias_kv=False, add_zero_attn=False, kdim=None, vdim=None, batch_first=False, device=None, dtype=None) [source]

Позволяет модели совместно обращать внимание на информацию из различных подпространств представлений, как описано в статье: Внимание — это всё, что вам нужно.

Многоголовое внимание определяется как:

$$\text{Многоголовое}(Q, K, V) = \text{Concat}(golova_1,\dots,golova_h)W^O$$

где $golova_i = \text{Внимание}(QW_i^Q, KW_i^K, VW_i^V)$.

nn.MultiHeadAttention будет использовать оптимизированные реализации scaled_dot_product_attention() по возможности.

В дополнение к поддержке новой функции scaled_dot_product_attention(), для ускорения вывода, MHA будет использовать быстрый вывод с поддержкой вложенных тензоров, если:

• вычисляется самовнимание (т.е., query, key, и value являются одним и тем же тензором).
• входные данные сгруппированы (3D) с batch_first==True
• либо вычисление градиента отключено (используя torch.inference_mode или torch.no_grad) или аргумент тензора requires_grad отсутствует
• обучение отключено (используя .eval())
• add_bias_kv равно False
• add_zero_attn равно False
• batch_first равно True и входные данные сгруппированы
• kdim и vdim равны embed_dim
• если передается вложенный тензор, то ни key_padding_mask, ни attn_mask не передаются
• автоматическое масштабирование отключено

Если используется оптимизированная реализация быстрого вывода, можно передать вложенный тензор для query/key/value для более эффективного представления заполнения, чем использование маски заполнения. В этом случае будет возвращен вложенный тензор, и можно ожидать дополнительного ускорения, пропорционального доле заполнения входных данных.

Параметры

• embed_dim – Общая размерность модели.
• num_heads – Количество параллельных головок внимания. Обратите внимание, что embed_dim будет разделен на num_heads (т.е. каждая головка будет иметь размер embed_dim // num_heads).
• dropout – Вероятность отбрасывания для attn_output_weights. По умолчанию: 0.0 (без отбрасывания).
• bias – Если указано, добавляет смещение к слоям проектирования входных / выходных данных. По умолчанию: True.
• add_bias_kv – Если указано, добавляет смещение к последовательностям ключей и значений по оси dim=0. По умолчанию: False.
• add_zero_attn – Если указано, добавляет новую группу нулей к последовательностям ключей и значений по оси dim=1. По умолчанию: False.
• kdim – Общее количество признаков для ключей. По умолчанию: None (использует kdim=embed_dim).
• vdim – Общее количество признаков для значений. По умолчанию: None (использует vdim=embed_dim).
• batch_first – Если True, то входные и выходные тензоры предоставляются как (пакет, последовательность, признак). По умолчанию: False (последовательность, пакет, признак).

Примеры:

>>> multihead_attn = nn.MultiheadAttention(embed_dim, num_heads)
>>> attn_output, attn_output_weights = multihead_attn(query, key, value)

forward(query, key, value, key_padding_mask=None, need_weights=True, attn_mask=None, average_attn_weights=True, is_causal=False) [source]

Параметры

• query (Tensor) – Векторы запросов формы $(L, E_q)$ для неразбитого входного сигнала, $(L, N, E_q)$ при batch_first=False или $(N, L, E_q)$ при batch_first=True, где $L$ — длина целевой последовательности, $N$ — размер пакета, и $E_q$ — размерность вложений запроса embed_dim. Запросы сравниваются с парами ключ-значение для получения результата. Более подробную информацию см. в статье «Attention Is All You Need».
• key (Tensor) – Векторы ключей формы $(S, E_k)$ для неразбитого входного сигнала, $(S, N, E_k)$ при batch_first=False или $(N, S, E_k)$ при batch_first=True, где $S$ — длина исходной последовательности, $N$ — размер пакета, и $E_k$ — размерность вложений ключей kdim. Более подробную информацию см. в статье «Attention Is All You Need».
• value (Tensor) – Векторы значений формы $(S, E_v)$ для неразбитого входного сигнала, $(S, N, E_v)$ при batch_first=False или $(N, S, E_v)$ при batch_first=True, где $S$ — длина исходной последовательности, $N$ — размер пакета, и $E_v$ — размерность вложений значений vdim. Более подробную информацию см. в статье «Attention Is All You Need».
• key_padding_mask (Optional[Tensor]) — При указании маска формы $(N, S)$, указывающая, какие элементы в key нужно игнорировать с целью внимания (т. е. обрабатывать как «заполнитель»). Для неразбитого query, форма должна быть $(S)$. Поддерживаются бинарные и числовые маски. Для бинарной маски значение True указывает, что соответствующее значение key будет проигнорировано для целей внимания. Для числовой маски оно будет непосредственно добавлено к соответствующему значению key.
• need_weights (bool) — При указании возвращает attn_output_weights в дополнение к attn_outputs. Установите need_weights=False для использования оптимизированного scaled_dot_product_attention и достижения наилучшей производительности для MHA. По умолчанию: True.
• attn_mask (Optional[Tensor]) — При указании 2D или 3D маска, предотвращающая внимание к определенным позициям. Должна иметь форму $(L, S)$ или $(N\cdot\text{num\_heads}, L, S)$, где $N$ — размер пакета, $L$ — длина целевой последовательности, и $S$ — длина исходной последовательности. 2D маска будет транслироваться по пакету, а 3D маска допускает различную маску для каждой записи в пакете. Поддерживаются бинарные и числовые маски. Для бинарной маски значение True указывает, что соответствующая позиция не допускает внимания. Для числовой маски значения маски будут добавлены к весу внимания. Если и attn_mask, и key_padding_mask указаны, их типы должны совпадать.
• average_attn_weights (bool) — Если True, указывает, что возвращаемое attn_weights должно быть усреднено по головам. В противном случае attn_weights предоставляются отдельно для каждой головы. Обратите внимание, что этот флаг действует только при need_weights=True. По умолчанию: True (т. е. усреднение весов по головам)

• is_causal (bool) – Если указано, применяется маскировка причинно-следственной связи как маска внимания. По умолчанию: False. Предупреждение: is_causal даёт подсказку, что attn_mask является маской причинно-следственной связи. Неправильные подсказки могут привести к неверному выполнению, включая обратную совместимость.

Тип возвращаемого значения

Tuple[Тензор, Optional[Тензор]]

Выходные данные:

• attn_output - Выходы внимания формы $(L, E)$ при невведённом размере пакета, $(L, N, E)$ при batch_first=False или $(N, L, E)$ при batch_first=True, где $L$ — длина целевой последовательности, $N$ — размер пакета, и $E$ — размерность встраивания embed_dim.
• attn_output_weights - Возвращается только при need_weights=True. Если average_attn_weights=True, возвращает средневзвешенные по головам веса внимания формы $(L, S)$ при невведённом размере пакета или $(N, L, S)$, где $N$ — размер пакета, $L$ — длина целевой последовательности, и $S$ — длина исходной последовательности. Если average_attn_weights=False, возвращает веса внимания на голову формы $(\text{num\_heads}, L, S)$ при невведённом размере пакета или $(N, \text{num\_heads}, L, S)$.

Примечание

batch_first аргумент игнорируется для невведённых входных данных.

merge_masks(attn_mask, key_padding_mask, query) [source]

Определить тип маски и объединить маски при необходимости. Если предоставлена только одна маска, эта маска и соответствующий тип маски будут возвращены. Если предоставлены обе маски, обе маски будут расширены до формы (batch_size, num_heads, seq_len, seq_len), объединены с логическим or и возвращён тип маски 2:param attn_mask: маска внимания формы (seq_len, seq_len), тип маски 0 :param key_padding_mask: маска заполнения формы (batch_size, seq_len), тип маски 1 :param query: векторы запроса формы (batch_size, seq_len, embed_dim)

Возвращаемое значение

объединённая маска mask_type: объединённый тип маски (0, 1 или 2)

Тип возвращаемого значения

merged_mask