大家好,我是你们的技术伙伴。👋

在2026年的今天,大语言模型(LLM)已经无处不在。然而,当我们惊叹于GPT等模型的强大时,往往容易忽略它最核心的基石——Transformer架构

很多同学看了无数遍“Attention is All You Need”的论文,也背熟了QKV(查询、键、值)的概念,但一旦落实到代码,往往两眼发黑。

今天,我不讲那些虚无缥缈的理论,我们来一场“显微镜”级的代码实战。我们将从最基础的词嵌入开始,一行代码一行代码地构建出一个完整的Transformer模型。

准备好了吗?让我们开始吧!🚀


🧱 第一章:基石——输入部分的双重奏(词嵌入与位置编码)

Transformer抛弃了RNN的循环结构,这意味着它无法天然地感知序列的顺序。为了解决这个问题,输入部分的设计至关重要,它由词嵌入(Word Embedding)位置编码(Positional Encoding)两部分组成。

1. 词嵌入层 (Word Embedding)

这是将离散的单词转化为连续向量的第一步。我们在代码中不仅实现了嵌入,还加入了一个关键的缩放技巧。

import torch
import torch.nn as nn
import math

class Embeddings(nn.Module):
    def __init__(self, vocab_size, d_model):
        super().__init__()
        self.vocab_size = vocab_size
        self.d_model = d_model
        # 核心嵌入层
        self.embed = nn.Embedding(vocab_size, d_model)
    
    def forward(self, x):
        # 关键点:乘以 sqrt(d_model)
        # 目的:为了平衡梯度,避免因维度变化导致的数值过大或过小,防止梯度消失或爆炸。
        return self.embed(x) * math.sqrt(self.d_model)

💡 原理解析:
为什么要乘以 sqrt(d_model)?这是为了在初始阶段保持向量的方差稳定,确保模型训练的平滑性。

2. 位置编码 (Positional Encoding)

这是Transformer的灵魂所在。我们使用正弦和余弦函数生成位置编码,让模型能够感知单词在序列中的位置。

class PositionEncoding(nn.Module):
    def __init__(self, d_model, dropout, max_len=60):
        super().__init__()
        self.dropout = nn.Dropout(p=dropout)
        
        # 1. 初始化位置编码矩阵 [max_len, d_model]
        pe = torch.zeros(max_len, d_model)
        
        # 2. 生成位置序列 [0, max_len)
        position = torch.arange(0, max_len).unsqueeze(1) # [max_len, 1]
        
        # 3. 核心公式:计算分母中的指数项
        # div_term = 1 / (10000^(2i/d_model))
        div_term = torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model))
        
        # 4. 偶数位置用sin, 奇数位置用cos
        pe[:, 0::2] = torch.sin(position * div_term) # 偶数维度
        pe[:, 1::2] = torch.cos(position * div_term) # 奇数维度
        
        # 5. 增加批次维度并注册为缓冲区(不更新参数)
        pe = pe.unsqueeze(0) # [1, max_len, d_model]
        self.register_buffer('pe', pe)
    
    def forward(self, x):
        # 将词嵌入与位置编码相加
        x = x + self.pe[:, :x.size(1)]
        return self.dropout(x)

可视化:
运行代码中的绘图函数,你会看到不同维度的位置编码呈现出不同的波形,这保证了每个位置都有独一无二的“指纹”。


🧠 第二章:核心——多头注意力机制 (Multi-Head Attention)

Attention机制是Transformer的“眼睛”。它让模型在处理某个词时,能够关注到句子中其他相关的词。

1. 缩放点积注意力 (Scaled Dot-Product Attention)

这是最底层的计算单元。注意其中的Mask机制,它用于遮挡未来的信息(在解码器中)或填充的无效信息。

def attention(query, key, value, mask=None, dropout=None):
    # 1. 计算Q和K的点积,并缩放
    d_k = query.size()[-1]
    scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)
    
    # 2. Mask机制:将mask位置填充为极小值(-1e9),确保softmax后概率为0
    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e9)
    
    # 3. 计算注意力权重
    p_attn = F.softmax(scores, dim=-1)
    
    # 4. 加权求和得到输出
    if dropout is not None:
        p_attn = dropout(p_attn)
        
    return torch.matmul(p_attn, value), p_attn
2. 多头注意力 (Multi-Head Attention)

单个注意力头可能只关注局部信息,多个头可以并行关注不同子空间的信息。

class MultiHeadAttention(nn.Module):
    def __init__(self, embed_dim, head, dropout_p=0.1):
        super().__init__()
        assert embed_dim % head == 0 # 确保能整除
        
        self.d_k = embed_dim // head
        self.head = head
        
        # 4个线性层:分别用于 Q, K, V 的投影 和 最终输出的投影
        self.linears = clones(nn.Linear(embed_dim, embed_dim), 4)
        self.dropout = nn.Dropout(dropout_p)
        self.atten = None

    def forward(self, query, key, value, mask=None):
        if mask is not None:
            mask = mask.unsqueeze(0) # 增加头维度
            
        batch = query.size(0)
        
        # 1. 线性变换 + 分头
        # 将 [Batch, Seq_Len, Embed_Dim] -> [Batch, Head, Seq_Len, d_k]
        query, key, value = [ 
            model(x).view(batch, -1, self.head, self.d_k).transpose(1, 2)
            for model, x in zip(self.linears, (query, key, value))
        ]
        
        # 2. 调用底层attention函数
        x, self.atten = attention(query, key, value, mask, self.dropout)
        
        # 3. 合并多头结果
        # [Batch, Head, Seq_Len, d_k] -> [Batch, Seq_Len, Embed_Dim]
        x = x.transpose(1, 2).contiguous().view(batch, -1, self.head * self.d_k)
        
        # 4. 最终线性投影
        return self.linears[-1](x)

