当前位置: 首页 > ds >正文

LLM学习笔记(七)注意力机制

文章目录

1. 传统序列标注的问题

image.png

传统的序列标注方法(如使用简单的全连接层FC或固定窗口的卷积)存在一些局限性:

  1. 全连接层 (FC) 的问题

    • 无法区分语境:对于一个序列中的同一个词,如果它在不同位置出现,全连接层通常会独立处理每个词的嵌入,无法根据上下文理解该词在特定语境下的具体含义。例如,“bank”可以是银行,也可以是河岸。FC层本身不具备捕捉这种上下文依赖的能力。
    • 固定权重:FC层的权重是固定的(在训练后),它对所有输入的处理方式相同,缺乏动态性。
  2. 固定窗口机制 (如CNN) 的问题

    • 上下文受限:通过固定大小的窗口来捕捉上下文信息,只能考虑到局部上下文,无法有效捕捉序列中距离较远的元素之间的依赖关系(长距离依赖)。
    • 过拟合风险:如果窗口大小选择不当,或者数据量不足,模型可能学习到一些仅在局部窗口内成立的模式,而这些模式不具有普适性,导致过拟合。

2. 自注意力机制 (Self-Attention)

自注意力机制旨在克服上述传统方法的局限性,特别是捕捉序列内部长距离依赖关系和根据上下文动态调整词表示。

image.png

核心特点

  1. 等长输出:输入序列包含多少个元素(词向量),自注意力机制就会输出多少个对应的新元素(新的词向量)。输入 a 1 , a 2 , . . . , a n a^1, a^2, ..., a^n a1,a2,...,an,输出 b 1 , b 2 , . . . , b n b^1, b^2, ..., b^n b1,b2,...,bn
  2. 可堆叠:自注意力层可以堆叠多层,每一层都在前一层输出的基础上进一步提炼信息,构建更复杂的表示。
  3. 全局上下文感知:每个输出元素 b i b^i bi 的计算都考虑了所有输入元素 a 1 , a 2 , . . . , a n a^1, a^2, ..., a^n a1,a2,...,an 的信息。这意味着 b i b^i bi 是整个输入序列的加权聚合,权重则体现了其他词与当前词的关联程度。
2.1 计算关联程度 α \alpha α

为了计算输出 b 1 b^1 b1,==我们需要知道输入序列中的每个向量 a i a^i ai (包括 a 1 a^1 a1 自身) 与 a 1 a^1 a1 的关联程度。==这个关联程度用数值 α 1 , i \alpha_{1,i} α1,i 来表示。
例如,计算 b 1 b^1 b1 时,需要计算 α 1 , 1 , α 1 , 2 , α 1 , 3 , α 1 , 4 \alpha_{1,1}, \alpha_{1,2}, \alpha_{1,3}, \alpha_{1,4} α1,1,α1,2,α1,3,α1,4
计算 b 2 b^2 b2 时,需要计算 α 2 , 1 , α 2 , 2 , α 2 , 3 , α 2 , 4 \alpha_{2,1}, \alpha_{2,2}, \alpha_{2,3}, \alpha_{2,4} α2,1,α2,2,α2,3,α2,4
通常,我们关注的是如何计算一个“查询”(query,比如 a 1 a^1 a1)与所有“键”(keys,比如 a 1 , a 2 , . . . , a n a^1, a^2, ..., a^n a1,a2,...,an)之间的关联度。

两种常见的计算 α \alpha α (原始注意力分数) 的方法

  1. 点积 (Dot-Product)
    image.png
    假设我们想计算 a 1 a^1 a1 (作为查询-K) 与 a i a^i ai (作为键-V) 之间的关联度 α 1 , i \alpha_{1,i} α1,i
    α 1 , i = q 1 ⋅ k i \alpha_{1,i} = q^1 \cdot k^i α1,i=q1ki
    在基础的自注意力中,查询向量 q 1 q^1 q1 和键向量 k i k^i ki 可以直接是输入向量 a 1 a^1 a1 a i a^i ai 经过线性变换得到的。
    更具体地,对于输入 a j a^j aj a i a^i ai:

    • q j = W Q a j q^j = W^Q a^j qj=WQaj (将 a j a^j aj 转换为查询向量)
    • k i = W K a i k^i = W^K a^i ki=WKai (将 a i a^i ai 转换为键向量)
      α j , i = ( q j ) T k i \alpha_{j,i} = (q^j)^T k^i αj,i=(qj)Tki
      在 “Attention Is All You Need” 论文中,还使用了一个缩放因子 d k \sqrt{d_k} dk (其中 d k d_k dk 是键向量的维度) 来防止点积结果过大导致梯度过小:
      α j , i = ( q j ) T k i d k \alpha_{j,i} = \frac{(q^j)^T k^i}{\sqrt{d_k}} αj,i=dk (qj)Tki(说明过程在附录中)
  2. 相加式/拼接式 (Additive Attention)
    image.png
    这种方法首先将查询向量 q j q^j qj 和键向量 k i k^i ki 拼接起来,然后通过一个带有非线性激活函数(通常是 tanh)的单层前馈网络,最后再通过一个权重向量 w a w_a wa 将其转换为标量:
    α j , i = w a T tanh ⁡ ( W q q j + W k k i + b ) \alpha_{j,i} = w_a^T \tanh(W_q q^j + W_k k^i + b) αj,i=waTtanh(Wqqj+Wkki+b)
    或者如图中所示,将 q j q^j qj k i k^i ki 分别乘以权重矩阵 W 1 W_1 W1 W 2 W_2 W2 (或者一个共享的 W W W 作用于拼接向量),然后相加,通过 tanh,再乘以 v T v^T vT
    e j , i = v T tanh ⁡ ( W 1 q j + W 2 k i ) e_{j,i} = v^T \tanh(W_1 q^j + W_2 k^i) ej,i=vTtanh(W1qj+W2ki) (这里的 e j , i e_{j,i} ej,i 就是 α j , i \alpha_{j,i} αj,i)

2.2 α \alpha α 与注意力机制的关联:Softmax归一化

image.png

计算得到的原始注意力分数 α 1 , i \alpha_{1,i} α1,i (这里指计算 b 1 b^1 b1 时, a 1 a^1 a1 与各个 a i a^i ai 的关联度) 需要被转换成权重,这些权重表示在生成 b 1 b^1 b1 时,应该给每个输入 a i a^i ai 分配多少“注意力”。

添加激活函数 Softmax
image.png
将针对特定查询(如 a 1 a^1 a1)计算出的所有原始分数 [ α 1 , 1 , α 1 , 2 , . . . , α 1 , n ] [\alpha_{1,1}, \alpha_{1,2}, ..., \alpha_{1,n}] [α1,1,α1,2,...,α1,n] 通过 Softmax 函数进行归一化,得到最终的注意力权重 α 1 , i ′ \alpha'_{1,i} α1,i
α 1 , i ′ = softmax ( α 1 , i ) = exp ⁡ ( α 1 , i ) ∑ k = 1 n exp ⁡ ( α 1 , k ) \alpha'_{1,i} = \text{softmax}(\alpha_{1,i}) = \frac{\exp(\alpha_{1,i})}{\sum_{k=1}^{n} \exp(\alpha_{1,k})} α1,i=softmax(α1,i)=k=1nexp(α1,k)exp(α1,i)

  • 为什么选择 Softmax?(附录有详细说明)
    1. 归一化:Softmax 将任意实数值的分数转换为总和为 1 的概率分布。这使得 α 1 , i ′ \alpha'_{1,i} α1,i 可以被解释为 a i a^i ai 对生成 b 1 b^1 b1 的贡献权重。
    2. 突出重要性:指数函数会放大数值之间的差异,使得较大的 α 1 , i \alpha_{1,i} α1,i 对应的 α 1 , i ′ \alpha'_{1,i} α1,i 更大,从而让模型更关注那些与当前查询更相关的输入。
    3. 非负性:权重都是正数。
    4. 可微性:Softmax 函数是可微的,这对于使用梯度下降进行模型训练至关重要。
2.3 计算输出 b i b^i bi

image.png

得到归一化的注意力权重 α j , i ′ \alpha'_{j,i} αj,i 后,输出向量 b j b^j bj 就是所有输入向量 a i a^i ai (更准确地说是它们对应的值向量 v i v^i vi) 的加权和:
v i = W V a i v^i = W^V a^i vi=WVai (将 a i a^i ai 转换为值向量)
b j = ∑ i = 1 n α j , i ′ v i b^j = \sum_{i=1}^{n} \alpha'_{j,i} v^i bj=i=1nαj,ivi

例如,对于 b 1 b^1 b1
b 1 = α 1 , 1 ′ v 1 + α 1 , 2 ′ v 2 + α 1 , 3 ′ v 3 + α 1 , 4 ′ v 4 b^1 = \alpha'_{1,1} v^1 + \alpha'_{1,2} v^2 + \alpha'_{1,3} v^3 + \alpha'_{1,4} v^4 b1=α1,1v1+α1,2v2+α1,3v3+α1,4v4
这意味着 b 1 b^1 b1 是通过聚合所有输入 a i a^i ai (转换后的 v i v^i vi) 的信息生成的,而聚合的权重 α 1 , i ′ \alpha'_{1,i} α1,i 则取决于 a 1 a^1 a1 (的查询形式 q 1 q^1 q1) 与各个 a i a^i ai (的键形式 k i k^i ki) 的相关性。

3. 矩阵乘法的角度 (Scaled Dot-Product Attention)

