一文理解透Transformer

你好，我是郭震

一、引言

"Attention Is All You Need"是一篇于2017年发表的开创性论文，首次介绍了Transformer模型。

这篇论文彻底改变了自然语言处理（NLP）领域的研究方向，为后续的众多NLP模型和应用奠定了基础。我们熟知的ChatGPT也是基于今天介绍的Transformer.

二、5个核心设计

Transformer模型的核心设计理念可以概括为以下几点：

1. 自注意力（Self-Attention）机制

• 核心概念：Transformer模型的基础是自注意力机制，它允许模型在处理序列（如文本）时，对序列中的每个元素计算其与序列中其他元素的关联度。这种机制使得模型能够捕捉到序列内长距离依赖关系。
• 优势：相比于之前的RNN和LSTM，自注意力机制能够在并行处理时有效地处理长距离依赖问题，显著提高了处理速度和效率。

2. 多头注意力（Multi-Head Attention）

• 设计：在自注意力的基础上，Transformer引入了多头注意力机制，通过将注意力机制“拆分”成多个头并行运行，模型可以从不同的子空间学习信息。
• 目的：这种设计使模型能够更好地理解语言的多种复杂关系，比如同义词和反义词关系、语法和语义关系等。

3. 位置编码（Positional Encoding）

• 问题：由于Transformer完全基于注意力机制，缺乏序列的位置信息。
• 解决方案：通过向输入序列的每个元素添加位置编码，模型能够利用这些信息来了解单词在句子中的位置关系。位置编码是与词嵌入相加的，以保留位置信息。

4. 编码器-解码器架构

• 架构：Transformer模型包含编码器和解码器两部分。编码器用于处理输入序列，解码器则基于编码器的输出和之前的输出生成目标序列。
• 特点：每个编码器和解码器层都包含多头注意力机制和前馈神经网络，通过残差连接和层归一化来优化训练过程。

5. 可扩展性和效率

• 并行处理：与RNN和LSTM等序列模型相比，Transformer的自注意力机制允许对整个序列进行并行处理，显著提高了训练和推理的速度。
• 适用范围：Transformer模型不仅适用于NLP任务，还被扩展到其他领域，如计算机视觉、音频处理等。

三、3个Q K V 向量

接下来，理解自注意力三个核心向量 Q  K  V：

Query（查询）

• Query代表当前单词或位置，是模型试图更好理解或对其编码时的焦点。在自注意力机制中，每个单词都会生成一个query向量，用于与其他单词的key向量进行匹配。

Key（键）

• Key与序列中的每个单词或位置相关联。它用于和query进行匹配，以确定每个单词对当前单词的重要性或"注意力"。基本上，key向量帮助模型了解它应该"关注"序列中的哪些部分。

Value（值）

• Value也与序列中的每个单词或位置相关联。一旦根据query和key的匹配计算出注意力分数，这些分数将用来加权对应的value向量。这意味着，每个单词的value向量被赋予的权重取决于其相对于当前焦点单词的重要性。

工作原理

• 在处理一个单词（或查询点）时，Transformer模型使用该单词的query向量去和其他所有单词的key向量进行点积操作，以此来计算一个注意力分数。这个分数决定了每个单词的value向量对当前单词的编码有多大影响。
• 然后，这些注意力分数会被标准化（通过softmax），并用来加权value向量。通过对所有加权的value向量求和，模型能够为当前单词生成一个加权的表示，这个表示考虑了整个序列中的上下文信息。
• 最终，这个过程为模型提供了一种动态的、基于内容的方法来决定在处理序列的每个部分时应该"关注"哪些信息。

针对句子“The cat sat on the mat”中的“sat”进行计算。为了简化，我们假设经过某种嵌入方法后，每个词的嵌入向量为：

• "The" -> [1, 0]
• "cat" -> [0, 1]
• "sat" -> [1, 1]
• "on" -> [0, 0]
• "the" -> [1, 0]
• "mat" -> [1, -1]

注意：为了演示目的，我们这里假设Query(Q)、Key(K)和Value(V)向量等同于词嵌入本身，而在实际的Transformer模型中，Q、K、V是通过词嵌入乘以不同的权重矩阵得到的。

步骤 1: 计算“sat”与所有单词的Key向量的点积得分

• 得分("sat", "The") = dot([1, 1], [1, 0]) = 1
• 得分("sat", "cat") = dot([1, 1], [0, 1]) = 1
• 得分("sat", "sat") = dot([1, 1], [1, 1]) = 2
• 得分("sat", "on") = dot([1, 1], [0, 0]) = 0
• 得分("sat", "the") = dot([1, 1], [1, 0]) = 1
• 得分("sat", "mat") = dot([1, 1], [1, -1]) = 0

