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

什么是注意力机制？

article 2024/10/27 5:31:09

一、注意力机制

注意力机制（Attention Mechanism）是一种在深度学习模型中模拟人类注意力的机制，它允许模型在处理信息时能够聚焦于当前任务最相关的部分。注意力机制的核心思想是计算输入数据中每个部分与当前任务的相关性，然后根据这些相关性来分配不同的权重。这些权重通常通过一个softmax函数进行归一化，使得所有权重的和为1，从而可以被解释为概率分布。公式见自注意力机制（没错，从公示表达上两者一模一样）。注意力机制特点如下：

查询和键来自不同来源：在传统的注意力机制中，查询（Query）和键（Key）通常来自不同的数据源。比如在机器翻译任务中，查询可能来自解码器的隐藏状态，而键和值（Value）来自编码器的输出。
依赖于外部上下文：注意力机制通常依赖于外部的上下文信息来计算注意力权重。这意味着模型在生成输出时，会考虑输入序列中的信息，但这些信息是通过外部的键和值来提供的。
应用场景：注意力机制广泛应用于各种序列到序列的任务中，如机器翻译、语音识别、图像标注等，其中模型需要同时考虑输入序列和输出序列的信息。

二、自注意力机制

自注意力（Self-Attention）是注意力机制的一种，它在单个序列的不同位置之间建立联系，以计算同一序列的表示。这种机制在自然语言处理（NLP）和计算机视觉等领域有广泛的应用，特别是在Transformer模型中发挥了重要作用。自注意力机制的核心思想是计算序列中每个元素对其他所有元素的注意力权重，然后根据这些权重对序列中的元素进行加权求和，得到最终的表示。

自注意力机制涉及三个主要组件：查询（Query）、键（Key）和值（Value）。这三个组件都是从输入数据中派生出来的，通常是通过与权重矩阵的矩阵乘法来实现的。

查询（Query）：代表当前元素想要关注的信息。
键（Key）：代表其他元素可以提供的信息。
值（Value）：代表关注后能够得到的信息。

自注意力机制通过计算查询和所有键之间的相似度得分，然后使用softmax函数对这些得分进行归一化，得到注意力权重。这些权重随后用于对值进行加权求和，得到最终的输出。

$eq?Attention%28Q%2CK%2CV%29%20%3D%20softmax%28%5Cfrac%7BQK%5E%7BT%7D%7D%7B%5Csqrt%7Bd_%7Bk%7D%7D%7D%29V$

其中，Q是查询矩阵。K是键矩阵。V是值矩阵。 $eq?d_%7Bk%7D$ 是键向量的维度，分母中的 $eq?%5Csqrt%7Bd_%7Bk%7D%7D$ 用于缩放点积，防止梯度消失或爆炸。设自注意力计算的输入向量是X，有：

$eq?Q%20%3D%20XW%5E%7BQ%7D$

$eq?K%20%3D%20XW%5E%7BK%7D$

$eq?V%20%3D%20XW%5E%7BV%7D$

其中，W是权重矩阵，通常是随机初始化的，以打破对称性并允许不同的层和节点从不同的起点开始学习。常见的初始化方法包括Glorot初始化（Xavier初始化）、He初始化和正态分布初始化等，这些方法旨在保持梯度流的稳定性，避免梯度消失或爆炸问题。自注意力机制特点如下：

查询、键和值来自同一来源：在自注意力机制中，查询、键和值都来自同一个数据源，通常是同一个输入序列。这意味着模型在处理序列的每个元素时，会考虑序列内部的依赖关系。
内部依赖关系：自注意力机制允许模型在序列内部建立依赖关系，这使得模型能够捕捉到序列中长距离的依赖关系。这对于理解句子的语法结构或图像中不同部分之间的关系特别有用。
应用场景：自注意力机制在处理需要内部依赖关系的任务中特别有效，如文本分类、句子嵌入、图像分割等。它也是Transformer模型的核心组成部分，后者在NLP领域取得了显著的成功。

三、多头注意力机制

多头注意力机制（Multi-Head Attention）是自注意力的一个扩展，它将输入序列投影到多个不同的子空间中，并在每个子空间中独立地执行自注意力计算。这样做的目的是让模型能够从不同的表示子空间中捕捉信息，增强模型的表达能力。多头注意力的计算过程可以表示为：

