transformer实战——多头注意力(Multi-head Attention)到底在干什么

zzz933
137 阅读8分钟

1 多头注意力(Multi-head Attention)

1.1 X的形状为[seq_len,d_model]

  1. 定义一组Q,K,V可以让一个词 attend to相关的词，我们可以定义多组 Q,K,V，让它们分别关注不同的上下文。计算 Q,K,V的过程还是一样，只不过线性变换的矩阵从一组 (WQW^Q ，WKW^K， WVW^V)变成了多组（W0QW_0^Q，W0KW_0^K，W0VW_0^V），（W1QW_1^Q，W1KW_1^K，W1VW_1^V）… 。

    • X的形状为[seq_len,d_model]
    • W0QW_0^Q的形状为[d_model,d_k]
    • W0KW_0^K 的形状为[d_model,d_k]
    • W0VW_0^V 的形状为[d_model,d_v]
    • Q 的形状为[seq_len,d_k]
    • K 的形状为[seq_len,d_k]
    • V 的形状为[seq_len,d_v]
    • ZiZ^i 的形状为[seq_len,d_v] image.png

  2. 如果heads=8,则对于输入矩阵X，每一组Q，K,V都可以由缩放点积注意力的方法得到一个输出矩阵Z。 image.png

  3. 接着我们先直接拼接Z0Z^0，Z1Z^1..Z7Z^7矩阵，然后与权重矩阵W0W^0相乘得到Z，矩阵就包含了所有注意力头从不同表示子空间学习到的信息，并且其维度与模型的其余部分（如残差连接、层归一化、前馈网络）是兼容的。它会被作为下一个模块的输入。

    • 拼接Z0Z^0，Z1Z^1..Z7Z^7矩阵之后的矩阵大小为[seq_len, 8*d_v]
    • 权重矩阵W0W^0的大小为[8*d_v,d_model]
    • 权重矩阵W0W^0的作用是将拼接后的多头注意力输出，重新投影回模型的原始维度 d_model
    • 最后Z矩阵的形状将是 [seq_len, d_model]，和一开始的 X 的形状相同
    • 所以我们一般设置 d_v = d_model / num_heads，拼接后的维度 num_heads * d_v 通常会等于 d_model，也就是W0W^0的形状是 [d_model, d_model] image.png

  4. 整体流程

    • 原始文本会经过词嵌入（Word Embedding）和位置编码（Positional Encoding），形成输入矩阵 X。
    • 除了第一个 Encoder 层，后续的 Encoder 层不再需要重新进行词嵌入和位置编码，它们直接接收前一个 Encoder 层的输出作为自己的输入。用 R 来表示这种输入，它实际上就是前一个 Encoder 层的输出。 image.png

1.2 X的形状为[batch_size,seq_len,d_model]

  1. 一开始以为WQW^Q 的形状为[batch_size,d_model,d_k]，这个是错误的，权重矩阵WQW^Q ，WKW^K， WVW^V本身并不包含 batch_size 维度。它们是模型学习到的固定参数，对于批次中的每个样本都是共享的。
  2. nn.Linear(in_features, out_features) 在内部创建的权重矩阵 weight 的形状是 [out_features, in_features]（这个形状和X不能相乘不要紧，因为不需要担心 PyTorch 内部的权重矩阵是转置存储的。只需要记住 nn.Linear(in_features, out_features) 就是在说：“我的输入向量有 in_features 这么多维，然后我通过线性变换后，希望得到一个 out_features 这么多维的输出向量。）
  3. 当 nn.Linear 接收到 [batch_size, seq_len, d_model] 这样的输入X 时，它会自动在内部对批次中的每个 [seq_len, d_model] 子矩阵执行相同的线性变换。正如Dot-Product Attention里写的，它会将 [batch_size, seq_len, d_model] 视为 (batch_size * seq_len) 个独立的 d_model 维向量进行处理。
  4. 所以权重矩阵的形状同一开始说的那样
    • W0QW_0^Q 的形状为[d_model,d_k]
    • W0KW_0^K 的形状为[d_model,d_k]
    • W0VW_0^V 的形状为[d_model,d_v]
  5. 当 X 经过这些线性层（例如 nn.Linear(d_model, d_k)）时，nn.Linear 会将变换应用到 d_model 维度，同时保留前面的 batch_size 和 seq_len 维度。
    • Q 的形状为[batch_size,seq_len,d_k]
    • K 的形状为[batch_size,seq_len,d_k]
    • V 的形状为[batch_size,seq_len,d_v]
  6. ZiZ^i 的形状为[batch_size,seq_len,d_v]
  7. 拼接ZiZ^i 后的形状为[batch_size,seq_len,heads*d_v]
  8. 最后Z的形状为[batch_size,seq_len,d_model]

