大模型基础——从零实现一个Transformer(1)
一、引言
上一章主要实现了一下Transformer里面的BPE算法和 Embedding模块定义
本章主要讲一下 Transformer里面的位置编码以及多头注意力
二、位置编码
2.1正弦位置编码(Sinusoidal Position Encoding)
其中:
pos:表示token在文本中的位置
i : i代表词向量具体的某一维度,即位置编码的每个维度对应一个波长不同的正弦或余弦波
d : d表示位置编码的最大维度,和词嵌入的维度相同,假设是512
对于位置0的编码为:
对于位置1的编码为:
2.2 正弦位置编码特性
- 相对位置关系:pos + k的位置编码可以被位置pos的位置编码线性表示
三角函数公式如下:
对于pos + k的位置编码:
根据式( 3 )和( 4 )整理上式有:
- 位置之间的相对距离
𝑃𝐸𝑝𝑜𝑠+𝑘∙𝑃𝐸𝑝𝑜𝑠 的内积:
位置之间内积的关系大小如下:
可以看到内积会随着相对位置的递增而减少,从而可以表示位置的相对距离。内积的结果是对称的,所以没有方向信息。
2.3 代码实现
import torch
from torch import nn,Tensor
import mathclass PositionalEmbedding(nn.Module):def __init__(self,d_model:int=512,dropout:float=0.1,max_positions:int=1024) -> None:''':param d_model: embedding向量的维度:param dropout::param max_positions: 最大长度'''super().__init__()self.dropout = nn.Dropout(p=dropout)# Position Embedding (max_positions,d_model)pe = torch.zeros(max_positions,d_model)# 创建position index列表 ,形状为:(max_positions, 1)position = torch.arange(0,max_positions).unsqueeze(1)# d_model 维度 偶数位是sin ,奇数位是cos# 计算除数,这里的除数将用于计算正弦和余弦的频率div_term = torch.exp(torch.arange(0,d_model,2) * -(math.log(10000.0) /d_model))# 对矩阵的偶数列(0,2,4...)进行正弦函数编码pe[:, 0::2] = torch.sin(position * div_term)# 对矩阵的奇数列(1,3,5...)进行余弦函数编码pe[:, 1::2] = torch.cos(position * div_term)# 扩展维度,增加batch_size: pe (1, max_positions, d_model)pe = pe.unsqueeze(0)# buffers will not be trainedself.register_buffer("pe", pe)def forward(self,x:Tensor) ->Tensor:"""Args:x (Tensor): (batch_size, seq_len, d_model) embeddingsReturns:Tensor: (batch_size, seq_len, d_model)"""# x.size(1)是指当前x的最大长度x = x + self.pe[:,:x.size(1)]return self.dropout(x)if __name__ == '__main__':seq_len = 128d_model = 512pe = PositionalEmbedding(d_model)x = torch.rand((1,100,d_model))print(pe(x).shape)
三、多头注意力
3.1 自注意力
公式如下:
- 假设一个矩阵X,分别乘上权重矩阵,,就得到了Q , K , V向量矩阵
- 然后除以 𝑑𝑘 进行缩放,再经过Softmax,得到注意力权重矩阵,接着乘以value向量矩阵V,就一次得到了所有单词的输出矩阵Z
3.2 多头注意力
将原来n_head分割乘Nx n_sub_head.对于每个头i,都有它自己不同的key,query和value矩阵: 𝑊𝑖𝐾,𝑊𝑖𝑄,𝑊𝑖𝑉 。在多头注意力中,key和query的维度是 𝑑𝑘 ,value嵌入的维度是 𝑑𝑣 (其中key,query和value的维度可以不同,Transformer里面一般设置的是相同的),这样每个头i,权重 𝑊𝑖𝑄∈𝑅𝑑×𝑑𝑘,𝑊𝑖𝐾∈𝑅𝑑×𝑑𝑘,𝑊𝑖𝑉∈𝑅𝑑×𝑑𝑣 ,然后与压缩到X中的输入相乘,得到 𝑄∈𝑅𝑁×𝑑𝑘,𝐾∈𝑅𝑁×𝑑𝑘,𝑉∈𝑅𝑁×𝑑𝑣 .
3.3 代码实现
import mathimport torch
from torch import nn,Tensor
from typing import *class MultiHeadAttention(nn.Module):def __init__(self,d_model: int = 512,n_heads: int=8,dropout: float = 0.1):''':param d_model: embedding大小:param n_heads: 多头个数:param dropout:'''super().__init__()assert d_model % n_heads == 0self.d_model = d_modelself.n_heads = n_headsself.d_key = d_model // n_headsself.q = nn.Linear(d_model,d_model)self.k = nn.Linear(d_model,d_model)self.k = nn.Linear(d_model,d_model)self.concat = nn.Linear(d_model,d_model)self.dropout = nn.Dropout(dropout)def split_heads(self,x:Tensor,is_key : bool = False) -> Tensor:'''分割向量为N个头,如果是key的话,softmax时候,key需要转置一下:param x::param is_key::return:'''batch_size = x.size(0)# x (batch_size,seq_len,n_heads,d_key)x = x.view(batch_size,-1,self.n_heads,self.d_key)if is_key:# (batch_size,n_heads,d_key,seq_len)return x.permute(0,2,3,1)# (batch_size,n_heads,seq_len,d_keyreturn x.transpose(1,2)def merge_heads(self,x: Tensor) -> Tensor:x = x.transpose(1,2).contigouse().view(x.size(0),-1,self.d_model)return xdef attention(self,query:Tensor,key:Tensor,value:Tensor,mask:Tensor = None,keep_attentions:bool = False):scores = torch.matmul(query,key) / math.sqrt(self.d_key)if mask is not None:scores = scores.masked_fill(mask == 0, -1e9)# weights (batch_size,n_heads,q_length,k_length)weights = self.dropout(torch.softmax(scores,dim=-1))# (batch_size,n_heads,q_length,k_length) x (batch_size,n_heads,v_length,d_key)# -> (batch_size,n_heads,q_length,d_key)# assert k_length == v_length# attn_output (batch_size, n_heads, q_length, d_key)atten_output = torch.matmul(weights,value)if keep_attentions:self.weights = weightselse:del weightsreturn atten_outputdef forward(self,query: Tensor,key: Tensor,value: Tensor,mask: Tensor = None,keep_attentions: bool = False)-> Tuple[Tensor,Tensor]:''':param query:(batch_size, q_length, d_model):param key:(batch_size, k_length, d_model):param value:(batch_size, v_length, d_model):param mask: mask for padding or decoder. Defaults to None.:param keep_attentions: whether keep attention weigths or not. Defaults to False.:return: (batch_size, q_length, d_model) attention output'''query = self.q(query)key = self.k(key)value = self.v(value)query,key,value = (self.split_heads(query),self.split_heads(key,is_key=True),self.split_heads(value))atten_output = self.attention(query,key,value,mask,keep_attentions)del querydel keydel value# concatconcat_output = self.merge_heads(atten_output)# the final liear# output (batch_size, q_length, d_model)output = self.concat(concat_output)return output