当前位置：首页 > article >正文

人工智能入门课【手写自注意力机制】

article 2025/2/3 2:19:07

原理

自注意力（Self-Attention）是一种强大的机制，广泛应用于自然语言处理、计算机视觉等领域，尤其是在Transformer架构中发挥了关键作用。它的核心思想是让模型能够动态地关注输入序列中不同位置之间的关系，从而更好地捕捉全局信息和长距离依赖。

在传统的序列处理模型如循环神经网络（RNN）中，信息是按时间步逐个传递的，模型在处理当前时刻的信息时，只能依赖于之前时刻的信息，这使得它在处理长序列时容易出现梯度消失或梯度爆炸的问题，难以捕捉长距离依赖关系。而自注意力机制通过计算序列中每个位置与其他所有位置的关联程度，直接在全局范围内进行信息交互，解决了这一问题。

具体来说，自注意力机制的实现过程可以分为以下几个步骤：

首先，输入序列会被转换为三个矩阵，分别是查询（Query）、键（Key）和值（Value）。这三个矩阵是通过输入序列与三个不同的可学习权重矩阵相乘得到的。
1. 具体公式： $\left\{\begin{matrix} query = W^q\cdot x\\ key = W^k\cdot x\\ value = W^v\cdot x \end{matrix}\right.$
2. 查询矩阵用于表示当前需要关注的内容，键矩阵用于表示序列中各个位置的特征，值矩阵则包含了序列中各个位置的实际信息。这三个矩阵的维度通常是相同的，且可以根据具体任务进行调整。
接下来，计算注意力分数。公式为 $attn\_score = \frac{QK^T}{\sqrt{d_k}}$ 。对于输入序列中的每个位置，都会将其查询向量与序列中所有位置的键向量进行点积运算，得到一个分数矩阵。
1. 为了使这些分数具有可比性，通常会对它们进行缩放，即将每个分数除以键向量维度的平方根。这样做的目的是防止点积结果过大而导致梯度消失或梯度爆炸。
然后，将分数矩阵通过softmax函数进行归一化，公式为 $attn\_weight = Softmax(attn\_score)$ ，得到注意力权重矩阵。softmax函数的作用是将分数矩阵中的每个元素转换为一个概率值，使得所有权重的和为1。这样，每个位置的权重就表示了它在全局范围内的重要性，权重越大，说明该位置与其他位置的关联越强。
最后，将注意力权重矩阵与值矩阵相乘，公式为 $output = attn\_weight\cdot V$ ，得到加权的值矩阵。这个加权的值矩阵就是自注意力机制的输出，它包含了输入序列中各个位置经过加权后的信息。通过这种方式，模型能够根据注意力权重动态地组合序列中的信息，从而更好地捕捉序列中的全局依赖关系。
总公式为： $output = Softmax(\frac{QK^T}{\sqrt{d_k}})\cdot V$

自注意力机制的一个重要特性是并行化计算。由于它不需要像RNN那样按顺序逐个处理序列中的元素，因此可以同时计算所有位置之间的注意力关系，大大提高了计算效率。这使得自注意力机制在处理大规模数据时具有显著的优势。

此外，自注意力机制还可以通过多头注意力（Multi-Head Attention）的方式进一步增强模型的表现能力。多头注意力的核心思想是将输入序列分成多个不同的“头”，每个头都独立地进行自注意力计算，然后将所有头的输出拼接在一起，再通过一个线性变换进行整合。这样做的好处是能够让模型从不同的角度捕捉序列中的信息，从而更好地理解序列的结构和语义。

单头自注意力

实现方法1

这段代码实现了一个简化版本的单头自注意力机制（Self-Attention），并展示了如何使用它处理输入数据。

在forward方法中

首先分别对输入数据 x 应用三个线性变换，得到查询（query）、键（key）和值（value）。它们的形状仍然是 [batch_size, seq_len, hidden_dim]。
计算query和key的点积
1. key.transpose(1, 2) 将 key 的形状从 [batch_size, seq_len, hidden_dim] 转置为 [batch_size, hidden_dim, seq_len]。
2. torch.matmul(query, key.transpose(1, 2)) 计算 query 和 key 的点积，得到形状为 [batch_size, seq_len, seq_len] 的注意力分数矩阵。
3. / math.sqrt(self.hidden_dim) 是一个缩放操作，用于防止点积结果过大导致的梯度消失或梯度爆炸问题。缩放因子是 hidden_dim 的平方根。
对注意力分数矩阵应用 Softmax 函数，将每个位置的分数转换为概率值，使得每一行的和为 1。这一步确保了注意力权重的合理性。
使用注意力权重矩阵 attn_weights 对值矩阵 value 进行加权求和，得到最终的自注意力输出。输出的形状为 [batch_size, seq_len, hidden_dim]。

import math
import torch
import torch.nn as nn


class SelfAttn1(nn.Module):     # 简化版本，单头
    def __init__(self, hidden_dim=728):
        super().__init__()
        self.hidden_dim = hidden_dim        # 隐藏层维度
        self.query_proj = nn.Linear(hidden_dim, hidden_dim)     # q的线性映射，映射到hidden_dim维
        self.key_proj = nn.Linear(hidden_dim, hidden_dim)       # k的线性映射，映射到hidden_dim维
        self.value_proj = nn.Linear(hidden_dim, hidden_dim)     # v的线性映射，映射到hidden_dim维

    def forward(self, x):
        # x: [batch_size, seq_len, hidden_dim]
        query = self.query_proj(x)    # 分别对输入数据 x 应用三个线性变换
        key = self.key_proj(x)
        value = self.value_proj(x)
        attn_weights = torch.matmul(query, key.transpose(1, 2)) / math.sqrt(self.hidden_dim)    # 计算注意力权重 [batch_size, seq_len, seq_len]
        attn_weights = torch.softmax(attn_weights, dim=-1)          # 对注意力分数矩阵应用 Softmax 函数，将每个位置的分数转换为概率值，使得每一行的和为 1。  [batch_size, seq_len, seq_len]
        attn_output = torch.matmul(attn_weights, value)         # 使用注意力权重矩阵 attn_weights 对值矩阵 value 进行加权求和，得到最终的自注意力输出。  [batch_size, seq_len, hidden_dim]  也可以使用@
        return attn_output


x = torch.randn(4, 10, 728)
self_attn = SelfAttn1()
attn_output = self_attn(x)
print(attn_output.shape)

实现方法2

这段代码实现了一个优化版本的单头自注意力机制（SelfAttn2），与之前的 SelfAttn1 代码相比，主要的区别在于对查询（Query）、键（Key）和值（Value）的计算方式进行了优化。主要区别如下：

在 SelfAttn2 中，使用了一个单一的线性变换层 self.proj，将输入数据映射到一个维度为 hidden_dim * 3 的空间（即query、key、value的三个映射矩阵被合并成一个）。这意味着输出张量的最后一个维度是输入维度的三倍。随后的forward()中，这个输出张量会被分割成三个部分，分别对应查询（Query）、键（Key）和值（Value）。
forward()
1. 首先，通过 self.proj(x) 将输入数据 x 映射到一个维度为 [batch_size, seq_len, hidden_dim * 3] 的张量 proj_output。
2. 然后，使用 torch.split 将 proj_output 按最后一个维度分割成三个部分，每部分的维度为 [batch_size, seq_len, hidden_dim]，分别对应查询（query）、键（key）和值（value）。这种分割方式确保了每个部分的维度与原始输入维度一致。
3. 接着同SelfAttn1()

import math
import torch
import torch.nn as nn



class SelfAttn2(nn.Module):     # 效率优化，qkv大矩阵运算
    def __init__(self, hidden_dim=728):
        super().__init__()
        self.hidden_dim = hidden_dim
        self.proj = nn.Linear(hidden_dim, hidden_dim * 3)       # qkv的线性映射，映射到hidden_dim * 3维

    def forward(self, x):
        # x: [batch_size, seq_len, hidden_dim]
        proj_output = self.proj(x)    # 将输入数据 x 映射到一个维度为 [batch_size, seq_len, hidden_dim * 3] 的张量 proj_output。
        query, key, value = torch.split(proj_output, self.hidden_dim, dim=-1)       # 使用 torch.split 将 proj_output 按最后一个维度分割成三个部分，分别对应查询（query）、键（key）和值（value）。[batch_size, seq_len, hidden_dim]
        attn_weights = torch.matmul(query, key.transpose(1, 2)) / math.sqrt(self.hidden_dim)
        attn_weights = torch.softmax(attn_weights, dim=-1)
        attn_output = torch.matmul(attn_weights, value)
        return attn_output


x = torch.randn(4, 10, 728)
self_attn = SelfAttn2()
attn_output = self_attn(x)
print(attn_output.shape)

在 SelfAttn1 中，分别使用了三个独立的线性变换层（self.query_proj、self.key_proj 和 self.value_proj）来计算查询、键和值。这种方式虽然直观，但计算效率较低，因为每次都需要对输入数据进行三次独立的线性变换。

但在 SelfAttn2 中，通过一个线性变换层将输入数据映射到一个更大的空间，然后一次性分割成查询、键和值。这种方式减少了线性变换的次数，从而提高了计算效率。具体来说，SelfAttn2 只需要一次矩阵乘法操作，而 SelfAttn1 需要三次矩阵乘法操作。SelfAttn2 的实现更加简洁，因为它将查询、键和值的计算合并到了一个线性变换层中，减少了代码的冗余性。

实现方法3

这段代码实现了一个更完整的单头自注意力机制（SelfAttn3），在之前的版本基础上增加了以下功能：

注意力掩码（attn_mask）：用于遮蔽某些位置的注意力权重，例如处理序列中的填充部分或防止解码器中的未来信息泄露
Dropout：在注意力权重上应用 Dropout，以防止模型过拟合。
输出映射（output_proj）：对自注意力的输出进行额外的线性变换，以调整输出的特征空间。

import math
import torch
import torch.nn as nn


class SelfAttn3(nn.Module):     # 加入dropout和attn_mask、以及output映射
    def __init__(self, hidden_dim=728, dropout_rate=0.1, *args, **kwargs):
        super().__init__()
        self.hidden_dim = hidden_dim
        self.proj = nn.Linear(hidden_dim, hidden_dim * 3)       # qkv的线性映射，映射到hidden_dim * 3维
        self.attention_dropout = nn.Dropout(dropout_rate)        # dropout

        self.output_proj = nn.Linear(hidden_dim, hidden_dim)

    def forward(self, x, attn_mask=None):
        # x: [batch_size, seq_len, hidden_dim]
        proj_output = self.proj(x)
        query, key, value = torch.split(proj_output, self.hidden_dim, dim=-1)       # [batch_size, seq_len, hidden_dim]
        attn_weights = torch.matmul(query, key.transpose(1, 2)) / math.sqrt(self.hidden_dim)
        if attn_mask is not None:    # 如果提供了注意力掩码 attn_mask，则将掩码中为 0 的位置对应的注意力权重设置为一个非常小的值（-1e20）。这样，在应用 Softmax 时，这些位置的权重会趋近于 0，从而实现遮蔽效果。
            attn_weights = attn_weights.masked_fill(attn_weights==0, float('-1e20'))
        attn_weights = torch.softmax(attn_weights, dim=-1)
        attn_weights = self.attention_dropout(attn_weights)
        attn_output = torch.matmul(attn_weights, value)

        attn_output = self.output_proj(attn_output)    # 对自注意力的输出进行额外的线性变换
        return attn_output


x = torch.randn(4, 10, 728)
mask = torch.randint(0, 2, (4, 10, 728))
self_attn = SelfAttn3()
attn_output = self_attn(x, mask)
print(attn_output.shape)

其中应用掩码部分：

如果提供了注意力掩码 attn_mask，则将掩码中为 0 的位置对应的注意力权重设置为一个非常小的值（-1e20）。这样，在应用 Softmax 时，这些位置的权重会趋近于 0，从而实现遮蔽效果。

if attn_mask is not None:
    attn_weights = attn_weights.masked_fill(attn_mask == 0, float('-1e20'))

输出映射：对自注意力的输出进行额外的线性变换，调整输出的特征空间。这一步可以进一步优化模型的表达能力。

attn_output = self.output_proj(attn_output)

多头自注意力

多头自注意力（Multi-Head Attention）是自注意力机制的一个扩展，它通过将输入数据分割成多个不同的“头”（heads），分别计算每个头的自注意力，然后将这些头的输出拼接起来，从而能够从多个不同的角度捕捉输入数据的特征。这种机制在Transformer架构中被广泛应用，极大地提升了模型对复杂数据结构的建模能力。

多头自注意力的原理

在单头自注意力中，输入数据首先被映射到查询（Query）、键（Key）和值（Value）三个矩阵，然后通过计算查询和键的点积得到注意力分数，再通过Softmax函数将这些分数转换为注意力权重，最后用这些权重对值进行加权求和，得到最终的输出。然而，单头自注意力只能从一个固定的视角来捕捉输入数据的特征，这在处理复杂的语言或图像数据时可能会限制模型的表现能力。

多头自注意力的核心思想是将输入数据分割成多个不同的“头”，每个头都独立地计算自注意力，从而能够从多个不同的角度捕捉输入数据的特征。具体来说，输入数据首先被分割成多个子空间，每个子空间对应一个“头”。在每个头中，分别计算查询、键和值，然后计算自注意力。最后，将所有头的输出拼接起来，再通过一个线性变换层进行整合，得到最终的输出。

多头自注意力的优势

多头自注意力的主要优势在于它能够从多个不同的角度捕捉输入数据的特征。在自然语言处理中，这意味着模型能够同时捕捉到句子中的局部信息和全局信息，从而更好地理解句子的语义。例如，一个头可能专注于捕捉句子中的主谓宾结构，而另一个头可能专注于捕捉句子中的修饰词和被修饰词之间的关系。通过将这些不同视角的信息融合在一起，模型能够得到一个更加丰富和全面的特征表示。

此外，多头自注意力还具有并行计算的优势。由于每个头的计算是独立的，因此可以并行进行，从而提高了计算效率。这使得多头自注意力机制在处理大规模数据时具有显著的优势。

这段代码实现了一个多头自注意力机制（MultiHeadSelfAttn），它是Transformer架构中的核心组件之一。代码中详细展示了如何将输入数据分割成多个“头”，分别计算每个头的自注意力，然后将这些头的输出拼接起来并进行整合。

import math
import torch
import torch.nn as nn


class MultiHeadSelfAttn(nn.Module):
    def __init__(self, hidden_dim=728, head_num=8, dropout_rate=0.1):
        super().__init__()
        self.hidden_dim = hidden_dim
        self.head_num = head_num
        self.head_dim = hidden_dim // head_num
        self.query_proj = nn.Linear(hidden_dim, hidden_dim)  # q的线性映射，映射到hidden_dim维
        self.key_proj = nn.Linear(hidden_dim, hidden_dim)  # k的线性映射，映射到hidden_dim维
        self.value_proj = nn.Linear(hidden_dim, hidden_dim)  # v的线性映射，映射到hidden_dim维

        self.attention_dropout = nn.Dropout(dropout_rate)
        self.output_proj= nn.Linear(hidden_dim, hidden_dim)

    def forward(self, x, attn_mask=None):
        # x: [batch_size, seq_len, hidden_dim]
        batch_size, seq_len, _ = x.size()
        query = self.query_proj(x)
        key = self.key_proj(x)
        value = self.value_proj(x)

        # [batch_size, seq_len, hidden_dim] -> [batch_size, seq_len, head_num, head_dim]
        query = query.view(batch_size, seq_len, self.head_num, self.head_dim).transpose(1, 2)       # [batch_size, head_num, seq_len, head_dim]
        key = key.view(batch_size, seq_len, self.head_num, self.head_dim).transpose(1, 2)
        value = value.view(batch_size, seq_len, self.head_num, self.head_dim).transpose(1, 2)
        attn_score = torch.matmul(query, key.transpose(-1, -2)) / math.sqrt(self.head_dim)      # [batch_size, head_num, seq_len, seq_len]
        if attn_mask is not None:
            attn_score = attn_score.masked_fill(attn_mask==0, float('-inf'))
        attn_score = torch.softmax(attn_score, dim=-1)
        attn_score = self.attention_dropout(attn_score)
        attn_output = torch.matmul(attn_score, value)
        attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, seq_len, self.hidden_dim)
        attn_output = self.output_proj(attn_output)
        return attn_output


