MultiheadAttention

class torch.nn.modules.activation.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) [исходный код]

Позволяет модели совместно учитывать информацию из различных подпространств представлений.

Этот слой MultiheadAttention реализует исходную архитектуру, описанную в статье Attention Is All You Need. Этот слой предназначен для использования в качестве эталонной реализации, помогающей понять основы, поэтому по сравнению с более новыми архитектурами он содержит лишь ограниченный набор функций. Учитывая быстрый темп развития архитектур на основе трансформеров, рекомендуем изучить этот учебный материал, чтобы создавать эффективные слои из базовых блоков ядра, или воспользоваться библиотеками более высокого уровня из экосистемы PyTorch.

Многоголовое внимание определяется следующим образом:

$$\text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1,\dots,\text{head}_h)W^O$$

где $\text{head}_i = \text{Attention}(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
• 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, входные и выходные тензоры имеют формат (batch, seq, feature). По умолчанию: False (seq, batch, feature).

Примеры:

>>> 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) [исходный код]

Вычисляет выходные значения внимания, используя эмбеддинги запроса, ключа и значения.

Поддерживает необязательные параметры для заполнения, масок и весов внимания.

Параметры:

• 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 (Tensor | None) – Если задана, маска формы $(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 (Tensor | None) – Если задана, двумерная или трехмерная маска, запрещающая внимание к определенным позициям. Ее форма должна быть $(L, S)$ или $(N\cdot\text{num\_heads}, L, S)$, где $N$ — размер пакета, $L$ — длина целевой последовательности, а $S$ — длина исходной последовательности. Двумерная маска будет распространена на весь пакет, а трехмерная маска позволяет задать отдельную маску для каждого элемента пакета. Поддерживаются двоичные маски и маски с плавающей точкой. Для двоичной маски значение True указывает, что соответствующей позиции запрещено учитывать внимание. Для маски с плавающей точкой значения маски будут прибавлены к весу внимания. Если заданы и attn_mask, и key_padding_mask, их типы должны совпадать.
• average_attn_weights (bool) – Если значение истинно, возвращаемые attn_weights усредняются по головам. В противном случае attn_weights возвращаются отдельно для каждой головы. Обратите внимание, что этот флаг действует только при need_weights=True. По умолчанию: True (то есть веса усредняются по головам)
• is_causal (bool) – Если задано, применяет причинную маску в качестве маски внимания. По умолчанию: False. Предупреждение: is_causal указывает, что attn_mask является причинной маской. Неверные подсказки могут привести к некорректному выполнению, включая проблемы с прямой и обратной совместимостью.

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

tuple[Tensor, Tensor | None]

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

• 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) [исходный код]

Определяет тип маски и при необходимости объединяет маски.

Если задана только одна маска, возвращаются эта маска и соответствующий ей тип. Если заданы обе маски, они обе расширяются до формы (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