从原理到实践：0基础实现用深度学习模型框架搭建Transformer模型

引言

2017年，Vaswani等人在《Attention Is All You Need》中提出的Transformer架构，彻底改变了自然语言处理（NLP）领域的技术范式。与传统RNN、LSTM等时序模型不同，Transformer完全基于自注意力机制，实现了并行化计算，大幅提升了训练效率，同时在长序列建模能力上也有显著优势。如今，BERT、GPT、LLaMA等主流大模型均以Transformer为核心架构。

本文将从底层原理出发，用PyTorch逐模块实现Transformer，不依赖nn.Transformer现成接口，深入解析每个组件的设计逻辑与数学原理，帮助读者真正理解Transformer的内部运作机制。

一、Transformer核心原理概述

Transformer架构分为编码器（Encoder） 和解码器（Decoder） 两部分，核心设计围绕以下3个关键点：

  1. 自注意力机制：捕捉序列内任意两个位置的依赖关系，实现“全局视野”
  2. 位置编码：弥补无时序结构的缺陷，注入词的位置信息
  3. 残差连接+层归一化：解决深层网络训练梯度消失问题，稳定训练过程

整体架构如下图所示（来源：原论文）：
在这里插入图片描述

二、基础模块实现与原理剖析

2.1 位置编码（Positional Encoding）

原理说明

Transformer无时序结构，无法像RNN那样通过迭代顺序获取位置信息。因此需要手动注入位置编码，核心要求：

  • 位置编码需与词嵌入维度一致（便于直接相加）
  • 能表达相对位置关系（即 P E p o s + k PE_{pos+k} PEpos+k​可由 P E p o s PE_{pos} PEpos​推导）
  • 对任意长度序列通用（支持训练外的长序列）

原论文采用正弦-余弦交替编码，公式如下：
P E ( p o s , 2 i ) = sin ⁡ ( p o s 1000 0 2 i / d m o d e l ) PE_{(pos, 2i)} = \sin\left( \frac{pos}{10000^{2i/d_{model}}} \right) PE(pos,2i)​=sin(100002i/dmodel​pos​)
P E ( p o s , 2 i + 1 ) = cos ⁡ ( p o s 1000 0 2 i / d m o d e l ) PE_{(pos, 2i+1)} = \cos\left( \frac{pos}{10000^{2i/d_{model}}} \right) PE(pos,2i+1)​=cos(100002i/dmodel​pos​)
其中：

  • p o s pos pos：词在序列中的位置（0,1,2,…）
  • i i i：词嵌入的维度索引（0,1,…,d_model/2-1）
  • d m o d e l d_{model} dmodel​：词嵌入维度（通常为512）

该设计的优势：

  1. 正弦函数的周期性天然支持相对位置（ sin ⁡ ( a + b ) = sin ⁡ a cos ⁡ b + cos ⁡ a sin ⁡ b \sin(a+b)=\sin a\cos b+\cos a\sin b sin(a+b)=sinacosb+cosasinb）
  2. 无需训练，固定计算，减少参数开销
  3. 可外推到训练时未见过的长序列
代码实现
import torch
import torch.nn as nn
import math

