中午网站做google广告好吗,微信公众号文章怎么转wordpress,百姓网招聘最新招聘信息,网站地址解析本文是github上的大模型教程LLMs-from-scratch的学习笔记#xff0c;教程地址#xff1a;教程链接 Chapter 3#xff1a; Attention Mechanism
本文首先从固定参数的注意力机制说起#xff0c;然后拓展到可以训练的注意力机制#xff0c;然后加入掩码mask#xff0c;最后… 本文是github上的大模型教程LLMs-from-scratch的学习笔记教程地址教程链接 Chapter 3 Attention Mechanism
本文首先从固定参数的注意力机制说起然后拓展到可以训练的注意力机制然后加入掩码mask最后拓展到多头注意力机制。 1. 注意力机制
一个句子中的每一个token都会受到其他token的影响这里先不考虑忽略未来的单词掩码的问题后面再说注意力机制可以让一个token收到其他token的影响生成一个最终我们想要的embedding。即每个token有一个原始的embedding通过注意力机制后得到了一个新的embedding这个embedding是结合了上下文语义得到的。
举个简单的例子我们直接使用tokens的embedding之间两两点乘得到互相之间的点乘结果然后将点乘结果归一化得到embeddings之间的注意力得分。
归一化一般使用softmax函数通过取指数除以求和得到归一化结果 torch.exp(x) / torch.exp(x).sum(dim0)
得到token之间的相关权重后我们就可以加权求和得到每一个token的最终embedding。 2. 可以训练的注意力头
在上面的例子中我们直接使用token对应的embedding来计算相关系数以及最终的加权求和这显然是不合理的如果这样的话那么我们只能训练token对应的词嵌入来学习模型或者是一些全连接层因此我们需要引入新的矩阵来学习到更多的参数这就是transformer的QKV矩阵。 QKV都是对原始的embedding做线性变换得到新的向量但是模型就可以通过训练QKV学习海量知识。 QKV的维度不固定可以与原始嵌入相同也可以不同。总之通过QKV三个矩阵我们将原始token的embedding转换成了3个新的向量。
Query vector: q ( i ) W q x ( i ) q^{(i)} W_q \,x^{(i)} q(i)Wqx(i)Key vector: k ( i ) W k x ( i ) k^{(i)} W_k \,x^{(i)} k(i)Wkx(i)Value vector: v ( i ) W v x ( i ) v^{(i)} W_v \,x^{(i)} v(i)Wvx(i)
可以使用矩阵乘法实现
keys inputs W_key
values inputs W_value然后我们计算K和Q之间的点积作为两两token之间的关联度。为什么要用两个不一样的矩阵我的猜测是如果使用的是一个矩阵计算相似度那么关于对角线对称的元素就会完全相同但是使用两个不同的矩阵计算就不会存在这样的情况可以学习到的内容更多。 我们使用K和Q的点积得到了两两之间的注意力得分同样使用softmax进行归一化得到最终的注意力权重。 注意到没有直接对注意力得分softmax而是除以维度的方根后再softmax这是因为在计算注意力权重时如果直接将Query和Key的点积结果用于softmax函数当Key的维度较高时点积的结果会变得非常大。这可能导致softmax函数在梯度下降过程中学习困难因为大的数值会使softmax的梯度变得非常小接近于0这在数值稳定性上是一个问题称为“梯度消失”。 最后一步不再使用原始的embedding加权我们使用V矩阵变换后的向量进行加权求和得到结果向量。 代码如下
class SelfAttention_v2(nn.Module):def __init__(self, d_in, d_out, qkv_biasFalse):super().__init__()self.W_query nn.Linear(d_in, d_out, biasqkv_bias)self.W_key nn.Linear(d_in, d_out, biasqkv_bias)self.W_value nn.Linear(d_in, d_out, biasqkv_bias)def forward(self, x):keys self.W_key(x)queries self.W_query(x)values self.W_value(x)attn_scores queries keys.Tattn_weights torch.softmax(attn_scores / keys.shape[-1]**0.5, dim-1)context_vec attn_weights valuesreturn context_vec3. 隐藏未来的单词 对语言任务来说在训练模型的时候不能使用未来的文本来预测之前的文本。因此我们需要屏蔽未见文本对先前文本的影响。在我们计算得到注意力权重后我们人为地将上三角矩阵的权重置为0。 有一种naive的方法就是将上注意力权重都置为0后重新对剩下的元素归一化。但是我们要介绍的是一般使用的方法 我们在计算出注意力得分后对右上角元素都赋值为负无穷大负无穷大在经过softmax后就变为了0。
mask torch.triu(torch.ones(context_length, context_length), diagonal1)
masked attn_scores.masked_fill(mask.bool(), -torch.inf)
attn_weights torch.softmax(masked / keys.shape[-1]**0.5, dim-1)最后为了防止过拟合一般会使用dropout对注意力权重矩阵进行随机丢弃加强模型泛化性能。
总结以上的所有内容我们现在就可以写出一个单头的注意力机制了并且加入了对batch输入的处理
class CausalAttention(nn.Module):def __init__(self, d_in, d_out, context_length,dropout, qkv_biasFalse):super().__init__()self.d_out d_outself.W_query nn.Linear(d_in, d_out, biasqkv_bias)self.W_key nn.Linear(d_in, d_out, biasqkv_bias)self.W_value nn.Linear(d_in, d_out, biasqkv_bias)self.dropout nn.Dropout(dropout) # Newself.register_buffer(mask, torch.triu(torch.ones(context_length, context_length), diagonal1)) # Newdef forward(self, x):b, num_tokens, d_in x.shape # New batch dimension bkeys self.W_key(x)queries self.W_query(x)values self.W_value(x)attn_scores queries keys.transpose(1, 2) # Changed transposeattn_scores.masked_fill_( # New, _ ops are in-placeself.mask.bool()[:num_tokens, :num_tokens], -torch.inf) # :num_tokens to account for cases where the number of tokens in the batch is smaller than the supported context_sizeattn_weights torch.softmax(attn_scores / keys.shape[-1]**0.5, dim-1)attn_weights self.dropout(attn_weights) # Newcontext_vec attn_weights valuesreturn context_vec4. 多头注意力机制 我们已经实现了单个头的注意力机制那么要实现多个头就是使用多个不同的注意力头各自对输入进行处理然后将各自得到的输出$z_i$拼接起来非常显而易见我们有第一个最直白的写法
class MultiHeadAttentionWrapper(nn.Module):def __init__(self, d_in, d_out, context_length, dropout, num_heads, qkv_biasFalse):super().__init__()self.heads nn.ModuleList([CausalAttention(d_in, d_out, context_length, dropout, qkv_bias) for _ in range(num_heads)])def forward(self, x):return torch.cat([head(x) for head in self.heads], dim-1)这是一种最简单直白的写法直接声明num_heads个注意力单元然后在前向传播的时候依次调用这num_heads个注意力头然后将输出拼接起来。(dim-1代表最后一维拼接 问题是这样的话需要循环num_heads次得到结果并且需要声明num_heads个注意力头相信熟悉线性代数的朋友已经想到了可以通过曾广矩阵来拓展注意力头。比如单个的注意力头是(d_in, d_out)那么有n个头的注意力机制就是(d_in, n*d_out)。 假设输入是(tokens, d_in)那么(tokens, d_in) (d_in, n*d_out) -- (tokens, n*d_out)输出的结果完美得到了n个头对应的输出我们只需要按照每d_out列拆开得到n个(tokens, d_out)的矩阵就能还原出n个头对应的结果进行后续的attention score计算。这样写起来虽然麻烦一些但是效率更高。
class MultiHeadAttention(nn.Module):def __init__(self, d_in, d_out, context_length, dropout, num_heads, qkv_biasFalse):super().__init__()assert (d_out % num_heads 0), \d_out must be divisible by num_headsself.d_out d_outself.num_heads num_headsself.head_dim d_out // num_heads # Reduce the projection dim to match desired output dimself.W_query nn.Linear(d_in, d_out, biasqkv_bias)self.W_key nn.Linear(d_in, d_out, biasqkv_bias)self.W_value nn.Linear(d_in, d_out, biasqkv_bias)self.out_proj nn.Linear(d_out, d_out) # Linear layer to combine head outputsself.dropout nn.Dropout(dropout)self.register_buffer(mask,torch.triu(torch.ones(context_length, context_length),diagonal1))def forward(self, x):b, num_tokens, d_in x.shapekeys self.W_key(x) # Shape: (b, num_tokens, d_out)queries self.W_query(x)values self.W_value(x)# We implicitly split the matrix by adding a num_heads dimension# Unroll last dim: (b, num_tokens, d_out) - (b, num_tokens, num_heads, head_dim)keys keys.view(b, num_tokens, self.num_heads, self.head_dim) values values.view(b, num_tokens, self.num_heads, self.head_dim)queries queries.view(b, num_tokens, self.num_heads, self.head_dim)# Transpose: (b, num_tokens, num_heads, head_dim) - (b, num_heads, num_tokens, head_dim)keys keys.transpose(1, 2)queries queries.transpose(1, 2)values values.transpose(1, 2)# Compute scaled dot-product attention (aka self-attention) with a causal maskattn_scores queries keys.transpose(2, 3) # Dot product for each head# Original mask truncated to the number of tokens and converted to booleanmask_bool self.mask.bool()[:num_tokens, :num_tokens]# Use the mask to fill attention scoresattn_scores.masked_fill_(mask_bool, -torch.inf)attn_weights torch.softmax(attn_scores / keys.shape[-1]**0.5, dim-1)attn_weights self.dropout(attn_weights)# Shape: (b, num_tokens, num_heads, head_dim)context_vec (attn_weights values).transpose(1, 2) # Combine heads, where self.d_out self.num_heads * self.head_dimcontext_vec context_vec.contiguous().view(b, num_tokens, self.d_out)context_vec self.out_proj(context_vec) # optional projectionreturn context_vec不同于前者将不同的注意力头分开计算第二种方法直接扩展querykey和value矩阵的列数将多个矩阵运算简化为一个矩阵运算计算完再更改维度还原成一个个注意力头效率更高。这样我们就完成了一个完整的注意力机制。
