掘金 人工智能 09月17日
深入解析Transformer架构:原理与PyTorch源码实现
index_new5.html
../../../zaker_core/zaker_tpl_static/wap/tpl_guoji1.html

 

本文深入剖析了Transformer模型的核心架构,从整体的编码器-解码器设计,到关键的位置编码机制,再到PyTorch源码中的具体实现细节。文章详细阐述了位置编码如何解决Transformer对位置信息不敏感的问题,并解释了其确定性、相对关系一致性和泛化能力。在PyTorch源码解析部分,作者详细讲解了Transformer类及其构成组件TransformerEncoderLayer和TransformerDecoderLayer的初始化与前向传播流程,并重点分析了多头自注意力机制、前馈神经网络以及层归一化等核心模块的实现。此外,文章还特别强调了注意力机制的计算原理,包括Scaled Dot-Product Attention和Multi-head Attention,并说明了在编码器和解码器中QKV的不同来源。最后,文章还拓展了Transformer与Mixture-of-Experts模型的比较。

💡 Transformer模型的核心设计在于其编码器-解码器架构,其中编码器负责理解源序列的语义,解码器则利用编码器输出和已生成序列来预测目标序列。这种分离使得模型能够并行处理输入(编码器)并在生成过程中保持自回归(解码器),是其高效性的关键。

🌟 位置编码是Transformer解决序列顺序信息缺失的创新之处。通过将位置信息以正弦和余弦函数的特定组合添加到词嵌入中,模型得以区分不同位置的词语,并具备处理任意长度序列的泛化能力。残差连接机制确保了位置信息在深层网络中得以保留。

⚙️ PyTorch的Transformer实现通过模块化设计,将编码器和解码器分解为多个TransformerEncoderLayer和TransformerDecoderLayer。每个Layer包含多头自注意力、前馈神经网络及层归一化等组件,通过串联这些Layer,构建了完整的Transformer模型,为研究者提供了灵活的接口。

🧠 注意力机制是Transformer的灵魂。Scaled Dot-Product Attention通过计算Query与Key的相似度来加权Value,实现了模型对输入序列中重要部分的聚焦。多头自注意力机制则通过并行计算多个注意力头,从不同表示子空间捕获更丰富的上下文信息,极大地提升了模型的表达能力。

本文较长,建议点赞收藏,以免遗失。更多AI大模型应用开发及AI大模型算法学习视频及资料,尽在聚客AI学院

在人工智能领域,Transformer架构无疑是大模型发展史上最重要的里程碑之一。它不仅构成了当前大模型处理任务的基础架构,更是深入理解现代大模型系统的关键。今天我将通过结合论文原理与PyTorch源码API,深度解析Transformer的设计思路与实现细节。如有遗漏,欢迎交流。

一、整体架构设计

Transformer模型分为两个主要部分:左侧的编码器(Encoder)和右侧的解码器(Decoder)。

编码器负责接受完整的源序列输入,将其转换为富含语义信息的表示序列。以机器翻译任务为例,编码器的作用类似于深度理解原文的专家,需要充分把握整个句子的含义、语法结构和上下文关系。

解码器承担着更为复杂的任务:它需要同时接受目标序列和编码器输出的表示序列,然后输出词汇/字符的概率分布。这好比翻译专家既要理解原文含义(通过编码器输出),又要根据已翻译内容决定下一个合适的词汇。

二、位置编码机制

Transformer模型本身对位置信息不敏感。例如"我爱你"和"你爱我"这两个句子,在没有位置信息的情况下,模型无法识别它们是语义完全不同的表达。这就像人类失去对词语顺序的感知能力,显然无法正确理解语言。

因此,需要引入带有位置信息的向量,将其添加到每个input embedding上,使不同位置获得不同的表征。这就是Positional Encoding模块的作用。

在设计位置编码时,遵循三个重要假设:

    确定性原则:每个位置的编码应该是确定的数字,不同序列中相同位置的编码必须一致。如果采用等分设计方法,将序列在0~1之间做均匀划分,那么序列长度不同时每个位置上的编码也会不同,这将违反确定性原则。相对关系一致性:不同句子中,任意两个位置之间的相对距离关系应该保持一致。这一设计目的是让模型学习通用的语言关系,例如"修饰词位于被修饰词前一个位置"这种通用模式。泛化能力:位置编码需要能够推广到训练时未见过的更长序列。即使测试集中的句子长度超过训练时的最大长度,模型仍能通过位置编码进行处理。

基于这些假设,Transformer采用正弦和余弦函数的组合来表征绝对位置信息:

