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": "\u3053\u308c\u306f\u3001PyTorch\u306e\u8ad6\u6587\u300c\u6ce8\u610f\u3055\u3048\u3042\u308c\u3070\u5341\u5206\u300d\u306e\u300c\u591a\u9762\u7684\u306a\u6ce8\u610f\u300d\u306e\u30c1\u30e5\u30fc\u30c8\u30ea\u30a2\u30eb/\u5b9f\u88c5\u3067\u3059\u3002\u5b9f\u88c5\u306f\u6ce8\u91c8\u4ed8\u304d\u30c8\u30e9\u30f3\u30b9\u30d5\u30a9\u30fc\u30de\u30fc\u304b\u3089\u7740\u60f3\u3092\u5f97\u3066\u3044\u307e\u3059
\u3002\n\u3053\u308c\u306f\u3001NLP\u81ea\u5df1\u56de\u5e30\u7528\u306eMHA\u3092\u5099\u3048\u305f\u57fa\u672c\u7684\u306a\u30c8\u30e9\u30f3\u30b9\u30d5\u30a9\u30fc\u30de\u30fc\u3092\u4f7f\u7528\u3059\u308b\u30c8\u30ec\u30fc\u30cb\u30f3\u30b0\u30b3\u30fc\u30c9\u3067\u3059\u3002
\n\n", "This method can be overridden for other variations like relative attention.
\n": "\u3053\u306e\u65b9\u6cd5\u306f\u3001\u76f8\u5bfe\u7684\u6ce8\u610f\u529b\u306a\u3069\u306e\u4ed6\u306e\u30d0\u30ea\u30a8\u30fc\u30b7\u30e7\u30f3\u3067\u306f\u30aa\u30fc\u30d0\u30fc\u30e9\u30a4\u30c9\u3067\u304d\u307e\u3059\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_^_0_^_\u4e0e\u3048\u3089\u308c\u305f\u30d9\u30af\u30c8\u30eb\u3084\u30d9\u30af\u30c8\u30eb\u306b\u5bfe\u3057\u3066\u3001\u30b9\u30b1\u30fc\u30ea\u30f3\u30b0\u3055\u308c\u305f\u30de\u30eb\u30c1\u30d8\u30c3\u30c9\u30fb\u30a2\u30c6\u30f3\u30b7\u30e7\u30f3\u3092\u8a08\u7b97\u3057\u307e\u3059\u3002_^_1_^_ _^_2_^_
\n_^_3_^_
\n\u7c21\u5358\u306b\u8a00\u3046\u3068\u3001\u30af\u30a8\u30ea\u306b\u4e00\u81f4\u3059\u308b\u30ad\u30fc\u3092\u898b\u3064\u3051\u3001\u305d\u308c\u3089\u306e\u30ad\u30fc\u306e\u5024\u3092\u53d6\u5f97\u3057\u307e\u3059\u3002
\n\u30af\u30a8\u30ea\u3068\u30ad\u30fc\u306e\u30c9\u30c3\u30c8\u7a4d\u304c\u3069\u306e\u7a0b\u5ea6\u4e00\u81f4\u3057\u3066\u3044\u308b\u304b\u3092\u793a\u3059\u6307\u6a19\u3068\u3057\u3066\u4f7f\u7528\u3057\u307e\u3059\u3002_^_4_^_\u64ae\u5f71\u524d\u306b\u30c9\u30c3\u30c8\u30d7\u30ed\u30c0\u30af\u30c8\u3092\u30b9\u30b1\u30fc\u30ea\u30f3\u30b0\u3057\u307e\u3059\u3002_^_5_^_\u3053\u308c\u306f\u3001\u30c9\u30c3\u30c8\u7a4d\u5024\u304c\u5927\u304d\u3044\u5834\u5408\u306b softmax \u306e\u30b0\u30e9\u30c7\u30fc\u30b7\u30e7\u30f3\u304c\u975e\u5e38\u306b\u5c0f\u3055\u304f\u306a\u308b\u539f\u56e0\u3068\u306a\u3089\u306a\u3044\u3088\u3046\u306b\u3059\u308b\u305f\u3081\u3067\u3059
\u3002_^_6_^_\nSoftmax \u306f\u3001\u30b7\u30fc\u30b1\u30f3\u30b9 (\u307e\u305f\u306f\u6642\u9593) \u306e\u8ef8\u306b\u6cbf\u3063\u3066\u8a08\u7b97\u3055\u308c\u307e\u3059\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\u3053\u306e\u30e2\u30b8\u30e5\u30fc\u30eb\u306f\u7dda\u5f62\u5909\u63db\u3092\u884c\u3044\u3001\u30d9\u30af\u30c8\u30eb\u3092\u6307\u5b9a\u3055\u308c\u305f\u6570\u306e\u30d8\u30c3\u30c9\u306b\u5206\u5272\u3057\u3066\u30de\u30eb\u30c1\u30d8\u30c3\u30c9\u30a2\u30c6\u30f3\u30b7\u30e7\u30f3\u3092\u884c\u3044\u307e\u3059\u3002\u3053\u308c\u306f\u3001\u30ad\u30fc\u3001\u30af\u30a8\u30ea\u3001\u304a\u3088\u3073\u5024\u306e\u30d9\u30af\u30c8\u30eb\u3092\u5909\u63db\u3059\u308b\u305f\u3081\u306b\u4f7f\u7528\u3055\u308c\u307e\u3059\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_^_\u306b\u306f\u5f62\u72b6\u304c\u3042\u308a_^_1_^_\u3001\u6700\u521d\u306e\u6b21\u5143\u306f\u30af\u30a8\u30ea\u6b21\u5143\u3067\u3059\u3002_^_2_^_\u30af\u30a8\u30ea\u30c7\u30a3\u30e1\u30f3\u30b7\u30e7\u30f3\u304c\u305d\u308c\u3068\u7b49\u3057\u3044\u5834\u5408\u306f\u30d6\u30ed\u30fc\u30c9\u30ad\u30e3\u30b9\u30c8\u3055\u308c\u307e\u3059
\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_^__^_2_^_\u304a\u3088\u3073\u306f\u3001\u30af\u30a8\u30ea\u3001\u30ad\u30fc\u3001\u304a\u3088\u3073\u5024\u306e\u30d9\u30af\u30c8\u30eb\u306e\u30b3\u30ec\u30af\u30b7\u30e7\u30f3\u3092\u683c\u7d0d\u3059\u308b\u30c6\u30f3\u30bd\u30eb\u3067\u3059\u3002\u5f62\u304c\u3042\u308a\u307e\u3059_^_3_^_\u3002
\n_^_4_^__^_5_^_\u5f62\u72b6\u304c\u3042\u308a\u3001\u30d0\u30c3\u30c1\u306e\u5834\u5408_^_7_^_\u3001_^_6_^__^_8_^_\u305d\u306e\u4f4d\u7f6e\u306e\u30af\u30a8\u30ea\u304c\u305d\u306e\u4f4d\u7f6e\u306e\u30ad\u30fc\u5024\u306b\u30a2\u30af\u30bb\u30b9\u3067\u304d\u308b\u304b\u3069\u3046\u304b\u3092\u793a\u3057\u307e\u3059\u3002_^_9_^_
\n", "_^_0_^_ attention along the key sequence dimension _^_1_^_
\n": "_^_0_^_\u30ad\u30fc\u30b7\u30fc\u30b1\u30f3\u30b9\u6b21\u5143\u306b\u6cbf\u3063\u3066\u6ce8\u76ee _^_1_^_
\n", "_^_0_^_, _^_1_^_ and _^_2_^_ have shape _^_3_^_
\n": "_^_0_^_\u3001_^_1_^__^_2_^_\u305d\u3057\u3066\u5f62\u304c\u3042\u308b _^_3_^_
\n", "Apply dropout
\n": "\u30c9\u30ed\u30c3\u30d7\u30a2\u30a6\u30c8\u3092\u9069\u7528
\n", "Apply mask
\n": "\u30de\u30b9\u30af\u3092\u9069\u7528
\n", "Calculate _^_0_^_ or _^_1_^_
\n": "_^_0_^_\u8a08\u7b97\u307e\u305f\u306f _^_1_^_
\n", "Compute attention scores _^_0_^_. This gives a tensor of shape _^_1_^_.
\n": "_^_0_^_\u30a2\u30c6\u30f3\u30b7\u30e7\u30f3\u30b9\u30b3\u30a2\u3092\u8a08\u7b97\u3057\u307e\u3059\u3002_^_1_^_\u3053\u308c\u306b\u3088\u308a\u5f62\u72b6\u306e\u30c6\u30f3\u30bd\u30eb\u304c\u5f97\u3089\u308c\u307e\u3059
\u3002\n", "Concatenate multiple heads
\n": "\u8907\u6570\u306e\u30d8\u30c3\u30c9\u3092\u9023\u7d50
\n", "Dropout
\n": "\u30c9\u30ed\u30c3\u30d7\u30a2\u30a6\u30c8
\n", "Input has shape _^_0_^_ or _^_1_^_. We apply the linear transformation to the last dimension and split that into the heads.
\n": "_^_0_^__^_1_^_\u5165\u529b\u306e\u5f62\u72b6\u306f\u307e\u305f\u306f\u3067\u3059\u3002\u7dda\u5f62\u5909\u63db\u3092\u6700\u5f8c\u306e\u6b21\u5143\u306b\u9069\u7528\u3057\u3001\u305d\u308c\u3092\u982d\u306b\u5206\u5272\u3057\u307e\u3059\u3002
\n", "Linear layer for linear transform
\n": "\u7dda\u5f62\u5909\u63db\u7528\u306e\u7dda\u5f62\u5c64
\n", "Linear transform
\n": "\u7dda\u5f62\u5909\u63db
\n", "Multiply by values _^_0_^_
\n": "\u5024\u306b\u3088\u308b\u4e57\u7b97 _^_0_^_
\n", "Number of dimensions in vectors in each head
\n": "\u5404\u30d8\u30c3\u30c9\u306e\u30d9\u30af\u30c8\u30eb\u306e\u6b21\u5143\u6570
\n", "Number of features per head
\n": "\u30d8\u30c3\u30c9\u3042\u305f\u308a\u306e\u6a5f\u80fd\u6570
\n", "Number of heads
\n": "\u30d8\u30c3\u30c9\u6570
\n", "Output has shape _^_0_^_ or _^_1_^_
\n": "_^_0_^_\u51fa\u529b\u306e\u5f62\u72b6\u304c\u3042\u308b\u304b _^_1_^_
\n", "Output layer
\n": "\u51fa\u529b\u30ec\u30a4\u30e4\u30fc
\n", "Prepare _^_0_^_, _^_1_^_ and _^_2_^_ for attention computation. These will then have shape _^_3_^_.
\n": "_^_0_^__^_1_^__^_2_^_\u6ce8\u610f\u529b\u8a08\u7b97\u306e\u6e96\u5099\u3092\u3057\u3066_^_3_^_\u3053\u308c\u3067\u5f62\u304c\u3067\u304d\u3042\u304c\u308a\u307e\u3059\u3002
\n", "Same mask applied to all heads.
\n": "\u3059\u3079\u3066\u306e\u982d\u306b\u540c\u3058\u30de\u30b9\u30af\u3092\u304b\u3051\u307e\u3057\u305f\u3002
\n", "Save attentions for any other calculations
\n": "\u4ed6\u306e\u8a08\u7b97\u306b\u6ce8\u610f\u3092\u5411\u3051\u3066\u304a\u304f
\n", "Save attentions if debugging
\n": "\u30c7\u30d0\u30c3\u30b0\u6642\u306e\u6ce8\u610f\u4e8b\u9805\u3092\u4fdd\u5b58
\n", "Scale scores _^_0_^_
\n": "\u30b9\u30b1\u30fc\u30eb\u30b9\u30b3\u30a2 _^_0_^_
\n", "Scaling factor before the softmax
\n": "\u30bd\u30d5\u30c8\u30de\u30c3\u30af\u30b9\u524d\u306e\u30b9\u30b1\u30fc\u30ea\u30f3\u30b0\u30d5\u30a1\u30af\u30bf\u30fc
\n", "Softmax for attention along the time dimension of _^_0_^_
\n": "\u6642\u9593\u8ef8\u306b\u6cbf\u3063\u305f\u6ce8\u76ee\u306e\u30bd\u30d5\u30c8\u30de\u30c3\u30af\u30b9 _^_0_^_
\n", "Split last dimension into heads
\n": "\u6700\u5f8c\u306e\u30c7\u30a3\u30e1\u30f3\u30b7\u30e7\u30f3\u3092\u30d8\u30c3\u30c9\u306b\u5206\u5272
\n", "These transform the _^_0_^_, _^_1_^_ and _^_2_^_ vectors for multi-headed attention.
\n": "\u3053\u308c\u3089\u306f_^_0_^_\u3001\u3001_^_1_^__^_2_^_\u306e\u30d9\u30af\u30c8\u30eb\u3092\u5909\u3048\u3066\u3001\u591a\u9762\u7684\u306a\u6ce8\u610f\u3092\u4fc3\u3057\u307e\u3059\u3002
\n", "We store attentions so that it can be used for logging, or other computations if needed
\n": "\u5fc5\u8981\u306b\u5fdc\u3058\u3066\u30ed\u30ae\u30f3\u30b0\u3084\u305d\u306e\u4ed6\u306e\u8a08\u7b97\u306b\u4f7f\u7528\u3067\u304d\u308b\u3088\u3046\u306b\u3001\u30a2\u30c6\u30f3\u30b7\u30e7\u30f3\u3092\u4fdd\u5b58\u3057\u307e\u3059
\n", "resulting mask has shape _^_0_^_
\n": "\u751f\u6210\u3055\u308c\u308b\u30de\u30b9\u30af\u306b\u306f\u5f62\u72b6\u304c\u3042\u308a\u307e\u3059 _^_0_^_
\n", "