BERT预训练模型公式解析与实战：从数学原理到高效实现

1次阅读

没有评论

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

自注意力机制是 BERT 模型的核心组件，其数学表达式为：

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

其中：

$Q$ (Query)、$K$ (Key)、$V$ (Value) 分别是通过线性变换得到的矩阵
$d_k$ 是 Key 的维度，用于缩放点积结果
softmax 函数将注意力权重归一化

这个公式实现了三个关键功能：

计算 Query 和 Key 的相似度（$QK^T$ 部分）
通过缩放因子 $\sqrt{d_k}$ 防止梯度消失
用 Value 加权求和得到最终输出

BERT 的前馈网络 (FFN) 由两个全连接层组成：

$$
\text{FFN}(x) = \text{max}(0, xW_1 + b_1)W_2 + b_2
$$

层归一化 (LayerNorm) 的公式为：

$$
\text{LayerNorm}(x) = \gamma \frac{x – \mu}{\sqrt{\sigma^2 + \epsilon}} + \beta
$$

其中 $\mu$ 和 $\sigma$ 是均值和方差，$\gamma$ 和 $\beta$ 是可学习参数。

原始 BERT 实现的主要性能问题：

O(n^2)复杂度：自注意力机制需要计算所有 token 对之间的关联，序列长度 n 较大时计算量剧增
显存占用高：存储中间激活值和梯度需要大量显存
长序列处理困难：超过 512token 时性能明显下降

使用 PyTorch 的 AMP(Automatic Mixed Precision)模块可以显著提升训练效率：

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()

当显存不足时，可以通过梯度累积模拟更大的 batch size：

accumulation_steps = 4

for i, (inputs, labels) in enumerate(train_loader):
    outputs = model(inputs)
    loss = criterion(outputs, labels)
    loss = loss / accumulation_steps
    loss.backward()

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