Multi-head self-attention

先验知识：Self-attention (SA)

通过阅读先验知识，我们知道实际上在单头自注意力（普通的 SA）中，

• 假设输入的特征 $\mathbf{X}$ 维度是 $d_{in}$ (可以通过变换改变维度至 $\mathbf{Q},\mathbf{K}$ 的维度，但是输入的维度这里就暂时这样表示)
• 输出的特征 $\mathbf{Y}$ 维度一般和 value $\mathbf{V}$ 的维度是一样的，我们用 $d_{out}$ 来表示

---

对于多头自注意力，上述维度的表示依然相同。

假设有 $N_h$ 个头，每个头都要对 $\mathbf{X}$ 进行一次单头 SA 操作，设每个 head 的输出维度是 $d_h$，如果 $d_h\times N_h=d_{out}$，那么皆大欢喜，我们直接将所有 head 的输出进行拼接即可。但是应用中很可能两者并不严格等于，此时我们应该额外添建一个 FC 层 (${\mathbf{W}}_{out}$) 用来对齐维度，使得最终的输出为 $d_{out}$，公式表示如下

$$\mathrm{MHSA}(\mathbf{X}):=\underset{h\in \left[N_h\right]}{\mathrm{concat}}[\text{ Self-Attention }(\mathbf{X})]{\mathbf{W}}_{\text{out }}+{\mathbf{b}}_{\text{out }}$$

[! Attention]

1. 实际应用中，我们常常指定 $N_h$, 然后计算 $d_h=d_{out}//N_h$
2. 每个 head 直接接受 $\mathbf{X}$ 作为输入，而不会对 $\mathbf{X}$ 做 Split (分成 $N_h$ 份)