大模型面试准备（五）：图解 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_dim
        query = query.transpose(1, 2)
        key = key.transpose(1, 2)
        value = value.transpose(1, 2)

        # 计算注意力得分
        scores = torch.matmul(query, key.transpose(-2, -1)) / self.head_dim
        if 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