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)[source]

Allows the model to jointly attend to information from different representation subspaces. See Attention Is All You Need

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

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

Parameters
• embed_dim – total dimension of the model.

• dropout – a Dropout layer on attn_output_weights. Default: 0.0.

• bias – add bias as module parameter. Default: True.

• add_bias_kv – add bias to the key and value sequences at dim=0.

• add_zero_attn – add a new batch of zeros to the key and value sequences at dim=1.

• kdim – total number of features in key. Default: None.

• vdim – total number of features in value. Default: None.

Note that if kdim and vdim are None, they will be set to embed_dim such that query, key, and value have the same number of features.

Examples:

>>> 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)[source]
Parameters
• key, value (query,) – map a query and a set of key-value pairs to an output. See “Attention Is All You Need” for more details.

• key_padding_mask – if provided, specified padding elements in the key will be ignored by the attention. When given a binary mask and a value is True, the corresponding value on the attention layer will be ignored. When given a byte mask and a value is non-zero, the corresponding value on the attention layer will be ignored

• need_weights – output attn_output_weights.

• attn_mask – 2D or 3D mask that prevents attention to certain positions. A 2D mask will be broadcasted for all the batches while a 3D mask allows to specify a different mask for the entries of each batch.

Shapes for inputs:
• query: $(L, N, E)$ where L is the target sequence length, N is the batch size, E is the embedding dimension.

• key: $(S, N, E)$ , where S is the source sequence length, N is the batch size, E is the embedding dimension.

• value: $(S, N, E)$ where S is the source sequence length, N is the batch size, E is the embedding dimension.

• key_padding_mask: $(N, S)$ where N is the batch size, S is the source sequence length. If a ByteTensor is provided, the non-zero positions will be ignored while the position with the zero positions will be unchanged. If a BoolTensor is provided, the positions with the value of True will be ignored while the position with the value of False will be unchanged.

• attn_mask: if a 2D mask: $(L, S)$ where L is the target sequence length, S is the source sequence length.

If a 3D mask: $(N\cdot\text{num\_heads}, L, S)$ where N is the batch size, L is the target sequence length, S is the source sequence length. attn_mask ensure that position i is allowed to attend the unmasked positions. If a ByteTensor is provided, the non-zero positions are not allowed to attend while the zero positions will be unchanged. If a BoolTensor is provided, positions with True is not allowed to attend while False values will be unchanged. If a FloatTensor is provided, it will be added to the attention weight.

Shapes for outputs:
• attn_output: $(L, N, E)$ where L is the target sequence length, N is the batch size, E is the embedding dimension.

• attn_output_weights: $(N, L, S)$ where N is the batch size, L is the target sequence length, S is the source sequence length.
{"extraction_info": {"found_math": true, "script_math_tex": 0, "script_math_asciimath": 0, "math_annotations": 10, "math_alttext": 0, "mathml": 0, "mathjax_tag": 0, "mathjax_inline_tex": 0, "mathjax_display_tex": 0, "mathjax_asciimath": 1, "img_math": 0, "codecogs_latex": 0, "wp_latex": 0, "mimetex.cgi": 0, "/images/math/codecogs": 0, "mathtex.cgi": 0, "katex": 0, "math-container": 0, "wp-katex-eq": 0, "align": 0, "equation": 0, "x-ck12": 0, "texerror": 0, "math_score": 0.5849102139472961, "perplexity": 2411.476693347769}, "config": {"markdown_headings": true, "markdown_code": true, "boilerplate_config": {"ratio_threshold": 0.18, "absolute_threshold": 10, "end_threshold": 15, "enable": true}, "remove_buttons": true, "remove_image_figures": true, "remove_link_clusters": true, "table_config": {"min_rows": 2, "min_cols": 3, "format": "plain"}, "remove_chinese": true, "remove_edit_buttons": true, "extract_latex": true}, "warc_path": "s3://commoncrawl/crawl-data/CC-MAIN-2021-04/segments/1610704843561.95/warc/CC-MAIN-20210128102756-20210128132756-00295.warc.gz"}