关于Transformer 架构的几点记录

Transformer Attention机制深度学习模型训练性能优化

作者: 醉月思发布日期: 2024-06-10 分类: 学习笔记

后来硬着头皮从头实现了一遍 Transformer，过程中踩了不少坑，也算是对这个架构有了些实际体会。

数学表达就是：

$$ \text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V $$

为什么要自己实现一遍 Transformer

现在直接用 Hugging Face Transformers 就能加载各种预训练模型，那为什么还要自己从头实现？

最直接的原因是理解。调 API 时你不知道模型内部在做什么，一出问题就不知道从哪查。自己实现一遍，至少知道每个组件大概在干什么，遇到问题时能有排查方向。

另一个原因是定制化。有些场景下现成的模型不太合适，需要修改架构。不自己实现一遍，改动起来心里没底。

当然，还有一个更现实的原因：面试官喜欢问 Transformer 的细节。

Attention 机制的数学原理

Attention 的核心思想很简单：给定一组查询（Query）、键（Key）和值（Value），通过计算 Query 和 Key 的相似度来决定对 Value 的关注程度。

数学表达就是：

$$ \text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V $$

这个公式里有几个关键点：

$QK^T$：计算 Query 和所有 Key 的相似度矩阵
除以 $\sqrt{d_k}$：缩放因子，防止数值过大导致 softmax 梯度消失
softmax：把相似度转换成概率分布
乘以 V：根据注意力权重对 Value 加权求和

为什么要有缩放因子

最开始我实现时没加这个 $\sqrt{d_k}$，训练到一半梯度就爆炸了。后来才知道，当 $d_k$ 很大时，$QK^T$ 的点积值会很大，导致 softmax 进入饱和区，梯度接近 0。

举个例子，假设 $d_k = 512$，$Q$ 和 $K$ 的分量均值是 0，标准差是 1，那么 $QK^T$ 的标准差就是 $\sqrt{512} \approx 22.6$。softmax 的输入方差这么大，梯度肯定出问题。

加上 $\sqrt{d_k}$ 后，点积的标准差被缩放回 1 左右，训练就稳定多了。

Self-Attention 的代码实现

用 PyTorch 实现 Self-Attention 并不复杂，关键是要处理好维度和批次。

import torch
import torch.nn as nn
import math

class SelfAttention(nn.Module):
    def __init__(self, embed_size, heads):
        super(SelfAttention, self).__init__()
        self.embed_size = embed_size
        self.heads = heads
        self.head_dim = embed_size // heads

        assert (self.head_dim * heads == embed_size), "Embed size needs to be div by heads"

        # 线性变换生成 Q, K, V
        self.values = nn.Linear(self.head_dim, self.head_dim, bias=False)
        self.keys = nn.Linear(self.head_dim, self.head_dim, bias=False)
        self.queries = nn.Linear(self.head_dim, self.head_dim, bias=False)

        # 输出投影
        self.fc_out = nn.Linear(heads * self.head_dim, embed_size)

    def forward(self, values, keys, query, mask):
        N = query.shape[0]
        value_len, key_len, query_len = values.shape[1], keys.shape[1], query.shape[1]

        # 分割多头
        values = values.reshape(N, value_len, self.heads, self.head_dim)
        keys = keys.reshape(N, key_len, self.heads, self.head_dim)
        query = query.reshape(N, query_len, self.heads, self.head_dim)

        # 线性变换
        values = self.values(values)
        keys = self.keys(keys)
        query = self.queries(query)

        # 计算 attention scores
        # einsum 解释：b=batch, h=heads, q=query_len, k=key_len, d=head_dim
        energy = torch.einsum("bhqd, bhkd -> bhqk", [query, keys])

        # 缩放
        energy = energy / (self.embed_size ** (1/2))

        # 应用 mask（可选）
        if mask is not None:
            energy = energy.masked_fill(mask == 0, float("-1e20"))

        # softmax 得到注意力权重
        attention = torch.softmax(energy, dim=3)

        # 加权求和
        out = torch.einsum("bhql, bhlv -> bhqv", [attention, values])

        # 拼接多头
        out = out.reshape(N, query_len, self.heads * self.head_dim)

        return self.fc_out(out)