x = torch.randn(4, 10, 728)
mask = torch.randint(0, 2, (4, 8, 10, 10))
self_attn = MultiHeadSelfAttn()
attn_output = self_attn(x, mask)
print(attn_output.shape)

核心部分代码讲解

首先分割多个头

将查询、键和值的形状从 [batch_size, seq_len, hidden_dim] 转换为 [batch_size, seq_len, head_num, head_dim]。
然后通过 transpose(1, 2) 将形状调整为 [batch_size, head_num, seq_len, head_dim]，以便每个头可以独立计算自注意力。

query = query.view(batch_size, seq_len, self.head_num, self.head_dim).transpose(1, 2)       # [batch_size, head_num, seq_len, head_dim]
key = key.view(batch_size, seq_len, self.head_num, self.head_dim).transpose(1, 2)
value = value.view(batch_size, seq_len, self.head_num, self.head_dim).transpose(1, 2)

计算多头版的查询和键的点积

计算每个头的查询和键的点积，得到注意力分数矩阵。key.transpose(-1, -2) 将键的形状从 [batch_size, head_num, seq_len, head_dim] 转置为 [batch_size, head_num, head_dim, seq_len]。
点积结果的形状为 [batch_size, head_num, seq_len, seq_len]。
通过除以 math.sqrt(self.head_dim) 进行缩放，防止点积结果过大导致的梯度消失或梯度爆炸问题。