1.3 在代码中的变化

1.3.1 计算QKV

  1. 从理论概念上讲，我们认为每个注意力头都需要一套独立的权重矩阵来将输入投影到其子空间。对于查询（Q）来说，这意味着：

    • 假设heads=8

    • head_0的W0QW_0^Q形状为 [d_model, d_k]。

    • head_1的W1QW_1^Q形状为 [d_model, d_k]。

    • head_i的WiQW_i^Q形状为 [d_model, d_k]。

    • image.png

  2. 手动代码实现,需要8行来算出每一个heads的Q 
    • # 假设 X 的形状是 [seq_len, d_model] 
Q_0 = input_word_vector @ WQ_0  # Q0 形状: [seq_len, d_k]
Q_1 = input_word_vector @ WQ_1  # Q1 形状: [seq_len, d_k]
Q_i = input_word_vector @ WQ_i  # Qi 形状: [seq_len, d_k]
  3. 代码实现上的技巧
    • 也就是使用一个 nn.Linear(d_model, d_model) 的层。
    • 该层内部的权重矩阵的形状是 [d_model, d_model]，可以同时完成W0QW_0^Q，W1QW_1^Q..W7QW_7^Q的工作
    • 这个矩阵可以被看作是由heads个[d_model, d_k]的小矩阵水平拼接而成的(之前说的d_k通常等于d_model/heads)：
    • image.png
    • 再用X乘这个矩阵，得到Q_0...Q_7，即结果形状[seq_len,d_model]是由heads个[seq_len, d_k]的小矩阵水平拼接 image.png

1.3.2 QKV的reshape

  1. Emb Sz就是d_model,Query Sz就是d_k,所以是从[seq,d_model] -> [seq,heads,d_k] -> [heads, seq, d_k]

image.png 2. 整体的Q的逻辑如下图，假设Heads=2

image.png

1.3.3 Z拼接的具体流程

假设heads=2,结果z0,z1的拼接结果的大小为从2个[seq_len,d_v]到[seq_len,d_model]（前提d_v = d_model / num_heads）。需要经过一次swap和reshape才能成功实现上述拼接。

image.png

1.4 代码

1.4.1 Multi-head Attention的全部流程

假设heads为2，Multi-head Attention的全部流程如下图 image.png

# 与之前Scaled Dot-Product Attention不同之处在于输入的q,k,v多了一个维度heads
# attention 函数本身不关心有多少个头，它只是对它接收到的多维张量在特定的维度上执行矩阵乘法和 Softmax 操作。
# 只要 MultiHeadedAttention 成功地将 query, key, value 重塑并转置成了 (batch_size, num_heads, sequence_length, d_k) 的形状
# attention 函数就能正确地对其进行并行计算。
def attention(query, key, value, mask=None, dropout=None):
    "Compute 'Scaled Dot Product Attention'"
    d_k = query.size(-1)
    scores = torch.matmul(query, key.transpose(-2, -1)) \
             / math.sqrt(d_k)
    # 如果提供了 mask，则将 mask 中值为 0 的位置对应的分数替换为一个非常小的负数 (-1e9)，
    # 在 Softmax 后使这些位置的注意力权重趋近于 0。
    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e9)
    p_attn = F.softmax(scores, dim = -1)
    # 如果提供了 dropout，则对注意力权重应用 dropout。
    if dropout is not None:
        p_attn = dropout(p_attn)
    return torch.matmul(p_attn, value), p_attn