class PositionalEncoding(nn.Module):
    def __init__(self, d_model: int, max_len: int = 5000, dropout: float = 0.1):
        super().__init__()
        # 初始化位置编码矩阵：[max_len, d_model]
        pe = torch.zeros(max_len, d_model)
        # 生成位置序列：[max_len, 1]
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        # 计算分母项：10000^(2i/d_model) → 转换为指数形式避免大数计算
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        
        # 偶数维度用正弦，奇数维度用余弦
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        
        # 扩展为[1, max_len, d_model]，适配批量数据（batch_size维度）
        pe = pe.unsqueeze(0)
        # 注册为缓冲区（不参与梯度更新）
        self.register_buffer('pe', pe)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: 词嵌入张量，shape=[batch_size, seq_len, d_model]
        Returns:
            注入位置信息后的词嵌入，shape不变
        """
        # 只取与输入序列长度匹配的位置编码
        x = x + self.pe[:, :x.size(1), :]
        return self.dropout(x)

2.2 缩放点积注意力（Scaled Dot-Product Attention）

原理说明

注意力机制的本质是“从查询（Query）出发，对键（Key）-值（Value）对进行加权求和”，核心公式如下：
Attention ( Q , K , V ) = softmax ( Q K T d k ) V \text{Attention}(Q, K, V) = \text{softmax}\left( \frac{QK^T}{\sqrt{d_k}} \right)V Attention(Q,K,V)=softmax(dk​ ​QKT​)V
其中：

  • Q Q Q：查询矩阵，shape=[batch_size, n_heads, seq_len_q, d_k]
  • K K K：键矩阵，shape=[batch_size, n_heads, seq_len_k, d_k]
  • V V V：值矩阵，shape=[batch_size, n_heads, seq_len_v, d_v]（通常seq_len_k=seq_len_v）
  • d k d_k dk​：单个注意力头的维度（ d k = d m o d e l / n h e a d s d_k = d_{model}/n_{heads} dk​=dmodel​/nheads​）

关键设计解析：

  1. 缩放因子 d k \sqrt{d_k} dk​ ​：当 d k d_k dk​增大时， Q K T QK^T QKT的元素值会增大，导致softmax后梯度趋近于0（梯度消失），缩放后可缓解该问题
  2. Softmax激活：将注意力权重归一化到[0,1]，确保权重之和为1
  3. 掩码（Mask）：用于屏蔽无效位置（如Padding）或未来信息（解码器），通过将无效位置的注意力得分设为 − ∞ -\infty −∞，使softmax后权重为0
代码实现
class ScaledDotProductAttention(nn.Module):
    def __init__(self, dropout: float = 0.1):
        super().__init__()
        self.dropout = nn.Dropout(dropout)

    def forward(
        self, 
        q: torch.Tensor, 
        k: torch.Tensor, 
        v: torch.Tensor, 
        mask: torch.Tensor = None
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """
        Args:
            q: 查询矩阵，shape=[batch_size, n_heads, seq_len_q, d_k]
            k: 键矩阵，shape=[batch_size, n_heads, seq_len_k, d_k]
            v: 值矩阵，shape=[batch_size, n_heads, seq_len_v, d_v]
            mask: 掩码矩阵，shape=[batch_size, 1, seq_len_q, seq_len_k]（1为n_heads维度广播）
        Returns:
            attn_output: 注意力输出，shape=[batch_size, n_heads, seq_len_q, d_v]
            attn_weights: 注意力权重，shape=[batch_size, n_heads, seq_len_q, seq_len_k]
        """
        d_k = q.size(-1)
        # 1. 计算注意力得分：Q*K^T / sqrt(d_k)
        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(d_k)  # [batch, n_heads, seq_q, seq_k]
        
        # 2. 应用掩码（无效位置设为-1e9）
        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)
        
        # 3. 计算注意力权重并 dropout
        attn_weights = torch.softmax(scores, dim=-1)  # [batch, n_heads, seq_q, seq_k]
        attn_weights = self.dropout(attn_weights)
        
        # 4. 加权求和得到注意力输出
        attn_output = torch.matmul(attn_weights, v)  # [batch, n_heads, seq_q, d_v]
        
        return attn_output, attn_weights

2.3 多头注意力（Multi-Head Attention）

原理说明

单头注意力只能捕捉单一维度的依赖关系，多头注意力通过将 Q , K , V Q, K, V Q,K,V投影到多个子空间（即“多头”），并行计算注意力后拼接，实现“多角度捕捉语义关联”。核心公式如下：
MultiHead ( Q , K , V ) = Concat ( head 1 , head 2 , . . . , head h ) W O \text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, \text{head}_2, ..., \text{head}_h)W^O MultiHead(Q,K,V)=Concat(head1​,head2​,...,headh​)WO
其中：
head i = Attention ( Q W i Q , K W i K , V W i V ) \text{head}_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V) headi​=Attention(QWiQ​,KWiK​,VWiV​)

  • W i Q , W i K , W i V W_i^Q, W_i^K, W_i^V WiQ​,WiK​,WiV​：投影矩阵，维度均为 [ d m o d e l , d k ] [d_{model}, d_k] [dmodel​,dk​]
  • W O W^O WO：输出投影矩阵，维度为 [ h ⋅ d k , d m o d e l ] [h \cdot d_k, d_{model}] [h⋅dk​,dmodel​]（ h h h为头数）
  • 约束： d m o d e l = h ⋅ d k d_{model} = h \cdot d_k dmodel​=h⋅dk​（确保拼接后维度与输入一致）
代码实现
class MultiHeadAttention(nn.Module):
    def __init__(self, d_model: int, n_heads: int, dropout: float = 0.1):
        super().__init__()
        # 校验：d_model必须能被n_heads整除
        assert d_model % n_heads == 0, f"d_model={d_model} must be divisible by n_heads={n_heads}"
        
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_k = d_model // n_heads  # 单个头的维度
        
        # 定义Q、K、V的投影矩阵（共享参数，而非每个头独立）
        self.w_q = nn.Linear(d_model, d_model)  # [d_model, d_model]
        self.w_k = nn.Linear(d_model, d_model)
        self.w_v = nn.Linear(d_model, d_model)
        
        # 定义输出投影矩阵
        self.w_o = nn.Linear(d_model, d_model)
        
        # 引入缩放点积注意力
        self.attention = ScaledDotProductAttention(dropout)

    def forward(
        self, 
        q: torch.Tensor, 
        k: torch.Tensor, 
        v: torch.Tensor, 
        mask: torch.Tensor = None
    ) -> torch.Tensor:
        """
        Args:
            q: 输入查询，shape=[batch_size, seq_len_q, d_model]
            k: 输入键，shape=[batch_size, seq_len_k, d_model]
            v: 输入值，shape=[batch_size, seq_len_v, d_model]
            mask: 掩码矩阵，shape=[batch_size, 1, seq_len_q, seq_len_k]
        Returns:
            多头注意力输出，shape=[batch_size, seq_len_q, d_model]
        """
        batch_size = q.size(0)
        
        # 1. 投影：Q、K、V分别通过线性层 → [batch, seq_len, d_model]
        q_proj = self.w_q(q)
        k_proj = self.w_k(k)
        v_proj = self.w_v(v)
        
        # 2. 分拆多头：将d_model维度拆分为n_heads×d_k，并调整维度顺序
        # 原shape: [batch, seq_len, d_model] → 拆分后: [batch, seq_len, n_heads, d_k] → 转置后: [batch, n_heads, seq_len, d_k]
        q_proj = q_proj.view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)  # [batch, n_heads, seq_q, d_k]
        k_proj = k_proj.view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)  # [batch, n_heads, seq_k, d_k]
        v_proj = v_proj.view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)  # [batch, n_heads, seq_v, d_k]
        
        # 3. 计算缩放点积注意力
        attn_output, _ = self.attention(q_proj, k_proj, v_proj, mask)  # [batch, n_heads, seq_q, d_k]
        
        # 4. 拼接多头：转置后合并n_heads×d_k为d_model
        attn_output = attn_output.transpose(1, 2).contiguous()  # [batch, seq_q, n_heads, d_k]
        # contiguous()确保内存连续，避免view报错
        attn_output = attn_output.view(batch_size, -1, self.d_model)  # [batch, seq_q, d_model]
        
        # 5. 输出投影
        output = self.w_o(attn_output)  # [batch, seq_q, d_model]
        
        return output

2.4 前馈网络（Position-Wise Feed-Forward Network）

原理说明

前馈网络对序列中每个位置的特征独立处理（与位置无关），作用是对注意力输出进行非线性变换和特征增强。原论文采用两层全连接网络，公式如下：
FFN ( x ) = max ⁡ ( 0 , x W 1 + b 1 ) W 2 + b 2 \text{FFN}(x) = \max(0, xW_1 + b_1)W_2 + b_2 FFN(x)=max(0,xW1​+b1​)W2​+b2​
其中：

  • W 1 W_1 W1​：第一层权重，维度 [ d m o d e l , d f f ] [d_{model}, d_{ff}] [dmodel​,dff​]
  • W 2 W_2 W2​：第二层权重，维度 [ d f f , d m o d e l ] [d_{ff}, d_{model}] [dff​,dmodel​]
  • d f f d_{ff} dff​：中间层维度（通常为 4 × d m o d e l 4 \times d_{model} 4×dmodel​，如2048）
  • max ⁡ ( 0 , ⋅ ) \max(0, \cdot) max(0,⋅)：ReLU激活函数（原论文也尝试过GELU，但ReLU更简洁）
代码实现
class PositionWiseFeedForward(nn.Module):
    def __init__(self, d_model: int, d_ff: int = 2048, dropout: float = 0.1):
        super().__init__()
        self.linear1 = nn.Linear(d_model, d_ff)  # [d_model, d_ff]
        self.linear2 = nn.Linear(d_ff, d_model)  # [d_ff, d_model]
        self.dropout = nn.Dropout(dropout)
        self.relu = nn.ReLU()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: 输入张量，shape=[batch_size, seq_len, d_model]
        Returns:
            前馈网络输出，shape不变
        """
        x = self.relu(self.linear1(x))  # [batch, seq_len, d_ff]
        x = self.dropout(x)
        x = self.linear2(x)  # [batch, seq_len, d_model]
        return x