这里用 einsum 处理张量运算，比手写索引更清晰。但我刚开始实现时不会用 einsum，写了一堆 reshape 和 permute，代码可读性很差，调 bug 时也费劲。

Multi-Head Attention 的好处

单头 Attention 只能在一个子空间里计算相似度，Multi-Head 让模型能在多个子空间里关注不同信息。

举个翻译的例子：有的头可能关注语法结构，有的头可能关注语义相似度，有的头可能关注位置关系。多头让模型能同时处理多种不同的注意力模式。

TransformerBlock 的实现

有了 Self-Attention，就可以组装 TransformerBlock 了。

class TransformerBlock(nn.Module):
    def __init__(self, embed_size, heads, dropout, forward_expansion):
        super(TransformerBlock, self).__init__()
        self.attention = SelfAttention(embed_size, heads)
        self.norm1 = nn.LayerNorm(embed_size)
        self.norm2 = nn.LayerNorm(embed_size)

        self.feed_forward = nn.Sequential(
            nn.Linear(embed_size, forward_expansion * embed_size),
            nn.ReLU(),
            nn.Linear(forward_expansion * embed_size, embed_size)
        )

        self.dropout = nn.Dropout(dropout)

    def forward(self, value, key, query, mask):
        attention = self.attention(value, key, query, mask)

        # Add & Norm
        x = self.dropout(self.norm1(attention + query))
        forward = self.feed_forward(x)

        # Add & Norm
        out = self.dropout(self.norm2(forward + x))
        return out

这里有两个关键设计：

残差连接（Add）：把原始输入加到输出上，帮助梯度流动，避免深层网络梯度消失
Layer Normalization（Norm）：在每个样本内部做归一化，比 BatchNorm 更适合序列数据

为什么用 LayerNorm 而不是 BatchNorm

最开始我用的是 BatchNorm，训练时效果还行，但推理时结果不稳定。后来才知道，BatchNorm 依赖批次统计量，而序列长度可能变化，导致归一化不稳定。

LayerNorm 在每个样本内部归一化，不依赖其他样本，更适合序列数据。而且 LayerNorm 在训练和推理时的行为一致，不会出现训练推理不一致的问题。

Encoder 和 Decoder 的实现

Transformer 的 Encoder 由多个 TransformerBlock 堆叠而成：

class Encoder(nn.Module):
    def __init__(self, src_vocab_size, embed_size, num_layers, heads,
                 device, forward_expansion, dropout, max_length):
        super(Encoder, self).__init__()
        self.embed_size = embed_size
        self.device = device

        # 词嵌入
        self.word_embedding = nn.Embedding(src_vocab_size, embed_size)
        # 位置编码
        self.position_embedding = nn.Embedding(max_length, embed_size)

        self.layers = nn.ModuleList([
            TransformerBlock(embed_size, heads, dropout, forward_expansion)
            for _ in range(num_layers)
        ])

        self.dropout = nn.Dropout(dropout)

    def forward(self, x, mask):
        N, seq_length = x.shape
        positions = torch.arange(0, seq_length).expand(N, seq_length).to(self.device)

        # 词嵌入 + 位置嵌入
        out = self.dropout(self.word_embedding(x) + self.position_embedding(positions))

        for layer in self.layers:
            out = layer(out, out, out, mask)

        return out

Decoder 稍微复杂一点，多了一个 Encoder-Decoder Attention 层：