#  创建 N 个 `module` 的深拷贝并放入 `nn.ModuleList` 中。
# 每个深拷贝都是一个独立的模块实例，拥有独立的参数。
# 返回一个 nn.ModuleList，其中包含 N 个独立的 nn.Module 实例
def clones(module, N):
    "Produce N identical layers."
    return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])


class MultiHeadedAttention(nn.Module):
    def __init__(self, h, d_model, dropout=0.1):
        super(MultiHeadedAttention, self).__init__()
        
        # h 代表的是 注意力头的数量
        # 断言模型的维度 `d_model` 必须能被注意力头的数量 `h` 整除。若不能，则立即报错。
        assert d_model % h == 0
        
        # 假设d_v 一直等于 d_k
        self.d_k = d_model // h 
        self.h = h 
        
        # 创建 4 个独立的 `nn.Linear` 层。每层的权重矩阵形状为 (d_model, d_model)
        # 前三个线性层（用于 Q、K、V 的投影）
        # 第四个线性层（用于最终的输出融合W_0)
        self.linears = clones(nn.Linear(d_model, d_model), 4)
        
        self.attn = None # 初始化 `attn` 属性。
        self.dropout = nn.Dropout(p=dropout) # 初始化 Dropout 层
            
    def forward(self, query, key, value, mask=None):
        # query, key, value 的输入形状： (batch_size, sequence_length, d_model)
        
        if mask is not None:
            # Same mask applied to all h heads.
            mask = mask.unsqueeze(1) # 将 mask 维度扩展
            # mask 的形状： (batch_size, 1, sequence_length, sequence_length)
        
        nbatches = query.size(0) # 获取当前批次的样本数量
            

        # query, key, value 进行线性投影和形状重塑
        # l(x): 对输入的 x (query/key/value) 应用线性变换
        # x 的形状： (nbatches, sequence_length, d_model)
        # l(x) 的输出形状： (nbatches, sequence_length, d_model)
        # view(nbatches, -1, self.h, self.d_k)： 重新整形为(nbatches, sequence_length, self.h, self.d_k)。
        # transpose(1, 2): 交换 `sequence_length` (索引 1) 和 `h` (索引 2) 维度。形状变为： (nbatches, self.h, sequence_length, self.d_k)
        query, key, value = \
            [l(x).view(nbatches, -1, self.h, self.d_k).transpose(1, 2)
             for l, x in zip(self.linears, (query, key, value))]
            

        # query, key, value 的形状均为： (nbatches, self.h, sequence_length, self.d_k)
        # attention`函数返回的 x是上下文向量，self.attn 是注意力权重
        # x的形状： (nbatches, self.h, sequence_length, self.d_k)
        # self.attn 的形状： (nbatches, self.h, sequence_length, sequence_length)
        x, self.attn = attention(query, key, value, mask=mask, 
                                 dropout=self.dropout)
            

        # 将多个注意力头的输出（x进行概念上的“拼接”，并通过一个最终的线性层进行融合
        # x.transpose(1, 2): 交换 h`(索引 1) 和 sequence_length`(索引 2) 维度。形状变为： (nbatches, sequence_length, self.h, self.d_k)
        # contiguous(): 确保张量在内存中是连续的
        # view(nbatches, -1, self.h * self.d_k): 将所有头的维度 self.h 和 self.d_k`拼接起来
        #  由于 self.h * self.d_k等于 d_model，形状变为： (nbatches, sequence_length, d_model)
        x = x.transpose(1, 2).contiguous() \
              .view(nbatches, -1, self.h * self.d_k)
        
        # self.linears[-1](x): 应用 linears列表中最后一个线性层（即 W_0 矩阵）。
        #   输入 x 的形状： (nbatches, sequence_length, d_model)。
        #   self.linears[-1]的权重矩阵形状： (d_model, d_model)。
        # 返回的最终输出 x的形状是： (nbatches, sequence_length, d_model)。
        return self.linears[-1](x)

参考

  1. Transformer 详解
  2. The Illustrated Transformer
  3. The Annotated Transformer
  4. Transformers Explained Visually (Part 3): Multi-head Attention, deep dive
avatar
文章
阅读
粉丝