手撕transformer（1）

hdy

本教程主要参考pytorch、pytorch-lighting文档以及3Blue1Brown中关于transformer部分的内容，并结合自己的理解写出的更基础的版本。Transformer 是一种基于自注意力机制（Self-Attention）的深度学习模型架构，由 Google 团队于 2017 年在论文《Attention Is All You Need》中提出。它彻底改变了自然语言处理（NLP）领域，并逐渐扩展至计算机视觉、语音识别等领域。Transformer 通过自注意力机制实现了对序列数据的全局建模，其并行化设计和模块化架构使其成为 AI 大模型的核心基础（ChatGPT，DeepSeek等等）。我们先拿出它的核心公式来看下：我们先暂且不管这个Q、K、V是如何从文字序列、图片、时序数据中得到的，先单拿出来整个流程来看，假设我们经过种种步骤将原始输入数据转换成了所需要的Q、K、V，那么接下来从计算图的角度来看是这样
按照公式来理解就是基于给定的三个矩阵或者向量Q,K,V进行运算，运算的顺序为首先Q和K的转置做点积，点积后的结果除以dk的开方。注：以下所指的点积均为矩阵乘法中的点积，请勿与其他概念混淆。Q、K、V这里分别指的是查询（Query）、键（Keys）、值（Values），Q是一个特征向量，它描述了我们在序列中要寻找的内容，即我们可能想要关注的内容。对于每个输入元素，我们都有一个键向量，它同样是一个特征向量。这个特征向量大致描述了该元素“提供”的内容，或者它可能重要的时刻。这些键的设计应使得我们能够根据查询来识别出需要关注的元素。对于每个输入元素，我们还有一个值向量，模型最终的目标正是通过值向量的加权组合，生成能够反映复杂上下文关系的表示，驱动下游任务（如翻译、分类）的预测结果。
也就是说我们先通过Q和K的点积得到注意力分数矩阵（Score Matrix），用于表示查询（Query）与键（Key）之间的相似性或关联程度。得到注意力分数矩阵之后下一步要做softmax，softmax的作用简单来讲就是做归一化，针对每一个经过softmax的输入都给出一个输出，并确保每个输出值在(0,1)区间内，且所有输出值之和为1。注意公式里的softmax针对的是点积后得到的矩阵，但其运算逻辑是取这个矩阵中的每一行做softmax，而不是把这个矩阵直接展平后做softmax。接下来是为什么要除以Q和K维度dk的开方。假设Q和K两者均采样于均值为0，方差为σ（假设为1）平方的正态分布。两个向量  和  的点积：

其方差为：

我们发现初始化σ为1时可以省去，那么点积后的方差和dk呈线性相关，当dk很大时，方差也会越大，代表其中的数据越不均衡，输入到softmax的“能量”过高，会导致softmax输出接近一个one-hot向量（最大值为1，其余接近0）。此时softmax的梯度几乎为零（因输出饱和），模型无法通过反向传播更新参数。解决方法就是除以dk的开方（缩放因子），使得缩放后的方差回到合理范围，此时的方差为缩放后方差若初始化时σ=1，则最后的方差仍为1，避免了数值爆炸和梯度消失。在注意力分数矩阵经过softmax 之后的结果称为注意力权重矩阵（Attention Weights Matrix），最后与值矩阵V做点积。其本质是将每个位置的注意力权重作为系数，对V中对应的值向量进行加权求和。下面我们用代码展示具体实现：先按照上述的步骤按顺序写成一个函数（忽略掩码mask部分，后续详解）
def scaled_dot_product(q, k, v, mask=None):    d_k = q.size()[-1] #获取特征维度，q与k应为相同维度    attn_logits = torch.matmul(q, k.transpose(-2, -1))#q与k的转置做点积    attn_logits = attn_logits / math.sqrt(d_k)#缩放因子    if mask is not None:        attn_logits = attn_logits.masked_fill(mask == 0, -9e15)    attention = F.softmax(attn_logits, dim=-1)#softmax归一化    values = torch.matmul(attention, v)#与值矩阵点积    return values, attention接下来我们生成一组Q,K,V来验证这个流程，这里我把随机生成的矩阵全部改为只保留一位小数的形式，从而方便手动验证。
seq_len, d_k = 3, 2pl.seed_everything(42)q = torch.randn(seq_len, d_k).round(decimals=1)  k = torch.randn(seq_len, d_k).round(decimals=1)  v = torch.randn(seq_len, d_k).round(decimals=1)  values, attention = scaled_dot_product(q, k, v)print("Q\n", q)print("K\n", k)print("V\n", v)print("Values\n", values)print("Attention\n", attention)输出如下


我们来手动验算一下这个过程，首先是就等于接下来除以Q和K维度dk的开方也就是根号2然后送入softmax函数，下面拆解softmax具体运算流程softmax的公式为其中按照softmax的公式对每行元素求和，得到分母 ：然后每行逐元素除以求和后的分母第一行为第二行为第三行为把得到的结果整理成矩阵形式，并和代码得到的矩阵比较是否一致（由于手动计算时我们选择的保留位数以及设备运算精度不同，可能会有稍许差异，我们这里是小数点后4位。得到attention后与值矩阵V做点积得到和自动计算的结果对比

到此注意力机制的计算流程就完成了，后续继续手撕transformer