步骤 2: 对得分应用Softmax进行归一化

假设得分经过Softmax函数归一化后（为简化计算，这里不展示Softmax的计算过程），得到的权重（假定值）为：

• 对于"The" -> 0.11
• 对于"cat" -> 0.11
• 对于"sat" -> 0.33
• 对于"on" -> 0.11
• 对于"the" -> 0.11
• 对于"mat" -> 0.11

步骤 3: 计算加权的Value向量和

根据上述得到的权重，我们现在计算加权的Value向量和（在本例中，Value向量和Key向量相同）：

• 加权和 = 0.11 * [1, 0] (The) + 0.11 * [0, 1] (cat) + 0.33 * [1, 1] (sat) + 0.11 * [0, 0] (on) + 0.11 * [1, 0] (the) + 0.11 * [1, -1] (mat)

进行计算：

• 加权和 = [0.11, 0] + [0, 0.11] + [0.33, 0.33] + [0, 0] + [0.11, 0] + [0.11, -0.11]
• 加权和 = [0.66, 0.33]

结果分析

加权和向量[0.66, 0.33]代表了“sat”这个词在考虑其它词的“关注”后的新表示。这个向量捕捉了句子中与“sat”相关性最高的信息。在这个简化的示例中，“sat”本身获得了最高的权重，这是有意义的，因为在自注意力机制中，当前处理的词往往对自身的表示贡献最大。

请注意，这个示例非常简化，实际上在Transformer模型中，词嵌入的维度会更大（例如，512维），并且Q、K、V向量是通过词嵌入与不同的权重矩阵相乘得到的。此外，还会应用多头注意力机制，进一步增强模型的能力。

在Transformer模型中，经过自注意力机制计算得到的加权和向量，如我们示例中的[0.66, 0.33]，会作为下一层或下一个处理步骤的输入。具体来说，这个向量会经过以下几个步骤：

1. 残差连接：首先，这个加权和向量会与原始输入向量（在我们的例子中是"sat"的嵌入[1, 1]）相加。这种操作称为残差连接（Residual Connection），它有助于防止深层网络中的梯度消失问题。
2. 层归一化：残差连接的结果随后会通过一个层归一化（Layer Normalization）步骤，以稳定训练过程。
3. 前馈网络：接着，这个向量会被送入一个前馈网络（Feed-forward Network），该网络对每个位置应用相同的操作，但是它是独立于其他位置的。这个前馈网络通常包含两个线性变换和一个ReLU激活函数。
4. 输出：前馈网络的输出可以被送入下一层的自注意力机制（如果有的话），作为下一层的输入。在Transformer模型中，这个过程会重复多次，每一层都会根据前一层的输出来计算新的加权和向量。
5. 最终输出：在最后一层之后，可能还会有额外的操作，如更多的层归一化、线性层等，最终产生模型的最终输出。在序列到序列的任务中，如机器翻译，这个输出会被送到解码器部分或直接用于生成预测结果。

因此，在下一次迭代（即下一层处理）时，加权和向量[0.66, 0.33]会经过残差连接、层归一化和前馈网络等一系列变换，然后可能成为下一层自注意力机制的输入。这是Transformer架构的核心设计之一，通过这种方式，模型能够捕获和整合序列中的信息，并在深层次上理解和处理文本。

四、从零实现一个Transformer

在PyTorch中实现注意力机制可以有多种方式，这里提供一个基本的自注意力（self-attention）实现示例。自注意力是Transformer网络中使用的一种注意力形式，它允许模型在序列的不同位置间加权聚合信息。

以下是一个简单的自注意力类的实现：

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

class SelfAttention(nn.Module):
    def __init__(self, embed_size, heads):
        super(SelfAttention, self).__init__()
        self.embed_size = embed_size
        self.heads = heads
        self.head_dim = embed_size // heads

        assert (
            self.head_dim * heads == embed_size
        ), "Embedding size needs to be divisible by heads"

        self.values = nn.Linear(self.head_dim, self.head_dim, bias=False)
        self.keys = nn.Linear(self.head_dim, self.head_dim, bias=False)
        self.queries = nn.Linear(self.head_dim, self.head_dim, bias=False)
        self.fc_out = nn.Linear(heads * self.head_dim, embed_size)

    def forward(self, values, keys, query, mask):
        N = query.shape[0]
        value_len, key_len, query_len = values.shape[1], keys.shape[1], query.shape[1]

        # Split the embedding into self.heads different pieces
        values = values.reshape(N, value_len, self.heads, self.head_dim)
        keys = keys.reshape(N, key_len, self.heads, self.head_dim)
        queries = query.reshape(N, query_len, self.heads, self.head_dim)

        values = self.values(values)
        keys = self.keys(keys)
        queries = self.queries(queries)

        # Einsum does matrix multiplication for query*keys for each training example
        # with each head
        attention = torch.einsum("nqhd,nkhd->nhqk", [queries, keys])

        if mask is not None:
            attention = attention.masked_fill(mask == 0, float("-1e20"))

        attention = F.softmax(attention / (self.embed_size ** (1/2)), dim=3)

        out = torch.einsum("nhqk,nvhd->nqhd", [attention, values]).reshape(
            N, query_len, self.heads * self.head_dim
        )

        out = self.fc_out(out)
        return out