为了高效计算,通常将上述操作表示为矩阵运算。
假设我们有一个输入序列包含 N N N 个词,每个词的嵌入向量维度为 d m o d e l d_{model} dmodel
输入矩阵 A A A (或称为 X X X) 的维度是 N × d m o d e l N \times d_{model} N×dmodel,其中每一行是一个词向量 a i a^i ai

  1. 生成 Query (Q), Key (K), Value (V) 矩阵
    通过将输入矩阵 A A A 分别乘以三个可学习的权重矩阵 W Q , W K , W V W^Q, W^K, W^V WQ,WK,WV 来得到 Q , K , V Q, K, V Q,K,V

    • W Q ∈ R d m o d e l × d k W^Q \in \mathbb{R}^{d_{model} \times d_k} WQRdmodel×dk
    • W K ∈ R d m o d e l × d k W^K \in \mathbb{R}^{d_{model} \times d_k} WKRdmodel×dk
    • W V ∈ R d m o d e l × d v W^V \in \mathbb{R}^{d_{model} \times d_v} WVRdmodel×dv
      (通常 d k = d v d_k = d_v dk=dv,在多头注意力中)

    Q = A W Q Q = A W^Q Q=AWQ (维度: N × d k N \times d_k N×dk)
    K = A W K K = A W^K K=AWK (维度: N × d k N \times d_k N×dk)
    V = A W V V = A W^V V=AWV (维度: N × d v N \times d_v N×dv)

    image.png
    (图中 X X X 即为我们的输入 A A A)

  2. 计算注意力分数 (Attention Scores)
    Scores = Q K T \text{Scores} = Q K^T Scores=QKT

    • K T K^T KT 的维度是 d k × N d_k \times N dk×N
    • Scores \text{Scores} Scores 矩阵的维度是 ( N × d k ) × ( d k × N ) = N × N (N \times d_k) \times (d_k \times N) = N \times N (N×dk)×(dk×N)=N×N
    • Scores j i \text{Scores}_{ji} Scoresji 表示第 j j j 个查询 q j q^j qj (即 Q Q Q 的第 j j j 行) 与第 i i i 个键 k i k^i ki (即 K K K 的第 i i i 行, K T K^T KT 的第 i i i 列) 之间的点积。

    image.png

  3. 缩放 (Scale)
    将分数除以 d k \sqrt{d_k} dk
    Scaled_Scores = Q K T d k \text{Scaled\_Scores} = \frac{Q K^T}{\sqrt{d_k}} Scaled_Scores=dk QKT

    image.png

  4. 应用 Softmax
    Scaled_Scores \text{Scaled\_Scores} Scaled_Scores 矩阵的每一行独立应用 Softmax 函数,得到注意力权重矩阵 P P P (Attention Weights)。
    P = softmax ( Q K T d k ) P = \text{softmax}(\frac{Q K^T}{\sqrt{d_k}}) P=softmax(dk QKT) (按行进行Softmax)

    • P P P 的维度仍然是 N × N N \times N N×N
    • P j i P_{ji} Pji 就是我们之前讨论的 α j , i ′ \alpha'_{j,i} αj,i,表示在计算第 j j j 个输出 b j b^j bj 时,第 i i i 个值向量 v i v^i vi 的权重。
    • 每一行的元素和为 1。

    image.png

  5. 计算输出 (Output)
    将注意力权重矩阵 P P P 与值矩阵 V V V 相乘。
    B = P V B = P V B=PV

    • B B B 的维度是 ( N × N ) × ( N × d v ) = N × d v (N \times N) \times (N \times d_v) = N \times d_v (N×N)×(N×dv)=N×dv
    • B B B 的第 j j j b j b^j bj ∑ i = 1 N P j i v i \sum_{i=1}^{N} P_{ji} v^i i=1NPjivi,这就是我们前面定义的输出向量。

    image.png

整体公式
Attention ( Q , K , V ) = softmax ( Q K T d k ) V \text{Attention}(Q, K, V) = \text{softmax}(\frac{Q K^T}{\sqrt{d_k}}) V Attention(Q,K,V)=softmax(dk QKT)V

3.1 矩阵计算展开示例 (假设 N=3, d k , d v d_k, d_v dk,dv 为某个值)