attn_score = torch.matmul(query, key.transpose(-1, -2)) / math.sqrt(self.head_dim)      # [batch_size, head_num, seq_len, seq_len]

计算注意力分数（这部分和之前的没区别）

对注意力分数应用 Softmax 函数，将分数转换为概率值。
对注意力权重应用 Dropout，随机丢弃一部分权重，以防止过拟合。
使用注意力权重对值进行加权求和，得到每个头的输出。输出的形状为 [batch_size, head_num, seq_len, head_dim]。

attn_score = torch.softmax(attn_score, dim=-1)
attn_score = self.attention_dropout(attn_score)
attn_output = torch.matmul(attn_score, value)

恢复原来的形状

将每个头的输出拼接起来，恢复到原始的形状 [batch_size, seq_len, hidden_dim]。
通过一个线性变换层 output_proj 对拼接后的输出进行整合，得到最终的多头自注意力输出。

attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, seq_len, self.hidden_dim)
attn_output = self.output_proj(attn_output)

Decoder

这段代码实现了一个基于Transformer架构的解码器（Decoder），它由多个解码器层（DecoderLayer）组成。每个解码器层包含多头自注意力机制和前馈神经网络（FFN）。代码还展示了如何使用这个解码器处理输入数据，并生成输出。

import math
import torch
import torch.nn as nn