class DecoderBlock(nn.Module):
    def __init__(self, embed_size, heads, forward_expansion, dropout, device):
        super(DecoderBlock, self).__init__()
        self.attention = SelfAttention(embed_size, heads)
        self.norm = nn.LayerNorm(embed_size)
        self.encoder_attention = SelfAttention(embed_size, heads)
        self.norm2 = nn.LayerNorm(embed_size)
        self.feed_forward = nn.Sequential(
            nn.Linear(embed_size, forward_expansion * embed_size),
            nn.ReLU(),
            nn.Linear(forward_expansion * embed_size, embed_size)
        )
        self.norm3 = nn.LayerNorm(embed_size)
        self.dropout = nn.Dropout(dropout)
        self.device = device

    def forward(self, x, value, key, src_mask, trg_mask):
        # Self-Attention
        attention = self.attention(x, x, x, trg_mask)
        query = self.dropout(self.norm(attention + x))

        # Encoder-Decoder Attention
        attention = self.encoder_attention(value, key, query, src_mask)
        query = self.dropout(self.norm2(attention + query))

        # Feed-Forward
        forward = self.feed_forward(query)
        out = self.dropout(self.norm3(forward + query))

        return out

这里的关键是 Decoder 有两个 Attention 层：

第一个 Self-Attention 处理目标序列，用的是 trg_mask（防止看到未来信息）
第二个 Attention 关注 Encoder 的输出，用的是 src_mask（处理源序列的 padding）

位置编码的实现

Transformer 本身没有循环或卷积结构，需要位置编码提供位置信息：

class PositionalEncoding(nn.Module):
    def __init__(self, embed_size, max_length, device):
        super(PositionalEncoding, self).__init__()
        self.device = device

        pe = torch.zeros(max_length, embed_size).to(device)
        position = torch.arange(0, max_length, dtype=torch.float).unsqueeze(1).to(device)
        div_term = torch.exp(torch.arange(0, embed_size, 2).float() * (-math.log(10000.0) / embed_size)).to(device)

        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)

        pe = pe.unsqueeze(0)
        self.register_buffer('pe', pe)

    def forward(self, x):
        return x + self.pe[:, :x.size(1)]

位置编码用正弦和余弦函数，不同维度对应不同频率，这样每个位置都有唯一的编码，而且能外推到比训练时更长的序列。

训练过程中的踩坑

1. 梯度爆炸

第一次训练时，没加梯度裁剪，训练到第 10 个 batch，梯度就变成了 NaN。

# 训练循环中加梯度裁剪
optimizer.zero_grad()
output = model(src, trg[:-1, :])
loss = criterion(output.reshape(-1, output.shape[2]), trg[1:, :].reshape(-1))
loss.backward()

# 梯度裁剪
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

optimizer.step()

加上梯度裁剪后，训练就稳定了。后来发现，除了梯度裁剪，调整学习率和初始化权重也能缓解这个问题。

2. 学习率过大

最开始用 0.001 的学习率，训练时 loss 一直在震荡，怎么都降不下来。后来改成 0.0001，loss 才开始正常下降。

Transformer 论文里用的是带预热的 Adam 学习率调度器：

class AdamWarmup:
    def __init__(self, model_size, warmup_steps, optimizer):
        self.model_size = model_size
        self.warmup_steps = warmup_steps
        self.optimizer = optimizer
        self.current_step = 0
        self.lr = 0

    def step(self):
        self.current_step += 1
        self.lr = self.model_size ** (-0.5) * min(
            self.current_step ** (-0.5),
            self.current_step * self.warmup_steps ** (-1.5)
        )

        for param_group in self.optimizer.param_groups:
            param_group['lr'] = self.lr

        self.optimizer.step()

这个调度器的逻辑是：先线性增加学习率，再按步数的平方根衰减。这样既能保证训练初期的稳定性，又能在后期收敛时精细调整。

3. 序列长度不一致

训练数据里序列长度不一致，处理起来很麻烦。最开始我直接 padding 到最大长度，但这样计算量太大，训练速度很慢。

后来学会了用动态 padding 和 pack_padded_sequence：

from torch.nn.utils.rnn import pad_sequence, pack_padded_sequence

# 动态 padding 到 batch 内最大长度
def collate_fn(batch):
    src_batch = [item['src'] for item in batch]
    trg_batch = [item['trg'] for item in batch]

    src_batch = pad_sequence(src_batch, batch_first=True, padding_idx=PAD_IDX)
    trg_batch = pad_sequence(trg_batch, batch_first=True, padding_idx=PAD_IDX)

    return {
        'src': src_batch,
        'trg': trg_batch
    }

