BERT（Bidirectional Encoder Representations from Transformers）是Google于2018年提出的预训练语言模型，通过双向Transformer编码器和掩码语言建模（MLM）任务，彻底改变了自然语言处理领域。本文将用Java实现一个简化版BERT的核心组件，帮助读者理解双向上下文编码与掩码预测机制。
核心原理
BERT的核心创新在于：
- 双向编码：传统语言模型只能从左到右或从右到左单向建模，而BERT通过Transformer的自注意力机制同时关注句子中每个词的左右上下文。
- 掩码语言建模（MLM）：随机遮挡输入序列中15%的token，让模型根据上下文预测被遮挡的词，从而学习深层语义表示。
- 位置编码：由于Transformer没有序列顺序感，BERT引入了可学习的位置嵌入来编码词的位置信息。
完整Java实现
import java.util.*;
/**
* 简化版BERT实现：包含词嵌入、位置编码、多头自注意力与前馈网络
*/
public class BERT {
// 超参数
private final int vocabSize; // 词表大小
private final int hiddenSize; // 隐藏层维度
private final int numHeads; // 注意力头数
private final int numLayers; // Transformer层数
private final int maxSeqLen; // 最大序列长度
private final double dropoutRate; // Dropout比率
// 可学习参数
private double[][] wordEmbedding; // 词嵌入矩阵 [vocabSize][hiddenSize]
private double[][] posEmbedding; // 位置嵌入矩阵 [maxSeqLen][hiddenSize]
private double[][] tokenTypeEmbedding; // 句子类型嵌入
// Transformer层参数（每层共享结构，独立参数）
private List<TransformerLayer> layers;
// 输出层：掩码语言模型预测头
private double[][] mlmWeight; // [hiddenSize][vocabSize]
private double[] mlmBias; // [vocabSize]
private Random random;
public BERT(int vocabSize, int hiddenSize, int numHeads,
int numLayers, int maxSeqLen, double dropoutRate) {
this.vocabSize = vocabSize;
this.hiddenSize = hiddenSize;
this.numHeads = numHeads;
this.numLayers = numLayers;
this.maxSeqLen = maxSeqLen;
this.dropoutRate = dropoutRate;
this.random = new Random(42);
initializeParameters();
}
/**
* 初始化所有可学习参数，使用Xavier初始化
*/
private void initializeParameters() {
double scale = Math.sqrt(2.0 / hiddenSize);
// 词嵌入
wordEmbedding = new double[vocabSize][hiddenSize];
for (int i = 0; i < vocabSize; i++) {
for (int j = 0; j < hiddenSize; j++) {
wordEmbedding[i][j] = random.nextGaussian() * scale;
}
}
// 位置嵌入
posEmbedding = new double[maxSeqLen][hiddenSize];
for (int i = 0; i < maxSeqLen; i++) {
for (int j = 0; j < hiddenSize; j++) {
posEmbedding[i][j] = random.nextGaussian() * scale;
}
}
// Token类型嵌入（用于区分句子A和句子B）
tokenTypeEmbedding = new double[2][hiddenSize];
for (int i = 0; i < 2; i++) {
for (int j = 0; j < hiddenSize; j++) {
tokenTypeEmbedding[i][j] = random.nextGaussian() * scale;
}
}
// 初始化Transformer层
layers = new ArrayList<>();
for (int i = 0; i < numLayers; i++) {
layers.add(new TransformerLayer(hiddenSize, numHeads, dropoutRate, random));
}
// MLM预测头
mlmWeight = new double[hiddenSize][vocabSize];
mlmBias = new double[vocabSize];
for (int i = 0; i < hiddenSize; i++) {
for (int j = 0; j < vocabSize; j++) {
mlmWeight[i][j] = random.nextGaussian() * scale;
}
}
}
/**
* 前向传播：输入token IDs，输出每个位置的隐藏表示
*/
public double[][] forward(int[] inputIds, int[] tokenTypeIds) {
int seqLen = inputIds.length;
double[][] hidden = new double[seqLen][hiddenSize];
// 嵌入层：词嵌入 + 位置嵌入 + 句子类型嵌入
for (int i = 0; i < seqLen; i++) {
for (int j = 0; j < hiddenSize; j++) {
hidden[i][j] = wordEmbedding[inputIds[i]][j]
+ posEmbedding[i][j]
+ tokenTypeEmbedding[tokenTypeIds[i]][j];
}
}
// Layer Norm + Dropout（简化版省略Dropout）
hidden = layerNorm(hidden);
// 通过所有Transformer层
for (TransformerLayer layer : layers) {
hidden = layer.forward(hidden);
}
return hidden;
}
/**
* 掩码语言模型预测：给定隐藏表示，预测被遮挡位置的词
*/
public double[][] mlmPredict(double[][] hidden) {
int seqLen = hidden.length;
double[][] logits = new double[seqLen][vocabSize];
for (int i = 0; i < seqLen; i++) {
for (int v = 0; v < vocabSize; v++) {
double sum = mlmBias[v];
for (int j = 0; j < hiddenSize; j++) {
sum += hidden[i][j] * mlmWeight[j][v];
}
logits[i][v] = sum;
}
}
return logits;
}
/**
* 计算MLM损失：交叉熵
*/
public double computeMLMLoss(double[][] logits, int[] labels, boolean[] mask) {
double loss = 0.0;
int count = 0;
for (int i = 0; i < logits.length; i++) {
if (!mask[i]) continue;
// Softmax
double maxLogit = Double.NEGATIVE_INFINITY;
for (double v : logits[i]) {
if (v > maxLogit) maxLogit = v;
}
double sumExp = 0.0;
double[] probs = new double[vocabSize];
for (int v = 0; v < vocabSize; v++) {
probs[v] = Math.exp(logits[i][v] - maxLogit);
sumExp += probs[v];
}
// 负对数似然
loss += -Math.log(probs[labels[i]] / sumExp);
count++;
}
return count > 0 ? loss / count : 0.0;
}
/**
* 掩码输入生成：随机遮挡15%的token
*/
public static MaskResult createMaskedInput(int[] inputIds, int maskTokenId,
int vocabSize, Random random) {
int[] maskedIds = inputIds.clone();
int[] labels = new int[inputIds.length];
boolean[] maskPositions = new boolean[inputIds.length];
Arrays.fill(labels, -1); // -1表示不参与loss计算
for (int i = 0; i < inputIds.length; i++) {
if (random.nextDouble() < 0.15) {
maskPositions[i] = true;
labels[i] = inputIds[i];
double r = random.nextDouble();
if (r < 0.8) {
maskedIds[i] = maskTokenId; // 80%替换为[MASK]
} else if (r < 0.9) {
maskedIds[i] = random.nextInt(vocabSize); // 10%替换为随机词
}
// 10%保持不变
}
}
return new MaskResult(maskedIds, labels, maskPositions);
}
private double[][] layerNorm(double[][] x) {
int n = x.length;
int d = x[0].length;
double[][] out = new double[n][d];
double eps = 1e-12;
for (int i = 0; i < n; i++) {
double mean = 0.0;
for (double v : x[i]) mean += v;
mean /= d;
double var = 0.0;
for (double v : x[i]) var += (v - mean) * (v - mean);
var /= d;
for (int j = 0; j < d; j++) {
out[i][j] = (x[i][j] - mean) / Math.sqrt(var + eps);
}
}
return out;
}
// ============ Transformer层 ============
static class TransformerLayer {
private final int hiddenSize;
private final int numHeads;
private final int headDim;
private final double dropoutRate;
// 多头注意力参数
private double[][] wq, wk, wv, wo;
private double[] bq, bk, bv, bo;
// 前馈网络参数
private double[][] w1, w2;
private double[] b1, b2;
private Random random;
public TransformerLayer(int hiddenSize, int numHeads,
double dropoutRate, Random random) {
this.hiddenSize = hiddenSize;
this.numHeads = numHeads;
this.headDim = hiddenSize / numHeads;
this.dropoutRate = dropoutRate;
this.random = random;
double scale = Math.sqrt(2.0 / hiddenSize);
// Q, K, V投影
wq = new double[hiddenSize][hiddenSize];
wk = new double[hiddenSize][hiddenSize];
wv = new double[hiddenSize][hiddenSize];
wo = new double[hiddenSize][hiddenSize];
bq = new double[hiddenSize];
bk = new double[hiddenSize];
bv = new double[hiddenSize];
bo = new double[hiddenSize];
initMatrix(wq, scale); initMatrix(wk, scale);
initMatrix(wv, scale); initMatrix(wo, scale);
// 前馈网络 (hiddenSize -> 4*hiddenSize -> hiddenSize)
w1 = new double[hiddenSize][4 * hiddenSize];
w2 = new double[4 * hiddenSize][hiddenSize];
b1 = new double[4 * hiddenSize];
b2 = new double[hiddenSize];
initMatrix(w1, scale); initMatrix(w2, scale);
}
private void initMatrix(double[][] m, double scale) {
for (double[] row : m) {
for (int i = 0; i < row.length; i++) {
row[i] = random.nextGaussian() * scale;
}
}
}
public double[][] forward(double[][] x) {
// 多头自注意力 + 残差连接
double[][] attnOut = multiHeadAttention(x);
double[][] normed1 = layerNorm(add(x, attnOut));
// 前馈网络 + 残差连接
double[][] ffOut = feedForward(normed1);
double[][] output = layerNorm(add(normed1, ffOut));
return output;
}
private double[][] multiHeadAttention(double[][] x) {
int seqLen = x.length;
// 线性投影得到Q, K, V
double[][] q = linear(x, wq, bq);
double[][] k = linear(x, wk, bk);
double[][] v = linear(x, wv, bv);
// 多头分割与注意力计算
double[][] attnOutput = new double[seqLen][hiddenSize];
for (int h = 0; h < numHeads; h++) {
int offset = h * headDim;
double[][] headOut = scaledDotProductAttention(
slice(q, offset, headDim),
slice(k, offset, headDim),
slice(v, offset, headDim)
);
// 合并头输出
for (int i = 0; i < seqLen; i++) {
System.arraycopy(headOut[i], 0, attnOutput[i], offset, headDim);
}
}
// 输出投影
return linear(attnOutput, wo, bo);
}
private double[][] scaledDotProductAttention(double[][] q, double[][] k, double[][] v) {
int seqLen = q.length;
double[][] scores = new double[seqLen][seqLen];
double scale = 1.0 / Math.sqrt(headDim);
// Q * K^T / sqrt(d)
for (int i = 0; i < seqLen; i++) {
for (int j = 0; j < seqLen; j++) {
double dot = 0.0;
for (int d = 0; d < headDim; d++) {
dot += q[i][d] * k[j][d];
}
scores[i][j] = dot * scale;
}
}
// Softmax
double[][] weights = softmax(scores);
// 加权求和
double[][] output = new double[seqLen][headDim];
for (int i = 0; i < seqLen; i++) {
for (int d = 0; d < headDim; d++) {
for (int j = 0; j < seqLen; j++) {
output[i][d] += weights[i][j] * v[j][d];
}
}
}
return output;
}
private double[][] feedForward(double[][] x) {
int seqLen = x.length;
// 第一层: hiddenSize -> 4*hiddenSize
double[][] h = linear(x, w1, b1);
// GELU激活（简化版使用ReLU）
for (int i = 0; i < seqLen; i++) {
for (int j = 0; j < h[i].length; j++) {
h[i][j] = Math.max(0, h[i][j]); // ReLU
}
}
// 第二层: 4*hiddenSize -> hiddenSize
return linear(h, w2, b2);
}
private double[][] linear(double[][] x, double[][] w, double[] b) {
int n = x.length;
int outDim = w[0].length;
double[][] out = new double[n][outDim];
for (int i = 0; i < n; i++) {
for (int j = 0; j < outDim; j++) {
double sum = b[j];
for (int k = 0; k < x[i].length; k++) {
sum += x[i][k] * w[k][j];
}
out[i][j] = sum;
}
}
return out;
}
private double[][] slice(double[][] x, int offset, int len) {
int n = x.length;
double[][] out = new double[n][len];
for (int i = 0; i < n; i++) {
System.arraycopy(x[i], offset, out[i], 0, len);
}
return out;
}
private double[][] add(double[][] a, double[][] b) {
int n = a.length;
int d = a[0].length;
double[][] out = new double[n][d];
for (int i = 0; i < n; i++) {
for (int j = 0; j < d; j++) {
out[i][j] = a[i][j] + b[i][j];
}
}
return out;
}
private double[][] layerNorm(double[][] x) {
// 同外层layerNorm
int n = x.length;
int d = x[0].length;
double[][] out = new double[n][d];
double eps = 1e-12;
for (int i = 0; i < n; i++) {
double mean = 0.0;
for (double v : x[i]) mean += v;
mean /= d;
double var = 0.0;
for (double v : x[i]) var += (v - mean) * (v - mean);
var /= d;
for (int j = 0; j < d; j++) {
out[i][j] = (x[i][j] - mean) / Math.sqrt(var + eps);
}
}
return out;
}
private double[][] softmax(double[][] x) {
int n = x.length;
double[][] out = new double[n][n];
for (int i = 0; i < n; i++) {
double max = Double.NEGATIVE_INFINITY;
for (double v : x[i]) if (v > max) max = v;
double sum = 0.0;
for (int j = 0; j < n; j++) {
out[i][j] = Math.exp(x[i][j] - max);
sum += out[i][j];
}
for (int j = 0; j < n; j++) out[i][j] /= sum;
}
return out;
}
}
static class MaskResult {
int[] maskedIds;
int[] labels;
boolean[] maskPositions;
MaskResult(int[] maskedIds, int[] labels, boolean[] maskPositions) {
this.maskedIds = maskedIds;
this.labels = labels;
this.maskPositions = maskPositions;
}
}
/**
* 主函数：演示BERT前向传播与MLM预测
*/
public static void main(String[] args) {
// 小型BERT配置
int vocabSize = 1000;
int hiddenSize = 128;
int numHeads = 4;
int numLayers = 2;
int maxSeqLen = 32;
int maskTokenId = 999; // [MASK] token
BERT bert = new BERT(vocabSize, hiddenSize, numHeads, numLayers, maxSeqLen, 0.1);
// 构造示例输入序列
int[] inputIds = {1, 45, 232, 87, 156, 333, 78, 2}; // [CLS] ... [SEP]
int[] tokenTypeIds = new int[inputIds.length]; // 单句，全为0
// 生成掩码输入
Random rand = new Random(123);
MaskResult maskResult = createMaskedInput(inputIds, maskTokenId, vocabSize, rand);
System.out.println("原始输入: " + Arrays.toString(inputIds));
System.out.println("掩码输入: " + Arrays.toString(maskResult.maskedIds));
// 前向传播
double[][] hidden = bert.forward(maskResult.maskedIds, tokenTypeIds);
System.out.println("隐藏表示维度: [" + hidden.length + ", " + hidden[0].length + "]");
// MLM预测
double[][] logits = bert.mlmPredict(hidden);
// 计算损失
double loss = bert.computeMLMLoss(logits, maskResult.labels, maskResult.maskPositions);
System.out.printf("MLM Loss: %.4f%n", loss);
// 展示掩码位置的预测结果
System.out.println("\n掩码位置预测:");
for (int i = 0; i < maskResult.maskPositions.length; i++) {
if (maskResult.maskPositions[i]) {
int predictedId = argmax(logits[i]);
System.out.printf(" 位置 %d: 真实=%d, 预测=%d%n", i, maskResult.labels[i], predictedId);
}
}
}
private static int argmax(double[] arr) {
int maxIdx = 0;
for (int i = 1; i < arr.length; i++) {
if (arr[i] > arr[maxIdx]) maxIdx = i;
}
return maxIdx;
}
}
关键设计解析
多头自注意力机制
每个注意力头独立学习不同的语义关系。通过将隐藏向量分割为numHeads份，模型可以同时关注句法结构、语义关联、指代关系等不同层面的信息。Scaled Dot-Product Attention中的1/sqrt(d)缩放因子防止内积过大导致softmax梯度消失。
掩码语言建模策略
BERT的掩码策略精心设计：80%概率替换为[MASK]标记，10%概率替换为随机词，10%概率保持不变。这种混合策略防止模型仅学习[MASK]位置的特定表示，迫使模型真正理解上下文语义。
层归一化与残差连接
Transformer层中，Layer Normalization对每个样本独立归一化，稳定深层网络训练。残差连接（x + F(x)）缓解梯度消失问题，使得BERT可以堆叠多达24层（BERT-Large）。
复杂度分析
| 组件 | 时间复杂度 | 空间复杂度 |
|---|---|---|
| 嵌入层 | O(n·d) | O(V·d + n·d) |
| 多头注意力 | O(n²·d) | O(n·d) |
| 前馈网络 | O(n·d²) | O(d²) |
| 总体（L层） | O(L·n²·d) | O(L·d²) |
其中n为序列长度，d为隐藏维度，V为词表大小。自注意力的O(n²)复杂度是BERT处理长文本时的主要瓶颈。
扩展方向
- NSP任务：实现下一句预测（Next Sentence Prediction），用于学习句子间关系。
- GELU激活：用高斯误差线性单元替代ReLU，平滑梯度流。
- AdamW优化器：实现权重衰减解耦的Adam优化器进行预训练。
- 分词器：集成WordPiece或BPE子词分词，处理未登录词。