class DecoderLayer(nn.Module):
    def __init__(self, hidden_dim, head_num, dropout_rate=0.1):
        super(DecoderLayer, self).__init__()
        self.hidden_dim = hidden_dim
        self.head_num = head_num
        self.head_dim = hidden_dim // head_num

        # multi-head self-attention
        self.query_proj = nn.Linear(hidden_dim, hidden_dim)  # q的线性映射，映射到hidden_dim维
        self.key_proj = nn.Linear(hidden_dim, hidden_dim)  # k的线性映射，映射到hidden_dim维
        self.value_proj = nn.Linear(hidden_dim, hidden_dim)  # v的线性映射，映射到hidden_dim维

        self.attention_dropout = nn.Dropout(dropout_rate)
        self.output_proj = nn.Linear(hidden_dim, hidden_dim)
        self.norm1 = nn.LayerNorm(hidden_dim, eps=1e-6)

        # ffn
        self.up_proj = nn.Linear(hidden_dim, hidden_dim * 4)    # 升维，gpt升维为4倍
        self.down_proj = nn.Linear(hidden_dim * 4, hidden_dim)
        self.act_fn = nn.GELU()
        self.drop_ffn = nn.Dropout(dropout_rate)
        self.norm2 = nn.LayerNorm(hidden_dim, eps=1e-6)

    def multi_head_self_attn(self, x, attn_mask=None):
        # x: [batch_size, seq_len, hidden_dim]
        batch_size, seq_len, _ = x.size()
        query = self.query_proj(x)
        key = self.key_proj(x)
        value = self.value_proj(x)

        # [batch_size, seq_len, hidden_dim] -> [batch_size, seq_len, head_num, head_dim]
        query = query.view(batch_size, seq_len, self.head_num, self.head_dim).transpose(1, 2)  # [batch_size, head_num, seq_len, head_dim]
        key = key.view(batch_size, seq_len, self.head_num, self.head_dim).transpose(1, 2)
        value = value.view(batch_size, seq_len, self.head_num, self.head_dim).transpose(1, 2)
        attn_score = torch.matmul(query, key.transpose(-1, -2)) / math.sqrt(self.head_dim)  # [batch_size, head_num, seq_len, seq_len]
        if attn_mask is not None:
            attn_mask = attn_mask.tril()
            attn_score = attn_score.masked_fill(attn_mask == 0, float('-inf'))
        else:
            attn_mask = torch.ones_like(attn_score).tril()
            attn_score = attn_score.masked_fill(attn_mask == 0, float('-inf'))
        attn_score = torch.softmax(attn_score, dim=-1)
        attn_score = self.attention_dropout(attn_score)
        attn_output = torch.matmul(attn_score, value)
        attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, seq_len, self.hidden_dim)
        attn_output = self.output_proj(attn_output)
        return attn_output

    def ffn(self, x):
        x = self.up_proj(x)
        x = self.act_fn(x)
        x = self.down_proj(x)
        x = self.drop_ffn(x)
        return x

    def forward(self, x, attn_mask=None):
        x = x + self.multi_head_self_attn(x, attn_mask)
        x = self.norm1(x)
        x = x + self.ffn(x)
        x = self.norm2(x)
        return x