设输入 A = ( − a 1 − − a 2 − − a 3 − ) A = \begin{pmatrix} - & a^1 & - \\ - & a^2 & - \\ - & a^3 & - \end{pmatrix} A= a1a2a3 (3行, d m o d e l d_{model} dmodel列, a i a^i ai是行向量)

  1. Q = ( − q 1 − − q 2 − − q 3 − ) Q = \begin{pmatrix} - & q^1 & - \\ - & q^2 & - \\ - & q^3 & - \end{pmatrix} Q= q1q2q3 , K = ( − k 1 − − k 2 − − k 3 − ) K = \begin{pmatrix} - & k^1 & - \\ - & k^2 & - \\ - & k^3 & - \end{pmatrix} K= k1k2k3 , V = ( − v 1 − − v 2 − − v 3 − ) V = \begin{pmatrix} - & v^1 & - \\ - & v^2 & - \\ - & v^3 & - \end{pmatrix} V= v1v2v3
    ( q i , k i q^i, k^i qi,ki 1 × d k 1 \times d_k 1×dk 的行向量; v i v^i vi 1 × d v 1 \times d_v 1×dv 的行向量)

  2. K T = ( ∣ ∣ ∣ k 1 k 2 k 3 ∣ ∣ ∣ ) K^T = \begin{pmatrix} | & | & | \\ k^1 & k^2 & k^3 \\ | & | & | \end{pmatrix} KT= k1k2k3 ( d k d_k dk 行, 3列)

  3. Scores = Q K T = ( − q 1 − − q 2 − − q 3 − ) ( ∣ ∣ ∣ ( k 1 ) T ( k 2 ) T ( k 3 ) T ∣ ∣ ∣ ) \text{Scores} = Q K^T = \begin{pmatrix} - & q^1 & - \\ - & q^2 & - \\ - & q^3 & - \end{pmatrix} \begin{pmatrix} | & | & | \\ (k^1)^T & (k^2)^T & (k^3)^T \\ | & | & | \end{pmatrix} Scores=QKT= q1q2q3 (k1)T(k2)T(k3)T
    (这里为了点积,我们应该将 k i k^i ki 视为列向量进行内积,或者 q i ( k i ) T q^i (k^i)^T qi(ki)T)
    Scores = ( q 1 ( k 1 ) T q 1 ( k 2 ) T q 1 ( k 3 ) T q 2 ( k 1 ) T q 2 ( k 2 ) T q 2 ( k 3 ) T q 3 ( k 1 ) T q 3 ( k 2 ) T q 3 ( k 3 ) T ) = ( α 11 α 12 α 13 α 21 α 22 α 23 α 31 α 32 α 33 ) \text{Scores} = \begin{pmatrix} q^1 (k^1)^T & q^1 (k^2)^T & q^1 (k^3)^T \\ q^2 (k^1)^T & q^2 (k^2)^T & q^2 (k^3)^T \\ q^3 (k^1)^T & q^3 (k^2)^T & q^3 (k^3)^T \end{pmatrix} = \begin{pmatrix} \alpha_{11} & \alpha_{12} & \alpha_{13} \\ \alpha_{21} & \alpha_{22} & \alpha_{23} \\ \alpha_{31} & \alpha_{32} & \alpha_{33} \end{pmatrix} Scores= q1(k1)Tq2(k1)Tq3(k1)Tq1(k2)Tq2(k2)Tq3(k2)Tq1(k3)Tq2(k3)Tq3(k3)T = α11α21α31α12α22α32α13α23α33
    (这里 α j i \alpha_{ji} αji q j q^j qj k i k^i ki 的原始点积)

  4. Scaled_Scores j i = α j i / d k \text{Scaled\_Scores}_{ji} = \alpha_{ji} / \sqrt{d_k} Scaled_Scoresji=αji/dk

  5. P = softmax ( Scaled_Scores ) P = \text{softmax}(\text{Scaled\_Scores}) P=softmax(Scaled_Scores) (按行softmax)
    P = ( α 11 ′ α 12 ′ α 13 ′ α 21 ′ α 22 ′ α 23 ′ α 31 ′ α 32 ′ α 33 ′ ) P = \begin{pmatrix} \alpha'_{11} & \alpha'_{12} & \alpha'_{13} \\ \alpha'_{21} & \alpha'_{22} & \alpha'_{23} \\ \alpha'_{31} & \alpha'_{32} & \alpha'_{33} \end{pmatrix} P= α11α21α31α12α22α32α13α23α33
    其中 α 11 ′ + α 12 ′ + α 13 ′ = 1 \alpha'_{11} + \alpha'_{12} + \alpha'_{13} = 1 α11+α12+α13=1, 等等。

  6. B = P V = ( α 11 ′ α 12 ′ α 13 ′ α 21 ′ α 22 ′ α 23 ′ α 31 ′ α 32 ′ α 33 ′ ) ( − v 1 − − v 2 − − v 3 − ) B = P V = \begin{pmatrix} \alpha'_{11} & \alpha'_{12} & \alpha'_{13} \\ \alpha'_{21} & \alpha'_{22} & \alpha'_{23} \\ \alpha'_{31} & \alpha'_{32} & \alpha'_{33} \end{pmatrix} \begin{pmatrix} - & v^1 & - \\ - & v^2 & - \\ - & v^3 & - \end{pmatrix} B=PV= α11α21α31α12α22α32α13α23α33 v1v2v3

    B = ( − ( α 11 ′ v 1 + α 12 ′ v 2 + α 13 ′ v 3 ) − − ( α 21 ′ v 1 + α 22 ′ v 2 + α 23 ′ v 3 ) − − ( α 31 ′ v 1 + α 32 ′ v 2 + α 33 ′ v 3 ) − ) = ( − b 1 − − b 2 − − b 3 − ) B = \begin{pmatrix} -& (\alpha'_{11}v^1 + \alpha'_{12}v^2 + \alpha'_{13}v^3) & - \\ -& (\alpha'_{21}v^1 + \alpha'_{22}v^2 + \alpha'_{23}v^3) & - \\ -& (\alpha'_{31}v^1 + \alpha'_{32}v^2 + \alpha'_{33}v^3) & - \end{pmatrix} = \begin{pmatrix} - & b^1 & - \\ - & b^2 & - \\ - & b^3 & - \end{pmatrix} B= (α11v1+α12v2+α13v3)(α21v1+α22v2+α23v3)(α31v1+α32v2+α33v3) = b1b2b3

    每一行 b j b^j bj 是一个 1 × d v 1 \times d_v 1×dv 的行向量,它是所有值向量 v i v^i vi 的加权和,权重由第 j j j 个查询与其他所有键的相关性决定。

image.png
输入 X X X (即我们的 A A A) 经过线性变换得到 Q , K , V Q, K, V Q,K,V。然后 Q Q Q K T K^T KT 相乘并缩放,经过 Softmax 得到注意力权重,最后与 V V V 相乘得到输出 Z Z Z (即我们的 B B B)。

3.2 具体代码实现

使用 PyTorch 实现的自注意力机制(Scaled Dot-Product Attention)的代码示

import torch
import torch.nn as nn
import torch.nn.functional as F
import mathclass SelfAttention(nn.Module):"""实现一个基础的自注意力机制 (Scaled Dot-Product Attention)。"""def __init__(self, embedding_dim, dk, dv):"""初始化自注意力层。参数:- embedding_dim (int): 输入词嵌入的维度 (d_model)。- dk (int): 查询 (Query) 和键 (Key) 向量的维度。- dv (int): 值 (Value) 向量的维度。通常 dk == dv。"""super(SelfAttention, self).__init__()self.embedding_dim = embedding_dimself.dk = dkself.dv = dv# 线性变换层,用于生成 Q, K, V# W_q, W_k, W_v 权重矩阵self.fc_q = nn.Linear(embedding_dim, dk)self.fc_k = nn.Linear(embedding_dim, dk)self.fc_v = nn.Linear(embedding_dim, dv)# 可选: 初始化权重 (这里为了简洁省略,PyTorch 默认有合理的初始化)# self._init_weights()# def _init_weights(self):#     nn.init.xavier_uniform_(self.fc_q.weight)#     nn.init.xavier_uniform_(self.fc_k.weight)#     nn.init.xavier_uniform_(self.fc_v.weight)#     if self.fc_q.bias is not None:#         nn.init.zeros_(self.fc_q.bias)#         nn.init.zeros_(self.fc_k.bias)#         nn.init.zeros_(self.fc_v.bias)def forward(self, x, mask=None):"""前向传播。参数:- x (torch.Tensor): 输入张量,形状为 (batch_size, seq_len, embedding_dim)。对应文中的 A (或 X) 矩阵,其中每一行是一个词向量 a_i。- mask (torch.Tensor, optional): 掩码张量,用于在 softmax 前屏蔽某些位置(例如 padding)。形状通常为 (batch_size, seq_len) 或 (batch_size, seq_len, seq_len)。值为 True 或 1 的位置表示要保留,False 或 0 表示要屏蔽。在 softmax 之前,被屏蔽的位置会被设置为一个很大的负数。返回:- output (torch.Tensor): 自注意力层的输出,形状为 (batch_size, seq_len, dv)。对应文中的 B (或 Z) 矩阵。- attention_weights (torch.Tensor): 注意力权重矩阵 P,形状为 (batch_size, seq_len, seq_len)。"""batch_size, seq_len, _ = x.shape# 1. 生成 Query (Q), Key (K), Value (V) 矩阵# x: (B, N, d_model)# q, k: (B, N, dk)# v: (B, N, dv)q = self.fc_q(x)  # Q = A * W_qk = self.fc_k(x)  # K = A * W_kv = self.fc_v(x)  # V = A * W_v# 2. 计算注意力分数 (Attention Scores): Q * K^T# q: (B, N, dk)# k.transpose(-2, -1): (B, dk, N)  (转置最后两个维度)# scores: (B, N, dk) @ (B, dk, N) -> (B, N, N)scores = torch.matmul(q, k.transpose(-2, -1))# 3. 缩放 (Scale)# scaled_scores = scores / sqrt(dk)# dk**0.5 is sqrt(dk)scaled_scores = scores / (self.dk ** 0.5)# 或者: scaled_scores = scores / math.sqrt(self.dk)# (可选) 应用掩码 (Masking)# 如果提供了掩码,将掩码中为 False (或0) 的位置的 scaled_scores 设置为一个非常小的数,# 这样在 softmax 后这些位置的权重会趋近于0。if mask is not None:# 假设 mask 的形状是 (B, N) 或 (B, 1, N) 或 (B, N, N)# 我们需要将其广播到 (B, N, N)if mask.dim() == 2: # (B, N) -> (B, 1, N) for broadcastingmask = mask.unsqueeze(1)# 将 False (或0) 替换为 -1e9 (一个很大的负数)scaled_scores = scaled_scores.masked_fill(mask == 0, -1e9) # 或者 -float('inf')# 4. 应用 Softmax# 对 scaled_scores 的每一行 (最后一个维度) 进行 softmax# attention_weights (P): (B, N, N)attention_weights = F.softmax(scaled_scores, dim=-1)# (可选) 应用 dropout 到注意力权重 (在 Transformer 论文中有提到)# attention_weights = F.dropout(attention_weights, p=dropout_rate)# 5. 计算输出 (Output): P * V# attention_weights: (B, N, N)# v: (B, N, dv)# output (B): (B, N, N) @ (B, N, dv) -> (B, N, dv)output = torch.matmul(attention_weights, v)return output, attention_weights# --- 使用示例 ---
if __name__ == '__main__':# 定义参数batch_size = 2seq_len = 4         # 序列长度 (N)embedding_dim = 6   # 输入词嵌入维度 (d_model)dk = 8              # Q, K 的维度dv = 8              # V 的维度 (通常 dk=dv)# 创建自注意力层实例self_attention_layer = SelfAttention(embedding_dim, dk, dv)# 创建一个随机输入张量 A (或 X)# 形状: (batch_size, seq_len, embedding_dim)input_tensor = torch.randn(batch_size, seq_len, embedding_dim)print(f"输入张量 A (X) 的形状: {input_tensor.shape}")# 进行前向传播output_b, attention_p = self_attention_layer(input_tensor)print(f"输出张量 B (Z) 的形状: {output_b.shape}")           # 预期: (batch_size, seq_len, dv)print(f"注意力权重 P 的形状: {attention_p.shape}")         # 预期: (batch_size, seq_len, seq_len)# 检查注意力权重是否归一化 (每行的和应为1)# 取第一个 batch 的第一个查询向量对应的权重print(f"\n第一个样本的注意力权重矩阵 P[0]:\n{attention_p[0]}")print(f"P[0] 每行权重之和:\n{attention_p[0].sum(dim=-1)}") # 应该接近全1向量print("\n--- 示例: 带掩码的自注意力 ---")# 创建一个掩码,例如屏蔽序列中的 padding token# 假设序列长度为4,但第二个样本只有前2个是有效token,后2个是padding# mask: 1表示有效,0表示padding/屏蔽# (batch_size, seq_len)input_mask = torch.tensor([[1, 1, 1, 1],  # 第一个样本全部有效[1, 1, 0, 0]   # 第二个样本后两个是padding], dtype=torch.bool) # 或者 torch.float32 用 0 和 1print(f"输入掩码的形状: {input_mask.shape}")# 进行带掩码的前向传播# 注意力机制中的 mask 通常是 (batch_size, seq_len, seq_len)# 这里我们用一个简单的 mask (batch_size, 1, seq_len) 来演示,# 它会屏蔽掉 K 矩阵中对应列的贡献。# 更复杂的 mask (如 decoder 中的 look-ahead mask) 会是 (seq_len, seq_len) 的上三角矩阵。# 对于 padding mask, 我们可以这样构造:# mask_for_scores (B, N, N): if mask[b,j]=0, then scores[b,:,j] = -inf# or if mask[b,i]=0, then scores[b,i,:] = -inf# 通常是让 Key 的 padding 位置不被 attend to:# pad_mask_for_scores (B, 1, N_k) -> expanded (B, N_q, N_k)# mask_for_scores = input_mask.unsqueeze(1).expand(-1, seq_len, -1)# print(f"扩展后用于 scaled_scores 的掩码形状: {mask_for_scores.shape}")# output_b_masked, attention_p_masked = self_attention_layer(input_tensor, mask=mask_for_scores)# 按照 `forward` 函数中的 mask 处理逻辑,我们直接传入 (B, N) 的 mask# 它会被 unsqueeze(1) 并用于 masked_filloutput_b_masked, attention_p_masked = self_attention_layer(input_tensor, mask=input_mask)print(f"\n带掩码的输出张量 B (Z) 的形状: {output_b_masked.shape}")print(f"带掩码的注意力权重 P 的形状: {attention_p_masked.shape}")print(f"\n第一个样本 (无padding) 的注意力权重矩阵 P_masked[0]:\n{attention_p_masked[0]}")print(f"P_masked[0] 每行权重之和:\n{attention_p_masked[0].sum(dim=-1)}")print(f"\n第二个样本 (后两个padding) 的注意力权重矩阵 P_masked[1]:\n{attention_p_masked[1]}")print(f"P_masked[1] 每行权重之和:\n{attention_p_masked[1].sum(dim=-1)}")# 观察 P_masked[1],最后两列的权重应该非常接近0 (因为 Key 的 padding 被屏蔽了)# 或者如果 mask 用于屏蔽 Query 的 padding 位置,则 P_masked[1] 的最后两行的分布会受影响。# 在我们的实现中,mask.unsqueeze(1) 会使得 mask 作用于 K 的维度。# 即对于一个 Query,它不会去关注那些被 mask 掉的 Key。# 所以 P_masked[b, i, j] 中,如果 input_mask[b, j] == 0,则 P_masked[b, i, j] 会接近0。

代码解释

  1. SelfAttention:
    • 继承自 torch.nn.Module,这是 PyTorch 中构建神经网络模块的标准方式。
    • __init__(self, embedding_dim, dk, dv):
      • embedding_dim ( d m o d e l d_{model} dmodel): 输入特征的维度。
      • dk: Query 和 Key 向量的维度。
      • dv: Value 向量的维度。通常 d k = d v d_k = d_v dk=dv
      • self.fc_q, self.fc_k, self.fc_v: 三个 nn.Linear 层,分别对应权重矩阵 W Q , W K , W V W^Q, W^K, W^V WQ,WK,WV。它们将输入 x (维度 embedding_dim) 转换为 Q, K, V 向量 (维度分别为 dk, dk, dv)。
    • forward(self, x, mask=None):
      • x: 输入张量,形状 (batch_size, seq_len, embedding_dim)
      • mask: 可选的掩码,用于处理例如 padding token 的情况。
      • 步骤 1 (生成 Q, K, V): q = self.fc_q(x), k = self.fc_k(x), v = self.fc_v(x)
      • 步骤 2 (计算分数): scores = torch.matmul(q, k.transpose(-2, -1))。这里 k.transpose(-2, -1) 是计算 K T K^T KT (对于每个 batch 内的矩阵)。矩阵乘法得到形状为 (batch_size, seq_len, seq_len) 的原始注意力分数。
      • 步骤 3 (缩放): scaled_scores = scores / (self.dk ** 0.5)。除以 d k \sqrt{d_k} dk
      • 步骤 (可选,掩码): 如果提供了 mask,则在 softmax 之前将 scaled_scores 中对应掩码为0(或False)的位置设置为一个非常小的数(如 -1e9),这样这些位置在 softmax 后权重会接近0。
      • 步骤 4 (Softmax): attention_weights = F.softmax(scaled_scores, dim=-1)dim=-1 表示对每个查询向量与所有键向量计算得到的分数进行 softmax,即对 scaled_scores 矩阵的每一行(最后一个维度)进行归一化。得到注意力权重矩阵 P P P
      • 步骤 5 (计算输出): output = torch.matmul(attention_weights, v)。将注意力权重 P P P 与值矩阵 V V V 相乘,得到最终的输出 B B B (或 Z Z Z),形状为 (batch_size, seq_len, dv)
        image.png

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

多头注意力机制是自注意力机制的一种扩展,它允许模型在不同的表示子空间中并行地学习来自输入序列不同部分的信息。这样做可以增强模型捕捉不同类型上下文依赖关系的能力。

image.png

核心思想

  1. 并行学习不同表示

    • 与其使用一组大的 W Q , W K , W V W^Q, W^K, W^V WQ,WK,WV 矩阵来计算单次注意力,多头注意力机制将输入(词嵌入)投影到 h h h (num_heads) 个不同的、低维度的查询(Query)、键(Key)和值(Value)子空间中。
    • 在每个子空间中,独立地执行一次缩放点积注意力(Scaled Dot-Product Attention)。
    • 这意味着对于每个输入词,我们会生成 h h h 组不同的 q , k , v q, k, v q,k,v 向量,每一组对应一个“头”。
  2. 独立计算注意力

    • 对于第 i i i-个头( i = 1 , . . . , h i=1, ..., h i=1,...,h),我们有自己的一组可学习的权重矩阵: W i Q , W i K , W i V W_i^Q, W_i^K, W_i^V WiQ,WiK,WiV
    • 输入 A A A (或 X X X) 会被分别乘以这些矩阵得到第 i i i-个头的 Q i , K i , V i Q_i, K_i, V_i Qi,Ki,Vi
    • 然后,对每个头独立计算注意力输出:
      head i = Attention ( Q i , K i , V i ) = softmax ( Q i K i T d k ) V i \text{head}_i = \text{Attention}(Q_i, K_i, V_i) = \text{softmax}\left(\frac{Q_i K_i^T}{\sqrt{d_k}}\right) V_i headi=Attention(Qi,Ki,Vi)=softmax(dk QiKiT)Vi
    • 注意,这里 d k d_k dk 通常是每个头的键向量维度,而不是原始嵌入维度 d m o d e l d_{model} dmodel。通常 d k = d m o d e l / h d_k = d_{model} / h dk=dmodel/h
  3. 拼接与最终线性变换

    • h h h 个头的注意力输出 head 1 , head 2 , . . . , head h \text{head}_1, \text{head}_2, ..., \text{head}_h head1,head2,...,headh 拼接(concatenate)起来。
    • 拼接后的结果再通过另一个可学习的线性变换(乘以权重矩阵 W O W^O WO)来产生最终的多头注意力输出。
      MultiHead ( Q , K , V ) = Concat ( head 1 , . . . , head h ) W O \text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, ..., \text{head}_h) W^O MultiHead(Q,K,V)=Concat(head1,...,headh)WO

