自注意力机制与多头注意力机制

44 阅读2分钟

1.引言

自注意力(self-attention)机制是Tramsformer架构的一个关键底层model,本篇文章-作为自用备忘录-将会着重讲解自注意力机制和多头注意力(multi head-attention)机制的思想和一般形式。

2.自注意力机制

2.1 原始数据输入

  对于文本翻译和文字生成类任务而言,输入通常是一个文本序列。将文本序列拆分为更小的词元(token)后,通过词嵌入(Word Embedding)的方式将词元投射到更高特征维度空间形成固定维度的稠密向量,这一步主要是为了提取出词语的高维特征。

2.2 算法核心思路

  自注意力层通过计算输入序列中每个位置与其他所有位置之间的相关性(由 Query 和 Key 决定),并将这些相关性归一化为权重,然后用这些权重对所有位置的 Value 向量进行加权求和,从而为每个位置生成一个融合了全局上下文信息的新表示。每个位置的输入向量aia_i都有对应一组qiq_i,kik_i和viv_i。下文将先以单个向量为例推导自注意力机制,再给出便于并行计算的矩阵形式
  对于输入序列中位置 ii 的向量 ai∈Rdmodela_i \in \mathbb{R}^{d_{\text{model}}},其对应的 Query、Key、Value 向量分别通过三个可学习的参数矩阵 WQW_Q、WKW_K、WVW_V 进行线性变换得到:

qi=aiWQ(WQ∈Rdmodel×dk)q_i = a_i W_Q \quad (W_Q \in \mathbb{R}^{d_{\text{model}} \times d_k})
ki=aiWK(WK∈Rdmodel×dk)k_i = a_i W_K \quad (W_K \in \mathbb{R}^{d_{\text{model}} \times d_k})
vi=aiWV(WV∈Rdmodel×dv)v_i = a_i W_V \quad (W_V \in \mathbb{R}^{d_{\text{model}} \times d_v})

  其中,dkd_k 和 dvd_v 分别是 Query/Key 和 Value 的维度,通常 dk=dvd_k = d_v。(在后续的多头注意力机制中 dk=dv=dmodel/hd_k = d_v = d_{\text{model}} / h)
  qiq_i,kik_i,viv_i均来自同一个向量aia_i,但是由于WQW_Q、WKW_K、WVW_V的不同被赋予了截然不同的角色。
  Query(查询) 表示当前词正在寻找的信息。每个位置都会用自己的 Query 去和其他位置的 Key 进行匹配,从而决定应该关注哪些位置。
  Key(键) 表示当前词能够提供什么匹配特征。它相当于一个标签,Query 与 Key 的点积结果反映了二者之间的相关程度,也就是注意力分数。
  Value(值) 表示当前词实际承载的语义内容。当 Query 与某个 Key 匹配后,模型就会从该位置取出对应的 Value,并将其融入当前位置的输出中。
  自注意力机制最重要的作用是能够让输出序列中的每个向量(包含自身)都融合所有位置的输入信息,而实现这一点的前提是先计算出当前位置输入aia_i与所有位置(包括自己在内)输入的相关性αi\alpha_i,再利用相关性权重αi\alpha_i计算得到的最终注意力权重计算输出。在《Attention is all you need》中相关性权重α\alpha的计算使用如下点积(Dot-product)的形式

  上图中展示的是点积计算相关性α\alpha的方法,qi,kjq_i, k_j均由位置输入ai,aja_i, a_j乘矩阵WQ,WkW_Q,W_k得到,再将qi,kjq_i, k_j点乘即可求得aia_i位置相对aja_j位置的相关性权重αi,j\alpha_{i,j}
  按照这样的方式,我们可以计算出aia_i相对所有输入位置的相关性(包括aia_i自己)

  在Transformer论文中,计算出的相关性权重还需要进行一层softmax层以获得最终的注意力权重α′\alpha'

  在得到了a1a_1相对所有位置的相关性权重后,就可以计算a1a_1位置的输出b1b_1。如下图所示,将v1,v2,v3,v4v_1,v_2,v_3, v_4分别按各自的相关性权重加权求和后便得到了b1b_1。按照同样的方法可以求得各位置的输出b2,b3,b4b_2,b_3,b_4

