Transformer中的掩码张量和注意力机制

好好学习Py

已于 2024-07-10 16:28:55 修改

阅读量1.3k

点赞数 11

文章标签： transformer 深度学习人工智能

于 2024-07-09 22:44:33 首次发布

本文链接：https://blog.csdn.net/weixin_74254879/article/details/140302374

版权

掩码张量

掩码张量在Transformer中的作用非常重要。它通常是一个与待处理张量具有相同尺寸的张量，其中的元素通常是0或1。这些元素决定了在应用注意力机制时哪些位置的信息应该被遮蔽或被替换。

在Transformer中，特别是在解码器部分，我们需要确保在生成输出时不会利用未来信息，因为解码器是逐步生成输出序列的。如果不使用掩码张量，模型可能会在计算注意力时利用未来的信息，导致结果不符合预期。

因此，掩码张量通过将某些位置的值设置为特定值（通常是一个很小的数或负无穷），有效地遮掩了那些不应该被当前计算利用的信息，确保了模型在训练和推理时的准确性和合理性。

import torch


def subsequent_mask(size):
    """生成向后遮掩的掩码张量, 参数 size 是掩码张量最后两个维度的大小。"""
    # 生成一个上三角矩阵，对角线及以上为0代表未遮蔽，下三角为1代表遮蔽
    mask_shape = (1, size, size)
    mask = 1 - torch.triu(torch.ones(mask_shape), diagonal=1).byte()
    return mask


if __name__ == "__main__":
    mask_size = 4
    subsequent_mask_tensor = subsequent_mask(mask_size)
    print("subsequent_mask_tensor:\n", subsequent_mask_tensor)

    import matplotlib.pyplot as plt
    plt.figure(figsize=(5, 5))
    plt.imshow(subsequent_mask(mask_size)[0])
    plt.show()

黄色部分被遮掩，紫色部分未被遮掩，每个位置只能注意到之前的位置，而不能看到之后的位置信息。

subsequent_mask_tensor:
tensor([[[1, 0, 0, 0],
[1, 1, 0, 0],
[1, 1, 1, 0],
[1, 1, 1, 1]]], dtype=torch.uint8)

注意力机制

当我们观察事物时，我们之所以能够迅速作出判断（即使可能是错误的判断），是因为大脑能够快速将注意力集中在事物最具有区分性的部分，而不是完整地逐个细节地观察。这种能力使得我们能够基于局部信息迅速做出整体判断，这也正是注意力机制的理论基础。

注意力机制通过接收三个特定输入：query，key，value，并应用特定的计算公式，生成表示query 在key和value作用下的注意力表示结果。

这里使用的计算规则： $\boldsymbol{\mathbf{}Attention(Q,K,V) = Softmax(\frac{Q\cdot K^{T}}{d_{model}}\mathbf{})\cdot V}$

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


def attention(query, key, value, mask=None, dropout=None):
    """
    注意力机制的实现.

    参数:
    - query: 查询张量.
    - key: 键张量.
    - value: 值张量.
    - mask: 掩码张量.
    - dropout: nn.Dropout层的实例化对象, 默认为None.
    """
    # 获取query张量的最后一维大小, 一般是词嵌入维度
    embedding_dim = query.size(-1)

    # 计算注意力得分张量scores
    scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(embedding_dim)

    # 如果提供了掩码张量, 对scores进行掩码处理
    if mask is not None:
        scores = scores.masked_fill(mask == 1, -1e9)

    # 对scores的最后一维进行softmax操作, 得到注意力张量attention_weights
    attention_weights = F.softmax(scores, dim=-1)

    # 如果提供了dropout, 对attention_weights进行dropout处理
    if dropout is not None:
        attention_weights = dropout(attention_weights)

    # 计算最终的注意力表示并返回
    return torch.matmul(attention_weights, value), attention_weights