image.png
这张图形象地展示了多头注意力的过程:

  1. 输入 X X X (即我们的 A A A) 经过多个并行的线性层,生成多组 Q i , K i , V i Q_i, K_i, V_i Qi,Ki,Vi
  2. 每组 Q i , K i , V i Q_i, K_i, V_i Qi,Ki,Vi 独立进行缩放点积注意力计算,得到 head i \text{head}_i headi
  3. 所有 head i \text{head}_i headi 拼接起来。
  4. 拼接结果再通过一个线性层( W O W^O WO)得到最终输出。

维度说明

  • 输入词嵌入维度: d m o d e l d_{model} dmodel
  • 注意力头的数量: h h h
  • 每个头的查询/键维度: d k = d m o d e l / h d_k = d_{model} / h dk=dmodel/h
  • 每个头的值维度: d v = d m o d e l / h d_v = d_{model} / h dv=dmodel/h (通常 d k = d v d_k = d_v dk=dv)
  • W i Q ∈ R d m o d e l × d k W_i^Q \in \mathbb{R}^{d_{model} \times d_k} WiQRdmodel×dk
  • W i K ∈ R d m o d e l × d k W_i^K \in \mathbb{R}^{d_{model} \times d_k} WiKRdmodel×dk
  • W i V ∈ R d m o d e l × d v W_i^V \in \mathbb{R}^{d_{model} \times d_v} WiVRdmodel×dv
  • 每个 head i \text{head}_i headi 的输出维度是 N × d v N \times d_v N×dv (假设输入序列长度为 N N N)。
  • 拼接后的张量维度是 N × ( h ⋅ d v ) N \times (h \cdot d_v) N×(hdv)。由于 h ⋅ d v = h ⋅ ( d m o d e l / h ) = d m o d e l h \cdot d_v = h \cdot (d_{model}/h) = d_{model} hdv=h(dmodel/h)=dmodel,所以拼接后的维度是 N × d m o d e l N \times d_{model} N×dmodel
  • W O ∈ R d m o d e l × d m o d e l W^O \in \mathbb{R}^{d_{model} \times d_{model}} WORdmodel×dmodel (因为 h ⋅ d v = d m o d e l h \cdot d_v = d_{model} hdv=dmodel,且通常希望输出维度与输入维度一致)。

优势

  • 多角度信息捕捉:不同的头可以学习关注输入序列的不同方面或不同类型的依赖关系。例如,一个头可能关注句法关系,另一个头可能关注语义相似性。
  • 更丰富的表示:通过在不同的子空间中进行投影和注意力计算,模型可以学习到更丰富的特征表示。
  • 稳定训练:将大的注意力计算分解为多个小的并行计算,有时可以帮助稳定训练过程。
  • 计算效率:虽然看起来计算量增加了 h h h 倍,但由于每个头的维度 d k , d v d_k, d_v dk,dv 被缩小了 h h h 倍,总的计算量与单头注意力(使用原始 d m o d e l d_{model} dmodel 维度)大致相当,并且可以高度并行化。
4.1 单头注意力 (Single-Head Attention) vs 多头注意力 (Multi-Head Attention)
特性单头注意力 (Scaled Dot-Product Attention)多头注意力 (Multi-Head Attention)
Query/Key/Value 生成一组 W Q , W K , W V W^Q, W^K, W^V WQ,WK,WV 矩阵,将 d m o d e l d_{model} dmodel 维输入投影到 d k , d k , d v d_k, d_k, d_v dk,dk,dv 维 (通常 d k = d v = d m o d e l d_k=d_v=d_{model} dk=dv=dmodel 或某个自定义值)。 h h h 组独立的 W i Q , W i K , W i V W_i^Q, W_i^K, W_i^V WiQ,WiK,WiV 矩阵。每组将 d m o d e l d_{model} dmodel 维输入投影到 d k = d m o d e l / h d_k = d_{model}/h dk=dmodel/h d v = d m o d e l / h d_v = d_{model}/h dv=dmodel/h 维的子空间。
注意力计算执行一次缩放点积注意力。并行执行 h h h 次独立的缩放点积注意力,每个头一次。
输出维度输出维度为 d v d_v dv (通常与 d m o d e l d_{model} dmodel 相同)。每个头的输出维度为 d v d_v dv。将 h h h 个头的输出拼接后,维度为 h ⋅ d v = d m o d e l h \cdot d_v = d_{model} hdv=dmodel,再通过 W O W^O WO 线性变换(通常保持 d m o d e l d_{model} dmodel 维)。
信息捕捉从一个角度或一个表示空间学习上下文依赖。 h h h 个不同的角度或 h h h 个不同的表示子空间学习上下文依赖。
模型复杂度参数量主要在 W Q , W K , W V W^Q, W^K, W^V WQ,WK,WV参数量在 h h h W i Q , W i K , W i V W_i^Q, W_i^K, W_i^V WiQ,WiK,WiV 以及 W O W^O WO。如果 d k , d v d_k, d_v dk,dv 按比例缩小,总参数量与单头(使用 d m o d e l d_{model} dmodel 维度)相似。
并行性单次计算。 h h h 个头的计算可以并行进行。
“查询”数量对于每个输出 b j b^j bj,本质上是使用一个“组合查询”(通过 W Q W^Q WQ 变换得到的 q j q^j qj)来聚合信息。对于每个输出 b j b^j bj,是通过 h h h 个不同的“子查询”(每个头的 q i j q_i^j qij)来并行聚合信息,然后组合结果。

关键区别的通俗理解

  • 单头注意力:就像一个人用一种方式(例如,只关注语法结构)去理解一句话中词与词之间的关系。

  • 多头注意力:就像 h h h 个人同时用 h h h 种不同的方式(例如,一个人关注语法,一个人关注语义,一个人关注词的位置等)去理解一句话中词与词之间的关系,然后把他们各自的理解综合起来,形成一个更全面的理解。

  • 在单头中,每个输入位置 a j a^j aj 通过 W Q W^Q WQ 得到一个 q j q^j qj

  • 在多头中,每个输入位置 a j a^j aj 通过 h h h 个不同的 W i Q W_i^Q WiQ (或一个大的 W Q W^Q WQ 然后分割) 得到 h h h 个不同的 q i j q_i^j qij (每个 q i j q_i^j qij 属于一个头)。每个 q i j q_i^j qij 都会独立地与该头对应的 K i K_i Ki 中的所有键进行交互,产生一个 head i \text{head}_i headi 的输出。

