Datewhale组队学习

从零实现注意力机制：NumPy 和 SciPy 的实践之旅

注意力机制，作为现代深度学习中的核心组件之一，早已在自然语言处理、计算机视觉等领域大放异彩。它的魅力在于能够动态地聚焦于输入数据的关键部分，从而让模型更“聪明”地处理信息。今天，我们将抛开复杂的框架，仅用 NumPy 和 SciPy，从零实现一个通用的注意力机制，深入理解其背后的数学逻辑。

1. 词向量：将词汇语义嵌入高维空间

一切从词向量开始。词向量是单词在嵌入空间中的数学表示，它将离散的符号转化为连续的向量，从而能够进行数学运算。这里，定义了四个单词的词向量，每个向量的维度为3：

import numpy as np
from scipy.special import softmax

word_1 = np.array([1, 0, 0])
word_2 = np.array([0, 1, 0])
word_3 = np.array([1, 1, 0])
word_4 = np.array([0, 0, 1])
words = np.array([word_1, word_2, word_3, word_4])

这些词向量被堆叠成一个矩阵 words，形状为 (4, 3)，每一行代表一个单词的嵌入。这一步的意义在于，将离散的符号转化为可计算的数学对象，为后续的注意力计算奠定基础。

2. 权重矩阵：映射到查询、键和值

注意力机制的核心在于查询（Query）、键（Key）和值（Value）的生成。这些概念听起来抽象，但本质上是将词向量映射到不同的空间中，以便计算它们之间的关系。通过随机初始化三个权重矩阵 W_Q、W_K 和 W_V，词向量被转换为查询、键和值：

np.random.seed(42)
W_Q = np.random.randint(3, size=(3, 3))
W_K = np.random.randint(3, size=(3, 3))
W_V = np.random.randint(3, size=(3, 3))

Q = words @ W_Q
K = words @ W_K
V = words @ W_V

查询、键和值的生成过程，本质上是矩阵乘法。这一步的意义在于，将原始的词向量映射到不同的语义空间中，为后续的相似度计算做准备。

3. 得分计算：量化相似度

接下来，通过查询和键的点积，计算得分矩阵。每个得分表示一个查询向量与一个键向量的相似度：

scores = Q @ K.T

得分的计算是注意力机制的关键步骤之一。点积越大，表示两个向量越相似。这一步的意义在于，量化每个查询向量与所有键向量之间的关系，为后续的权重分配提供依据。

4. 权重分配：Softmax 概率化

得分矩阵虽然量化了相似度，但还需要将其转化为概率分布。这里，softmax 函数登场了。通过对得分进行缩放并应用 softmax，得到了注意力权重：

weights = softmax(scores / np.sqrt(K.shape[1]), axis=1)

softmax 的作用是将得分转化为概率分布，使得每个查询向量对所有键向量的权重之和为1。缩放操作则是为了保持数值稳定性，防止得分过大或过小导致计算溢出或下溢。

5. 注意力输出：加权求和

最后，通过加权求和的方式，计算注意力输出。每个注意力输出是所有值向量的加权和，权重由 softmax 函数计算得到：

attention = weights @ V
print(attention)

注意力输出的计算是注意力机制的最终步骤。通过加权求和，模型能够动态地聚焦于输入序列中最相关的部分，从而生成更准确的输出。

多头注意力机制：从单头到多头的进化

注意力机制的核心思想是让模型能够动态地关注输入序列中的不同部分，从而更好地捕捉上下文信息。然而，单头注意力机制有一个明显的局限性：它只能从一个角度捕捉输入序列的信息。为了克服这一限制，多头注意力机制（Multi-Head Attention）应运而生。通过并行地计算多个注意力头，模型能够从不同的子空间中提取信息，从而更全面地理解输入数据。

今天，我们将深入探讨多头注意力机制的实现，并通过 PyTorch 实现一个完整的 MultiHeadAttention 类。

1. 多头注意力机制的核心思想

多头注意力机制的核心在于将输入数据映射到多个子空间中，并在每个子空间中独立地计算注意力。具体来说，输入数据会被线性变换为多个查询（Query）、键（Key）和值（Value），然后在每个头上分别计算注意力。最后，所有头的输出会被拼接起来，并通过一个线性层映射回原始维度。

这种设计的好处在于：

• 多视角捕捉信息：每个头可以关注输入序列的不同部分，从而捕捉更丰富的特征。
• 并行计算：多个头的计算可以并行进行，提高了计算效率。
• 灵活性：通过调整头的数量，可以控制模型的复杂度和表达能力。