这种设计的关键优势在于:pe(pos+k)可以表示为pe(k)的线性组合(利用三角函数公式sin(A+B)=sin(A)cos(B)+cos(A)sin(B))。这意味着即使测试集中出现pos+k这种未见过的位置,也可以表示为训练集中已见过位置的线性组合,从而保证了对长句子的推广能力。

关于位置信息在深层网络中可能丢失的担忧,通过残差连接机制得到解决。假设有N层神经网络,输入x₀包含位置编码,那么:

第1层: x₁ = x₀ + F₁(x₀)

第2层: x₂ = x₁ + F₂(x₁) = x₀ + F₁(x₀) + F₂(x₁)

...

第N层: xₙ = x₀ + Σᵢ₌₁ⁿ Fᵢ(xᵢ₋₁)

初始位置信息x₀始终存在于每一层的输出中,确保了位置信息不会随网络深度增加而消失。

三、PyTorch实现解析

PyTorch中的Transformer实现位于/pytorch/torch/nn/modules/transformer.py,本文基于v2.5.1版本进行分析。

PyTorch定义了顶层的Transformer类,可通过torch.nn.Transformer调用:

transformer_model = nn.Transformer(d_model=512, nhead=8, num_encoder_layers=6)src = torch.rand((10, 32, 512))tgt = torch.rand((20, 32, 512))out = transformer_model(src, tgt)

在Transformer的__init__函数中,包含5个核心参数:

    d_model:整个Transformer的特征维度,原论文设置为512。这一维度需要足够大以承载丰富语义信息,同时不能过大导致计算复杂度过高。nhead:Multi-head attention的头数。多头设计使模型能够捕捉更多位置与位置之间的关系。num_encoder_layers:编码器encoder的block数量,每个block包含多头自注意力机制和前馈神经网络,默认6个。num_decoder_layers:解码器decoder的block数量,每个block包含多头自注意力机制、交叉注意力机制以及前馈神经网络。dim_feedforward:前馈神经网络层中间的特征维度,默认2048。Multihead attention输出首先映射到2048维特征空间,再映射回512维空间,以保证输出维度与输入一致,便于残差连接。

init函数的作用是实例化模块,首先实例化encoder模块。

encoder通过TransformerEncoder类实现实例化,需要传入encoder_layer参数。TransformerEncoderLayer类实现了Multihead self attention的调用、残差连接、层归一化、全连接网络,这些组件共同构成一个encoder_layer。

解码器部分同样需要传入decodeLayer参数,该layer包含自注意力机制、交叉注意力机制以及前馈神经网络。

总体而言,Transformer源码由四个核心class构成:

在forward函数中,Transformer的计算流程清晰明确:

编码器输入源句子及padding_mask。编码器中的注意力机制不需要因果掩码,但需要处理样本长度差异。训练时序列长度不一致,短样本的后续位置无效,通过softmax中将无效位置值转为负无穷,使这些位置概率归零。

解码器输入包括:目标句子target、编码器输出memory、目标掩码target mask和内存掩码memory mask。target mask是考虑因果关系的上三角矩阵,确保输出每次只根据当前预测词汇左侧的所有词汇进行预测,而不使用右侧词汇信息。

Transformer本质上是自回归的解码过程,非并行预测输出,而是每次预测一个输出,通过不断解码预测完整目标句子。

四、编码器实现细节

单个编码器的实现集中在TransformerEncoderLayer类中。在__init__函数中,需要实例化四个关键组件:

init函数参数与transformer中一致:d_model为512维特征维度,nhead为多头注意力头数。

多头设计目的是让模型捕捉更多位置间关系。多头分为多组query、key和value,每组独立计算attention上下文向量,最后拼接这些向量并通过FFN得到最终输出。原始512维embedding分为8个头后,每个头向量维度变为64。每个头独立计算注意力后,得到8个64维输出向量,通过拼接操作重新组合成512维向量,再经输出线性变换层得到最终512维输出。

dim_feedforward是前馈神经网络FFN的维度,设置两个全连接层实现从512到2048再到512的变换。前馈神经网络对每个位置独立建模,不同位置参数共享,类似于1×1的pointwise卷积,对图像中每个像素位置的特征向量独立进行变换。参数共享旨在让模型学会"如何处理特征"的通用能力,而非"如何处理第x个位置特征"的特定能力。

FFN实现embedding相同位置不同维度间的融合,注意力机制负责位置间信息交流。

init函数中实例化的组件包括:

    Multi-head attention模块FFN前馈神经网络中的两个Linear层层归一化layer normDropout模块,提高网络泛化能力

forward函数中,编码器层的调用流程为:首先通过self attention block得到表征,加上残差连接后经层归一化处理;然后将输出送入feed forward block,经过每个位置独立的全连接层,再次加上残差连接后经层归一化得到最终输出。

原始论文设计采用后归一化方式。

