LLM - Transformer 的 Q/K/V 详解

目录

一.引言

二.传统 Q/K/V

三.Transformer Q/K/V

- Input Query

- Q/K/V 获取

- Q/K 相似度计算

- 注意力向量

- Multi Head

四.代码测试

- 初始化

- Attention

- Main

五.总结


一.引言

Transformer 的输入是我们的一个 query 句子,例如 "我爱中国",但是 Transformer 处理时却 1 生 3 得到了 Q/K/V,下面我们从传统机器学习和 Transformer 两个角度看一下 Q/K/V 从哪里来,去哪里去。

二.传统 Q/K/V

LLM - Transformer 的 Q/K/V 详解_第1张图片

对于传统机器学习的 Attention 而言,Q/K 分别代表 Query 和 Keys 即我们候选物品与另一个物品的向量,通过 matmul(Q,K) 即可计算二者的相似度,最后 softmax 归一化乘上 valuse 即可得到。MatMul 后的 Scala 是除以当前 sqrt(dim) 的缩放操作,防止乘积结果太大。

最常见的就是电商场景,例如我们浏览了一件羽绒服,这个就是我们的 Q,而我们历史浏览过多个商品例如鞋子、包包、大衣,这些都是 K,通过 matmul(Q,K) 可以计算出羽绒服与我们浏览的 N 个商品每个商品的相似度,归一化后的值有大有小,从而代表当前商品 K 对 Q 向量的贡献度衡量,进行加权求和即可得到候选集 Q 在浏览记录 K 下的注意力向量表征。

LLM - Transformer 的 Q/K/V 详解_第2张图片

三.Transformer Q/K/V

Transformer 场景下,我们的输入只有 text = "我爱中国",此时不存在上面类似候选商品 Q 和历史浏览商品 K 的关系,怎门办,也好办,Q/K 都是自己就可以,也就是 Self-Attention,下面我们看下如何操作。

- Input Query

对于给定的句子 text,其长度为 L,以 "我想吃酸菜鱼" 为例,其 seq_len = 6, dim = [L x 768]

LLM - Transformer 的 Q/K/V 详解_第3张图片

这里我们可以构建 Embedding 层,将 text 使用 Tokenizer 分词后,得到其对应的 Embedding 表征。除此之外,我们定义 Attention Module,同时根据传入的 Head 数计算拆分后的 head_dim。

class BitDDDAttention(nn.Module):

    def __init__(self, embed_dim, num_heads):
        super(BitDDDAttention, self).__init__()

        self.embed_dim = embed_dim  # embedding 维度
        self.num_heads = num_heads  # head 数量
        assert embed_dim % num_heads == 0
        self.head_dim = embed_dim // num_heads

        # 构建 Q/K/V 向量以及最后的全连接 MLP
        self.query = nn.Linear(embed_dim, embed_dim)
        self.key = nn.Linear(embed_dim, embed_dim)
        self.value = nn.Linear(embed_dim, embed_dim)
        self.fc_out = nn.Linear(embed_dim, embed_dim)

- Q/K/V 获取

LLM - Transformer 的 Q/K/V 详解_第4张图片

这里 Q/K/V 都是通过一层 Linear 映射得到的,torch 的话就是 nn.Linear,映射后的 Q/K/V 维度不变,依然是 L x 768:

LLM - Transformer 的 Q/K/V 详解_第5张图片

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

        # Split the embedding into num_heads and reshape to (batch_size, num_heads, seq_len, head_dim)
        query = self.query(x).view(batch_size, seq_len, self.num_heads, self.head_dim).permute(0, 2, 1, 3)
        key = self.key(x).view(batch_size, seq_len, self.num_heads, self.head_dim).permute(0, 2, 1, 3)
        value = self.value(x).view(batch_size, seq_len, self.num_heads, self.head_dim).permute(0, 2, 1, 3)

初始化 Q/K/V 的转换 Linear 矩阵,并根据 head_num 与 head_dim 转换,转换后的维度为:

(bsz, num_heads, seq_len, head_dim)

- Q/K 相似度计算

Attention(Q,K_i,V_i) = softmax(\frac{Q^TK_i}{\sqrt{d_k}})V_i

相似度计算考虑如上 Scale-Dot-Attention 公式,Q 的维度为 L x 768,K 转置后为 768 x L,二者 MatMul 得到 L x L 的方阵,其中 (0,0) 位置代表 "我想吃酸菜鱼" 中 "我" 对 V 中 "我" 字向量的注意力权重,以此类推,整个结果表示 "我想吃酸菜鱼" 中每个字对其他字的注意力,当然也包括自己,所以这种操作也被称为自注意力。

LLM - Transformer 的 Q/K/V 详解_第6张图片

        # Compute the attention scores
        attention_scores = torch.matmul(query, key.permute(0, 1, 3, 2)) / self.head_dim ** 0.5
        attention_probs = torch.softmax(attention_scores, dim=-1)

Q Dim: (bsz, num_heads, seq_len, head_dim)

K.T Dim: (bsz, num_heads, head_dim, seq_len)

相乘后得到 (bsz, num_heads, seq_len, seq_len) 的注意力权重。

Tips:

这里做了除以 sqrt(dim) 的操作,主要是为了缩小点积范围,确保 softmax 梯度稳定性,而为什么需要做 softmax 一是为了保证权重的非负性,同时增加非线性操作。第 row 行归一化的权重,就是第 row 个元素对其他元素的注意力权重。

