This is a tutorial/implementation of multi-headed attention from paper Attention Is All You Need in PyTorch. The implementation is inspired from Annotated Transformer.
\nHere is the training code that uses a basic transformer with MHA for NLP auto-regression.
\nHere is an experiment implementation that trains a simple transformer.
\n": "\u8fd9\u662f\u8bba\u6587\u300a Attention is All You Need \u300b\u4e2d\u591a\u5934\u6ce8\u610f\u529b\u7684PyTorch\u6559\u7a0b/\u5b9e\u73b0\u3002\u8be5\u5b9e\u73b0\u7684\u7075\u611f\u6765\u81ea\u300a\u5e26\u6ce8\u91ca\u7684 Transformer \u300b\u3002
\u8fd9\u662f\u4f7f\u7528\u57fa\u7840 Transformer \u548c MHA \u8fdb\u884c NLP \u81ea\u56de\u5f52\u7684\u8bad\u7ec3\u4ee3\u7801\u3002
\u8fd9\u662f\u4e00\u4e2a\u8bad\u7ec3\u7b80\u5355 Transformer \u7684\u4ee3\u7801\u5b9e\u73b0\u3002
\n", "This method can be overridden for other variations like relative attention.
\n": "\u8fd9\u79cd\u65b9\u6cd5\u53ef\u4ee5\u540c\u6837\u9002\u7528\u4e8e\u5176\u4ed6\u53d8\u4f53\uff0c\u5982\u76f8\u5bf9\u6ce8\u610f\u529b\u3002
\n", "\nThis computes scaled multi-headed attention for given _^_0_^_, _^_1_^_ and _^_2_^_ vectors.
\n_^_3_^_
\nIn simple terms, it finds keys that matches the query, and gets the values of those keys.
\nIt uses dot-product of query and key as the indicator of how matching they are. Before taking the _^_4_^_ the dot-products are scaled by _^_5_^_. This is done to avoid large dot-product values causing softmax to give very small gradients when _^_6_^_ is large.
\nSoftmax is calculated along the axis of of the sequence (or time).
\n": "\n\u8fd9\u5c06\u8ba1\u7b97\u7ed9\u51fa\u7684_^_1_^_\u3001_^_2_^_\u548c_^_0_^_\u5411\u91cf\u7f29\u653e\u540e\u7684\u591a\u5934\u6ce8\u610f\u529b\u3002
\n_^_3_^_
\n\u7b80\u5355\u6765\u8bf4\uff0c\u5b83\u4f1a\u627e\u5230\u4e0e\u67e5\u8be2 (Query) \u5339\u914d\u7684\u952e (key)\uff0c\u5e76\u83b7\u53d6\u8fd9\u4e9b\u952e (Key) \u7684\u503c (Value) \u3002
\n\u5b83\u4f7f\u7528\u67e5\u8be2\u548c\u952e\u7684\u70b9\u79ef\u4f5c\u4e3a\u8861\u91cf\u5b83\u4eec\u4e4b\u95f4\u5339\u914d\u7a0b\u5ea6\u7684\u6307\u6807\u3002\u5728\u8fdb\u884c_^_4_^_\u4e4b\u524d\uff0c\u70b9\u79ef\u4f1a\u4e58\u4ee5_^_5_^_\u3002\u8fd9\u6837\u505a\u662f\u4e3a\u4e86\u907f\u514d\u5f53_^_6_^_\u8f83\u5927\u65f6\uff0c\u5927\u7684\u70b9\u79ef\u503c\u5bfc\u81f4 Softmax \u64cd\u4f5c\u8f93\u51fa\u975e\u5e38\u5c0f\u7684\u68af\u5ea6\u3002
\nSoftmax \u662f\u6cbf\u5e8f\u5217\uff08\u6216\u65f6\u95f4\uff09\u8f74\u8ba1\u7b97\u7684\u3002
\n", "\nThis module does a linear transformation and splits the vector into given number of heads for multi-head attention. This is used to transform key, query, and value vectors.
\n": "\n\u8be5\u90e8\u5206\u6267\u884c\u7ebf\u6027\u53d8\u6362\uff0c\u5e76\u5c06\u5411\u91cf\u5206\u5272\u6210\u7ed9\u5b9a\u6570\u91cf\u7684\u5934\u4ee5\u83b7\u5f97\u591a\u5934\u6ce8\u610f\u529b\u3002\u8fd9\u7528\u4e8e\u952e\u3001\u67e5\u8be2\u548c\u503c\u5411\u91cf\u3002
\n", "_^_0_^_ has shape _^_1_^_, where first dimension is the query dimension. If the query dimension is equal to _^_2_^_ it will be broadcasted.
\n": "_^_0_^_\u7684\u5f62\u72b6\u4e3a_^_1_^_\uff0c\u5176\u4e2d\u7b2c\u4e00\u7ef4\u662f\u67e5\u8be2\u7ef4\u5ea6\u3002\u5982\u679c\u67e5\u8be2\u7ef4\u5ea6\u7b49\u4e8e_^_2_^_\uff0c\u5219\u4f1a\u8fdb\u884c\u5e7f\u64ad\u3002
\n", "_^_0_^_, _^_1_^_ and _^_2_^_ are the tensors that store collection of query, key and value vectors. They have shape _^_3_^_.
\n_^_4_^_ has shape _^_5_^_ and _^_6_^_ indicates whether for batch _^_7_^_, query at position _^_8_^_ has access to key-value at position _^_9_^_.
\n": "_^_0_^_\u3001_^_1_^_\u548c_^_2_^_\u662f\u5b58\u50a8\u67e5\u8be2\u3001\u952e\u548c\u503c\u5411\u91cf\u96c6\u5408\u7684\u5f20\u91cf\u3002\u5b83\u4eec\u7684\u5f62\u72b6\u4e3a_^_3_^_\u3002
\n_^_4_^_\u7684\u5f62\u72b6\u4e3a_^_5_^_\uff0c_^_6_^_\u8868\u793a\u6279\u6b21_^_7_^_\uff0c\u5728\u4f4d\u7f6e_^_8_^_\u5904\u67e5\u8be2\u662f\u5426\u6709\u6743\u8bbf\u95ee\u4f4d\u7f6e_^_9_^_\u5904\u7684\u952e\u503c\u5bf9\u3002
\n", "_^_0_^_ attention along the key sequence dimension _^_1_^_
\n": "\u5bf9 Key \u5e8f\u5217\u7ef4\u5ea6\u4e0a\u7684\u6ce8\u610f\u529b\u8fdb\u884c_^_0_^_\u64cd\u4f5c\uff0c_^_1_^_
\n", "_^_0_^_, _^_1_^_ and _^_2_^_ have shape _^_3_^_
\n": "_^_0_^_\uff0c_^_1_^_\u548c_^_2_^_\u7684\u5f62\u72b6\u4e3a_^_3_^_
\n", "Apply dropout
\n": "\u5e94\u7528 Dropout
\n", "Apply mask
\n": "\u5e94\u7528\u63a9\u7801
\n", "Calculate _^_0_^_ or _^_1_^_
\n": "\u8ba1\u7b97_^_0_^_\u6216_^_1_^_
\n", "Compute attention scores _^_0_^_. This gives a tensor of shape _^_1_^_.
\n": "\u8ba1\u7b97\u6ce8\u610f\u529b\u5206\u6570_^_0_^_\uff0c\u8fd9\u5c06\u5f97\u5230\u4e00\u4e2a\u5f62\u72b6\u4e3a_^_1_^_\u7684\u5f20\u91cf\u3002
\n", "Concatenate multiple heads
\n": "\u8fde\u63a5\u591a\u4e2a\u5934
\n", "Dropout
\n": "Dropout
\n", "Input has shape _^_0_^_ or _^_1_^_. We apply the linear transformation to the last dimension and split that into the heads.
\n": "\u8f93\u5165\u7684\u5f62\u72b6\u4e3a_^_0_^_\u6216_^_1_^_\u3002\u6211\u4eec\u5bf9\u6700\u540e\u4e00\u7ef4\u5e94\u7528\u7ebf\u6027\u53d8\u6362\uff0c\u5e76\u5c06\u5176\u5206\u4e3a\u591a\u4e2a\u5934\u3002
\n", "Linear layer for linear transform
\n": "\u7ebf\u6027\u5c42\u7528\u4e8e\u7ebf\u6027\u53d8\u6362
\n", "Linear transform
\n": "\u7ebf\u6027\u53d8\u6362
\n", "Multiply by values _^_0_^_
\n": "\u4e58\u4ee5\u6570\u503c_^_0_^_
\n", "Number of dimensions in vectors in each head
\n": "\u6bcf\u4e2a\u5934\u90e8\u4e2d\u5411\u91cf\u7684\u7ef4\u5ea6\u6570\u91cf
\n", "Number of features per head
\n": "\u6bcf\u4e2a\u5934\u90e8\u7684\u7279\u5f81\u6570\u91cf
\n", "Number of heads
\n": "\u6ce8\u610f\u529b\u5934\u6570
\n", "Output has shape _^_0_^_ or _^_1_^_
\n": "\u8f93\u51fa\u5177\u6709\u5f62\u72b6_^_0_^_\u6216_^_1_^_
\n", "Output layer
\n": "\u8f93\u51fa\u5c42
\n", "Prepare _^_0_^_, _^_1_^_ and _^_2_^_ for attention computation. These will then have shape _^_3_^_.
\n": "\u4e3a\u6ce8\u610f\u529b\u8ba1\u7b97\u51c6\u5907\u5411\u91cf_^_0_^_\uff0c_^_1_^_\u5e76_^_2_^_\u5b83\u4eec\u7684\u5f62\u72b6\u5c06\u53d8\u4e3a_^_3_^_\u3002
\n", "Same mask applied to all heads.
\n": "\u6240\u6709\u7684\u5934\u90e8\u4f7f\u7528\u76f8\u540c\u7684\u63a9\u7801\u3002
\n", "Save attentions for any other calculations
\n": "\u4e3a\u5176\u4ed6\u8ba1\u7b97\u4fdd\u5b58\u6ce8\u610f\u529b\u4fe1\u606f
\n", "Save attentions if debugging
\n": "\u8c03\u8bd5\u65f6\u4fdd\u5b58\u6ce8\u610f\u529b\u4fe1\u606f
\n", "Scale scores _^_0_^_
\n": "\u7f29\u653e\u5206\u6570_^_0_^_
\n", "Scaling factor before the softmax
\n": "Softmax \u4e4b\u524d\u7684\u7f29\u653e\u7cfb\u6570
\n", "Softmax for attention along the time dimension of _^_0_^_
\n": "\u5728\u952e\uff08 Key \uff09\u7684\u65f6\u95f4\u7ef4\u5ea6\u4e0a\u8fdb\u884c\u6ce8\u610f\u529b Softmax_^_0_^_
\n", "Split last dimension into heads
\n": "\u5c06\u6700\u540e\u4e00\u4e2a\u7ef4\u5ea6\u5206\u6210\u591a\u4e2a\u5934\u90e8
\n", "These transform the _^_0_^_, _^_1_^_ and _^_2_^_ vectors for multi-headed attention.
\n": "\u8fd9\u4e9b\u5c06\u5bf9\u591a\u5934\u6ce8\u610f\u529b\u7684\u5411\u91cf_^_0_^_\u3001_^_1_^_\u548c_^_2_^_\u8fdb\u884c\u8f6c\u6362\u3002
\n", "We store attentions so that it can be used for logging, or other computations if needed
\n": "\u5b58\u50a8\u6ce8\u610f\u529b\u4fe1\u606f\uff0c\u4ee5\u4fbf\u5728\u9700\u8981\u65f6\u7528\u4e8e\u8bb0\u5f55\u6216\u5176\u4ed6\u8ba1\u7b97\u3002
\n", "resulting mask has shape _^_0_^_
\n": "\u751f\u6210\u7684\u63a9\u7801\u5f62\u72b6\u4e3a_^_0_^_
\n", "