BERT预训练模型架构图解析：从原理到工程实践

1次阅读

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

BERT（Bidirectional Encoder Representations from Transformers）的核心创新在于双向 Transformer 编码器结构。与单向语言模型不同，BERT 通过掩码语言模型（MLM）和下一句预测（NSP）任务实现双向上下文建模。其架构图包含以下关键组件：

输入层 ：Token 嵌入（WordPiece 分词）+ 位置编码 + 段嵌入（Segment Embeddings）
Transformer 编码器堆叠 ：通常由 12/24 层组成，每层包含多头自注意力机制和前馈神经网络
输出层 ：根据任务类型（如分类 / 序列标注）适配不同输出结构

（注：此处应为实际架构图 URL）

预训练阶段
计算资源消耗：需数十 GB 显存和数百小时 GPU 训练时间
数据需求：需要大规模无监督文本（如 Wikipedia+BookCorpus）
超参数敏感：学习率、warmup 步数等对效果影响显著
微调阶段
小样本适应：下游任务数据量不足时易过拟合
领域迁移：医疗 / 法律等专业领域效果下降
部署开销：模型体积大导致推理延迟高

特性	Transformer	RNN
并行计算	全序列并行	时序串行
长距离依赖	自注意力直达	依赖梯度传播
内存占用	O(n²) 注意力矩阵	O(n) 隐状态
典型应用	BERT/GPT	LSTM/GRU

# PyTorch 实现示例
class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super().__init__()
        assert d_model % num_heads == 0
        self.d_k = d_model // num_heads
        self.num_heads = num_heads
        self.q_linear = nn.Linear(d_model, d_model)
        self.k_linear = nn.Linear(d_model, d_model)
        self.v_linear = nn.Linear(d_model, d_model)
        self.out_linear = nn.Linear(d_model, d_model)

    def forward(self, q, k, v, mask=None):
        # 线性变换并分头 [batch, seq_len, d_model] -> [batch, num_heads, seq_len, d_k]
        q = self.q_linear(q).view(q.size(0), -1, self.num_heads, self.d_k).transpose(1,2)
        k = self.k_linear(k).view(k.size(0), -1, self.num_heads, self.d_k).transpose(1,2)
        v = self.v_linear(v).view(v.size(0), -1, self.num_heads, self.d_k).transpose(1,2)

        # 计算注意力得分
        scores = torch.matmul(q, k.transpose(-2,-1)) / math.sqrt(self.d_k)
        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)
        attn_weights = F.softmax(scores, dim=-1)

        # 加权求和并合并头
        output = torch.matmul(attn_weights, v)
        output = output.transpose(1,2).contiguous().view(output.size(0), -1, self.num_heads * self.d_k)
        return self.out_linear(output)

def positional_encoding(max_seq_len, d_model):
    position = torch.arange(max_seq_len).unsqueeze(1)
    div_term = torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model))
    pe = torch.zeros(max_seq_len, d_model)
    pe[:, 0::2] = torch.sin(position * div_term)
    pe[:, 1::2] = torch.cos(position * div_term)
    return pe  # [max_seq_len, d_model]

import torch
import torch.nn as nn
import math

class BERT(nn.Module):
    def __init__(self, vocab_size, max_len=512, d_model=768, n_layers=12, n_heads=12, dropout=0.1):
        super().__init__()
        self.token_embed = nn.Embedding(vocab_size, d_model)
        self.pos_embed = nn.Parameter(torch.zeros(1, max_len, d_model))
        self.segment_embed = nn.Embedding(2, d_model)
        self.layers = nn.ModuleList([TransformerLayer(d_model, n_heads, dropout) for _ in range(n_layers)])

    def forward(self, input_ids, segment_ids, attention_mask):
        # 嵌入求和
        token_emb = self.token_embed(input_ids)
        pos_emb = self.pos_embed[:, :input_ids.size(1), :]
        seg_emb = self.segment_embed(segment_ids)
        x = token_emb + pos_emb + seg_emb

        # Transformer 编码
        for layer in self.layers:
            x = layer(x, attention_mask)
        return x