【pytorch系列】 nn.MultiheadAttention 详解

关于MultiheadAttention ：

一种注意力机制，常置于Transformer的开头。

Transformer自2017年推出之后，已经横扫NLP领域，成为当之无愧的state-of-the-art。原始paper “Attention is All you Need”中对attention提出了通用的query/key/value抽象。

新型的网络结构： Transformer，里面所包含的注意力机制称之为 self-attention。这套 Transformer 是能够计算 input 和 output 的 representation 而不借助 RNN 的的 model，所以作者说有 attention 就够了。

模型：同样包含 encoder 和 decoder 两个 stage，encoder 和 decoder 都是抛弃 RNN，而是用堆叠起来的 self-attention，和 fully-connected layer 来完成，模型的架构如下：

所谓Multi-Head Attention其实是把QKV的计算并行化，

1. 原始attention计算d_model维的向量，而Multi-Head Attention则是将d_model维向量先经过一个Linear Layer，
2. 再分解为h个Head计算attention，
3. 最终将这些attention向量连在一起后再经过一层Linear Layer输出。

所以在整个过程中需要4个输入和输出维度都是d_model的Linear Layer，而整个Model的输入是(batch_size, seq_length, d_model)，输出也是(batch_size, seq_length, d_model)。

multi-head attention 由多个 scaled dot-product attention 这样的基础单元经过 stack 而成。

那么 Q、K、V 到底是什么？encoder 里的 attention 叫 self-attention，顾名思义，就是自己和自己做 attention。在传统的 seq2seq 中的 encoder 阶段，我们得到 n 个时刻的 hidden states 之后，可以用每一时刻的 hidden state hi，去分别和任意的 hidden state hj,j=1,2,…,n 计算 attention，这就有点 self-attention 的意思。回到当前的模型，由于抛弃了 RNN，encoder 过程就没了 hidden states，那拿什么做 self-attention 来自嗨呢？

可以想到，假如作为 input 的 sequence 共有 n 个 word，那么我可以先对每一个 word 做 embedding 吧？就得到 n 个 embedding，然后我就可以用 embedding 代替 hidden state 来做 self-attention 了。所以 Q 这个矩阵里面装的就是全部的 word embedding，K、V 也是一样。

所以为什么管 Q 叫query？就是你每次拿一个 word embedding，去“查询”其和任意的 word embedding 的 match 程度（也就是 attention 的大小），你一共要做 n 轮这样的操作。

我们记 word embedding 的 dimension 为 dmodel ，所以 Q 的 shape 就是 n*dmodel， K、V 也是一样，第 i 个 word 的 embedding 为 vi，所以该 word 的 attention 应为：

nn.MultiheadAttention 用法：

self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
1

params:
d_model:是每一个单词本来的词向量长度；
nhead是我们MultiheadAttention的head的数量。

进行forward操作

attn_output, attn_output_weights = multihead_attn(query, key, value)
1