本文主要是介绍大模型面试准备（五）：图解 Transformer 最关键模块 MHA，希望对大家解决编程问题提供一定的参考价值，需要的开发者们随着小编来一起学习吧！
节前，我们组织了一场算法岗技术&面试讨论会，邀请了一些互联网大厂朋友、参加社招和校招面试的同学，针对大模型技术趋势、大模型落地项目经验分享、新手如何入门算法岗、该如何备战、面试常考点分享等热门话题进行了深入的讨论。
合集在这里：《大模型面试宝典》(2024版) 正式发布！
Transformer 原始论文中的模型结构如下图所示：
上一篇文章讲解了 Transformer 的关键模块 Positional Encoding（大家可以自行翻阅），本篇文章讲解一下 Transformer 的最重要模块 Multi-Head Attention（MHA），毕竟 Transformer 的论文名称就叫 《Attention Is All You Need》。
Transformer 中的 Multi-Head Attention 可以细分为3种，Multi-Head Self-Attention（对应上图左侧Multi-Head Attention模块），Multi-Head Cross-Attention（对应上图右上Multi-Head Attention模块），Masked Multi-Head Self-Attention（对应上图右下Masked Multi-Head Attention模块）。
其中 Self 和 Cross 的区分是对应的 Q和 K、 V是否来自相同的输入。是否Mask的区分是是否需要看见全部输入和预测的输出，Encoder需要看见全部的输入问题，所以不能Mask；而Decoder是预测输出，当前预测只能看见之前的全部预测，不能看见之后的预测，所以需要Mask。
本篇文章主要通过图解的方式对 Multi-Head Attention 的核心思想和计算过程做讲解，喜欢本文记得收藏、点赞、关注。技术和面试交流，文末加入我们
MHA核心思想
MHA过程图解
注意力计算公式如下：
图示过程图下：
多头注意力
MHA通过多个头的方式，可以增强自注意力机制聚合上下文信息的能力，以关注上下文的不同侧面，作用类似于CNN的多个卷积核。下面我们就通过一张图来完成MHA的解析：
单头注意力
知道了多头注意力的实现方式后，那如果是通过单头注意力完成同样的计算，矩阵形式是什么样的呢？下面我还是以一图胜千言的方式来回答这个问题：
通过单头注意力的比较，相信大家对多头注意力（MHA）应该有了更好的理解。我们可以发现多头注意力就是将一个单头进行了切分计算，最后又将结果进行了合并，整个过程中的整体维度和计算量基本是不变的，但提升了模型的学习能力。
最后附上一份MHA的实现和Transformer的构建代码：
import torch
import torch.nn as nn# 定义多头自注意力层
class MultiHeadAttention(nn.Module):def __init__(self, d_model, n_heads):super(MultiHeadAttention, self).__init__()self.n_heads = n_heads # 多头注意力的头数self.d_model = d_model # 输入维度（模型的总维度）self.head_dim = d_model // n_heads # 每个注意力头的维度assert self.head_dim * n_heads == d_model, "d_model必须能够被n_heads整除" # 断言，确保d_model可以被n_heads整除# 线性变换矩阵，用于将输入向量映射到查询、键和值空间self.wq = nn.Linear(d_model, d_model) # 查询（Query）的线性变换self.wk = nn.Linear(d_model, d_model) # 键（Key）的线性变换self.wv = nn.Linear(d_model, d_model) # 值（Value）的线性变换# 最终输出的线性变换，将多头注意力结果合并回原始维度self.fc_out = nn.Linear(d_model, d_model) # 输出的线性变换def forward(self, query, key, value, mask):# 将嵌入向量分成不同的头query = query.view(query.shape[0], -1, self.n_heads, self.head_dim)key = key.view(key.shape[0], -1, self.n_heads, self.head_dim)value = value.view(value.shape[0], -1, self.n_heads, self.head_dim)# 转置以获得维度 batch_size, self.n_heads, seq_len, self.head_dimquery = query.transpose(1, 2)key = key.transpose(1, 2)value = value.transpose(1, 2)# 计算注意力得分scores = torch.matmul(query, key.transpose(-2, -1)) / self.head_dimif mask is not None:scores = scores.masked_fill(mask == 0, -1e9)attention = torch.nn.functional.softmax(scores, dim=-1)out = torch.matmul(attention, value)# 重塑以恢复原始输入形状out = out.transpose(1, 2).contiguous().view(query.shape[0], -1, self.d_model)out = self.fc_out(out)return out# 定义Transformer编码器层
class TransformerEncoderLayer(nn.Module):def __init__(self, d_model, n_heads, dim_feedforward, dropout):super(TransformerEncoderLayer, self).__init__()# 多头自注意力层，接收d_model维度输入，使用n_heads个注意力头self.self_attn = MultiHeadAttention(d_model, n_heads)# 第一个全连接层，将d_model维度映射到dim_feedforward维度self.linear1 = nn.Linear(d_model, dim_feedforward)# 第二个全连接层，将dim_feedforward维度映射回d_model维度self.linear2 = nn.Linear(dim_feedforward, d_model)# 用于随机丢弃部分神经元，以减少过拟合self.dropout = nn.Dropout(dropout)# 第一个层归一化层，用于归一化第一个全连接层的输出self.norm1 = nn.LayerNorm(d_model)# 第二个层归一化层，用于归一化第二个全连接层的输出self.norm2 = nn.LayerNorm(d_model)def forward(self, src, src_mask):# 使用多头自注意力层处理输入src，同时提供src_mask以屏蔽不需要考虑的位置src2 = self.self_attn(src, src, src, src_mask)# 残差连接和丢弃：将自注意力层的输出与原始输入相加，并应用丢弃src = src + self.dropout(src2)# 应用第一个层归一化src = self.norm1(src)# 经过第一个全连接层，再经过激活函数ReLU，然后进行丢弃src2 = self.linear2(self.dropout(torch.nn.functional.relu(self.linear1(src))))# 残差连接和丢弃：将全连接层的输出与之前的输出相加，并再次应用丢弃src = src + self.dropout(src2)# 应用第二个层归一化src = self.norm2(src)# 返回编码器层的输出return src# 实例化模型
vocab_size = 10000 # 词汇表大小（根据实际情况调整）
d_model = 512 # 模型的维度
n_heads = 8 # 多头自注意力的头数
num_encoder_layers = 6 # 编码器层的数量
dim_feedforward = 2048 # 全连接层的隐藏层维度
max_seq_length = 100 # 最大序列长度
dropout = 0.1 # 丢弃率# 创建Transformer模型实例
model = Transformer(vocab_size, d_model, n_heads, num_encoder_layers, dim_feedforward, max_seq_length, dropout)
最后的最后再贴上一张非常不错的 Transformer 手绘吧！
技术交流群
前沿技术资讯、算法交流、求职内推、算法竞赛、面试交流(校招、社招、实习)等、与 10000+来自港科大、北大、清华、中科院、CMU、腾讯、百度等名校名企开发者互动交流~
我们建了算法岗技术与面试交流群， 想要进交流群、需要源码&资料、提升技术的同学，可以直接加微信号：mlc2040。加的时候备注一下：研究方向 +学校/公司+CSDN，即可。然后就可以拉你进群了。
方式①、微信搜索公众号：机器学习社区，后台回复：加群
方式②、添加微信号：mlc2040，备注：技术交流
用通俗易懂方式讲解系列
- 《大模型面试宝典》(2024版) 正式发布！
- 《大模型实战宝典》（2024版）正式发布！
- 大模型面试准备（一）：LLM主流结构和训练目标、构建流程
- 大模型面试准备（二）：LLM容易被忽略的Tokenizer与Embedding
- 大模型面试准备（三）：聊一聊大模型的幻觉问题
- 大模型面试准备（四）：大模型面试必会的位置编码（绝对位置编码sinusoidal，旋转位置编码RoPE，以及相对位置编码ALiBi）
参考文献：
参考资料：
[1] https://jalammar.github.io/illustrated-transformer/
[2] https://zhuanlan.zhihu.com/p/264468193
[3] https://zhuanlan.zhihu.com/p/662777298
这篇关于大模型面试准备（五）：图解 Transformer 最关键模块 MHA的文章就介绍到这儿，希望我们推荐的文章对编程师们有所帮助！
