深度学习代码|Multi-Headed Attention (MHA)多头注意力机制的代码实现

相关文章

李沐《动手学深度学习》注意力机制

---

一、导入相关库

import math
from typing import Optional, List# 从 typing 模块中导入 Optional 和 List 类型，用于类型提示

import torch
from torch import nn

from labml import tracker #跟踪实验中的指标和损失值等信息

二、准备工作

（一）理论基础

在多头自注意力机制中，每个注意力头都需要一个查询（Query）、一个键（Key）和一个值（Value）向量。这些向量通过线性变换从输入特征中提取而来，然后用于计算注意力权重和加权求和。具体来说，给定输入特征张量 $X$，我们首先通过三个线性变换来计算查询$Q$、键$K$和值向量$V$：
$$\begin{gathered} Q=X\cdot W_Q \\ K=X\cdot W_K \\ V=X\cdot W_V \end{gathered}$$

其中，$W_Q$、$W_K$ 和 $W_V$ 是学习到的权重矩阵，用于将输入特征 $X$ 映射到查询、键和值向量的空间中。这些线性变换可以通过 PyTorch 中的 nn.Linear 层来实现。

（二）定义PrepareForMultiHeadAttention模块

该模块用于准备多头自注意力机制中的查询、键和值向量：

• 定义多头自注意力机制中的线性变换操作（在自注意力机制中，需要将输入的特征向量通过线性变换映射到不同的空间中，以便进行多头注意力的计算。）
• 将向量拆分为给定数量的头部，以获得多头注意。

class PrepareForMultiHeadAttention(nn.module):
	'''
	d_model：模型输入的特征维度；
	heads：注意力机制中的头数；
	d_k：每个头部中以向量表示的维度数；
	bias：是否使用偏置项
	'''
	def __init__(self,d_model:int,heads:int,d_k:int,bias:bool):
		super().__init__()
		#线性变换的线性层，输入为d_model，输出为heads*d_k
		self.linear=nn.Linear(d_model,heads*d_k,bias=bias) 
		self.heads=heads 
		self.d_k=d_k 
	
	'''
	x的形状为：seq_len,batch_size,d_model或batch_size,d_model
	输出形状：seq_len,batch_size,heads,d_k或batch_size,heads,d_k
	'''
	
	def forward(self,x:torch.Tensor):
		#获取输入张量 x 的形状，去掉最后一个维度，得到一个形状为 (seq_len, batch_size) 或 (batch_size,) 的元组
		head_shape=x.shape[:-1]
		x=self.linear(x)
		#将线性变换后的张量进行重塑操作，将最后一个维度拆分为heads个头部，每个头部的维度为d_k
		x=x.view(*head_shape,self.heads,self.d_k)
		return x

pytorch中的view方法：用于对张量进行重塑（reshape）。其作用是将张量的形状变换为指定的形状，但是要求变换后的形状与原始形状的元素数量保持一致。

* 表示解包（unpacking）操作符。在函数调用或函数定义中，*args 表示将传入的参数打包成一个元组，而在函数调用时，*args 则表示将元组解包为独立的参数。

三、多头注意模块

（一）理论基础

通过计算查询$Q$和键$K$之间的点积，并应用缩放因子 $\frac{1}{\sqrt{d_k}}$，得到注意力权重：
$$\alpha =softmax(\frac{Q{K}^T}{\sqrt{d_k}})$$

利用注意力权重对值向量$V$进行加权求和，得到最终的输出向量
$$Attention(Q,K,V)=\alpha V=softmax(\frac{Q{K}^T}{\sqrt{d_k}})V$$
这样，对于每个注意力头，都可以得到一个输出向量，在多头自注意力机制中，会并行地进行多个这样的注意力头的计算，最后将他们的输出向量连接起来，形成最终的输出。

（二）创建MultiHeadAttention模块

• __init__：初始化对象的属性
• get_score：计算查询和键之间的分数$S_{ijbh}={\sum }_d{Q}_{ibhd}{K}_{jbhd}$
• prepare_mask：对注意力遮罩进行处理，使其与后续计算注意力权重的张量形状相匹配
• forward：前向计算过程

注意力遮罩（Attention Mask）用于指示哪些位置的信息是有效的，哪些位置是无效的。在自注意力机制中，有时需要对注意力权重进行调整，以便在计算注意力时忽略某些位置的信息，或者对某些位置的信息赋予特定的权重。
注意力遮罩主要应用在查询和键之间的相似度计算过程中，用于调整或者限制查询和键之间的关系。