代码解析:
这里的 clones 函数使用了 nn.ModuleList 深拷贝多个线性层,这是为了确保参数不共享。


⚙️ 第三章:骨架——编码器与解码器

有了底层积木,现在我们来搭建模型的主体。

1. 编码器层 (Encoder Layer)

编码器层由两个子层组成:多头自注意力机制前馈全连接网络(FFN)。每个子层都使用了残差连接(Residual Connection)层规范化(Layer Normalization)

class EncoderLayer(nn.Module):
    def __init__(self, d_model, self_attn, feed_forward, dropout=0.1):
        super().__init__()
        self.d_model = d_model
        self.self_attn = self_attn
        self.feed_forward = feed_forward
        # 克隆两个子层连接结构(包含残差连接和规范化)
        self.sublayer = clones(SublayerConnection(d_model, dropout), 2)

    def forward(self, x, mask):
        # 第1个子层:多头自注意力
        x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, mask))
        # 第2个子层:前馈全连接网络
        x = self.sublayer[1](x, lambda x: self.feed_forward(x))
        return x
2. 解码器层 (Decoder Layer)

解码器比编码器更复杂,它多了一个编码器-解码器注意力层,用于关注源序列的信息。

class DecoderLayer(nn.Module):
    def __init__(self, d_model, self_attn, src_attn, feed_forward, dropout=0.1):
        super().__init__()
        self.d_model = d_model
        self.self_attn = self_attn      # 掩码多头注意力(防止看到未来信息)
        self.src_attn = src_attn        # 编码器-解码器注意力
        self.feed_forward = feed_forward
        # 3个子层连接
        self.layers = clones(SublayerConnection(d_model, dropout), mit 3)

    def forward(self, x, encoder_output, source_mask, target_mask):
        # 第1子层:掩码自注意力(Q=K=V=x)
        x = self.layers[0](x, lambda x: self.self_attn(x, x, x, target_mask))
        
        # 第2子层:编码器-解码器注意力
        # Q来自解码器,K和V来自编码器输出
        x = self.layers[1](x, lambda x: self.src_attn(x, encoder_output, encoder_output, source_mask))
        
        # 第3子层:前馈网络
        x = self.layers[2](x, self.feed_forward)
        
        return x

🏁 第四章:输出——生成最终结果

模型的最后一步是将解码器输出的特征向量转化为具体的词汇概率分布。

1. 生成器 (Generator)

使用一个线性层将特征映射到词汇表大小,然后通过 LogSoftmax 转化为对数概率。

class Generator(nn.Module):
    def __init__(self, d_model, vocab_size):
        super().__init__()
        # 线性层:将解码器输出的特征转为词汇表大小的分数
        self.linear = nn.Linear(d_model, vocab_size)

    def forward(self, x):
        # LogSoftmax + NLLLoss 是分类任务的经典组合
        return F.log_softmax(self.linear(x), dim=-1)
2. 组装完整的Transformer

最后,我们将所有组件拼装起来,形成一个完整的EncoderDecoder类。

class EncoderDecoder(nn.Module):
    def __init__(self, source_embed, encoder, target_embed, decoder, generator):
        super().__init__()
        self.source_embed = source_embed
        self.encoder = encoder
        self.target_embed = target_embed
        self.decoder = decoder
        self.generator = generator

    def forward(self, source_x, target_y, source_mask, target_mask):
        # 1. 编码
        encoder_result = self.encode(source_x, source_mask)
        # 2. 解码
        decoder_result = self.decode(target_y, encoder_result, source_mask, target_mask)
        # 3. 生成最终输出
        return self.generator(decoder_result)

    def encode(self, source_x, source_mask):
        embed_x = self.source_embed(source_x)
        return self.encoder(embed_x, source_mask)

    def decode(self, target_y, encoder_output, source_mask, target_mask):
        embed_y = self.target_embed(target_y)
        return self.decoder(embed_y, encoder_output, source_mask, target_mask)

🚀 第五章:实战测试——让模型跑起来

我们在make_model函数中实例化所有组件,并进行前向传播测试。

核心参数配置:

  • d_model (词向量维度): 512
  • Head (多头数量): 8
  • d_ff (前馈网络隐藏层): 2048
  • N (编码器/解码器层数): 6

运行测试代码,如果看到控制台输出类似 result.shape: [2, 4, vocab_size] 的信息,恭喜你,你已经成功复现了Transformer的核心架构!


📝 总结与展望

通过这篇文章,我们完成了一次从理论到代码的完整穿越:

  1. 输入端:掌握了词嵌入与位置编码的融合。
  2. 核心端:亲手实现了多头注意力机制与掩码逻辑。
  3. 结构端:构建了带有残差连接和层规范化的编码器与解码器。
  4. 输出端:完成了最终的分类生成。

最后的叮嘱:
虽然现在的大模型看似复杂,但其核心架构依然是Transformer。希望这篇2026年的实战指南能为你打下坚实的基础。

如果你觉得这篇文章硬核且有用,请务必点赞、收藏,并关注我。有任何关于深度学习的问题,欢迎在评论区留言,我会一一解答。💬

Logo

AtomGit 是由开放原子开源基金会联合 CSDN 等生态伙伴共同推出的新一代开源与人工智能协作平台。平台坚持“开放、中立、公益”的理念,把代码托管、模型共享、数据集托管、智能体开发体验和算力服务整合在一起,为开发者提供从开发、训练到部署的一站式体验。

更多推荐