从零实现BERT预训练模型：核心架构与工程实践

1次阅读

共计 2564 个字符，预计需要花费 7 分钟才能阅读完成。

自然语言处理（NLP）领域近年来取得了显著的进展，其中预训练模型扮演了至关重要的角色。BERT（Bidirectional Encoder Representations from Transformers）作为其中的佼佼者，通过双向 Transformer 编码器和大规模无监督预训练，显著提升了各种 NLP 任务的性能。BERT 的核心创新点包括：

双向上下文建模：传统的语言模型（如 GPT）仅从左到右或从右到左单向建模，而 BERT 通过掩码语言模型（MLM）任务实现了双向上下文理解。
Transformer 架构：BERT 基于 Transformer 编码器堆叠而成，利用自注意力机制捕捉长距离依赖关系。
预训练 + 微调范式：先在大量无标注数据上预训练，再通过简单微调适配下游任务，显著降低了标注数据需求。

BERT 的核心是 Transformer 编码器，其实现细节如下：

自注意力机制：
计算 Query、Key、Value 矩阵，通过点积得到注意力权重
使用多头注意力（Multi-Head Attention）并行处理不同子空间的特征
公式：$Attention(Q,K,V)=softmax(\frac{QK^T}{\sqrt{d_k}})V$
位置编码：
Transformer 本身不具备序列位置信息，需通过位置编码注入
BERT 使用可学习的位置嵌入（Position Embeddings）而非固定正弦函数
层归一化与残差连接：
每个子层（注意力 / 前馈网络）后接 LayerNorm 和残差连接
缓解深层网络梯度消失问题

以下是一个规范的 BERT 实现框架核心代码（简化版）：

import torch
import torch.nn as nn
from torch.nn import functional as F

class BertSelfAttention(nn.Module):
    def __init__(self, hidden_size=768, num_heads=12):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = hidden_size // num_heads

        self.query = nn.Linear(hidden_size, hidden_size)
        self.key = nn.Linear(hidden_size, hidden_size)
        self.value = nn.Linear(hidden_size, hidden_size)

    def forward(self, x, attention_mask=None):
        batch_size = x.size(0)

        # 线性变换并分头
        q = self.query(x).view(batch_size, -1, self.num_heads, self.head_dim)
        k = self.key(x).view(batch_size, -1, self.num_heads, self.head_dim)
        v = self.value(x).view(batch_size, -1, self.num_heads, self.head_dim)

        # 注意力分数计算
        scores = torch.einsum('bnih,bnjh->bnij', q, k) / math.sqrt(self.head_dim)
        if attention_mask is not None:
            scores = scores.masked_fill(attention_mask == 0, -1e9)

        # 注意力权重与输出
        attn = F.softmax(scores, dim=-1)
        output = torch.einsum('bnij,bnjh->bnih', attn, v)

        return output.view(batch_size, -1, self.num_heads * self.head_dim)

针对大规模预训练的显存挑战，推荐以下优化方案：

梯度检查点（Gradient Checkpointing）：
只保留部分层的激活值，其余层在前向时重新计算
牺牲 30% 计算时间换取显存下降 50% 以上
混合精度训练：
使用 torch.cuda.amp 自动管理 FP16/FP32 转换

典型配置：

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()
with autocast():
    outputs = model(inputs)
    loss = criterion(outputs, labels)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()