PyTorch——自注意力（self-attention）机制实现（代码详解）

cqu_shuai

113846人浏览 · 2021-03-28 16:24:21

cqu_shuai · 2021-03-28 16:24:21 发布

参考链接

https://www.bilibili.com/video/BV1JE411g7XF?p=54
https://arxiv.org/abs/1706.03762
https://blog.csdn.net/qq_36653505/article/details/83375160

简述自注意力机制（self-attention）

self-attention可以视为一个特征提取层，给定输入特征 $a1,a2,⋅⋅⋅ana^{1},a^{2},\cdot \cdot \cdot a^{n}$ ，经过self-attention layer，融合每个输入特征，得到新的特征 $b1,b2,⋅⋅⋅bnb^{1},b^{2},\cdot \cdot \cdot b^{n}$ 。具体如下：

设输入特征为 $I$ ，分别将其乘以三个矩阵 $W^{q}$ 、 $W^{k}$ 和 $W^{v}$ 得到 $Q$ （query）、 $K$ （key）和 $V$ （value）三个矩阵；接下来使用矩阵 $Q$ 和 $K$ 的乘积得到注意力矩阵 $A$ ，归一化得到 $A^\hat{A}$ ；最后，将归一化后的注意力矩阵 $A^\hat{A}$ 乘上 $V$ ，得到最后的输出特征 $O$ 。
在这里插入图片描述

多头自注意力机制（multi-head self-attention）

上述的self-attention中，每个输入特征 $a^{i}$ 乘上矩阵 $W^{q}$ 、 $W^{k}$ 和 $W^{v}$ 后，分别得到一个向量 $q^{i}$ 、 $k^{i}$ 和 $v^{i}$ ，称为单头自注意力机制。如果将这些向量 $q^{i}$ 、 $k^{i}$ 和 $v^{i}$ 分裂为 $n$ 个就得到 $n$ 头自注意力机制了。公认多头自注意力机制的效果好于单头的，因为前者可以捕获更多维度的信息。示意图如下：
在这里插入图片描述

代码实现

设超参数num_attention_heads为自注意力机制的头数，如此，计算出每个头的维度attention_head_size。

self.num_attention_heads = num_attention_heads
self.attention_head_size = int(hidden_size / num_attention_heads)
self.all_head_size = hidden_size

定义 $W^{q}$ 、 $W^{k}$ 和 $W^{v}$ 三个矩阵。

self.query = nn.Linear(input_size, self.all_head_size)
self.key = nn.Linear(input_size, self.all_head_size)
self.value = nn.Linear(input_size, self.all_head_size)

下面开始逐步计算，需要主要的是计算过程中张量维度的变化。
将输入特征乘以三个矩阵 $W^{q}$ 、 $W^{k}$ 和 $W^{v}$ ，输出的张量此时还没有区分出多个头。维度变化为：input_tensor $(batch,n,input_size)\left ( batch,n,input\_size\right )$ 到mixed_query_layer $(batch,n,all_head_size)\left ( batch,n,all\_head\_size\right )$

mixed_query_layer = self.query(input_tensor)
mixed_key_layer = self.key(input_tensor)
mixed_value_layer = self.value(input_tensor)

切分为num_attention_heads个头，并变换维度。维度变化为：mixed_query_layer $(batch,n,all_head_size)\left ( batch,n,all\_head\_size\right )$ 到query_layer $(batch,num_attention_heads,n,attention_head_size)\left ( batch,num\_attention\_heads,n,attention\_head\_size\right )$

def transpose_for_scores(self, x):
   new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
   x = x.view(*new_x_shape)
   return x.permute(0, 2, 1, 3)

query_layer = self.transpose_for_scores(mixed_query_layer)
key_layer = self.transpose_for_scores(mixed_key_layer)
value_layer = self.transpose_for_scores(mixed_value_layer)

矩阵 $Q$ 和 $K$ 相乘，得到注意力矩阵，并除以向量的维度的开方，防止注意力分数随维度增大而增大。维度变化为：query_layer $(batch,num_attention_heads,n,attention_head_size)\left ( batch,num\_attention\_heads,n,attention\_head\_size\right )$ 到attention_scores $(batch,num_attention_heads,n,n)\left ( batch,num\_attention\_heads,n,n\right )$

attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))

attention_scores = attention_scores / math.sqrt(self.attention_head_size)

注意力矩阵归一化。维度变化为：attention_scores $(batch,num_attention_heads,n,n)\left ( batch,num\_attention\_heads,n,n\right )$ 到attention_probs $(batch,num_attention_heads,n,n)\left ( batch,num\_attention\_heads,n,n\right )$

attention_probs = nn.Softmax(dim=-1)(attention_scores)

将注意力矩阵乘以矩阵 $V$ 。维度变化为：ttention_probs $(batch,num_attention_heads,n,n)\left ( batch,num\_attention\_heads,n,n\right )$ 乘以value_layer $(batch,num_attention_heads,n,attention_head_size)\left ( batch,num\_attention\_heads,n,attention\_head\_size\right )$ 到context_layer $(batch,num_attention_heads,n,attention_head_size)\left ( batch,num\_attention\_heads,n,attention\_head\_size\right )$ 。

context_layer = torch.matmul(attention_probs, value_layer)

