Многоголовое внимание

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]

Позволяет модели совместно обращать внимание на информацию из разных подпространств представлений, как описано в статье: Attention Is All You Need.

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

$$\text{MultiHead}(Q, K, V) = \text{Concat}(head_1,\dots,head_h)W^O$$

где $head_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V)$.

forward() будет использовать специальную оптимизированную реализацию, если выполнены все следующие условия:

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

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

Параметры:

• embed_dim – Общая размерность модели.
• num_heads – Количество параллельных головок внимания. Обратите внимание, что embed_dim будет разделено по num_heads (т.е. каждая головка будет иметь размер embed_dim // num_heads).
• dropout – Вероятность дропаута для attn_output_weights. По умолчанию: 0.0 (дропаут отсутствует).
• bias – Если указано, добавляет смещение к слоям проекции входных/выходных данных. По умолчанию: True.
• add_bias_kv – Если указано, добавляет смещение к последовательностям ключей и значений в измерении 0. По умолчанию: False.
• add_zero_attn – Если указано, добавляет новый пакет нулей к последовательностям ключей и значений в измерении 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) [source]

Параметры:

• query (Тензор) – Векторы запросов формы $(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 (Тензор) – Векторы ключей формы $(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 (Тензор) – Векторы значений формы $(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 (Необязательный[Тензор]) – Если указано, маска формы $(N, S)$, указывающая, какие элементы внутри key следует игнорировать для целей внимания (т.е. рассматривать как «заполнение»). Для неразмеченного query, форма должна быть $(S)$. Поддерживаются двоичные и байтовые маски. Для двоичной маски значение True указывает, что соответствующее значение key будет проигнорировано для целей внимания. Для плавающей маски она будет напрямую добавлена к соответствующему значению key.
• need_weights (логическое) – Если указано, возвращает attn_output_weights помимо attn_outputs. По умолчанию: True.
• attn_mask (Необязательный[Тензор]) – Если указано, 2D или 3D маска, предотвращающая внимание к определённым позициям. Должна иметь форму $(L, S)$ или $(N\cdot\text{num\_heads}, L, S)$, где $N$ — размер пакета, $L$ — длина целевой последовательности, а $S$ — длина последовательности источника. 2D маска будет транслироваться по пакету, а 3D маска позволит иметь различную маску для каждого элемента пакета. Поддерживаются двоичные, байтовые и плавающие маски. Для двоичной маски значение True указывает, что соответствующая позиция не может участвовать во внимании. Для байтовой маски ненулевое значение указывает, что соответствующая позиция не может участвовать во внимании. Для плавающей маски значения маски будут добавлены к весу внимания.

• average_attn_weights (bool) – Если значение равно true, указывает, что возвращаемые attn_weights должны быть усреднены по всем головкам. В противном случае attn_weights предоставляются отдельно для каждой головки. Обратите внимание, что этот флаг действует только при need_weights=True. Значение по умолчанию: True (т.е. усреднение весов по всем головкам)

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

Tuple[Tensor, Optional[Tensor]]

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

• 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 аргумент игнорируется для непакетированных входов.