2. 实现细节：从代码中理解多头注意力

以下是一个完整的 MultiHeadAttention 类的实现，我们将逐行解析其关键部分。

import math
import torch
import torch.nn as nn
import torch.nn.functional as F

class MultiHeadAttention(nn.Module):
    def __init__(self, heads, d_model, dropout=0.1):
        super().__init__()
        self.d_model = d_model  # 模型的总维度
        self.d_k = d_model // heads  # 每个头的维度
        self.h = heads  # 头的数量

        # 线性层，用于生成 Q、K、V
        self.q_linear = nn.Linear(d_model, d_model)
        self.k_linear = nn.Linear(d_model, d_model)
        self.v_linear = nn.Linear(d_model, d_model)

        # Dropout 层，用于正则化
        self.dropout = nn.Dropout(dropout)

        # 输出线性层，用于将多头输出映射回模型维度
        self.out = nn.Linear(d_model, d_model)

    def attention(self, q, k, v, mask=None):
        # 计算 Q 和 K 的点积，并除以 sqrt(d_k) 进行缩放
        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)

        # 如果提供了掩码，则将掩码对应的位置设置为负无穷
        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)

        # 应用 softmax 计算注意力权重
        scores = F.softmax(scores, dim=-1)

        # 应用 dropout
        scores = self.dropout(scores)

        # 将注意力权重与 V 相乘，得到输出
        output = torch.matmul(scores, v)
        return output

    def forward(self, q, k, v, mask=None):
        batch_size = q.size(0)

        # 将输入通过线性层，并调整形状以进行多头计算
        q = self.q_linear(q).view(batch_size, -1, self.h, self.d_k).transpose(1, 2)
        k = self.k_linear(k).view(batch_size, -1, self.h, self.d_k).transpose(1, 2)
        v = self.v_linear(v).view(batch_size, -1, self.h, self.d_k).transpose(1, 2)

        # 计算注意力输出
        scores = self.attention(q, k, v, mask)

        # 将多头输出拼接，并调整形状以匹配模型维度
        concat = scores.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model)

        # 通过输出线性层
        output = self.out(concat)
        return output

3. 关键步骤解析

1. 初始化参数

• d_model：模型的总维度。
• d_k：每个头的维度，等于 d_model // heads。
• heads：头的数量。
• 线性层 q_linear、k_linear、v_linear：用于将输入映射到查询、键和值。
• out：输出线性层，用于将多头输出映射回模型维度。

2. 注意力计算

• 点积计算：通过矩阵乘法计算查询和键的点积，并除以 sqrt(d_k) 进行缩放，以防止点积过大。
• 掩码处理：如果提供了掩码，则将掩码对应的位置设置为负无穷，这样在 softmax 后这些位置的值为0。
• Softmax：将得分转换为概率分布。
• Dropout：对注意力权重进行随机丢弃，防止过拟合。
• 加权求和：将注意力权重与值相乘，得到输出。

3. 前向传播

• 线性变换：将输入通过线性层，生成查询、键和值。
• 形状调整：将输入调整为 (batch_size, heads, seq_len, d_k)，以便进行多头计算。
• 注意力计算：在每个头上独立计算注意力。
• 拼接输出：将所有头的输出拼接起来，并通过输出线性层映射回模型维度。

4. 测试与验证

以下代码用于测试 MultiHeadAttention 类的功能：

if __name__ == "__main__":
    heads = 4
    d_model = 128  # d_model 必须是 heads 的整数倍
    dropout = 0.1

    model = MultiHeadAttention(heads, d_model, dropout)

    batch_size = 2
    seq_len = 5
    q = torch.rand(batch_size, seq_len, d_model)  # Query
    k = torch.rand(batch_size, seq_len, d_model)  # Key
    v = torch.rand(batch_size, seq_len, d_model)  # Value

    output = model(q, k, v)
    print("Output shape:", output.shape)  # 应该是 [batch_size, seq_len, d_model]

    loss = output.mean()
    loss.backward()
    print("Backward pass completed.")

输出结果：

Output shape: torch.Size([2, 5, 128])
Backward pass completed.

欢迎各位大佬在评论区交流&批评指正！

Datewhale组队学习

参考链接

1. Datawhale AI+X 共学活动 Fun-Transformer —— Task5：项目实践
2. Attention Is All You Need