变换context_layer维度，为了后面将各头得到的结果拼接。这里的contiguous()是将tensor的内存变成连续的，为后面的view()做准备。维度变化为：context_layer $(batch,num_attention_heads,n,attention_head_size)\left ( batch,num\_attention\_heads,n,attention\_head\_size\right )$ 到context_layer $(batch,n,num_attention_heads,attention_head_size)\left ( batch,n,num\_attention\_heads,attention\_head\_size\right )$

context_layer = context_layer.permute(0, 2, 1, 3).contiguous()

将各头的结果拼接起来。维度变化为：context_layer $(batch,n,num_attention_heads,attention_head_size)\left ( batch,n,num\_attention\_heads,attention\_head\_size\right )$ 到context_layer $(batch,n,all_head_size)\left ( batch,n,all\_head\_size\right )$

new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
context_layer = context_layer.view(*new_context_layer_shape)

完整代码

class LayerNorm(nn.Module):
    def __init__(self, hidden_size, eps=1e-12):
        """Construct a layernorm module in the TF style (epsilon inside the square root).
        """
        super(LayerNorm, self).__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.bias = nn.Parameter(torch.zeros(hidden_size))
        self.variance_epsilon = eps

    def forward(self, x):
        u = x.mean(-1, keepdim=True)
        s = (x - u).pow(2).mean(-1, keepdim=True)
        x = (x - u) / torch.sqrt(s + self.variance_epsilon)
        return self.weight * x + self.bias
        
class SelfAttention(nn.Module):
    def __init__(self, num_attention_heads, input_size, hidden_size, hidden_dropout_prob):
        super(SelfAttention, self).__init__()
        if hidden_size % num_attention_heads != 0:
            raise ValueError(
                "The hidden size (%d) is not a multiple of the number of attention "
                "heads (%d)" % (hidden_size, num_attention_heads))
        self.num_attention_heads = num_attention_heads
        self.attention_head_size = int(hidden_size / num_attention_heads)
        self.all_head_size = hidden_size

        self.query = nn.Linear(input_size, self.all_head_size)
        self.key = nn.Linear(input_size, self.all_head_size)
        self.value = nn.Linear(input_size, self.all_head_size)

        self.attn_dropout = nn.Dropout(attention_probs_dropout_prob)

        # 做完self-attention 做一个前馈全连接 LayerNorm 输出
        self.dense = nn.Linear(hidden_size, hidden_size)
        self.LayerNorm = LayerNorm(hidden_size, eps=1e-12)
        self.out_dropout = nn.Dropout(hidden_dropout_prob)

    def transpose_for_scores(self, x):
        new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
        x = x.view(*new_x_shape)
        return x.permute(0, 2, 1, 3)

    def forward(self, input_tensor):
        mixed_query_layer = self.query(input_tensor)
        mixed_key_layer = self.key(input_tensor)
        mixed_value_layer = self.value(input_tensor)

        query_layer = self.transpose_for_scores(mixed_query_layer)
        key_layer = self.transpose_for_scores(mixed_key_layer)
        value_layer = self.transpose_for_scores(mixed_value_layer)

        # Take the dot product between "query" and "key" to get the raw attention scores.
        attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))

        attention_scores = attention_scores / math.sqrt(self.attention_head_size)
        # Apply the attention mask is (precomputed for all layers in BertModel forward() function)
        # [batch_size heads seq_len seq_len] scores
        # [batch_size 1 1 seq_len]

        # attention_scores = attention_scores + attention_mask

        # Normalize the attention scores to probabilities.
        attention_probs = nn.Softmax(dim=-1)(attention_scores)
        # This is actually dropping out entire tokens to attend to, which might
        # seem a bit unusual, but is taken from the original Transformer paper.
        # Fixme
        attention_probs = self.attn_dropout(attention_probs)
        context_layer = torch.matmul(attention_probs, value_layer)
        context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
        new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
        context_layer = context_layer.view(*new_context_layer_shape)
        hidden_states = self.dense(context_layer)
        hidden_states = self.out_dropout(hidden_states)
        hidden_states = self.LayerNorm(hidden_states + input_tensor)

        return hidden_states

GitCode 开源社区

旨在为数千万中国开发者提供一个无缝且高效的云端环境，以支持学习、使用和贡献开源项目。

更多推荐

SCI论文快速写作指南：沁言学术实用方法

在SCI论文写作领域，尤其针对CSDN社区内的计算机科学、AI和算法优化研究者，高效方法的应用已成为提升投稿成功率的关键。AI工具的整合能显著优化从检索到润色的流程。本文基于公开调研数据和用户案例，剖析SCI写作难点，并以沁言学术为例，探讨实用方法与技术机制。重点在于算法深度与效率insights，旨在为CSDN用户提供可操作的指导，帮助科研人员加速SCI论文产出。

GitCode 开源社区

仓颉共学创作营正式启动！邀您共探新兴技术，解锁仓颉创作红利

GitCode 开源社区

跨学科研究者如何高效整理文献

AI阅读与摘要提取是高效整理的关键，沁言学术集成OCR和Transformer模型，支持深度解析。阅读模块自动生成结构化摘要，捕捉核心概念，如从一篇“生物信息学中的神经网络”论文中提取算法细节与生物关联。摘要提取利用注意力机制，优先处理跨学科元素：例如，分析中英混用文献时，AI创建可视化摘要云和关联图，突出“深度学习模型”与“基因序列分析”的融合。这降低阅读门槛，特别在新兴交叉如“AI伦理与医学”