跟李沐学AI：Transformer

Transformer架构

(图源:10.7. Transformer — 动手学深度学习 2.0.0 documentation)

基于编码器-解码器架构来处理序列对

与使用注意力的seq2seq不同，Transformer纯基于注意力

多头注意力（Multi-Head Attention）

(图源:10.5. 多头注意力 — 动手学深度学习 2.0.0 documentation)

对于同一QKV，我们希望抽取不同的信息。多头注意力机制相比于标准的注意力机制，允许模型捕捉一个序列的不同方面并对不同方面生成不同权重。

假设输入序列为： I love palying football.

假设embedding_size=3并且使用2个注意力头，序列可表示为一个(4, 3)的矩阵(sequence_length, embed_size)。

[ I ]         [ love ]       [ playing ]    [ football ]
[[0.1, 0.2, 0.3], [0.4, 0.5, 0.6], [0.7, 0.8, 0.9], [1.0, 1.1, 1.2]]  # Shape: (4, 3)

首先我们需要把embedding投影为Q，K，V。

权重矩阵的形状由输入特征数和头数共同决定，不同头的权重矩阵的形状是相同的。

$W_Q^{i} \epsilon R^{d_{feature}\times d_{head}}$

$W_K^{i} \epsilon R^{d_{feature}\times d_{head}}$

$W_V^{i} \epsilon R^{d_{feature}\times d_{head}}$

$d_{feature}$ 是输入的embedding_size，在本例中为3。

$d_{head}=d_{feature}/h$ 。其中h是注意力头的数量。

故本例中权重矩阵的形状为(3, 1)，其中1为向下取整得到。

对于头1：

Q1 = Wq1 * Embedding
K1 = Wk1 * Embedding
V1 = Wv1 * Embedding

假设头1的权重为：

$W_Q^{head1} =[[0.1],[0.2],[0.3]]$

$W_K^{head1} =[[0.4],[0.5],[0.6]]$

$W_V^{head1} =[[0.7],[0.8],[0.9]]$

那么我们可以列出输入embedding变换至Q的计算过程

$Q_I^{\text{head1}} = [0.1, 0.2, 0.3] * [[0.1], [0.2], [0.3]] = [0.14]$

$Q_{love}^{\text{head1}} = [0.32]$

$Q_{playing}^{\text{head1}} = [0.5]$

$Q_{football}^{\text{head1}} = [0.68]$

矩阵表示为：

$Q^{\text{head1}} = [[0.14], [0.32], [0.5], [0.68]]$

同理可得：

$K^{\text{head1}} = [[0.32], [0.77], [1.22], [1.67]]$

$V^{\text{head1}} = [[0.38], [0.92], [1.46], [2.0]]$

对于头2，我们需要不同的权重，计算过程相同

随后对于每个头，对Q，K进行点积并进行softmax计算注意力权重，将权重与V相乘计算加权和。本例中，QK点积的结果为一个(4, 4)的矩阵。将注意力权重与V相乘得到(4, 1)的注意力输出。

计算了每个头的注意力后，对输出进行拼接(concatenate)操作。再将拼接后的结果输入另一个线性层以得到最后的输出。本例中两个注意力头拼接后得到(4, 2)的输出。随后使用一个线性层将结果转换为需要的维度(4, 3)。

过程总结：

输入维度: (4, 3) (4 词元, 嵌入维度为 3).
每个头的Q、K、V： (4, 1) (4 tokens, 1 dimension per head).
每个头的输出：(4, 1).
每个头拼接后：(4, 2).
线性映射后的输出：(4, 3).

基于位置的前馈网络FFN

将输入形状由(b, n, d)变为(bn, d)

输出形状由(bn, d)变回(b, n, d)

FFN等价于两层kernel_size为1的一维卷积层

将输入由(b, n, d)变为(bn, d)有如下原因：

全连接层要求输入为2维，一维为batch_size，二维表示特征。
将(b, n, d)变为(bn, d)而不是(b, nd)是为了保留每个元素在序列中的位置信息（位置编码提供的位置信息）。
可以有效地将所有批次序列中的每个位置视为具有d特征的单个样本。这允许密集层对序列中的每个位置应用相同的权重集合。
同时保障了并行计算，(bn, d)可以并行处理多个序列，可以同时对整个序列进行操作，而不是按顺序处理每个元素，这大大加快了计算速度并减少了整体训练时间。