$eq?MultiHead%28Q%2CK%2CV%29%20%3D%20concat%28head_%7B1%7D%2Chead_%7B2%7D%2C...%2Chead_%7Bn%7D%29W%5E%7BO%7D$

其中，每个head的计算为：

$eq?head_%7Bi%7D%20%3D%20Attention%28Q_%7Bi%7D%2CK_%7Bi%7D%2CV_%7Bi%7D%29$

由于多头注意力机制是对查询Q、键K和值V通过不同的线性变换（矩阵乘法），将它们分别投影到n个不同的子空间（即不同的头）。则Q、K、V的计算有：

$eq?Q_%7Bi%7D%20%3D%20XW_%7Bi%7D%5E%7BQ%7D$

$eq?K_%7Bi%7D%20%3D%20XW_%7Bi%7D%5E%7BK%7D$

$eq?V_%7Bi%7D%20%3D%20XW_%7Bi%7D%5E%7BV%7D$

其中，W即为当前头的权重矩阵， $eq?W_%7Bi%7D$ 为当前头的投影矩阵。将拼接后的结果通过另一个线性变换 $eq?W%5E%7BO%7D$ 进行投影，得到最终的多头注意力输出。

四、注意力机制示例

1、自注意力机制

以下是一个简单的自注意力机制的计算，可以发现其实这玩意儿并没有多神秘高深。但实际上的应用却又非常有效。

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

class SelfAttention(nn.Module):
    def __init__(self, embed_dim):
        super(SelfAttention, self).__init__()
        self.embed_dim = embed_dim

        self.query_linear = nn.Linear(embed_dim, embed_dim)
        self.key_linear = nn.Linear(embed_dim, embed_dim)
        self.value_linear = nn.Linear(embed_dim, embed_dim)

        self.final_linear = nn.Linear(embed_dim, embed_dim)

    def forward(self, x):
        # x: (batch_size, seq_len, embed_dim)
        batch_size, seq_len, embed_dim = x.size()

        # Calculate queries, keys, and values
        queries = self.query_linear(x)  # (batch_size, seq_len, embed_dim)
        keys = self.key_linear(x)      # (batch_size, seq_len, embed_dim)
        values = self.value_linear(x)    # (batch_size, seq_len, embed_dim)

        # Compute attention scores
        attention_scores = torch.matmul(queries, keys.transpose(-2, -1)) / math.sqrt(self.embed_dim)
        attention_weights = F.softmax(attention_scores, dim=-1)

        # Compute attention output
        attention_output = torch.matmul(attention_weights, values)

        # Apply final linear layer
        output = self.final_linear(attention_output)

        return output

# Example usage
embed_dim = 64  # Embedding dimension
seq_len = 10    # Sequence length
batch_size = 1  # Batch size

# Dummy input tensor with shape (batch_size, seq_len, embed_dim)
x = torch.randn(batch_size, seq_len, embed_dim)

self_attention = SelfAttention(embed_dim)
attention_output = self_attention(x)

print(attention_output.shape)  # Should be (batch_size, seq_len, embed_dim)

2、多头注意力机制

公式上看比较费解，但实际上就是向量的拆分和线性变换。

class MultiHeadAttention(nn.Module):
    def __init__(self, embed_dim, num_heads):
        super(MultiHeadAttention, self).__init__()
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads

        assert self.head_dim * num_heads == embed_dim, "embed_dim must be divisible by num_heads"

        self.q_linear = nn.Linear(embed_dim, embed_dim)
        self.k_linear = nn.Linear(embed_dim, embed_dim)
        self.v_linear = nn.Linear(embed_dim, embed_dim)
        self.out_linear = nn.Linear(embed_dim, embed_dim)

    def forward(self, x):
        batch_size = x.size(0)

        # 从单头到多头的转换，当然，实际上还有更复杂的做法，这里是简单的实现
        q = self.q_linear(x).view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)
        k = self.k_linear(x).view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)
        v = self.v_linear(x).view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)

        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)
        attention_weights = F.softmax(scores, dim=-1)

        output = torch.matmul(attention_weights, v).transpose(1, 2).contiguous()
        output = output.view(batch_size, -1, self.embed_dim)
        output = self.out_linear(output)

        return output