if __name__ == "__main__":
    # 设置参数
    test_embedding_dim = 512
    test_vocab_size = 10000
    test_max_len = 100

    # 文本嵌入层
    text_embeddings = TextEmbeddings(test_vocab_size, test_embedding_dim)
    test_input_tensor = torch.LongTensor([[1, 2, 3, 4], [4, 3, 2, 1]])
    text_embeddings_output = text_embeddings(test_input_tensor)

    # 添加位置编码
    positional_encoding = PositionalEncoding(test_embedding_dim, dropout=0.1,
                                             max_sequence_length=test_max_len)
    positional_encoded_output = positional_encoding(text_embeddings_output)

    # 未设置mask
    query_tensor = key_tensor = value_tensor = positional_encoded_output
    attention_output, attention_weights_output = attention(query_tensor, key_tensor, value_tensor)
    print("未设置mask的attention_output:\n", attention_output)
    print("未设置mask的attention_weights:\n", attention_weights_output)

    # 设置mask
    query_tensor = key_tensor = value_tensor = positional_encoded_output
    test_mask = torch.oness(2, 4, 4)    # 令mask为一个2x4x4的全1张量
    attention_output, attention_weights_output = attention(query_tensor, key_tensor, value_tensor,
                                                           mask=test_mask)
    print("设置mask的attention_output:\n", attention_output)
    print("设置mask的attention_weights:\n", attention_weights_output)

未设置mask的注意力每个位置只关注自身，没有其他位置的权重信息。设置mask的注意力权重是一个均匀分布，每个query向量平均地关注了所有key向量。

未设置mask的attention_output:
tensor([[[-59.3522, 42.8812, 11.9322, ..., -15.6505, -11.0933, 25.4808],
[ 38.4859, 10.4008, 25.0720, ..., -12.1591, -22.3926, -1.2321],
[-27.6069, 0.0000, 44.6558, ..., 54.9970, -40.9909, -5.4670],
[ -7.0379, 21.8327, -50.7755, ..., -38.0019, 0.0000, -6.7024]],

[[ -7.1947, 24.0438, -51.0394, ..., -38.0019, 37.1588, -6.7024],
[-27.6823, -8.4545, 44.5322, ..., 54.9970, -40.9911, -5.4670],
[ 38.5613, 9.3381, 25.1956, ..., -12.1591, 0.0000, -1.2321],
[-59.1954, 0.0000, 12.1960, ..., -15.6505, -11.0926, 25.4808]]],
grad_fn=<UnsafeViewBackward0>)
未设置mask的attention_weights:
tensor([[[1., 0., 0., 0.],
[0., 1., 0., 0.],
[0., 0., 1., 0.],
[0., 0., 0., 1.]],

[[1., 0., 0., 0.],
[0., 1., 0., 0.],
[0., 0., 1., 0.],
[0., 0., 0., 1.]]], grad_fn=<SoftmaxBackward0>)
设置mask的attention_output:
tensor([[[-13.8778, 18.7787, 7.7211, ..., -2.7036, -18.6192, 3.0198],
[-13.8778, 18.7787, 7.7211, ..., -2.7036, -18.6192, 3.0198],
[-13.8778, 18.7787, 7.7211, ..., -2.7036, -18.6192, 3.0198],
[-13.8778, 18.7787, 7.7211, ..., -2.7036, -18.6192, 3.0198]],

[[-13.8778, 6.2319, 7.7211, ..., -2.7036, -3.7312, 3.0198],
[-13.8778, 6.2319, 7.7211, ..., -2.7036, -3.7312, 3.0198],
[-13.8778, 6.2319, 7.7211, ..., -2.7036, -3.7312, 3.0198],
[-13.8778, 6.2319, 7.7211, ..., -2.7036, -3.7312, 3.0198]]],
grad_fn=<UnsafeViewBackward0>)
设置mask的attention_weights:
tensor([[[0.2500, 0.2500, 0.2500, 0.2500],
[0.2500, 0.2500, 0.2500, 0.2500],
[0.2500, 0.2500, 0.2500, 0.2500],
[0.2500, 0.2500, 0.2500, 0.2500]],

[[0.2500, 0.2500, 0.2500, 0.2500],
[0.2500, 0.2500, 0.2500, 0.2500],
[0.2500, 0.2500, 0.2500, 0.2500],
[0.2500, 0.2500, 0.2500, 0.2500]]], grad_fn=<SoftmaxBackward0>)

多头注意力机制

多头注意力机制通过线性变换对查询query、key、value进行映射，每组变换矩阵都是方阵且不改变张量尺寸。每个头在注意力计算中独立地关注输入张量的不同部分，通过分割最后一维的词嵌入向量来实现。这种设计使得每个注意力机制可以专注于优化词汇不同特征的表示，从而平衡单一注意力机制可能产生的偏差，增强词义的多元表达能力。实验证明，这种结构能有效提升模型的性能。

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


# 定义克隆函数，用于生成相同网络层的克隆
def clones(module, N):
    """生成多个相同网络层的克隆函数"""
    return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])