这段代码实现了一个简单的多头自注意力机制，它包括以下步骤：

1. 初始化线性层来转换输入的值、键和查询。
2. 在前向传播中，将输入的值、键和查询分别通过对应的线性层。
3. 使用einsum进行矩阵乘法，以计算查询和键之间的注意力分数。
4. 可选地，应用一个掩码（mask）来避免在注意力分数上关注某些特定位置。
5. 应用softmax函数来获取注意力权重。
6. 用einsum将注意力权重应用于值，获得加权的值。
7. 最后，将结果通过另一个线性层进行可能的尺寸调整。
   上面代码SelfAttention类 实现下面过程：
   
   

要使用上面的自注意力机制，你需要将其整合到你的神经网络模型中。以下是一个如何在一个简单的序列处理任务中使用自注意力模块的示例：

import torch
import torch.nn as nn

# 假设我们有一个特定大小的嵌入层和自注意力层
embed_size = 256
heads = 8
sequence_length = 100  # 输入序列的长度
batch_size = 32
vocab_size = 10000  # 假设的词汇大小

class TransformerBlock(nn.Module):
    def __init__(self, embed_size, heads):
        super(TransformerBlock, self).__init__()
        self.attention = SelfAttention(embed_size, heads)
        self.norm1 = nn.LayerNorm(embed_size)
        self.norm2 = nn.LayerNorm(embed_size)

        self.feed_forward = nn.Sequential(
            nn.Linear(embed_size, 2 * embed_size),
            nn.ReLU(),
            nn.Linear(2 * embed_size, embed_size)
        )

    def forward(self, value, key, query, mask):
        attention = self.attention(value, key, query, mask)
        x = self.norm1(attention + query)
        forward = self.feed_forward(x)
        out = self.norm2(forward + x)
        return out

class Transformer(nn.Module):
    def __init__(self, embed_size, heads, vocab_size, sequence_length):
        super(Transformer, self).__init__()
        self.word_embedding = nn.Embedding(vocab_size, embed_size)
        self.position_embedding = nn.Embedding(sequence_length, embed_size)
        self.layers = nn.ModuleList([TransformerBlock(embed_size, heads) for _ in range(6)])
        self.fc_out = nn.Linear(embed_size, vocab_size)

    def forward(self, x, mask):
        N, seq_length = x.shape
        positions = torch.arange(0, seq_length).expand(N, seq_length).to(x.device)
        out = self.word_embedding(x) + self.position_embedding(positions)

        for layer in self.layers:
            out = layer(out, out, out, mask)

        out = self.fc_out(out)

        return out

# 创建一个模型实例
model = Transformer(embed_size, heads, vocab_size, sequence_length)

# 创建一个随机输入序列
input_seq = torch.randint(0, vocab_size, (batch_size, sequence_length))

# 创建一个mask(这个例子中为None)
mask = None

# 前向传播
output = model(input_seq, mask)

TransformerBlock类实现下面图：

这里，Transformer类实现了一个包含6个TransformerBlock层的简单Transformer模型，TransformerBlock包含了自注意力层（SelfAttention中的多头自注意力中 多头此处等于heads 为 8）和前馈神经网络。模型的输入是一个整数序列，这些整数代表词汇表中的索引，然后模型输出一个相同长度的序列，其中的每个元素是对应的词汇表大小的向量，表示概率分布。

mask 参数是可选的，它可以用来掩盖序列中的某些部分，例如在处理变长输入或者防止模型在解码时看到未来的信息时非常有用。

在实际应用中，需要根据具体任务调整模型结构和超参数，比如词嵌入的大小、注意力头的数量、序列的长度等，并且可能需要添加额外的层和功能，比如词嵌入层、位置嵌入层和最终的输出层。此外，还需要准备训练数据，定义损失函数和优化器，并执行训练循环。

五、结果分析