class Decoder(nn.Module):
    def __init__(self, vocab_size, hidden_dim, head_num, dropout_rate=0.1):
        super(Decoder, self).__init__()
        self.layer_list = nn.ModuleList([DecoderLayer(hidden_dim, head_num, dropout_rate) for _ in range(6)])
        self.emb = nn.Embedding(vocab_size, hidden_dim)
        self.out = nn.Linear(hidden_dim, vocab_size)

    def forward(self, x, attn_mask=None):
        x = self.emb(x)
        for layer in self.layer_list:
            x = layer(x, attn_mask)
        x = self.out(x)
        return torch.softmax(x, dim=-1)


x = torch.randint(low=0, high=12, size=(3, 4))
mask = (torch.tensor([[1, 1, 1, 1], [1, 1, 0, 0], [1, 1, 1, 0]]).unsqueeze(1).unsqueeze(2).repeat(1, 8, 4, 1))      # [3,4]->[3,8,4,4
print(mask.shape)
decoder = Decoder(vocab_size=12, hidden_dim=64, head_num=8)
output = decoder(x, mask)
print(output.shape)

DecoderLayer解释

参数解析

hidden_dim：隐藏层维度，表示输入数据的特征维度。
head_num：头的数量，表示将输入数据分割成多少个子空间。
head_dim：每个头的维度，计算公式为 hidden_dim // head_num。
query_proj、key_proj 和 value_proj：三个线性变换层，分别用于将输入数据映射到查询（Query）、键（Key）和值（Value）空间。
attention_dropout：Dropout层，用于在注意力权重上应用 Dropout，防止过拟合。
output_proj：一个额外的线性变换层，用于对多头自注意力的输出进行整合。
norm1 和 norm2：两个LayerNorm层，用于在自注意力和前馈网络之后进行归一化。
up_proj 和 down_proj：前馈网络的升维和降维线性变换层。
act_fn：激活函数，这里使用了GELU。
drop_ffn：Dropout层，用于在前馈网络中防止过拟合。

