注意力机制详解

注意力机制(Attention Mechanism)是现代深度学习中最核心的技术之一,广泛应用于自然语言处理、计算机视觉等领域。本文将深入解析注意力机制的输入输出及其运算过程。

什么是注意力机制

注意力机制的核心思想源于人类的视觉注意力:人类在观察事物时,会聚焦于关键部分而忽略无关信息。同样地,神经网络中的注意力机制允许模型动态地关注输入的不同部分,从而更好地捕捉重要信息。

输入与输出

输入

注意力机制的输入通常包括三个向量序列:

  • Query (Q): 查询向量,表示当前关注点
  • Key (K): 键向量,用于和查询进行匹配
  • Value (V): 值向量,包含实际的信息内容

假设我们有一个长度为 nn 的输入序列,每个位置的向量维度为 dkd_k,则:

  • QRm×dkQ \in \mathbb{R}^{m \times d_k} (通常 mm 是目标序列长度)
  • KRn×dkK \in \mathbb{R}^{n \times d_k}
  • VRn×dvV \in \mathbb{R}^{n \times d_v}

输出

注意力机制的输出是一个新的向量序列,其形状为:

  • OutputRm×dvOutput \in \mathbb{R}^{m \times d_v}

输出向量是所有 Value 向量的加权和,权重由 Query 和 Key 的相似度决定。

核心运算过程

第一步:计算注意力分数

首先计算 Query 和 Key 之间的相似度分数。最常用的是点积注意力:

Score(Q,K)=QKTScore(Q, K) = QK^T

对于单个查询向量 qq 和所有键向量 KK,分数计算为:

score(q,ki)=qkiTscore(q, k_i) = q \cdot k_i^T

这会产生一个分数向量 scoresRnscores \in \mathbb{R}^n

第二步:缩放

为了防止点积过大导致 softmax 函数梯度消失,通常会对分数进行缩放:

scoresscaled=QKTdkscores_{scaled} = \frac{QK^T}{\sqrt{d_k}}

其中 dkd_k 是键向量的维度。缩放因子的引入是因为:假设 qqkk 的元素独立且均值为 0、方差为 1,则它们的点积均值为 0、方差为 dkd_k。除以 dk\sqrt{d_k} 可以将方差归一化为 1。

第三步:Softmax 归一化

将缩放后的分数通过 softmax 函数转换为概率分布:

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

对于单个查询,注意力权重为:

αi=exp(scorei)j=1nexp(scorej)\alpha_i = \frac{exp(score_i)}{\sum_{j=1}^{n} exp(score_j)}

其中 αi\alpha_i 表示第 ii 个位置的注意力权重,满足 i=1nαi=1\sum_{i=1}^{n} \alpha_i = 1

第四步:加权求和

最后,用注意力权重对 Value 向量进行加权求和:

output=i=1nαivioutput = \sum_{i=1}^{n} \alpha_i v_i

写成矩阵形式:

Output=Attention(Q,K,V)VOutput = Attention(Q, K, V) \cdot V

其中 Attention(Q,K,V)Rm×nAttention(Q, K, V) \in \mathbb{R}^{m \times n} 是注意力权重矩阵。

完整公式

将上述步骤整合,注意力机制的完整计算公式为:

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

自注意力机制(Self-Attention)

自注意力是注意力机制的一种特殊形式,其中 Query、Key、Value 都来自同一个输入序列 XX

Q=XWQ,K=XWK,V=XWVQ = XW^Q, \quad K = XW^K, \quad V = XW^V

其中 WQRdmodel×dkW^Q \in \mathbb{R}^{d_{model} \times d_k}WKRdmodel×dkW^K \in \mathbb{R}^{d_{model} \times d_k}WVRdmodel×dvW^V \in \mathbb{R}^{d_{model} \times d_v} 是可学习的参数矩阵。

自注意力使得序列中的每个位置都能关注到序列中的其他所有位置,从而捕捉序列内部的依赖关系。

多头注意力(Multi-Head Attention)

为了从不同表示子空间捕获不同类型的信息,多头注意力机制将上述过程并行执行多次:

MultiHead(Q,K,V)=Concat(head1,...,headh)WOMultiHead(Q, K, V) = Concat(head_1, ..., head_h)W^O

其中:

headi=Attention(QWiQ,KWiK,VWiV)head_i = Attention(QW_i^Q, KW_i^K, VW_i^V)

每个 head 都有独立的参数矩阵,允许模型在不同的位置关注来自不同表示子空间的信息。

掩码注意力(Masked Attention)

在某些场景(如解码器)中,需要限制注意力只关注特定位置。通过掩码矩阵 MM 实现:

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

掩码矩阵中需要屏蔽的位置设为 -\infty,允许关注的位置设为 00。这样在 softmax 后,被屏蔽位置的权重将接近于 0。

计算复杂度分析

对于序列长度 nn 和维度 dd

  • 时间复杂度: O(n2d)O(n^2 \cdot d),主要由 QKTQK^T 的矩阵乘法决定
  • 空间复杂度: O(n2+nd)O(n^2 + n \cdot d),需要存储 n×nn \times n 的注意力矩阵

这也是注意力机制在处理长序列时的主要瓶颈。

实际应用示例

假设我们有一个句子:"The cat sat on the mat",经过词嵌入后得到序列 XR6×dX \in \mathbb{R}^{6 \times d}

在自注意力计算中,对于单词 "cat":

  1. 生成 qcat=xcatWQq_{cat} = x_{cat}W^Q,以及所有位置的 KKVV
  2. 计算 "cat" 与所有单词的注意力分数
  3. 可能得到高权重对:cat→sat(动词关系),cat→the(修饰关系)
  4. 输出向量融合了这些相关单词的信息

这种机制使得模型能够动态理解单词在上下文中的含义和关系。

总结

注意力机制通过 Query-Key-Value 的设计,实现了信息的动态选择和聚合:

组件作用
Query定义"我在寻找什么"
Key定义"我是什么"
Value定义"我包含什么信息"
注意力权重衡量相关性,决定信息流

这种优雅的设计使得模型能够自适应地关注最重要的信息,成为现代深度学习架构的基础组件。

参考

  1. Vaswani, A., et al. "Attention Is All You Need." NeurIPS 2017.
  2. Bahdanau, D., et al. "Neural Machine Translation by Jointly Learning to Align and Translate." ICLR 2015.
  3. https://arxiv.org/abs/1706.03762