三、编码器与解码器实现

3.1 编码器层（Encoder Layer）

原理说明

编码器层是编码器的基本单元，由“多头自注意力”和“前馈网络”两个子层构成，每个子层后均添加残差连接和层归一化（LayerNorm）。结构如下：
LayerNorm ( x + MultiHead ( x , x , x ) ) → LayerNorm ( x + FFN ( x ) ) \text{LayerNorm}(x + \text{MultiHead}(x, x, x)) \rightarrow \text{LayerNorm}(x + \text{FFN}(x)) LayerNorm(x+MultiHead(x,x,x))→LayerNorm(x+FFN(x))

  • 自注意力： Q = K = V = x Q=K=V=x Q=K=V=x，捕捉序列内的依赖关系
  • 残差连接： x + SubLayer ( x ) x + \text{SubLayer}(x) x+SubLayer(x)，缓解梯度消失
  • 层归一化：对特征维度归一化（ LayerNorm ( x ) = x − E [ x ] Var [ x ] + ϵ γ + β \text{LayerNorm}(x) = \frac{x - \mathbb{E}[x]}{\sqrt{\text{Var}[x] + \epsilon}} \gamma + \beta LayerNorm(x)=Var[x]+ϵ ​x−E[x]​γ+β），加速训练
代码实现
class EncoderLayer(nn.Module):
    def __init__(self, d_model: int, n_heads: int, d_ff: int = 2048, dropout: float = 0.1):
        super().__init__()
        # 子层1：多头自注意力
        self.self_attn = MultiHeadAttention(d_model, n_heads, dropout)
        # 子层2：前馈网络
        self.feed_forward = PositionWiseFeedForward(d_model, d_ff, dropout)
        # 层归一化
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        # Dropout（子层输出后应用）
        self.dropout = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor, mask: torch.Tensor = None) -> torch.Tensor:
        """
        Args:
            x: 输入张量，shape=[batch_size, seq_len, d_model]
            mask: 填充掩码，shape=[batch
# 深度学习 # transformer # 人工智能
Logo
码道开发者社区

码道开发者社区，聚焦华为云码道 CodeArts 代码智能体，沉淀 Agent、Skill、鸿蒙开发实战内容，供开发者查阅资料、交流技术、分享工程实践

更多推荐

cover

不想先学一堆 React？用 Codex 帮你完成第一支 Remotion 动效视频

avatar 码道开发者社区
cover

突发！英伟达限制员工使用Claude；程序员一条代码删掉89TB数据获刑：被判五年十个月；曝iPhone Duo有望支持微信双开 | 极客头条

avatar 码道开发者社区
cover

每日热评｜学术牛马的救星还是幻觉放大器？47k星学术研究Agent技能包深度评测

avatar 码道开发者社区