这个方法实现了多头自注意力机制。

def multi_head_self_attn(self, x, attn_mask=None):
    # x: [batch_size, seq_len, hidden_dim]    x 的形状为 [batch_size, seq_len, hidden_dim]
    batch_size, seq_len, _ = x.size()
    query = self.query_proj(x)        # 使用三个线性变换层分别计算查询（query）、键（key）和值（value）。
    key = self.key_proj(x)
    value = self.value_proj(x)

    # [batch_size, seq_len, hidden_dim] -> [batch_size, seq_len, head_num, head_dim]    将查询、键和值的形状从 [batch_size, seq_len, hidden_dim] 转换为 [batch_size, head_num, seq_len, head_dim]。
    query = query.view(batch_size, seq_len, self.head_num, self.head_dim).transpose(1, 2)
    key = key.view(batch_size, seq_len, self.head_num, self.head_dim).transpose(1, 2)
    value = value.view(batch_size, seq_len, self.head_num, self.head_dim).transpose(1, 2)
    attn_score = torch.matmul(query, key.transpose(-1, -2)) / math.sqrt(self.head_dim)    # 计算每个头的查询和键的点积，得到注意力分数矩阵。
    if attn_mask is not None:    # 如果提供了注意力掩码 attn_mask，则将掩码中为 0 的位置对应的注意力分数设置为负无穷（-inf）。
        attn_mask = attn_mask.tril()
        attn_score = attn_score.masked_fill(attn_mask == 0, float('-inf'))
    else:    # 如果没有提供掩码，则生成一个下三角矩阵作为掩码，以防止解码器中的未来信息泄露。
        attn_mask = torch.ones_like(attn_score).tril()
        attn_score = attn_score.masked_fill(attn_mask == 0, float('-inf'))
    attn_score = torch.softmax(attn_score, dim=-1)    # 对注意力分数应用 Softmax 函数，将分数转换为概率值。
    attn_score = self.attention_dropout(attn_score)    # 对注意力权重应用 Dropout。
    attn_output = torch.matmul(attn_score, value)    # 使用注意力权重对值进行加权求和，得到每个头的输出。
    attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, seq_len, self.hidden_dim)    # 将所有头的输出拼接起来，恢复到原始的形状 [batch_size, seq_len, hidden_dim]。
    attn_output = self.output_proj(attn_output)    # 通过一个线性变换层 output_proj 对拼接后的输出进行整合。
    return attn_output

