机器学习周报（transformer学习2）

摘要

本文实现了Transformer模型中的关键组件，包括多头注意力机制、编码器层以及位置编码。首先 ，通过手写实现了多头注意力机制，利用多个注意力头并行计算，提升了信息的表达能力。然后，设计了自定义的编码器模块，结合了多头注意力、前馈神经网络和层归一化。最后，构建了完整的Transformer架构，为序列处理任务提供了一个基础框架。

Abstract

This paper implements key components of the Transformer model, including the multi-head attention mechanism, encoder layer, and positional encoding. First, the multi-head attention mechanism is manually implemented to enhance information representation by computing with multiple attention heads in parallel. Next, a custom encoder module is designed, integrating multi-head attention, feedforward neural networks, and layer normalization. Finally, a complete Transformer architecture is constructed, providing a foundational framework for sequence processing tasks.

1 手写多头注意力机制

多头注意力机制原理：

多头注意力的核心思想是：通过多个注意力头并行计算，使得模型能够从不同的子空间学习到更多的信息。每个注意力头都会对输入的Query、Key、Value进行加权处理，最终的结果会被合并。具体来说，计算过程包括：

1. Q、K、V 变换： 首先，将输入的词向量通过线性变换分别映射到Query、Key、Value空间中。

2. Scaled Dot-Product Attention： 对于每个头，计算Q与K的点积，并除以$\sqrt{d_k}$（即Key的维度的平方根）进行缩放，然后使用Softmax进行归一化，得到每个位置的权重。

3. 头组合： 对每个头的结果加权平均，最终将所有头的结果拼接起来并通过一个线性变换生成最终输出。

通过该机制，我们能够通过多个注意力头并行地对输入序列进行处理，从而捕获不同的子空间信息。

代码部分

import torch
from torch import nn
import torch.functional as f
import math
#%%
# 测试数据
X = torch.randn( 128, 64, 512) # Batch,Time,Dimension
print(X.shape)
#%%
# 设置multihead_attention基本参数
d_model = 512  # 映射到Q，K，V空间中有多少位
n_head = 8 # 有多少个头
#%%
class multi_head_attention(nn.Module):
    def __init__(self, d_model,n_head) -> None:
        super(multi_head_attention,self).__init__()
 
        self.n_head = n_head
        self.d_model = d_model
        self.w_q = nn.Linear(d_model, d_model)  # 线性层映射函数，把初始向量映射到Q,K,V（query,key,value）
        # 简单来说就是去寻找一些query去跟key，问他（key）哪些数据是跟我匹配的上的，匹配上之后，key所对应的value值进行加权组合，最终得到attention的输出
        self.w_k = nn.Linear(d_model, d_model)
        self.w_v = nn.Linear(d_model, d_model)
        self.w_combine = nn.Linear(d_model, d_model) # 由于是多头注意力，所以要在最后做一个组合映射(多写一个w_combine的线性映射)
        self.softmax = nn.Softmax(dim=-1)
 
    def forward(self, q, k, v):
        batch, time, dimension = q.shape
        n_d = self.d_model // self.n_head   # 得到新维度
        q, k, v = self.w_q(q), self.w_k(k), self.w_v(v)  # 把qkv分别丢到上面定义的三个线性映射层中，就可以得到qkv空间中的一个表示
        # 对空间表示进行切分，对我们需要得到几个头进行切分
        q = q.view(batch, time, self.n_head, n_d).permute(0, 2, 1, 3) # 把q进行维度划分，一维是batch，二维是time， 三维是n.head（分成几个头），四维是n.d（分完头之后的维度）
        k = k.view(batch, time, self.n_head, n_d).permute(0, 2, 1, 3) # 也可以说把最后一维拆成了n.head和n.d两个维度的乘积
        v = v.view(batch, time, self.n_head, n_d).permute(0, 2, 1, 3) # 做attention操作的时候head维是不能放在最后的，对最后两个维度进行处理，所以要用permute指令做一个维度变换
        # 原先的维度是0，1，2，3现在则是0，2，1，3
 
        score = q @ k.transpose(2, 3) / math.sqrt(n_d) # q乘以k的转置除以它的维度开根号（让方差变小） @是矩阵乘法
        # torch.tril命令-生成下三角矩阵（左下角都是1，右上角都是0）
        mask = torch.tril(torch.ones(time, time, dtype=bool))
        score = score.masked_fill(mask == 0, float("-inf")) # 把mask等于0的地方都填充为负无穷
        # 填充为负无穷的原因:softmax操作时e^-inf就是0，就相当于我们不去care后面部分的信息
        score = self.softmax(score) @ v
 
        # 最后把得分的格式变回来（因为之前把time维和self.n_head维度进行了旋转，现在则是要旋转回来），然后再过一个连续性函数
        score = score.permute(0, 2, 1, 3).contiguous().view(batch, time