# DataLoader 里用 collate_fn
train_loader = DataLoader(train_dataset, batch_size=32, collate_fn=collate_fn)

这样每个 batch 的 padding 长度不同，减少了无效计算，训练速度明显提升。

4. 掩码（Mask）写错

Decoder 的 Self-Attention 需要用三角掩码防止看到未来信息，我一开始把掩码写反了，导致模型能看到未来信息，训练结果一塌糊涂。

def create_trg_mask(trg):
    N, trg_len = trg.shape
    trg_mask = torch.tril(torch.ones((trg_len, trg_len))).expand(
        N, 1, trg_len, trg_len
    )

    return trg_mask

# trg_mask 的样子（trg_len=4）
# [[[[1, 0, 0, 0],
#    [1, 1, 0, 0],
#    [1, 1, 1, 0],
#    [1, 1, 1, 1]]]]

注意这里用的是 torch.tril（下三角矩阵），不是 torch.triu（上三角矩阵）。写反了就会导致模型能看到未来信息。

性能优化经验

1. 混合精度训练

Transformer 的计算量很大，用混合精度训练能显著提速：

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()

for batch in train_loader:
    optimizer.zero_grad()

    with autocast():
        output = model(src, trg[:-1, :])
        loss = criterion(output.reshape(-1, output.shape[2]), trg[1:, :].reshape(-1))

    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

混合精度训练用 FP16 做大部分计算，只在必要时用 FP32，既减少了显存占用，又提升了计算速度。

2. 梯度累积

显存不够时，可以用梯度累积模拟更大的 batch size：

accumulation_steps = 4

for i, batch in enumerate(train_loader):
    with autocast():
        output = model(src, trg[:-1, :])
        loss = criterion(output.reshape(-1, output.shape[2]), trg[1:, :].reshape(-1))
        loss = loss / accumulation_steps

    scaler.scale(loss).backward()

    if (i + 1) % accumulation_steps == 0:
        scaler.step(optimizer)
        scaler.update()
        optimizer.zero_grad()

这样每 accumulation_steps 个 batch 才更新一次参数，等效于 batch size 扩大了 accumulation_steps 倍。

3. 模型并行

单卡显存不够时，可以把模型拆到多张卡上：

import torch.distributed as dist
import torch.multiprocessing as mp

def train(rank, world_size):
    dist.init_process_group("nccl", rank=rank, world_size=world_size)
    torch.cuda.set_device(rank)

    model = Transformer(...)
    model = model.to(rank)
    model = nn.parallel.DistributedDataParallel(model, device_ids=[rank])

    # 训练循环
    for batch in train_loader:
        output = model(src.to(rank), trg.to(rank))
        loss = criterion(...)
        loss.backward()
        optimizer.step()

模型并行适合大模型训练，但通信开销大，小模型时可能还不如单卡快。

一些实际应用场景

文本分类

用 Transformer 的 Encoder 部分做文本分类很简单：

class TransformerClassifier(nn.Module):
    def __init__(self, vocab_size, embed_size, num_layers, heads, num_classes, device):
        super(TransformerClassifier, self).__init__()
        self.encoder = Encoder(vocab_size, embed_size, num_layers, heads, device, 4, 0.1, 100)
        self.fc = nn.Linear(embed_size, num_classes)

    def forward(self, x, mask):
        out = self.encoder(x, mask)
        # 用 [CLS] token 的输出或者 mean pooling
        out = out.mean(dim=1)
        return self.fc(out)

文本分类任务只需要 Encoder 部分，因为不需要生成序列，只需要提取特征。

机器翻译

完整的 Transformer 用于机器翻译：

