Multi-head Self-Attention Machanism

3. 多头自注意力机制Multi-head Self-Attention Machanism多头注意力机制是在自注意力机制的基础上发展起来的是自注意力机制的变体旨在增强模型的表达能力和泛化能力。它通过使用多个独立的注意力头分别计算注意力权重并将它们的结果进行拼接或加权求和从而获得更丰富的表示。在自注意力机制中每个单词或者字都仅仅只有一个 q、k、v与其对应如下图所示为什么需要多个head在做self-attention的时候是用q去找相关的k但是相关有很多不同形式和定义不同q负责不同种类的相关性。多头注意力机制则是在 ai 乘以一个q,k,v后会再分配多个 q,k,v这里以2个 q,k,v 为例如下图所示:3.1q、k操作在多头注意力机制中ai会先乘q矩阵qiWq*ai;其次会为其多分配两个head以q为例包括qi,1,qi,2;同样地k和 v 也是一样的操作。那么下面就是 q和 k 的点乘操作了在多头注意力机制中有多个q 和 k究竟应该选择哪个进行操作呢其实很简单就是看下标如下图所示。qi,1,会和 ki,1和kj,1进行点乘,再进行softmax操作3.2v操作将对应的attention分数乘对应的v 比如head1的qi,1和 ki,1和kj,1进行点乘dot_product后得到相应的注意力分数归一化后分别和 vi,1和vj,1相乘再求和得到注意力输出head2也是同时这样进行的得到bi,1和bi,2.3.3注意力输出将head1和head2得到的输出拼接起来进行transpose就得到了注意力输出bi3.4single-head与multi-head对比single-head-attention中从每个a中提取出q、k、v,他们的维度为d-model每个token的维度multi-head-attention,与singlehead一样先得出q、k、v,不同的是每个q、k、v还要分出多个head将多头得到的bii先按照列再按行con-cat可以得到mxd_model的矩阵m为token个数d_model为每个token的维度3.5流程总结3.6代码实现import torch from torch import nn as nn #dim_in:输入中每个token的维度也就是输入x的最后一个维度 #d_model:single-head-attention情况下q、k、v总的向量长度 #num_heads:head个数 class MultiHead_SelfAttention(nn.Module): def __init__(self,input_dim,num_heads): super().__init__() self.num_headsnum_heads self.head_diminput_dim//num_heads#head的维度为输入维度除以head个数方便后面拼接 assert input_dim%num_heads0 ,Input dimension must be divisible by the number of heads. ​ # Linear layers for the query, key, and value projections for each head self.querynn.Linear(input_dim,input_dim) self.keynn.Linear(input_dim,input_dim) self.valuenn.Linear(input_dim,input_dim) ​ self.output_linearnn.Linear(input_dim,input_dim) ​ def forward(self,x): batch_size,seq_len,input_dimx.size() #输入数据shape[batch_size,token个数token长度] #将输出向量经过矩阵乘法后拆分为多头 queryself.query(x).view(batch_size,seq_len,self.num_heads,self.head_dim) #输入数据shape[batch_size,token个数head数head维度] keyself.query(x).view(batch_size, seq_len, self.num_heads, self.head_dim) valueself.query(x).view(batch_size, seq_len, self.num_heads, self.head_dim) #对调序列的长度和head个数(batch_size, seq_len, num_heads, head_dim) to (batch_size, num_heads, seq_len, head_dim) #方便后续矩阵乘法和不同头部的注意力计算 queryquery.transpose(1,2)#(batch_size, num_heads, seq_len, head_dim) keykey.transpose(1,2) valuevalue.transpose(1,2) #计算注意力分数和权重matmul:最后两个维度做矩阵乘法 attention_scorestorch.matmul(query,key.transpose(-2,-1))/torch.sqrt(torch.tensor(self.head_dim,dtypetorch.float)) #query:(batch_size, num_heads, seq_len, head_dim) * key(batch_size,num_heads,head_dim,seq_len) #attention_scores:(batch_size, num_heads, seq_len, seq_len) attention_weightstorch.softmax(attention_scores,dim-1) #注意力加权求和 attentiontorch.matmul(attention_weights,value) # attention_scores:(batch_size, num_heads, seq_len, seq_len)* value(batch_size, num_heads, seq_len, head_dim) #attention:(batch_size, num_heads, seq_len, head_dim) #连接和线性变换 attentionattention.transpose(1,2).contiguous().view(batch_size,seq_len,input_dim)#contiguos深拷贝不改变原数据 #(batch_size,num_heads, seq_len , head_dim) to (batch_size, seq_len, num_heads, head_dim) to (batch_size,seq_len,input_dim) outputself.output_linear(attention) ​ return output ​ #定义多头自注意力机制模型 class MultiHead_SelfAttention_Classifier(nn.Module): def __init__(self,input_dim,num_heads,hidden_dim,num_classes): super().__init__() self.attentionMultiHead_SelfAttention(input_dim,num_heads) self.fc1nn.Linear(input_dim,hidden_dim) self.fc2nn.Linear(hidden_dim,num_classes) self.relunn.ReLU() def forward(self,x): xself.attention(x) xx.mean(dim1)#(batch_size, seq_len, input_dim) to (batch_size, input_dim) xself.fc1(x) xself.relu(x) xself.fc2(x) return x ​ if __name____main__: modelMultiHead_SelfAttention_Classifier(input_dim4,num_heads2,hidden_dim32,num_classes2) xtorch.randn((2,4,4))#batch4个token每个token长度input_dim outputmodel(x) print(output) ​输出tensor([[-0.0598, -0.0333], [-0.2237, -0.0573]], grad_fnAddmmBackward0)