4.2 多头注意力机制的代码实现

下面是使用 PyTorch 实现多头注意力机制的代码:

import torch
import torch.nn as nn
import torch.nn.functional as F
import mathclass MultiHeadAttention(nn.Module):"""实现多头注意力机制。"""def __init__(self, embedding_dim, num_heads):"""初始化多头注意力层。参数:- embedding_dim (int): 输入词嵌入的维度 (d_model)。- num_heads (int): 注意力头的数量 (h)。"""super(MultiHeadAttention, self).__init__()assert embedding_dim % num_heads == 0, "embedding_dim 必须能被 num_heads 整除"self.embedding_dim = embedding_dimself.num_heads = num_headsself.head_dim = embedding_dim // num_heads # dk 和 dv 都设为 head_dim# 线性层用于 Q, K, V 的整体变换,之后再分割到各个头# 也可以为每个头创建独立的 fc_q, fc_k, fc_v,但这种方式更常见且高效self.fc_q = nn.Linear(embedding_dim, embedding_dim) # 输出维度是 d_model (h * head_dim)self.fc_k = nn.Linear(embedding_dim, embedding_dim) # 输出维度是 d_modelself.fc_v = nn.Linear(embedding_dim, embedding_dim) # 输出维度是 d_model# 最终的输出线性层 W_oself.fc_o = nn.Linear(embedding_dim, embedding_dim)# 可选: 初始化权重# self._init_weights()# def _init_weights(self):#     nn.init.xavier_uniform_(self.fc_q.weight)#     nn.init.xavier_uniform_(self.fc_k.weight)#     nn.init.xavier_uniform_(self.fc_v.weight)#     nn.init.xavier_uniform_(self.fc_o.weight)#     if self.fc_q.bias is not None:#         nn.init.zeros_(self.fc_q.bias)#         nn.init.zeros_(self.fc_k.bias)#         nn.init.zeros_(self.fc_v.bias)#         nn.init.zeros_(self.fc_o.bias)def _split_heads(self, x, batch_size):"""将最后一个维度分割成 (num_heads, head_dim),并调整形状以进行并行计算。x: (batch_size, seq_len, embedding_dim)返回: (batch_size, num_heads, seq_len, head_dim)"""x = x.view(batch_size, -1, self.num_heads, self.head_dim)return x.transpose(1, 2) # (B, h, N, head_dim)def _combine_heads(self, x, batch_size):"""将多个头的输出合并回原始形状。x: (batch_size, num_heads, seq_len, head_dim)返回: (batch_size, seq_len, embedding_dim)"""x = x.transpose(1, 2).contiguous() # (B, N, h, head_dim)return x.view(batch_size, -1, self.embedding_dim) # (B, N, d_model)def scaled_dot_product_attention(self, q, k, v, mask=None):"""执行缩放点积注意力。q, k, v: (batch_size, num_heads, seq_len, head_dim)mask: (batch_size, 1, 1, seq_len_k) or (batch_size, 1, seq_len_q, seq_len_k) for broadcasting"""# q: (B, h, N_q, head_dim)# k.transpose(-2, -1): (B, h, head_dim, N_k)# scores: (B, h, N_q, N_k)scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)if mask is not None:scores = scores.masked_fill(mask == 0, -1e9)attention_weights = F.softmax(scores, dim=-1)# attention_weights: (B, h, N_q, N_k)# v: (B, h, N_v, head_dim) where N_v = N_k# output: (B, h, N_q, head_dim)output = torch.matmul(attention_weights, v)return output, attention_weightsdef forward(self, query, key, value, mask=None):"""前向传播。在自注意力中,query, key, value 都是同一个输入 x。在编码器-解码器注意力中,query 来自解码器,key 和 value 来自编码器。参数:- query (torch.Tensor): 查询张量,形状 (batch_size, seq_len_q, embedding_dim)。- key (torch.Tensor):   键张量,形状 (batch_size, seq_len_k, embedding_dim)。- value (torch.Tensor): 值张量,形状 (batch_size, seq_len_v, embedding_dim)。(通常 seq_len_k == seq_len_v)- mask (torch.Tensor, optional): 掩码张量。返回:- output (torch.Tensor): 多头注意力的输出,形状 (batch_size, seq_len_q, embedding_dim)。- attention_weights (torch.Tensor): 注意力权重,形状 (batch_size, num_heads, seq_len_q, seq_len_k)。"""batch_size = query.size(0)# 1. 线性变换并分割成多个头# q, k, v 初始形状: (B, N, d_model)q = self.fc_q(query) # (B, N_q, d_model)k = self.fc_k(key)   # (B, N_k, d_model)v = self.fc_v(value) # (B, N_v, d_model)q = self._split_heads(q, batch_size) # (B, h, N_q, head_dim)k = self._split_heads(k, batch_size) # (B, h, N_k, head_dim)v = self._split_heads(v, batch_size) # (B, h, N_v, head_dim)# 2. 对每个头独立计算缩放点积注意力# mask 需要能广播到 (B, h, N_q, N_k)# 例如,如果 mask 是 (B, N_k),需要变为 (B, 1, 1, N_k)# 如果 mask 是 (B, N_q, N_k),需要变为 (B, 1, N_q, N_k)if mask is not None:# 假设 mask 作用于 key 的 padding# e.g., mask (B, N_k) -> (B, 1, 1, N_k)if mask.dim() == 2: # (B, N_k)mask = mask.unsqueeze(1).unsqueeze(2) # (B, 1, 1, N_k)elif mask.dim() == 3: # (B, N_q, N_k) like causal maskmask = mask.unsqueeze(1) # (B, 1, N_q, N_k)# 确保 mask 数据类型正确mask = mask.to(dtype=q.dtype) # to match q, k, v dtype for masked_fillattention_output, attention_weights = self.scaled_dot_product_attention(q, k, v, mask)# attention_output: (B, h, N_q, head_dim)# attention_weights: (B, h, N_q, N_k)# 3. 合并多个头的输出attention_output = self._combine_heads(attention_output, batch_size) # (B, N_q, d_model)# 4. 最终的线性变换output = self.fc_o(attention_output) # (B, N_q, d_model)return output, attention_weights# --- 使用示例 ---
if __name__ == '__main__':batch_size = 2seq_len = 4embedding_dim = 12 # d_model, e.g., 512num_heads = 3      # h, e.g., 8. embedding_dim % num_heads == 0 (12 % 3 == 0)assert embedding_dim % num_heads == 0multi_head_attention_layer = MultiHeadAttention(embedding_dim, num_heads)# 随机输入 (对于自注意力,query, key, value 是相同的)input_tensor = torch.randn(batch_size, seq_len, embedding_dim)print(f"输入张量 (query/key/value) 的形状: {input_tensor.shape}")# 前向传播output_mha, attn_weights_mha = multi_head_attention_layer(input_tensor, input_tensor, input_tensor)print(f"多头注意力输出的形状: {output_mha.shape}")             # 预期: (batch_size, seq_len, embedding_dim)print(f"多头注意力权重的形状: {attn_weights_mha.shape}")     # 预期: (batch_size, num_heads, seq_len, seq_len)print("\n--- 示例: 带掩码的多头自注意力 ---")# mask for padding# (batch_size, seq_len_k)input_mask = torch.tensor([[1, 1, 1, 1],[1, 1, 0, 0]], dtype=torch.bool) # True for valid, False for padprint(f"输入掩码的形状: {input_mask.shape}")output_mha_masked, attn_weights_mha_masked = multi_head_attention_layer(input_tensor, input_tensor, input_tensor, mask=input_mask)print(f"\n带掩码的多头注意力输出形状: {output_mha_masked.shape}")print(f"带掩码的多头注意力权重形状: {attn_weights_mha_masked.shape}")print(f"\n第二个样本 (后两个padding) 的注意力权重 attn_weights_mha_masked[1]:\n")# 打印第二个样本的每个头的注意力权重for head_idx in range(num_heads):print(f"Head {head_idx+1}:\n{attn_weights_mha_masked[1, head_idx]}")print(f"Head {head_idx+1} 每行权重之和:\n{attn_weights_mha_masked[1, head_idx].sum(dim=-1)}")# 观察 attn_weights_mha_masked[1, head_idx, :, 2:] 这些列的权重应该非常接近0

代码解释

  1. MultiHeadAttention:
    • __init__:
      • embedding_dim ( d m o d e l d_{model} dmodel), num_heads ( h h h).
      • self.head_dim = embedding_dim // num_heads: 计算每个头的维度 ( d k , d v d_k, d_v dk,dv).
      • self.fc_q, self.fc_k, self.fc_v: 线性层,将输入投影到 d m o d e l d_{model} dmodel 维。这些输出随后会被分割成 h h h 个头。这种实现方式(先整体变换再分割)比为每个头创建独立的 W i Q , W i K , W i V W_i^Q, W_i^K, W_i^V WiQ,WiK,WiV (每个输出 d k d_k dk d v d_v dv 维) 在参数上是等价的,但在实现上可能更简洁和高效,因为可以利用大的矩阵乘法。
      • self.fc_o: 最终的输出线性层 W O W^O WO
    • _split_heads: 辅助函数,将形状 (B, N, d_model) 的张量重塑并转置为 (B, h, N, head_dim),以便每个头可以独立计算。
    • _combine_heads: 辅助函数,将 _split_heads 的操作逆转,将 (B, h, N, head_dim) 合并回 (B, N, d_model)
    • scaled_dot_product_attention: 这是一个静态方法或可以在类内部实现的辅助函数,执行实际的缩放点积注意力计算。它接收已经分割好的多头 Q, K, V。
    • forward:
      • 接收 query, key, valuemask。对于自注意力,query, key, value 都是同一个输入 x
      • 步骤 1 (线性变换和分割):
        • q = self.fc_q(query), k = self.fc_k(key), v = self.fc_v(value): 计算整体的 Q, K, V。
        • q = self._split_heads(q, batch_size) 等:将 Q, K, V 分割成多个头。
      • 步骤 2 (独立注意力计算): 调用 self.scaled_dot_product_attention(q, k, v, mask)。注意掩码 mask 需要正确地广播到与 scores 矩阵 (B, h, N_q, N_k) 兼容的形状。
      • 步骤 3 (合并头): attention_output = self._combine_heads(attention_output, batch_size)
      • 步骤 4 (最终线性变换): output = self.fc_o(attention_output)