class Transformer(nn.Module):
    def __init__(self, src_vocab_size, trg_vocab_size, embed_size, num_layers, heads,
                 device, forward_expansion, dropout, max_length):
        super(Transformer, self).__init__()
        self.encoder = Encoder(src_vocab_size, embed_size, num_layers, heads,
                               device, forward_expansion, dropout, max_length)
        self.decoder = Decoder(trg_vocab_size, embed_size, num_layers, heads,
                               forward_expansion, dropout, device, max_length)
        self.device = device

    def forward(self, src, trg):
        src_mask = create_src_mask(src)
        trg_mask = create_trg_mask(trg)

        enc_src = self.encoder(src, src_mask)
        out = self.decoder(trg, enc_src, src_mask, trg_mask)

        return out

机器翻译需要完整的 Encoder-Decoder 架构，Encoder 处理源语言，Decoder 生成目标语言。

生成文本

用训练好的模型生成文本：

def generate_text(model, src, max_length, device):
    model.eval()
    with torch.no_grad():
        # 编码源序列
        enc_src = model.encoder(src, None)

        # 初始化目标序列为 <SOS> token
        trg = torch.ones(1, 1).fill_(SOS_IDX).long().to(device)

        for i in range(max_length - 1):
            # 创建 target mask
            trg_mask = model.decoder.create_trg_mask(trg)

            # 前向传播
            output = model.decoder(trg, enc_src, None, trg_mask)

            # 取最后一个 token 的输出
            next_token = output.argmax(2)[:, -1].item()

            # 拼接到目标序列
            trg = torch.cat([trg, torch.ones(1, 1).fill_(next_token).long().to(device)], dim=1)

            # 如果生成了 <EOS> token，停止生成
            if next_token == EOS_IDX:
                break

    return trg

生成文本时用贪婪解码，也可以用 beam search 提升质量。

最后一些反思

从头实现 Transformer 确实能学到很多东西，但投入产出比不高。除非你真的需要定制模型架构，否则直接用现成的框架更划算。

Hugging Face Transformers 已经做了大量优化，包括：

混合精度训练
梯度检查点
模型并行
丰富的预训练权重
完善的 API

这些优化自己实现一遍太费时间，而且效果可能还不如现成的。

但自己实现一遍的价值在于理解。当你知道每个组件在做什么，遇到问题时就能有排查方向，而不是只知道调参数。

技术的价值不在复杂度，而在解决实际问题的能力。Transformer 复杂，但它的核心思想（注意力机制）其实很简单。把简单的事情做好，比追求复杂更有意义。

那次把 Transformer 跑通后，我发给导师看。导师回了句：“代码写得还行，但你知道为什么这么多头吗？”

我当时愣了一下，才发现自己只是把公式翻译成了代码，并没有真正理解设计背后的思考。

这大概就是学和做的区别吧。

版权声明: 本文首发于指尖魔法屋-关于Transformer 架构的几点记录（https://blog.thinkmoon.cn/post/165-transformer-architecture-attention-mechanism-practice/）转载或引用必须申明原指尖魔法屋来源及源地址！

数据库连接池折腾手记

先复现，再谈优化。很多人一上来就讲数据库连接池的全景图；我更想先把这次卡住的点说清楚。

📁 技术实践

📅 2021-03-04

分布式缓存：这次怎么落地的

服务从单实例扩到 5 个之后，本地缓存反而把数据库查询次数抬高了。

📁 技术实践

📅 2021-05-08

GraphQL实战踩坑记录

这次搞 GraphQL，起因很直接：团队内部前端反馈查用户详情接口返回了 30 多个字段，实际只用了 5 个；另一个接口要三次请求才能拿到完整数…

📁 技术实践

📅 2021-05-28

关于Transformer 架构的几点记录

为什么要自己实现一遍 Transformer

Attention 机制的数学原理

为什么要有缩放因子

Self-Attention 的代码实现

Multi-Head Attention 的好处

TransformerBlock 的实现

为什么用 LayerNorm 而不是 BatchNorm

Encoder 和 Decoder 的实现

位置编码的实现

训练过程中的踩坑

1. 梯度爆炸

2. 学习率过大

3. 序列长度不一致

4. 掩码（Mask）写错

性能优化经验

1. 混合精度训练

2. 梯度累积

3. 模型并行

一些实际应用场景

文本分类

机器翻译

生成文本

最后一些反思

评论

相关阅读