class MultiHeadAttention(nn.Module):
	'''
	heads:头的数量。
	d_model:query 、key 和value 向量中的要素数
	d_k:每头特征数
	
	'''
	def __init__(self,heads:int,d_model:int,dropout_prob:float=0.1,bias=True):
		super().__init__()
		self.d_k=d_model // heads #计算每个头部的查询、键和值的维度
		self.heads=heads
		
		self.query=PrepareForMultiHeadAttention(d_model,heads,self.d_k,bias=bias)
		self.key=PrepareForeMultiHeadAttention(d_model,heads,self.d_k,bias=bias)
		self.value=PrepareForMultiHeadAttention(d_model,heads,self.d_k,bias=True)
		
		#创建一个 Softmax 层，用于计算注意力权重，dim=1 表示在时间维度上进行 Softmax 计算。
		self.softmax=nn.Softmax(dim=1)
		#创建一个 Dropout 层，用于在训练过程中进行随机失活
		self.dropout=nn.Dropout(dropout_prob)
		#计算缩放因子，这里使用了倒数的平方根进行缩放
		self.scale=1/math.sqrt(self.d_k)
		#初始化一个属性 attn，用于存储注意力权重
		self.attn=None
	
	def get_scores(self,query:torch.Tensor,key:torch.Tensor):
		return torch.einsum('ibhd,jbhd -> ijbh',query,key)
	
	'''
	mask: 输入的注意力遮罩，形状为(seq_len_q, seq_len_k, batch_size)。
	query_shape: 查询张量的形状，其中包含序列长度和批次大小。
	key_shape: 键张量的形状，其中包含序列长度和批次大小。
	'''
	def prepare_mask(self,mask:torch.Tensor,query_shape:List[int],key_shape:List[int]):
		
		#确保输入的注意力遮罩的第一个维度的大小与查询张量的序列长度维度大小匹配或者为1
		assert mask.shape[0] == 1 or mask.shape[0] == query_shape[0]
		#确保输入的注意力遮罩的第二个维度大小与键张量的序列长度维度大小匹配
		assert mask.shape[1] == key_shape[0]
		#确保输入的注意力遮罩的第三个维度的大小与查询张量的批次大小维度大小匹配或者为1
		assert mask.shape[2] == 1 or mask.shape[2] == query_shape[1]
		
		#将注意力遮罩的最后一个维度扩展一个维度
		#形状从(seq_len_q, seq_len_k, batch_size)扩展为(seq_len_q, seq_len_k, batch_size, heads)
		mask = mask.unsqueeze(-1)
		return mask

	'''
	query、key、value的形状：[seq_len, batch_size, d_model]
	mask的形状：[seq_len, seq_len, batch_size] 
	'''
	def forward(self,*,
						query:torch.Tensor,
						key:torch.Tensor,
						value:torch.Tensor,
						mask:Optional[torch.Tensor] = None):
						
		seq_len,batch_size,_ = query.shape
		if mask is not None:
			mask=self.prepare_mask(mask,query.shape,key.shape)
		
		#query、key、value的形状经过处理后变为：[seq_len,batch_size,heads,d_k]
		query=self.query(query)
		key=self.key(key)
		value=self.value(value)
		
		#计算注意力分数
		scores=self.get_scores(query,key)
		#应用缩放因子
		scores *= self.scale

		#应用mask
		if mask is not None:
			#将遮罩中值为 0 的位置替换为负无穷，这样在计算 Softmax 时对应位置的注意力权重将为 0
			scores=scores.masked_fill(mask==0,float('-inf'))
		#计算注意力权重
		attn=self.softmax(scores)
		#在调试过程中输出注意力权重 attn 的信息
		tracker.debug('attn',attn)
		#应用dropout
		attn=self.dropout(attn)
		#加权求和
		x=torch.einsum("ijbh,jbhd->ibhd",attn,value)
		
		#将注意力权重从计算图中分离出来，以避免在反向传播过程中对注意力权重的梯度进行更新
		self.attn=attn.detach()
		
		#连接多个头：将经过加权求和得到的向量按照特定形状重新排列
		x=x.reshape(seq_len,batch_size,-1)
		return self.output(x)

参考：https://github.com/labmlai/annotated_deep_learning_paper_implementations?tab=readme-ov-file