2.3 算法的矩阵形式

  假设输入序列长度为 nn,模型维度为 dmodeld_{\text{model}},输入矩阵为:

X∈Rn×dmodelX \in \mathbb{R}^{n \times d_{\text{model}}}

  其中每一行对应一个位置的输入向量。

  通过三个可学习的参数矩阵 WQW_Q、WKW_K、WVW_V,分别得到 Query、Key、Value 矩阵:

Q=XWQ(WQ∈Rdmodel×dk)Q = X W_Q \quad (W_Q \in \mathbb{R}^{d_{\text{model}} \times d_k})
K=XWK(WK∈Rdmodel×dk)K = X W_K \quad (W_K \in \mathbb{R}^{d_{\text{model}} \times d_k})
V=XWV(WV∈Rdmodel×dv)V = X W_V \quad (W_V \in \mathbb{R}^{d_{\text{model}} \times d_v})

得到的矩阵形状分别为:

Q∈Rn×dkQ \in \mathbb{R}^{n \times d_k}
K∈Rn×dkK \in \mathbb{R}^{n \times d_k}
V∈Rn×dvV \in \mathbb{R}^{n \times d_v}

其中:

  • nn 是序列长度;

  • dkd_k 是 Query 和 Key 的维度,二者必须相同,因为需要计算 QKTQK^T;

  • dvd_v 是 Value 的维度,可以不同于 dkd_k,它决定了注意力输出的维度。

    自注意力层的输出可以表示为如下形式:

Attention(Q,K,V)=softmax(QKTdk)V\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{Q K^T}{\sqrt{d_k}}\right) V

上式与先前推导的唯一不同就是在softmax前除以dk\sqrt{d_k}以排除因维度dkd_k过大而导致的softmax饱和,保证梯度能够有效传播(也可以理解为让注意力分数的方差与维度dkd_k无关)

3 多头注意力机制(Multi-Head Attention)

多头注意力机制的公式为:

MultiHead(Q,K,V)=Concat(head1,head2,…,headh)WO\text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, \text{head}_2, \dots, \text{head}_h) W_O

其中每个头的计算为:

headi=Attention(QWQi, KWKi, VWVi)\text{head}_i = \text{Attention}(Q W_Q^i,\ K W_K^i,\ V W_V^i)

Attention 函数为缩放点积注意力:

Attention(Q,K,V)=softmax(QKTdk)V\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{Q K^T}{\sqrt{d_k}}\right) V

参数矩阵与维度:

WQi∈Rdmodel×dk,WKi∈Rdmodel×dk,WVi∈Rdmodel×dvW_Q^i \in \mathbb{R}^{d_{\text{model}} \times d_k},\quad W_K^i \in \mathbb{R}^{d_{\text{model}} \times d_k},\quad W_V^i \in \mathbb{R}^{d_{\text{model}} \times d_v}
WO∈Rhdv×dmodelW_O \in \mathbb{R}^{h d_v \times d_{\text{model}}}

通常设 dk=dv=dmodel/hd_k = d_v = d_{\text{model}} / h,其中 hh 是注意力头的数量。

多头注意力机制可以理解为:将输入分别投影到n个低维子空间,在每一个空间分别独立进行一次维度为dmodel/hd_{\text{model}} / h的注意力计算,将所有头的输出在特征维度上拼接,并最终乘以一个线性变换矩阵得到最终输出。

4 参考文献

【Transformer论文逐段精读【论文精读】】 www.bilibili.com/video/BV1pu…

【(2025版)李宏毅机器学习深度学习系列课程全集,公认体验感最好的入门课程!--人工智能/机器学习/深度学习】 www.bilibili.com/video/BV1TA…

Vaswani, Ashish et al. “Attention is All you Need.” Neural Information Processing Systems (2017).