层归一化

什么是层归一化？

对每个样本的所有feature做归一化。意思是对每个样本独立地计算均值和方差。

假设输入 $x=[[1, 2], [3, 4], [5, 6]]$ 为(3, 2)的矩阵，样本归一化对每行计算均值和方差：

对第一个样本：均值为(1+2)/2=1.5，方差为 $(1-1.5)^2+(2-1.5)^2=0.25$ 。

对每行计算并归一化后得到输出 $\hat{x}=[[-1, 1], [-1, 1], [-1, 1]]$

对比BatchNorm

Batch normalization对输入的所有batch的每个feature批量计算均值和方差并进行归一化。对每个feature独立地计算均值和方差。

同样假设 $x=[[1, 2], [3, 4], [5, 6]]$ ，对于第一个特征计算均值和方差：

$mean=(1+3+5)/3=3$ ， $variance=(1-3)^2+(3-3)^3+(5-3)^2=4$

对每列计算并归一化后的到输出： $\hat{x}=[[-1, -1], [0, 0], [1, 1]]$

为什么是有layer norm而不是batch norm？

Transform的输入一般为一个序列，但是每个序列的有效长度(valid length)有所不同，会导致batch norm不稳定，故不适合batch norm。

对每个样本的元素做归一化，在长度变化时更加稳定。

信息传递

编码器中的输出 $y_1,\dots,y_n$ 作为解码器中第i个Transformer块中多头注意力的K和V，解码器的Q来自目标序列。

说明编码器和解码器中的Block个数和输出维度是相同的。

预测

预测时，解码器的掩码注意力和多头注意力的Q、K和V的来源分别如下：

Masked Attention:
- Q: 解码器在时间步 t 的状态（之前生成的tokens的表示）。
- K: 解码器在时间步 t 的状态（之前生成的tokens的表示）。
- V: 解码器在时间步 t 的状态（之前生成的tokens的表示）。
Multi-Head Attention（交叉注意力）:
- Q: 解码器在时间步 ttt 的状态。
- K: 编码器的最终输出（输入序列的上下文表示）。
- V: 编码器的最终输出（输入序列的上下文表示）。

在训练时，解码器的Masked Attention和Multi-Head Attention的Q、K和V分别由以下组成：

Masked Attention:
- Q: 解码器当前时间步的输入向量（包括所有已生成的tokens的表示）。
- K: 解码器当前时间步的输入向量（包括所有已生成的tokens的表示）。
- V: 解码器当前时间步的输入向量（包括所有已生成的tokens的表示）。
Multi-Head Attention（交叉注意力）:
- Q: 解码器当前时间步的输入向量（表示已生成的tokens）。
- K: 编码器的最终输出（输入序列的上下文表示）。
- V: 编码器的最终输出（输入序列的上下文表示）。

基于Pytorch的Transformer

MultiHeadAttention