5. 位置编码 (Positional Encoding)

位置编码补充

问题背景:为何需要位置编码?

Transformer 模型的核心是自注意力机制。自注意力机制在处理输入序列时,是不区分词语顺序的。也就是说,对于一个词,它会计算与序列中所有其他词的注意力权重,但这个计算过程本身并不包含词语在序列中的位置信息。例如,“猫追老鼠” 和 “老鼠追猫”,对于自注意力机制来说,如果仅输入词嵌入,它可能无法区分这两种顺序,因为每个词看到的上下文词集合是相同的,只是顺序不同。

然而,在自然语言处理中,词语的顺序至关重要,它决定了句子的含义。因此,我们需要一种方法将词语的位置信息注入到模型中。

解决方案:位置编码

位置编码就是为了解决这个问题而提出的。它的核心思想是:为输入序列中的每个词嵌入向量添加一个额外的向量,这个向量专门用来表示该词在序列中的绝对或相对位置。

方式
如上图所示,位置编码向量 p i p^i pi 被直接加到对应的词嵌入向量 x i x^i xi 上:
InputEmbeddingWithPosition = WordEmbedding ( x i ) + PositionalEncoding ( p i ) \text{InputEmbeddingWithPosition} = \text{WordEmbedding}(x^i) + \text{PositionalEncoding}(p^i) InputEmbeddingWithPosition=WordEmbedding(xi)+PositionalEncoding(pi)

这个相加后的向量将作为 Transformer 编码器或解码器堆栈的实际输入。

位置编码的特性要求

  1. 唯一性:对于序列中的每个位置,它应该产生一个唯一的位置编码。
  2. 确定性:对于给定的位置,其编码应该是固定的,不应随训练而改变(至少在原始 Transformer 论文中是这样设计的,也有可学习的位置编码变体)。
  3. 泛化性:模型应该能够泛化到比训练时遇到的序列更长的序列。位置编码的设计应能支持这一点。
  4. 距离感知:理想情况下,不同位置之间的编码差异应该能反映它们之间的距离关系,即对于任意固定的偏移量 k k k P E p o s + k PE_{pos+k} PEpos+k 应该可以由 P E p o s PE_{pos} PEpos 线性表示,这样模型可以学习到相对位置信息。

“Attention Is All You Need” 论文中使用的正弦/余弦位置编码

论文中提出了一种使用不同频率的正弦和余弦函数来生成位置编码的方法。对于位置 p o s pos pos 和编码向量的维度索引 i i i (从0到 d m o d e l − 1 d_{model}-1 dmodel1),位置编码 P E ( p o s , i ) PE_{(pos, i)} PE(pos,i) 的计算方式如下:

  • i = 2 k i = 2k i=2k (偶数维度):
    P E ( p o s , 2 k ) = sin ⁡ ( p o s 1000 0 2 k / d m o d e l ) PE_{(pos, 2k)} = \sin\left(\frac{pos}{10000^{2k/d_{model}}}\right) PE(pos,2k)=sin(100002k/dmodelpos)
  • i = 2 k + 1 i = 2k+1 i=2k+1 (奇数维度):
    P E ( p o s , 2 k + 1 ) = cos ⁡ ( p o s 1000 0 2 k / d m o d e l ) PE_{(pos, 2k+1)} = \cos\left(\frac{pos}{10000^{2k/d_{model}}}\right) PE(pos,2k+1)=cos(100002k/dmodelpos)

其中:

  • p o s pos pos 是词在序列中的位置索引(从0开始)。
  • d m o d e l d_{model} dmodel 是词嵌入和位置编码向量的维度。
  • 2 k 2k 2k 2 k + 1 2k+1 2k+1 表示编码向量中的维度索引。

为什么选择这种函数形式?

  1. 唯一性:每个位置 p o s pos pos 都会产生一个独特的编码向量。

  2. 周期性带来的相对位置信息
    对于任意固定的偏移量 k ′ k' k, P E p o s + k ′ PE_{pos+k'} PEpos+k 可以表示为 P E p o s PE_{pos} PEpos 的线性函数。具体来说:
    sin ⁡ ( A + B ) = sin ⁡ ( A ) cos ⁡ ( B ) + cos ⁡ ( A ) sin ⁡ ( B ) \sin(A+B) = \sin(A)\cos(B) + \cos(A)\sin(B) sin(A+B)=sin(A)cos(B)+cos(A)sin(B)
    cos ⁡ ( A + B ) = cos ⁡ ( A ) cos ⁡ ( B ) − sin ⁡ ( A ) sin ⁡ ( B ) \cos(A+B) = \cos(A)\cos(B) - \sin(A)\sin(B) cos(A+B)=cos(A)cos(B)sin(A)sin(B)
    A = p o s wavelength A = \frac{pos}{\text{wavelength}} A=wavelengthpos B = k ′ wavelength B = \frac{k'}{\text{wavelength}} B=wavelengthk,其中 wavelength = 1000 0 2 k / d m o d e l \text{wavelength} = 10000^{2k/d_{model}} wavelength=100002k/dmodel
    这意味着 P E ( p o s + k ′ , 2 k ) PE_{(pos+k', 2k)} PE(pos+k,2k) P E ( p o s + k ′ , 2 k + 1 ) PE_{(pos+k', 2k+1)} PE(pos+k,2k+1) 可以由 P E ( p o s , 2 k ) , P E ( p o s , 2 k + 1 ) , P E ( k ′ , 2 k ) , P E ( k ′ , 2 k + 1 ) PE_{(pos, 2k)}, PE_{(pos, 2k+1)}, PE_{(k', 2k)}, PE_{(k', 2k+1)} PE(pos,2k),PE(pos,2k+1),PE(k,2k),PE(k,2k+1) (实际上是 sin ⁡ ( B ) \sin(B) sin(B) cos ⁡ ( B ) \cos(B) cos(B)) 线性组合得到。这允许模型很容易地学习到相对位置信息,因为不同位置之间的关系可以通过一个固定的旋转矩阵来表示。

  3. 平滑变化:位置编码随着位置的变化是平滑的。

  4. 能够处理不同长度的序列:由于是基于三角函数,理论上可以生成任意长度序列的位置编码,即使序列长度超过训练时见过的最大长度(尽管泛化能力仍可能受限)。

  5. 波长变化:分母中的 1000 0 2 k / d m o d e l 10000^{2k/d_{model}} 100002k/dmodel 项使得波长从 2 π 2\pi 2π (当 2 k / d m o d e l ≈ 0 2k/d_{model} \approx 0 2k/dmodel0) 变化到 10000 ⋅ 2 π 10000 \cdot 2\pi 100002π (当 2 k / d m o d e l ≈ 1 2k/d_{model} \approx 1 2k/dmodel1)。这意味着编码向量的不同维度使用不同频率的正弦/余弦波。低维度的编码变化频率高(对应短波长),高维度的编码变化频率低(对应长波长)。这种设计允许模型同时关注不同尺度的位置信息。

实现细节

  • 位置编码矩阵 P E PE PE 的形状通常是 (max_seq_len, d_model),其中 max_seq_len 是模型能处理的最大序列长度。
  • 在实际使用时,会取这个预先计算好的 P E PE PE 矩阵的前 current_seq_len行,然后加到词嵌入上。
  • 这个 P E PE PE 矩阵通常在模型初始化时一次性计算好,并且在训练过程中保持不变(非可学习)。

位置编码在模型中的位置
位置编码通常在输入嵌入层之后,进入 Transformer 的第一个编码器层(或解码器层)之前添加。

image.png
这张图清晰地展示了:

  • x 1 , x 2 , . . . , x t x^1, x^2, ..., x^t x1,x2,...,xt 是原始的词序列(或它们的ID)。
  • 通过 “Input Embedding” 层得到词嵌入向量。
  • p 1 , p 2 , . . . , p t p^1, p^2, ..., p^t p1,p2,...,pt 是对应位置的位置编码向量。
  • 词嵌入和位置编码向量按元素相加。
  • 相加后的结果作为后续 Transformer 层的输入。

其他类型的位置编码
虽然正弦/余弦位置编码是经典 Transformer 中的标准做法,但后续也出现了一些其他的位置编码方法:

  • 可学习的位置编码 (Learned Positional Embeddings):将位置编码视为模型参数,在训练过程中学习得到。例如 BERT 使用这种方式。这种方法简单直接,但可能在泛化到未见过的序列长度时表现不如三角函数编码。
  • 相对位置编码 (Relative Positional Embeddings):直接在自注意力计算过程中引入相对位置信息,而不是在输入端添加绝对位置编码。例如 Transformer-XL, T5, DeBERTa 等模型采用了不同形式的相对位置编码。这种方法更关注词与词之间的相对距离,而不是它们在序列中的绝对位置。
  • 旋转位置编码 (Rotary Positional Embedding, RoPE):由 RoFormer 提出,通过对查询和键向量进行旋转来注入相对位置信息,具有良好的内外插能力。LLaMA 等模型采用了这种方式。

总结
位置编码是 Transformer 模型能够处理序列顺序信息的关键机制。通过将位置信息编码成向量并融入词嵌入中,模型可以在进行自注意力计算时感知到词语的顺序和相对位置。正弦/余弦位置编码因其数学特性和良好的泛化潜力而被广泛使用。

5.1 代码实现 (正弦/余弦位置编码)
import torch
import torch.nn as nn
import math
import matplotlib.pyplot as plt
import numpy as npclass PositionalEncoding(nn.Module):"""实现Transformer论文中描述的正弦和余弦位置编码。"""def __init__(self, d_model, max_len=5000, dropout=0.1):super(PositionalEncoding, self).__init__()self.dropout = nn.Dropout(p=dropout)pe = torch.zeros(max_len, d_model)position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))pe[:, 0::2] = torch.sin(position * div_term)pe[:, 1::2] = torch.cos(position * div_term)self.register_buffer('pe', pe.unsqueeze(0))def forward(self, x):x = x + self.pe[:, :x.size(1), :]return self.dropout(x)if __name__ == '__main__':d_model = 256  # 嵌入维度 (建议偶数,例如256或512)max_seq_len_for_pe = 200 # 位置编码预计算的最大长度 (Y轴范围)dropout_rate = 0.0 # 关闭dropout以便于观察PE本身pos_encoder = PositionalEncoding(d_model, max_seq_len_for_pe, dropout_rate)current_seq_len = max_seq_len_for_pe # 获取PE矩阵并移除batch维度,转换为numpype_matrix = pos_encoder.pe.squeeze(0).cpu().numpy() # Shape: (max_seq_len, d_model)# 分离sin和cos部分pe_sin = pe_matrix[:current_seq_len, 0::2] # Shape: (current_seq_len, d_model/2)pe_cos = pe_matrix[:current_seq_len, 1::2] # Shape: (current_seq_len, d_model/2)# --- 可视化代码 ---# 调整figsize使图像更宽,更接近您提供的示例# 两个子图垂直排列,共享X轴,无间距fig, axes = plt.subplots(2, 1, figsize=(8, 4.5), sharex=True, gridspec_kw={'hspace': 0.05})# 绘制 Sin 部分 (上半部分)im_sin = axes[0].imshow(pe_sin, aspect='auto', cmap='viridis', vmin=-1, vmax=1, interpolation='nearest')axes[0].set_yticks([])  # 移除Y轴刻度和标签axes[0].set_ylabel("")# 绘制 Cos 部分 (下半部分)im_cos = axes[1].imshow(pe_cos, aspect='auto', cmap='viridis', vmin=-1, vmax=1, interpolation='nearest')axes[1].set_yticks([])  # 移除Y轴刻度和标签axes[1].set_xticks([])  # 移除X轴刻度和标签axes[1].set_xlabel("")axes[1].set_ylabel("")# 添加一个横向的、位于底部的共享颜色条# 调整subplots_adjust的bottom值和cbar_ax的位置/大小以匹配您的图像plt.subplots_adjust(left=0.05, right=0.95, top=0.95, bottom=0.15) # [left, bottom, width, height] for the colorbar axes in figure coordinatescbar_ax = fig.add_axes([0.25, 0.06, 0.5, 0.03]) cb = fig.colorbar(im_sin, cax=cbar_ax, orientation='horizontal')cb.set_ticks([-1, 1]) # 只显示 -1 和 1 刻度cb.ax.tick_params(labelsize=8) # 调整刻度标签大小plt.show()# --- 原始的检查代码 ---print("--- Original Checks ---")batch_size = 4seq_len_test = 20 # 用于测试的短序列长度word_embeddings = torch.randn(batch_size, seq_len_test, d_model)print(f"原始词嵌入的形状: {word_embeddings.shape}")pos_encoder.eval() # 确保dropout关闭 (因为我们设置dropout_rate=0.0, 这步不是必须的,但好习惯)embeddings_with_pe = pos_encoder(word_embeddings)print(f"添加位置编码后的嵌入形状: {embeddings_with_pe.shape}")print(f"\n位置0的编码 (前5维): {pos_encoder.pe[0, 0, :5].cpu().numpy()}")print(f"位置1的编码 (前5维): {pos_encoder.pe[0, 1, :5].cpu().numpy()}")print(f"位置0和位置1的编码是否相同: {np.allclose(pos_encoder.pe[0, 0, :5].cpu().numpy(), pos_encoder.pe[0, 1, :5].cpu().numpy())}")

