大模型LLMs组件系列：深入浅出解析注意力Attention机制

1 简单self-Attention：注意力机制的原理与优势

1.1 最简单、单Query的、无Batch_size维度的Attention机制图解

如上图，假设我们的输入的batch_size=1、seq_len=3的token序列input_ids在经过embedding，前归一化层，位置编码后，以input_embed张量的形式进入Attention机制。在这个假设中，input_embed.shape=torch.Size([3, 8])，分别对应token序列长度和embedding空间维数hidden_size（直观起见没有batch_size维度），简化的代码如下。

inputs = "你好啊"
input_ids = tokenizer( inputs )               # torch.tensor( [8, 4, 7] )
input_embed = self.embedding( input_ids )     # torch.shape = torch.Size([3, 8])
'''
在经过前归一化层、位置编码后，input_embed = tensor([ 
  [ 0.23, 0.81, 0.45, 0.33, 0.44, 0.12, 0.17, 0.88 ],
  [ 0.90, 0.78, 0.32, 0.55, 0.28, 0.31, 0.49, 0.73 ],
  [ 0.12, 0.31, 0.56, 0.71, 0.02, 0.80, 0.34, 0.36 ],
])
'''

注意力机制的核心就是让所有token都用自己的状态向量与包括自己在内的所有token进行点积，相互打分。所以让input_embed中的三个长度为8的张量两两相乘？不行，太简单了，还会使得结果方阵对角线下一半是镜像。那么怎么办？

答：先用两个形状一样但值不一样的可学习权重矩阵对input_embed进行投影变换，再两两相乘，这下不仅在Attention机制中加入了神经网络权重参数，还解决了input_embed自乘带来的镜像问题。这两个权重矩阵就是投影矩阵q_proj和k_proj，为了消除量纲，还需要除以一个基于head_dim的缩放因子，简化代码如下：

self.q_proj = torch.nn.Linear( hidden_size, head_dim )
self.k_proj = torch.nn.Linear( hidden_size, head_dim )
# d_q和d_k分别是两个投影变换矩阵的输出维度，相等，均为head_dim，因为两者的输出结果要进行点积。再上图例子中，两者值都等于6。
query_states = self.q_proj( input_embed )
key_states = self.k_proj( input_embed )
attn_weights = torch.matmul(query_states, key_states.transpose(0, 2)) / math.sqrt(self.head_dim)
# attn_weights.shape = torch.Size([3, 3])

这样就够了吗？当然不行，这里的运算结果将是一个seq_len*seq_len的评分方阵，输入token的隐藏状态不见了，那么怎么办？

答：要想再次获得token的状态信息，最直接的方式是再用这个评分矩阵乘以输入input_embed，理论上也可以这样做，只不过是效果如何的问题。Transformer机制给出的方案是再来一个投影矩阵v_proj对input_embed进行变换，然后再用上述评分矩阵乘以这个变换结果，不论如何，在计算效率允许的范围内，加一个可学习层总共是利大于弊的吧？代码简化如下：

self.v_proj = torch.nn.Linear( hidden_size, d_v )
# d_v和head_dim（d_q和d_k）可以不一样，上图中就设置为4，但很多大模型都设为一样，可能是为了计算方便。
value_states = self.v_proj( input_embed )
attn_weights = nn.functional.softmax(attn_weights, dim=-1)
attn_output = torch.matmul(attn_weights, value_states