class MultiHeadAttention(nn.Module):
    # num_hiddens: 输入的特征数
    def __init__(self, key_size, query_size, value_size, num_hiddens, num_heads, dropout, bias=False, **kwargs):
        super(MultiHeadAttention, self).__init__(**kwargs)
        # heads的数量
        self.num_heads = num_heads
        # 点积注意力
        self.attention = d2l.DotProductAttention(dropout)
        # 计算Q的权重矩阵
        self.W_q = nn.Linear(query_size, num_hiddens, bias=bias)
        # 计算K的权重矩阵
        self.W_k = nn.Linear(key_size, num_hiddens, bias=bias)
        # 计算V的权重矩阵
        self.W_v = nn.Linear(value_size, num_hiddens, bias=bias)
        # 对注意力输出进行变换的矩阵
        self.W_o = nn.Linear(num_hiddens, num_hiddens, bias=bias)

    def forward(self, queries, keys, values, valid_lens):
        # QKV的shape：(batch_size, nums, num_hiddens)
        queries = transpose_qkv(self.W_q(queries), self.num_heads)
        keys = transpose_qkv(self.W_k(keys), self.num_heads) 
        values = transpose_qkv(self.W_v(values), self.num_heads)
        
        # valid_lens的shape：(batch_size, ) or (batch_size, num_query)
        if valid_lens is not None:
            # 每个头都需要valid_lens的信息，因此需要复制num_heads次
            # repeats=self.num_heads, dim=0 表示沿着第一个维度（即 batch 维度）将 valid_lens 的每一项都重复 num_heads
            # 由于每个头都需要知道它正在处理的序列的有效长度，因此我们需要将 valid_lens 对应地复制给每一个头。
            valid_lens = torch.repeat_interleave(valid_lens, repeats=self.num_heads, dim=0)

        # output的shape：(batch_size*num_heads, num_query, num_hiddens/num_heads)
        # 将不同头的QKV合并成为一个大矩阵进行计算，提升计算并行度。即K=K1 concat K2 concat .. Kn
        output = self.attention(queries, keys, values, valid_lens)
        # output_concat的形状：(batch_size, num_query, num_hiddens)
        # 将不同头的结果拼接
        output_concat = transpose_output(output, self.num_heads)
        # 对注意力输出进行线性变换后输出
        return self.W_o(output_concat)

# 两个辅助函数
#@save
def transpose_qkv(X, num_heads):
    # 为了多注意头并行计算而变换形状
    # X的输入为：(batch_size, num_query, num_hiddens)
    # 输出为：(batch_szie, num_query, num_heads, num_hiddnes / num_heads)
    X = X.reshape(X.shape[0], X.shape[1], num_heads, -1)
    # 交换heads维和num_query维
    # 输出为：(batch_size, num_heads, num_query, num_hiddens / num_heads)
    X = X.permute(0, 2, 1, 3)
    # 将输出变为(batch_size * num_heads, num_query, num_hiddens / num_heads)用于并行计算
    return X.reshape(-1, X.shape[2], X.shape[3])

#@save
def transpose_output(X, num_heads):
    # 将合并的X还原num_heads维度
    # 输入为：(batch_size, num_heads, num_query, num_hiddens/num_heads)
    X = X.reshape(-1, num_heads, X.shape[1], X.shape[2])
    # 转换为：(batch_size, num_query, num_heads, num_hiddens / num_heads)
    X = X.permute(0, 2, 1, 3)
    # 转换为原格式：(batch_size, num_query, num_hiddens)
    return X.reshape(X.shape[0], X.shape[1], -1)

num_hiddens, num_heads = 100, 5
attention = MultiHeadAttention(num_hiddens, num_hiddens, num_hiddens,
                               num_hiddens, num_heads, 0.5)
attention.eval()

batch_size, num_queries = 2, 4
# valid_lens表示第一个batch的有效长度为3，第二个为2
num_kvpairs, valid_lens =  6, torch.tensor([3, 2])
'''
在编码器-解码器架构中，比如在Transformer模型里，解码器部分的每个时间步都会产生一个查询向量去与编码器的输出进行交互。
查询的数量通常是解码器当前时间步的数目，或者是整个目标序列的长度（如果一次性处理整个序列的话）。
对于自注意力机制（self-attention）(编码器中)，每个位置上的元素都作为一个查询去与其他所有元素交互，所以这里的查询数量就是输入序列的长度。
'''
X = torch.ones((batch_size, num_queries, num_hiddens))
'''
在编码器中的自注意力层，键-值对的数量等同于输入序列的长度。
在解码器中，键-值对可以来自于编码器的输出，这时它们的数量就是编码器输出序列的长度；或者来自解码器自身的先前输出，
这时它们的数量将是解码器已经产生的序列长度。
'''
Y = torch.ones((batch_size, num_kvpairs, num_hiddens))
attention(X, Y, Y, valid_lens).shape
'''
Q与KV数量不相同，如何进行注意力计算？
Q.shape = (batch_size, num_queries, d_model) 
K.shape = (batch_size, num_kvpairs, d_model) 
Q @ K =  (batch_size, num_queries, num_kvpairs) = A
将结果进行softmax后再与V相乘
V.shape = (batch_size, num_kvpairs, d_model) 
A @ K = (batch_size, num_queries, d_model) 
'''