# 实现多头注意力机制的类
class MultiHeadedAttention(nn.Module):
    def __init__(self, num_heads, embedding_dim, dropout=0.1):
        """初始化多头注意力机制

        参数:
        - num_heads: 头数
        - embedding_dim: 词嵌入的维度
        - dropout: dropout 丢弃率，默认为 0.1
        """
        super(MultiHeadedAttention, self).__init__()

        # 检查 embedding_dim 是否可以被 num_heads 整除
        assert embedding_dim % num_heads == 0

        # 计算每个头得到的分割词向量维度
        self.head_dim = embedding_dim // num_heads

        # 头数
        self.num_heads = num_heads

        # 初始化线性层列表，包括 Q、K、V 和最后的线性变换层
        self.linears = clones(nn.Linear(embedding_dim, embedding_dim), 4)

        # 注意力权重
        self.attn_weight = None

        # Dropout
        self.dropout = nn.Dropout(p=dropout)

    def forward(self, query, key, value, mask=None):
        """前向逻辑函数

        参数:
        - query: 查询张量
        - key: 键张量
        - value: 值张量
        - mask: 掩码张量，默认为 None
        """
        # 如果存在掩码张量 mask，则进行维度扩展
        if mask is not None:
            mask = mask.unsqueeze(0)

        # 获取 batch_size
        batch_size = query.size(0)

        # 对输入的 Q、K、V 分别进行线性变换，并分割成多个头
        query, key, value = \
            [model(x).view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)
             for model, x in zip(self.linears, (query, key, value))]

        # 进行注意力计算，调用 attention 函数
        attn, self.attn_weight = attention(query, key, value, mask=mask, dropout=self.dropout)

        # 将多头注意力计算结果重新整合并线性变换得到最终输出
        attn = attn.transpose(1, 2).contiguous().view(batch_size, -1, self.num_heads * self.head_dim)

        return self.linears[-1](attn)


if __name__ == "__main__":
    # 设置参数
    test_embedding_dim = 512
    test_vocab_size = 10000
    test_max_len = 100
    test_heads = 8
    test_dropout = 0.2

    # 文本嵌入层
    text_embeddings = TextEmbeddings(test_vocab_size, test_embedding_dim)
    test_input_tensor = torch.LongTensor([[1, 2, 3, 4], [4, 3, 2, 1]])
    text_embeddings_output = text_embeddings(test_input_tensor)

    # 添加位置编码
    positional_encoding = PositionalEncoding(test_embedding_dim, dropout=0.1,
                                             max_sequence_length=test_max_len)
    positional_encoded_output = positional_encoding(text_embeddings_output)

    # 多头自注意力机制计算
    query_tensor = key_tensor = value_tensor = positional_encoded_output
    test_mask = torch.zeros(8, 4, 4)  # 令mask为一个8x4x4的零张量
    self_mha = MultiHeadedAttention(test_heads, test_embedding_dim, test_dropout)
    self_mha_result = self_mha(query_tensor, key_tensor, value_tensor, test_mask)
    print("Multi-Head Attention Result:\n", self_mha_result)
    print("Shape of Multi-Head Attention Result:", self_mha_result.shape)

输出的具体数值代表了经过多头注意力就算后的每个位置的表示。这些值可以包含有关每个位置与其他位置的关联信息，以及输入序列中不同部分的重要性。

Multi-Head Attention Result:
tensor([[[ 3.7060, 2.9021, -7.1416, ..., -7.1680, -3.0814, 9.1765],
[ -5.5620, -9.7481, -0.9496, ..., -1.4335, -6.4775, 15.4199],
[-10.6869, 1.0085, -14.1587, ..., 0.5768, 4.0398, -8.2267],
[ 17.4881, -14.9494, -7.7935, ..., 15.6512, 9.8766, 1.4503]],

[[ 17.9870, -12.7330, -10.7763, ..., 18.9070, 9.4836, -3.9758],
[-16.8575, 4.4356, -5.4201, ..., 7.8603, 11.2810, -12.5162],
[ -5.3459, -5.8356, 3.7956, ..., 5.0795, 1.3836, 14.9227],
[ -0.3273, -7.7487, -1.8062, ..., -13.1506, -9.4419, -0.4178]]],
grad_fn=<ViewBackward0>)
Shape of Multi-Head Attention Result: torch.Size([2, 4, 512])