Transformer架构自2017年”Attention Is All You Need”论文提出以来，彻底改变了自然语言处理和计算机视觉领域。与RNN/CNN不同，Transformer完全依赖注意力机制捕捉序列间的依赖关系，实现了高度并行化训练。本文将用Java从零实现一个简化版Transformer，涵盖多头自注意力、缩放点积注意力、位置编码与前馈网络。
Transformer核心组件
缩放点积注意力
给定查询矩阵 $Q$、键矩阵 $K$ 和值矩阵 $V$，注意力输出为：
$$ ext{Attention}(Q, K, V) = ext{softmax}\left( rac{QK^T}{\sqrt{d_k}}
ight)V$$
缩放因子 $\sqrt{d_k}$ 防止点积值过大导致softmax梯度消失。
多头注意力
将 $Q, K, V$ 投影到 $h$ 个不同的子空间，分别计算注意力后拼接：
$$ ext{MultiHead}(Q, K, V) = ext{Concat}( ext{head}_1, …, ext{head}_h)W^O$$
位置编码
由于Transformer没有循环或卷积结构，需要通过位置编码注入位置信息。使用正弦/余弦函数：
$$PE_{(pos, 2i)} = \sin(pos / 10000^{2i/d_{model}})$$
$$PE_{(pos, 2i+1)} = \cos(pos / 10000^{2i/d_{model}})$$
Java实现
import java.util.*;
/**
* 简化版Transformer实现
* 包含：多头自注意力、前馈网络、位置编码、层归一化
*/
public class Transformer {
static final Random rand = new Random(42);
// 矩阵工具类
static class Matrix {
double[][] data;
int rows, cols;
Matrix(int r, int c) { data = new double[r][c]; rows = r; cols = c; }
Matrix(double[][] d) { data = d; rows = d.length; cols = d[0].length; }
static Matrix zeros(int r, int c) { return new Matrix(r, c); }
static Matrix randn(int r, int c) {
Matrix m = new Matrix(r, c);
for (int i = 0; i < r; i++)
for (int j = 0; j < c; j++)
m.data[i][j] = rand.nextGaussian() * 0.02;
return m;
}
Matrix add(Matrix b) {
Matrix r = new Matrix(rows, cols);
for (int i = 0; i < rows; i++)
for (int j = 0; j < cols; j++)
r.data[i][j] = data[i][j] + b.data[i][j];
return r;
}
Matrix scale(double s) {
Matrix r = new Matrix(rows, cols);
for (int i = 0; i < rows; i++)
for (int j = 0; j < cols; j++)
r.data[i][j] = data[i][j] * s;
return r;
}
static Matrix matmul(Matrix a, Matrix b) {
Matrix r = new Matrix(a.rows, b.cols);
for (int i = 0; i < a.rows; i++)
for (int k = 0; k < a.cols; k++)
for (int j = 0; j < b.cols; j++)
r.data[i][j] += a.data[i][k] * b.data[k][j];
return r;
}
Matrix transpose() {
Matrix r = new Matrix(cols, rows);
for (int i = 0; i < rows; i++)
for (int j = 0; j < cols; j++)
r.data[j][i] = data[i][j];
return r;
}
Matrix addBias(double[] bias) {
Matrix r = new Matrix(rows, cols);
for (int i = 0; i < rows; i++)
for (int j = 0; j < cols; j++)
r.data[i][j] = data[i][j] + bias[j];
return r;
}
}
// 辅助函数
static double[] softmax(double[] x) {
double max = x[0];
for (double v : x) if (v > max) max = v;
double sum = 0;
double[] exp = new double[x.length];
for (int i = 0; i < x.length; i++) {
exp[i] = Math.exp(x[i] - max);
sum += exp[i];
}
for (int i = 0; i < x.length; i++) exp[i] /= sum;
return exp;
}
static Matrix softmaxRows(Matrix m) {
Matrix r = new Matrix(m.rows, m.cols);
for (int i = 0; i < m.rows; i++)
r.data[i] = softmax(m.data[i]);
return r;
}
static Matrix relu(Matrix m) {
Matrix r = new Matrix(m.rows, m.cols);
for (int i = 0; i < m.rows; i++)
for (int j = 0; j < m.cols; j++)
r.data[i][j] = Math.max(0, m.data[i][j]);
return r;
}
static Matrix layerNorm(Matrix x) {
Matrix r = new Matrix(x.rows, x.cols);
double eps = 1e-6;
for (int i = 0; i < x.rows; i++) {
double mean = 0, var = 0;
for (double v : x.data[i]) mean += v;
mean /= x.cols;
for (double v : x.data[i]) var += (v - mean) * (v - mean);
var /= x.cols;
double std = Math.sqrt(var + eps);
for (int j = 0; j < x.cols; j++)
r.data[i][j] = (x.data[i][j] - mean) / std;
}
return r;
}
// 位置编码
static class PositionalEncoding {
final int dModel, maxLen;
final double[][] pe;
PositionalEncoding(int dModel, int maxLen) {
this.dModel = dModel; this.maxLen = maxLen;
pe = new double[maxLen][dModel];
for (int pos = 0; pos < maxLen; pos++) {
for (int i = 0; i < dModel; i += 2) {
double div = Math.pow(10000, (2.0 * i) / dModel);
pe[pos][i] = Math.sin(pos / div);
if (i + 1 < dModel) pe[pos][i + 1] = Math.cos(pos / div);
}
}
}
Matrix encode(Matrix x) {
Matrix r = new Matrix(x.rows, x.cols);
for (int i = 0; i < x.rows; i++)
for (int j = 0; j < x.cols; j++)
r.data[i][j] = x.data[i][j] + pe[i][j];
return r;
}
}
// 多头自注意力
static class MultiHeadAttention {
final int dModel, numHeads, dHead;
final Matrix Wq, Wk, Wv, Wo;
final double[] bq, bk, bv, bo;
MultiHeadAttention(int dModel, int numHeads) {
this.dModel = dModel; this.numHeads = numHeads;
this.dHead = dModel / numHeads;
Wq = Matrix.randn(dModel, dModel); Wk = Matrix.randn(dModel, dModel);
Wv = Matrix.randn(dModel, dModel); Wo = Matrix.randn(dModel, dModel);
bq = new double[dModel]; bk = new double[dModel];
bv = new double[dModel]; bo = new double[dModel];
}
Matrix forward(Matrix x) {
int seqLen = x.rows;
// 线性投影
Matrix Q = Matrix.matmul(x, Wq).addBias(bq);
Matrix K = Matrix.matmul(x, Wk).addBias(bk);
Matrix V = Matrix.matmul(x, Wv).addBias(bv);
// 分头计算注意力
Matrix out = Matrix.zeros(seqLen, dModel);
double scale = 1.0 / Math.sqrt(dHead);
for (int h = 0; h < numHeads; h++) {
// 提取当前头的 Q_h, K_h, V_h
double[][] Qh = new double[seqLen][dHead];
double[][] Kh = new double[seqLen][dHead];
double[][] Vh = new double[seqLen][dHead];
for (int i = 0; i < seqLen; i++) {
for (int j = 0; j < dHead; j++) {
Qh[i][j] = Q.data[i][h * dHead + j];
Kh[i][j] = K.data[i][h * dHead + j];
Vh[i][j] = V.data[i][h * dHead + j];
}
}
// scores = Qh * Kh^T
double[][] scores = new double[seqLen][seqLen];
for (int i = 0; i < seqLen; i++)
for (int k = 0; k < seqLen; k++)
for (int j = 0; j < dHead; j++)
scores[i][k] += Qh[i][j] * Kh[k][j];
// 缩放 + softmax
for (int i = 0; i < seqLen; i++) {
for (int j = 0; j < seqLen; j++) scores[i][j] *= scale;
scores[i] = softmax(scores[i]);
}
// attention = scores * Vh
for (int i = 0; i < seqLen; i++)
for (int j = 0; j < dHead; j++) {
double sum = 0;
for (int k = 0; k < seqLen; k++) sum += scores[i][k] * Vh[k][j];
out.data[i][h * dHead + j] = sum;
}
}
// 输出投影
return Matrix.matmul(out, Wo).addBias(bo);
}
}
// 前馈网络
static class FeedForward {
final int dModel, dFF;
final Matrix W1, W2;
final double[] b1, b2;
FeedForward(int dModel, int dFF) {
this.dModel = dModel; this.dFF = dFF;
W1 = Matrix.randn(dModel, dFF); W2 = Matrix.randn(dFF, dModel);
b1 = new double[dFF]; b2 = new double[dModel];
}
Matrix forward(Matrix x) {
Matrix hidden = Matrix.matmul(x, W1).addBias(b1);
hidden = relu(hidden);
return Matrix.matmul(hidden, W2).addBias(b2);
}
}
// Transformer编码器层
static class EncoderLayer {
final MultiHeadAttention mha;
final FeedForward ff;
EncoderLayer(int dModel, int numHeads, int dFF) {
mha = new MultiHeadAttention(dModel, numHeads);
ff = new FeedForward(dModel, dFF);
}
Matrix forward(Matrix x) {
// 子层1：多头自注意力 + 残差连接 + 层归一化
Matrix attn = mha.forward(x);
Matrix norm1 = layerNorm(x.add(attn));
// 子层2：前馈网络 + 残差连接 + 层归一化
Matrix ffn = ff.forward(norm1);
return layerNorm(norm1.add(ffn));
}
}
// 简化Transformer编码器
static class TransformerEncoder {
final PositionalEncoding pe;
final EncoderLayer[] layers;
final int dModel;
TransformerEncoder(int dModel, int numHeads, int dFF, int numLayers, int maxLen) {
this.dModel = dModel;
pe = new PositionalEncoding(dModel, maxLen);
layers = new EncoderLayer[numLayers];
for (int i = 0; i < numLayers; i++)
layers[i] = new EncoderLayer(dModel, numHeads, dFF);
}
Matrix encode(Matrix inputEmbeddings) {
Matrix x = pe.encode(inputEmbeddings);
for (EncoderLayer layer : layers)
x = layer.forward(x);
return x;
}
}
public static void main(String[] args) {
int dModel = 16, numHeads = 4, dFF = 64, numLayers = 2, seqLen = 5;
// 随机输入嵌入（实际应用中应为词嵌入）
Matrix input = Matrix.randn(seqLen, dModel);
TransformerEncoder encoder = new TransformerEncoder(dModel, numHeads, dFF, numLayers, 100);
Matrix output = encoder.encode(input);
System.out.println("输入维度: " + input.rows + " x " + input.cols);
System.out.println("输出维度: " + output.rows + " x " + output.cols);
System.out.println("输出示例（第一行前5个值）：");
for (int j = 0; j < Math.min(5, output.cols); j++)
System.out.printf("%.4f ", output.data[0][j]);
System.out.println();
}
}
算法复杂度分析
- 自注意力计算：对于序列长度 $L$ 和模型维度 $d$，单次注意力计算为 $O(L^2 \cdot d)$，主要来自 $QK^T$ 的矩阵乘法。
- 多头并行：$h$ 个头分别计算，总复杂度仍为 $O(L^2 \cdot d)$（假设 $d = h \cdot d_{head}$）。
- 前馈网络：$O(L \cdot d \cdot d_{ff})$，其中 $d_{ff}$ 通常为 $4d$。
- 整体编码器：$N$ 层叠加后为 $O(N \cdot L^2 \cdot d + N \cdot L \cdot d \cdot d_{ff})$。
- 关键优势：相比RNN的 $O(L \cdot d^2)$ 串行复杂度，Transformer的 $L^2$ 项可以高度并行，在GPU上效率极高。
扩展方向
- 掩码注意力：在解码器中添加因果掩码，确保位置 $i$ 只能关注到位置 $\le i$ 的信息。
- KV缓存：推理时缓存已计算的Key和Value，避免重复计算，将复杂度从 $O(L^2)$ 降至 $O(L)$。
- 位置编码变体：使用可学习位置嵌入或旋转位置编码（RoPE）替代正弦编码。
总结
Transformer通过自注意力机制实现了序列元素间的全局交互，其并行化特性使其能够高效处理长序列。Java实现展示了多头注意力的分头计算、位置编码的注入方式以及残差连接与层归一化的稳定训练技巧。理解Transformer的底层计算逻辑，是掌握现代大语言模型（如GPT、BERT）的必经之路。