image.png
image.png

附录

一、注意力机制中的缩放因子 1 d k \frac{1}{\sqrt{d_k}} dk 1:为何及如何作用

在Transformer等模型中广泛使用的缩放点积注意力(Scaled Dot-Product Attention)机制中,有一个关键的步骤是对查询向量 q \mathbf{q} q 和键向量 k \mathbf{k} k 的点积结果进行缩放,即除以 d k \sqrt{d_k} dk (其中 d k d_k dk 是键向量的维度)。这一操作对于模型的稳定训练和性能至关重要。

1. 内积的方差与维度 d k d_k dk 的关系

在点积注意力中,注意力分数是通过查询向量 q \mathbf{q} q 和键向量 k \mathbf{k} k 的内积 q ⋅ k \mathbf{q} \cdot \mathbf{k} qk 计算得到的。假设 q \mathbf{q} q k \mathbf{k} k 的每个元素是独立同分布的随机变量,例如均值为 0、方差为 1,那么内积 q ⋅ k \mathbf{q} \cdot \mathbf{k} qk 的方差会随着向量维度 d k d_k dk 的增加而增大。具体来说:

  • 内积 q ⋅ k = ∑ i = 1 d k q i k i \mathbf{q} \cdot \mathbf{k} = \sum_{i=1}^{d_k} q_i k_i qk=i=1dkqiki d k d_k dk 个独立项的和。
  • 假设 E [ q i ] = 0 , Var ( q i ) = 1 \mathbb{E}[q_i] = 0, \text{Var}(q_i) = 1 E[qi]=0,Var(qi)=1 E [ k i ] = 0 , Var ( k i ) = 1 \mathbb{E}[k_i] = 0, \text{Var}(k_i) = 1 E[ki]=0,Var(ki)=1,且 q i , k i q_i, k_i qi,ki 独立。
    • 那么 E [ q i k i ] = E [ q i ] E [ k i ] = 0 × 0 = 0 \mathbb{E}[q_i k_i] = \mathbb{E}[q_i]\mathbb{E}[k_i] = 0 \times 0 = 0 E[qiki]=E[qi]E[ki]=0×0=0
    • Var ( q i k i ) = E [ ( q i k i ) 2 ] − ( E [ q i k i ] ) 2 = E [ q i 2 ] E [ k i 2 ] − 0 = Var ( q i ) Var ( k i ) = 1 × 1 = 1 \text{Var}(q_i k_i) = \mathbb{E}[(q_i k_i)^2] - (\mathbb{E}[q_i k_i])^2 = \mathbb{E}[q_i^2]\mathbb{E}[k_i^2] - 0 = \text{Var}(q_i)\text{Var}(k_i) = 1 \times 1 = 1 Var(qiki)=E[(qiki)2](E[qiki])2=E[qi2]E[ki2]0=Var(qi)Var(ki)=1×1=1 (因为 E [ X 2 ] = Var ( X ) + ( E [ X ] ) 2 \mathbb{E}[X^2] = \text{Var}(X) + (\mathbb{E}[X])^2 E[X2]=Var(X)+(E[X])2)。
  • 因此, Var ( q ⋅ k ) = Var ( ∑ i = 1 d k q i k i ) = ∑ i = 1 d k Var ( q i k i ) = ∑ i = 1 d k 1 = d k \text{Var}(\mathbf{q} \cdot \mathbf{k}) = \text{Var}\left(\sum_{i=1}^{d_k} q_i k_i\right) = \sum_{i=1}^{d_k} \text{Var}(q_i k_i) = \sum_{i=1}^{d_k} 1 = d_k Var(qk)=Var(i=1dkqiki)=i=1dkVar(qiki)=i=1dk1=dk

这意味着,当 d k d_k dk 较大时(例如在Transformer中常见的512或1024),内积 q ⋅ k \mathbf{q} \cdot \mathbf{k} qk 的数值的方差会很大,导致其绝对值可能变得非常大,分布范围过宽。

2. Softmax 函数的饱和问题

注意力机制中,内积 q ⋅ k \mathbf{q} \cdot \mathbf{k} qk (或其缩放版本) 会被输入到 softmax 函数中,用于计算注意力权重:
softmax ( x i ) = e x i ∑ j e x j \text{softmax}(x_i) = \frac{e^{x_i}}{\sum_j e^{x_j}} softmax(xi)=jexjexi
然而,softmax 函数对输入值的大小非常敏感:

  • 当某些 x i x_i xi 变得非常大时:对应的 e x i e^{x_i} exi 会变得极大,导致 softmax 的输出在该项上接近 1,而在其他项上接近 0。
  • 当所有 x i x_i xi 都非常大或非常小时 (或者差异极大时):softmax 函数可能会进入“饱和区”。

在饱和区,softmax 函数的梯度 ∂ softmax ( x i ) ∂ x k \frac{\partial \text{softmax}(x_i)}{\partial x_k} xksoftmax(xi) (对于 k ≠ i k \neq i k=i k = i k=i k=i) 会变得非常小(接近 0)。这会导致梯度消失问题,阻碍模型通过反向传播有效更新参数,严重影响训练过程和模型的收敛。

3. 除以 d k \sqrt{d_k} dk 的作用:控制方差,避免饱和

为了解决上述问题,缩放点积注意力将内积除以 d k \sqrt{d_k} dk ,即计算:
q ⋅ k d k \frac{\mathbf{q} \cdot \mathbf{k}}{\sqrt{d_k}} dk qk
这样做的核心效果是:

  • 控制方差
    如前所述, Var ( q ⋅ k ) = d k \text{Var}(\mathbf{q} \cdot \mathbf{k}) = d_k Var(qk)=dk
    根据方差的性质 Var ( c X ) = c 2 Var ( X ) \text{Var}(cX) = c^2 \text{Var}(X) Var(cX)=c2Var(X),我们有:
    Var ( q ⋅ k d k ) = ( 1 d k ) 2 Var ( q ⋅ k ) = 1 d k × d k = 1 \text{Var}\left( \frac{\mathbf{q} \cdot \mathbf{k}}{\sqrt{d_k}} \right) = \left( \frac{1}{\sqrt{d_k}} \right)^2 \text{Var}(\mathbf{q} \cdot \mathbf{k}) = \frac{1}{d_k} \times d_k = 1 Var(dk qk)=(dk 1)2Var(qk)=dk1×dk=1
    这意味着,通过除以 d k \sqrt{d_k} dk ,我们将点积的方差从 d k d_k dk 缩放回了 1 (在上述理想假设下)。这使得输入到 softmax 函数的值的尺度大致保持一致,而不会随着 d k d_k dk 的增大而剧烈变化。

  • 避免饱和区:通过将内积的数值范围缩小,可以确保 softmax 函数的输入不会过大或过小,从而有效避免其进入饱和区。这使得 softmax 函数能够在其梯度较大的区域工作。