TransformerEncoder类的作用是将多个编码器层串联起来,将上一层输出作为下一层输入,经多层处理得到最终编码器输出。

init函数主要传入两个参数:encoder_layer(TransformerEncoderLayer实例)和num_layers(transformer encoder层数)。

五、解码器实现细节

解码器实现比编码器更复杂,包含三个子模块,需要处理更多交互。

在TransformerDecoderLayer中,需要实例化三套组件(自注意力+交叉注意力+前馈神经网络):

init参数中,d_model表示模型特征大小(默认512),nhead为解码器多头自注意力机制头数,dim_feedforward为解码器中FFN前馈神经网络维度。

解码器与编码器的不同之处在于需要实例化两个Multihead attention模块:

    自注意力机制:处理解码器输入序列(目标句子embedding)的自身表征交叉注意力机制:建立解码器多头注意力输出与编码器输出状态的关联性,跨越encoder和decoder两个不同序列

同样,解码器实现两个Linear层,将交叉自注意力机制输出投射到2048维空间后再降维至512维。由于解码器有三个模块,需要实现3个归一化和3个dropout层。

解码器的forward函数体现三个模块的协同工作:

    目标句子序列x和target mask输入到self._sa_block,对目标句子做自注意力计算,结果经残差网络和层归一化上一模块输出与编码器输出memory进行交叉注意力计算,新表征经残差网络和归一化输出上一模块输出送入FFN前馈神经网络,再次经残差网络和归一化得到解码器最终输出

_sa_block和_mha_block均调用Multihead attention,但query、key、value来源不同。

自注意力机制中query、key、value均来自目标序列,计算自身相关性。交叉注意力机制中,query来自解码器输出,key和value来自编码器输出(memory)。

多个DecoderLayer构成TransformerDecoder,实现方式与TransformerEncoder类似。

六、注意力机制核心

PyTorch的实际实现进行了更多优化,但核心思想与论文版本一致:

直观理解,attention函数建立query与key-value对之间的连接,最终产生输出。以百度搜索为例:搜索词条为query,数据库中的词条信息为key,具体内容为value。通过query,百度返回搜索结果,这就是注意力机制基于query和key+value计算上下文的过程。

注意力计算结果是对Value的加权求和,权重基于Query和Key的相似度计算。先计算Query与每个Key的相似度,经Softmax归一化得到权重,再将权重与对应Value加权求和。

Transformer采用"Scaled Dot-Product Attention",公式中QK相乘后除以√d_k,目的是使Softmax输入分布更稳定,方差更小。

注意力计算由Q、K、V三个向量构成。首先进行query和key的矩阵相乘(单样本为向量内积,批量为矩阵相乘),然后除以√d_K进行缩放,得到每个位置上的内积结果后进行归一化处理。Softmax输出结果总和为1,每个值在0-1之间,表示概率分布。将此概率与每个位置value加权求和,得到attention输出。

Multi-head self attention同时计算多个自注意力机制,将每个结果拼接后得到最终输出。

代码实现中,注意力计算输入包括query、key和value。

首先将q与k的转置进行矩阵相乘,得到向量后添加mask。mask将等于0的位置填充为极小的负数(负无穷),经Softmax归一化后概率变为0,使不重要位置概率赋值为0。得到注意力概率分布后,与value进行加权求和,得到自注意力机制输出。多头注意力将多个单头输出拼接为最终输出。

在Transformer中,不同注意力机制的QKV来源和映射方式不同:

    编码器层:QKV均由word embedding加position encoding后,通过三个独立线性映射得到解码器自注意力层:同样由目标句子embedding加position encoding后,经三个独立线性映射得到QKV交叉注意力层:query来自解码器输出的线性映射,key和value分别来自编码器输出memory的两个不同映射

本期拓展知识点

那么,你知道在大语言模型中,Transformer与专家混合模型之间的区别吗?感兴趣的粉丝朋友可以查阅我之前整理的技术文档:《Transformer与专家混合模型的比较》

好了,本次的分享也花了不少时间去整理,如果对你有所帮助,记得给个小红心,我们下期见。

Fish AI Reader

Fish AI Reader

AI辅助创作,多种专业模板,深度分析,高质量内容生成。从观点提取到深度思考,FishAI为您提供全方位的创作支持。新版本引入自定义参数,让您的创作更加个性化和精准。

FishAI

FishAI

鱼阅,AI 时代的下一个智能信息助手,助你摆脱信息焦虑

联系邮箱 441953276@qq.com

相关标签

Transformer AI大模型 深度学习 PyTorch 自然语言处理 位置编码 注意力机制 Encoder-Decoder NLP Transformer Architecture AI Machine Learning Deep Learning Positional Encoding Attention Mechanism
相关文章