LLM - Transformer 的 Q/K/V 详解_第7张图片

- 注意力向量

经过 softmax 后的 LxL 的注意力权重矩阵分别与 LxDim 的 Value 相乘,最终得到自注意力后的矩阵。第一行代表 "我" 对 "我想吃酸菜鱼" 中每个注意力权重在 "我" 向量上的加权平均结果,以此类推,最后得到 seq_len 即 L 个加权平均的表征。

LLM - Transformer 的 Q/K/V 详解_第8张图片

        attention_output = torch.matmul(attention_probs, value)
                                .permute(0, 2, 1, 3)
                                .contiguous()
                                .view(batch_size, seq_len, self.embed_dim)                                                                                                                          

计算后通过 view 将多个 head 的向量合并,恢复为原始的 embed_dim,此时维度依然为 

(bsz, seq_len, head_num x head_dim)

- Multi Head

LLM - Transformer 的 Q/K/V 详解_第9张图片

实践场景中很多都使用 multi-head Attention 即多头注意力,我们可以理解为把原始向量进行拆分,将其多个部分分别进行上面的 Attention 操作。比如还是 "我想吃酸菜鱼",我们得到的 Q/K/V 都是 L x 768 维,此时令 head = 4,则 L x 768 会切分为 4 份 L x 192,分别将 L x 192 做 Self-Attention 后再 Concat,即可再次得到 L x 768。Multi-Head 的思想意在学习向量不同位置对语义的不同表征。

四.代码测试

- 初始化

#!/usr/bin/python
# -*- coding: UTF-8 -*-

import torch
import torch.nn as nn
from transformers import BertModel, BertTokenizer


class BitDDDAttention(nn.Module):

    def __init__(self, embed_dim, num_heads):
        super(BitDDDAttention, self).__init__()

        self.embed_dim = embed_dim  # embedding 维度
        self.num_heads = num_heads  # head 数量
        assert embed_dim % num_heads == 0
        self.head_dim = embed_dim // num_heads

        # 构建 Q/K/V 向量以及最后的全连接 MLP
        self.query = nn.Linear(embed_dim, embed_dim)
        self.key = nn.Linear(embed_dim, embed_dim)
        self.value = nn.Linear(embed_dim, embed_dim)
        self.fc_out = nn.Linear(embed_dim, embed_dim)

定义 head_num、head_dim 以及 Q/K/V 的 Linear 线性层。

- Attention

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

        # Split the embedding into num_heads and reshape to (batch_size, num_heads, seq_len, head_dim)
        query = self.query(x).view(batch_size, seq_len, self.num_heads, self.head_dim).permute(0, 2, 1, 3)
        key = self.key(x).view(batch_size, seq_len, self.num_heads, self.head_dim).permute(0, 2, 1, 3)
        value = self.value(x).view(batch_size, seq_len, self.num_heads, self.head_dim).permute(0, 2, 1, 3)

        # Compute the attention scores
        attention_scores = torch.matmul(query, key.permute(0, 1, 3, 2)) / self.head_dim ** 0.5
        attention_probs = torch.softmax(attention_scores, dim=-1)

        # Apply the attention weights to the value
        attention_output = torch.matmul(attention_probs, value).permute(0, 2, 1, 3).contiguous().view(batch_size,
                                                                                                      seq_len,
                                                                                                      self.embed_dim)

        # Apply a linear layer to the output
        x = self.fc_out(attention_output)
        return x

执行 Scale-Dot-Attention,Multi-Head 可以看作是分开做了多个 Self-Attention。

- Main

if __name__ == '__main__':
    tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
    pretrained_bert = BertModel.from_pretrained('bert-base-uncased')

    input_texts = ["This is a test sentence.", "Here is another test sentence."]
    input_ids = [tokenizer.encode(text, add_special_tokens=True, max_length=10, padding='max_length', truncation=True,
                                  return_tensors='pt') for text in input_texts]
    input_ids = torch.cat(input_ids, dim=0)  # Concatenate and add batch dimension

    with torch.no_grad():
        embedded_output = pretrained_bert(input_ids)[0]  # Get the output of the BERT model
        print(embedded_output.size())  # Output shape should be (2, 10, embedding_dim)

        embed_dim = embedded_output.size(-1)
        num_heads = 4
        model = BitDDDAttention(embed_dim, num_heads)

        output = model(embedded_output)
        print(output.size())  # Output shape should be (2, 10, embed_dim)

输入两条样本,batch 为 2,seq_len 为 10,最终调用 Attention 进行 forward 即可得到一次 Attention 的结果,维度为: torch.Size([2, 10, 768])。

五.总结

本节内容基于 Torch 实现了 MultiHeadAttention,由于 Q/K/V 都是来自于同一句话,故其名为自注意力即 Self-Attention,如果想要类似推荐算法中的不同商品比较,则 Q 是一句话,KV 换成另一句话即可。这里有一点需要注意,我们在 K/V 中调整任意两个字的位置,其对最终的计算没有影响,即在注意力机制中是没有位置信息的,这也是为什么 LLM 要引入 Position Embedding 的原因。

你可能感兴趣的:(Python,深度学习,transformer,深度学习,人工智能)