大模型基础——从零实现一个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 math


class 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 trained
        self.register_buffer("pe", pe)

    def forward(self,x:Tensor) ->Tensor:
        """

                Args:
                    x (Tensor): (batch_size, seq_len, d_model) embeddings

                Returns:
                    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 = 128
    d_model = 512

    pe = 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 math

import 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 == 0
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_key = d_model // n_heads

        self.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_key
        return 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 x

    def 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 = weights
        else:
            del weights

        return atten_output

    def 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 query
        del key
        del value

        # concat
        concat_output = self.merge_heads(atten_output)

        # the final liear
        # output (batch_size, q_length, d_model)
        output = self.concat(concat_output)
        
        return output