4. 保障梯度稳定性

缩放后的注意力权重定义为:
α i j = softmax ( q i ⋅ k j d k ) \alpha_{ij} = \text{softmax}\left( \frac{\mathbf{q}_i \cdot \mathbf{k}_j}{\sqrt{d_k}} \right) αij=softmax(dk qikj)
在反向传播过程中,梯度(例如 ∂ L ∂ ( q i ⋅ k j ) \frac{\partial L}{\partial (\mathbf{q}_i \cdot \mathbf{k}_j)} (qikj)L,其中 L L L 是损失函数)的计算会涉及到 softmax 的输出。如果不进行缩放,当 d k d_k dk 很大时,内积的绝对值可能过大,导致 softmax 梯度过小,使得参数(如 W Q , W K W^Q, W^K WQ,WK)的更新变得非常缓慢或停滞。而除以 d k \sqrt{d_k} dk 后,softmax 输入的尺度得到控制,其梯度的大小不会因 d k d_k dk 的增加而急剧减小,从而保持了数值稳定性和训练的稳定性。

总结

除以 d k \sqrt{d_k} dk 的核心目的是在高维空间中控制点积 q ⋅ k \mathbf{q} \cdot \mathbf{k} qk 的数值范围 (特别是其方差),防止其因维度 d k d_k dk 增大而导致数值过大或过小。这不仅避免了 softmax 函数的饱和问题,还确保了梯度的稳定性,使得模型能够更高效地学习和收敛。这一看似简单的缩放技巧是缩放点积注意力机制在 Transformer 模型中成功应用的关键因素之一。


二、为什么选择Softmax函数?—— 兼论与其他激活函数的对比

在机器学习中,尤其是在多分类任务和注意力机制中,softmax函数被广泛应用。它的选择并非偶然,而是基于其优良的数学特性和实际效果。

1. 将任意实数转换为有效的概率分布

神经网络的原始输出(通常称为logits)可以是任意实数值,包括正数、负数或零。Softmax函数能够将这些原始分数转换成一个有效的概率分布。
softmax ( z i ) = e z i ∑ j = 1 K e z j \text{softmax}(z_i) = \frac{e^{z_i}}{\sum_{j=1}^{K} e^{z_j}} softmax(zi)=j=1Kezjezi
其中 z i z_i zi 是第 i i i 个类别的logit, K K K 是类别总数。

  • 处理负值与确保非负性:指数函数 e z i e^{z_i} ezi 确保了所有转换后的值都是正数。
  • 归一化到0-1之间:通过除以所有指数项的总和,确保每个输出值 softmax ( z i ) \text{softmax}(z_i) softmax(zi) 都在0和1之间,符合概率的定义。
  • 对比简单归一化
    若尝试 normalize ( z i ) = z i ∑ j = 1 K z j \text{normalize}(z_i) = \frac{z_i}{\sum_{j=1}^{K} z_j} normalize(zi)=j=1Kzjzi,当 z i z_i zi 中存在负值时,可能会导致“概率”为负或分母为零(或接近零导致数值不稳定),这在概率解释上是不可接受的。Softmax通过指数变换优雅地解决了这个问题。
2. 保证输出概率总和为1

对于互斥的多分类问题(一个样本只属于一个类别),所有类别的概率之和必须为1。Softmax函数天然满足这一特性:
∑ i = 1 K softmax ( z i ) = ∑ i = 1 K e z i ∑ j = 1 K e z j = ∑ i = 1 K e z i ∑ j = 1 K e z j = 1 \sum_{i=1}^{K} \text{softmax}(z_i) = \sum_{i=1}^{K} \frac{e^{z_i}}{\sum_{j=1}^{K} e^{z_j}} = \frac{\sum_{i=1}^{K} e^{z_i}}{\sum_{j=1}^{K} e^{z_j}} = 1 i=1Ksoftmax(zi)=i=1Kj=1Kezjezi=j=1Kezji=1Kezi=1

  • 对比独立Sigmoid函数
    Sigmoid函数 σ ( z ) = 1 1 + e − z \sigma(z) = \frac{1}{1 + e^{-z}} σ(z)=1+ez1 主要用于二分类问题或多标签分类问题(一个样本可属于多个类别)。若在多分类问题中对每个logit独立应用Sigmoid,得到的输出虽然都在0-1之间,但它们的总和不一定为1。例如,[0.7, 0.8, 0.6] 就不构成一个有效的多类别概率分布。Softmax则强制这种互斥性。
3. 放大差异,突出最相关的选项

Softmax函数中的指数运算 e z i e^{z_i} ezi 具有放大输入值之间差异的特性。

  • 如果某个logit z k z_k zk 远大于其他logits,那么 e z k e^{z_k} ezk 将会比其他的 e z j e^{z_j} ezj 大得多,导致 softmax ( z k ) \text{softmax}(z_k) softmax(zk) 接近1,而其他项接近0。
  • 这种“赢者通吃”(winner-takes-all-like)的效应使得模型在预测时更有区分度,能够更自信地指向最可能的类别或注意力最应集中的部分。
  • 对比简单归一化
    简单归一化(即使是针对正数)通常会保留原始值的相对比例,而不会像softmax那样显著放大最大值的影响,可能导致概率分布过于平滑,决策不够果断。
4. 与交叉熵损失结合时具有简洁高效的梯度

在训练分类模型时,Softmax通常与交叉熵损失(Cross-Entropy Loss)函数配合使用。
L = − ∑ i = 1 K y i log ⁡ ( softmax ( z i ) ) L = -\sum_{i=1}^{K} y_i \log(\text{softmax}(z_i)) L=i=1Kyilog(softmax(zi))
其中 y i y_i yi 是真实标签(通常为one-hot编码, y k = 1 y_k=1 yk=1 表示真实类别为 k k k y i = 0 y_i=0 yi=0 for i ≠ k i \neq k i=k)。
在这种组合下,损失函数对logit z j z_j zj 的偏导数具有非常简洁的形式:
∂ L ∂ z j = softmax ( z j ) − y j \frac{\partial L}{\partial z_j} = \text{softmax}(z_j) - y_j zjL=softmax(zj)yj
这个梯度形式简单、计算高效,并且在数值上相对稳定,非常有利于深度神经网络的优化过程。

5. 坚实的概率理论基础

Softmax函数是多项逻辑回归(Multinomial Logistic Regression)的核心组成部分,后者是逻辑回归在多类别问题上的自然推广。它为模型的输出提供了一个清晰的概率解释,使其在统计学和概率模型的框架下具有坚实的理论基础。

6. 在注意力机制中的关键作用

在Transformer等模型的自注意力或注意力机制中,Softmax用于将原始的注意力得分(通常是点积结果)转换为注意力权重。这些权重表示在构建一个特定元素的表示时,应该给予输入序列中其他元素多少“关注度”。确保这些权重总和为1是至关重要的,因为它代表了注意力的分配。Softmax的放大效应也有助于模型集中关注最相关的部分。

总结

综上所述,Softmax函数因其以下关键特性而被广泛选择:

  1. 有效概率转换:能将任意实数logits转换为非负、0-1之间且总和为1的概率分布。
  2. 互斥类别建模:适合对互斥的多分类结果进行建模。
  3. 增强决策区分度:通过指数运算放大差异,使模型预测更明确。
  4. 优化友好:与交叉熵损失结合时梯度计算简洁高效。
  5. 理论完备性:具有良好的概率模型理论支持。
  6. 注意力权重分配:在注意力机制中完美契合权重分配的需求。
http://www.xdnf.cn/news/6837.html

相关文章:

  • C# NX二次开发-实体离散成点
  • 使用pyinstaller生成exe时,如何指定生成文件名字
  • Linux!启动~
  • WHAT - 前端同构 Isomorphic Javascript
  • Ubuntu系统安装VsCode
  • UAI 2025重磅揭晓:录取数据公布(附往届数据)
  • Python字符串常用内置函数详解
  • 独立开发者利用AI工具快速制作产品MVP
  • Qt功能区:Ribbon使用
  • Linux复习笔记(六)shell编程
  • 实现书签-第一部分
  • 中大型水闸安全监测系统建设实施方案
  • 在服务器上安装AlphaFold2遇到的问题(2)
  • 【C++】 —— 笔试刷题day_30
  • 【C++ | 内存管理】C++ weak_ptr 详解:成员函数、使用例子与实现原理
  • 力扣654题:最大二叉树(递归)
  • 实时技术方案对比:SSE vs WebSocket vs Long Polling
  • Java Set系列集合详解:HashSet、LinkedHashSet、TreeSet底层原理与使用场景
  • 产品经理入门——认识产品经理
  • OCCT知识笔记之Poly_Triangulation详解
  • YOLOv7训练时4个类别只出2个类别
  • vue使用Fabric和pdfjs完成合同签章及批注
  • 第八节第三部分:认识枚举、枚举的作用和应用场景
  • DeepSearch:WebThinker开启AI搜索研究新纪元!
  • 游戏站的几种形式
  • redis数据结构-11(了解 Redis 持久性选项:RDB 和 AOF)
  • STM32H743IIT6_ADC采集误差分析与ADC_DMA
  • 【论信息系统项目的整合管理】
  • leetcode 2900. 最长相邻不相等子序列 I 简单
  • 【LeetCode 热题 100】搜索插入位置 / 搜索旋转排序数组 / 寻找旋转排序数组中的最小值