这个方法实现了前馈神经网络（FFN）。

首先通过 up_proj 将输入数据升维到 hidden_dim * 4。
应用激活函数 GELU。
通过 down_proj 将数据降维回 hidden_dim。
应用 Dropout 防止过拟合。

def ffn(self, x):
    x = self.up_proj(x)    # 首先通过 up_proj 将输入数据升维到 hidden_dim * 4。
    x = self.act_fn(x)    # 应用激活函数 GELU。
    x = self.down_proj(x)    # 通过 down_proj 将数据降维回 hidden_dim。
    x = self.drop_ffn(x)
    return x

这是解码器层的前向传播方法。

def forward(self, x, attn_mask=None):
    x = x + self.multi_head_self_attn(x, attn_mask)    # 将输入数据 x 与多头自注意力的输出相加，然后通过 LayerNorm 层 norm1 进行归一化。
    x = self.norm1(x)    # 通过 LayerNorm 层 norm1 进行归一化。
    x = x + self.ffn(x)    # 将归一化后的数据与前馈神经网络的输出相加，然后通过 LayerNorm 层 norm2 进行归一化。
    x = self.norm2(x)
    return x

Decoder解释

vocab_size：词汇表大小，表示输入数据的词汇量。
hidden_dim：隐藏层维度。
head_num：头的数量。
dropout_rate：Dropout比率。
layer_list：一个包含 6 个解码器层的 ModuleList。
emb：一个嵌入层，用于将输入的词汇索引映射到隐藏层维度的向量。
out：一个线性变换层，用于将解码器的输出映射到词汇表大小的维度

forward

def forward(self, x, attn_mask=None):
    x = self.emb(x)    # 使用嵌入层 emb 将输入的词汇索引映射到隐藏层维度的向量。
    for layer in self.layer_list:    # 依次将数据通过 6 个解码器层。
        x = layer(x, attn_mask)
    x = self.out(x)    # 最后，通过线性变换层 out 将解码器的输出映射到词汇表大小的维度。
    return torch.softmax(x, dim=-1)    # 应用 Softmax 函数，将输出转换为概率分布。

总结

总的来说，自注意力机制是一种强大的神经网络架构组件，用于动态地衡量输入序列中不同位置之间的关联程度。它通过计算查询（Query）、键（Key）和值（Value）之间的点积，生成注意力权重，再利用这些权重对值进行加权求和，从而实现对输入数据的全局信息捕捉。这种机制允许模型在处理每个元素时，同时考虑整个序列的信息，有效解决了传统序列模型难以捕捉长距离依赖的问题。自注意力机制的核心优势在于其并行计算能力和对全局信息的高效利用，使其在自然语言处理和计算机视觉等领域得到了广泛应用。

如果你喜欢我的内容，别忘了点赞、关注和收藏哦！你的每一个支持都是我不断进步的动力，也让我更有信心继续创作更多有价值的内容。感谢你的陪伴，让我们一起在知识的海洋里探索更多！❤️✨

查看全文

http://www.kler.cn/a/529195.html