引言
Transformer是Google于2017年在论文《Attention Is All You Need》中提出的革命性模型架构，它完全摒弃了循环和卷积结构，仅依靠自注意力机制（Self-Attention）就实现了卓越的序列建模效果。Transformer奠定了现代大语言模型（如GPT、BERT）的基础，是当前人工智能领域最重要的算法突破之一。本文将用Java实现Transformer的核心组件——自注意力机制、多头注意力和位置编码，并通过一个简化模型展示其序列转换能力。
算法原理
自注意力机制
自注意力通过计算序列中每个位置与其他所有位置的关联权重来捕捉全局依赖关系。核心计算如下：
Query、Key、Value变换：
$$Q = XW_Q, \quad K = XW_K, \quad V = XW_V$$
注意力分数：
$$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$$
其中 $d_k$ 是Key向量的维度，除以 $\sqrt{d_k}$ 是为了防止点积过大导致Softmax梯度消失。
多头注意力
多头注意力将查询空间划分为 $h$ 个头，每个头独立学习不同的注意力模式：
$$\text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, …, \text{head}_h)W_O$$
$$\text{head}_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V)$$
位置编码
由于Transformer没有循环结构，需要显式注入位置信息。使用正弦/余弦位置编码：
$$PE_{(pos, 2i)} = \sin\left(\frac{pos}{10000^{2i/d_{model}}}\right)$$
$$PE_{(pos, 2i+1)} = \cos\left(\frac{pos}{10000^{2i/d_{model}}}\right)$$
Java实现
1. 矩阵运算工具类
import java.util.Random;
/**
* 矩阵运算工具类：支持Transformer所需的基本线性代数操作
*/
public class MatrixOps {
/** 生成随机矩阵（Xavier初始化） */
public static double[][] randomMatrix(int rows, int cols, long seed) {
Random rand = new Random(seed);
double[][] matrix = new double[rows][cols];
double scale = Math.sqrt(2.0 / (rows + cols));
for (int i = 0; i < rows; i++) {
for (int j = 0; j < cols; j++) {
matrix[i][j] = rand.nextGaussian() * scale;
}
}
return matrix;
}
/** 矩阵乘法：A × B */
public static double[][] matMul(double[][] A, double[][] B) {
int m = A.length, n = B[0].length, p = B.length;
double[][] C = new double[m][n];
for (int i = 0; i < m; i++) {
for (int j = 0; j < n; j++) {
double sum = 0;
for (int k = 0; k < p; k++) {
sum += A[i][k] * B[k][j];
}
C[i][j] = sum;
}
}
return C;
}
/** 矩阵转置 */
public static double[][] transpose(double[][] matrix) {
int rows = matrix.length, cols = matrix[0].length;
double[][] result = new double[cols][rows];
for (int i = 0; i < rows; i++) {
for (int j = 0; j < cols; j++) {
result[j][i] = matrix[i][j];
}
}
return result;
}
/** 矩阵逐元素相加 */
public static double[][] add(double[][] A, double[][] B) {
double[][] result = new double[A.length][A[0].length];
for (int i = 0; i < A.length; i++) {
for (int j = 0; j < A[0].length; j++) {
result[i][j] = A[i][j] + B[i][j];
}
}
return result;
}
/** 矩阵与标量相乘 */
public static double[][] scale(double[][] matrix, double scalar) {
double[][] result = new double[matrix.length][matrix[0].length];
for (int i = 0; i < matrix.length; i++) {
for (int j = 0; j < matrix[0].length; j++) {
result[i][j] = matrix[i][j] * scalar;
}
}
return result;
}
/** Softmax函数：按行计算 */
public static double[][] softmax(double[][] matrix) {
double[][] result = new double[matrix.length][matrix[0].length];
for (int i = 0; i < matrix.length; i++) {
double max = Double.NEGATIVE_INFINITY;
for (double v : matrix[i]) max = Math.max(max, v);
double sum = 0;
for (int j = 0; j < matrix[0].length; j++) {
result[i][j] = Math.exp(matrix[i][j] - max);
sum += result[i][j];
}
for (int j = 0; j < matrix[0].length; j++) {
result[i][j] /= sum;
}
}
return result;
}
/** 打印矩阵（用于调试） */
public static void printMatrix(double[][] matrix, String name) {
System.out.println(name + " [" + matrix.length + "x" + matrix[0].length + "]");
for (int i = 0; i < Math.min(4, matrix.length); i++) {
System.out.print(" ");
for (int j = 0; j < Math.min(6, matrix[0].length); j++) {
System.out.printf("%8.4f ", matrix[i][j]);
}
if (matrix[0].length > 6) System.out.print("...");
System.out.println();
}
if (matrix.length > 4) System.out.println(" ...");
}
}
2. 正弦位置编码
/**
* 正弦位置编码：为序列中的每个位置生成唯一的位置向量
*/
public class PositionalEncoding {
private final double[][] encoding; // [maxLen, dModel]
public PositionalEncoding(int maxLen, int dModel) {
this.encoding = new double[maxLen][dModel];
for (int pos = 0; pos < maxLen; pos++) {
for (int i = 0; i < dModel; i += 2) {
double angle = pos / Math.pow(10000, (2.0 * i) / dModel);
encoding[pos][i] = Math.sin(angle);
if (i + 1 < dModel) {
encoding[pos][i + 1] = Math.cos(angle);
}
}
}
}
/** 获取指定位置编码向量 */
public double[] getEncoding(int position) {
return encoding[position].clone();
}
/** 为输入序列添加位置编码 */
public double[][] addPositionalEncoding(double[][] input) {
double[][] result = new double[input.length][input[0].length];
for (int i = 0; i < input.length; i++) {
for (int j = 0; j < input[0].length; j++) {
result[i][j] = input[i][j] + encoding[i][j];
}
}
return result;
}
}
3. 多头注意力层
/**
* 多头注意力层：实现Transformer核心自注意力机制
*/
public class MultiHeadAttention {
private final int dModel; // 模型维度
private final int numHeads; // 注意力头数
private final int dHead; // 每个头的维度
// 投影权重矩阵
private final double[][] Wq, Wk, Wv, Wo;
public MultiHeadAttention(int dModel, int numHeads, long seed) {
if (dModel % numHeads != 0) {
throw new IllegalArgumentException("dModel必须能被numHeads整除");
}
this.dModel = dModel;
this.numHeads = numHeads;
this.dHead = dModel / numHeads;
// 初始化投影权重
this.Wq = MatrixOps.randomMatrix(dModel, dModel, seed);
this.Wk = MatrixOps.randomMatrix(dModel, dModel, seed + 1);
this.Wv = MatrixOps.randomMatrix(dModel, dModel, seed + 2);
this.Wo = MatrixOps.randomMatrix(dModel, dModel, seed + 3);
}
/**
* 前向传播：计算多头注意力输出
* @param X 输入序列 [seqLen, dModel]
* @return 注意力输出 [seqLen, dModel]
*/
public double[][] forward(double[][] X) {
int seqLen = X.length;
// 1. 线性投影得到Q、K、V
double[][] Q = MatrixOps.matMul(X, Wq); // [seqLen, dModel]
double[][] K = MatrixOps.matMul(X, Wk); // [seqLen, dModel]
double[][] V = MatrixOps.matMul(X, Wv); // [seqLen, dModel]
// 2. 分头处理：将dModel维度分割为numHeads个头
double[][] output = new double[seqLen][dModel];
for (int h = 0; h < numHeads; h++) {
// 提取当前头的Qh, Kh, Vh
double[][] Qh = extractHead(Q, h); // [seqLen, dHead]
double[][] Kh = extractHead(K, h); // [seqLen, dHead]
double[][] Vh = extractHead(V, h); // [seqLen, dHead]
// 3. 计算注意力分数：Qh × Kh^T / sqrt(dHead)
double[][] KhT = MatrixOps.transpose(Kh);
double[][] scores = MatrixOps.matMul(Qh, KhT); // [seqLen, seqLen]
scores = MatrixOps.scale(scores, 1.0 / Math.sqrt(dHead));
// 4. Softmax归一化
double[][] attnWeights = MatrixOps.softmax(scores); // [seqLen, seqLen]
// 5. 加权求和：attnWeights × Vh
double[][] headOutput = MatrixOps.matMul(attnWeights, Vh); // [seqLen, dHead]
// 6. 将头输出拼接回dModel维度
mergeHead(output, headOutput, h);
}
// 7. 最终线性投影
return MatrixOps.matMul(output, Wo); // [seqLen, dModel]
}
/** 从dModel维度中提取第headIndex个头 [seqLen, dHead] */
private double[][] extractHead(double[][] matrix, int headIndex) {
int seqLen = matrix.length;
double[][] head = new double[seqLen][dHead];
int startCol = headIndex * dHead;
for (int i = 0; i < seqLen; i++) {
System.arraycopy(matrix[i], startCol, head[i], 0, dHead);
}
return head;
}
/** 将单头输出合并到结果矩阵中 */
private void mergeHead(double[][] output, double[][] headOutput, int headIndex) {
int startCol = headIndex * dHead;
for (int i = 0; i < output.length; i++) {
for (int j = 0; j < dHead; j++) {
output[i][startCol + j] = headOutput[i][j];
}
}
}
}
4. 前馈网络与层归一化
/**
* 前馈神经网络：Transformer中的FFN子层
*/
public class FeedForwardNetwork {
private final int dModel;
private final int dFF; // 前馈隐藏层维度（通常为4*dModel）
private final double[][] W1, W2;
public FeedForwardNetwork(int dModel, int dFF, long seed) {
this.dModel = dModel;
this.dFF = dFF;
this.W1 = MatrixOps.randomMatrix(dModel, dFF, seed);
this.W2 = MatrixOps.randomMatrix(dFF, dModel, seed + 1);
}
/** ReLU激活函数 */
private double relu(double x) {
return Math.max(0, x);
}
/** 前向传播：FFN(x) = max(0, xW1 + b1)W2 + b2 */
public double[][] forward(double[][] X) {
// 第一层线性变换 + ReLU
double[][] hidden = MatrixOps.matMul(X, W1); // [seqLen, dFF]
for (int i = 0; i < hidden.length; i++) {
for (int j = 0; j < hidden[0].length; j++) {
hidden[i][j] = relu(hidden[i][j]);
}
}
// 第二层线性变换
return MatrixOps.matMul(hidden, W2); // [seqLen, dModel]
}
}
/**
* 层归一化：稳定训练过程
*/
public class LayerNormalization {
private final int features;
private final double epsilon = 1e-6;
public LayerNormalization(int features) {
this.features = features;
}
public double[][] normalize(double[][] X) {
double[][] result = new double[X.length][X[0].length];
for (int i = 0; i < X.length; i++) {
// 计算均值
double mean = 0;
for (double v : X[i]) mean += v;
mean /= features;
// 计算方差
double var = 0;
for (double v : X[i]) var += (v - mean) * (v - mean);
var /= features;
// 归一化
for (int j = 0; j < features; j++) {
result[i][j] = (X[i][j] - mean) / Math.sqrt(var + epsilon);
}
}
return result;
}
}
5. 完整Transformer编码器层
/**
* Transformer编码器层：多头注意力 + 前馈网络 + 残差连接 + 层归一化
*/
public class TransformerEncoderLayer {
private final MultiHeadAttention selfAttention;
private final FeedForwardNetwork ffn;
private final LayerNormalization norm1, norm2;
public TransformerEncoderLayer(int dModel, int numHeads, int dFF, long seed) {
this.selfAttention = new MultiHeadAttention(dModel, numHeads, seed);
this.ffn = new FeedForwardNetwork(dModel, dFF, seed + 10);
this.norm1 = new LayerNormalization(dModel);
this.norm2 = new LayerNormalization(dModel);
}
public double[][] forward(double[][] X) {
// 子层1：多头自注意力 + 残差连接 + 层归一化
double[][] attnOut = selfAttention.forward(X);
double[][] norm1Out = norm1.normalize(MatrixOps.add(X, attnOut));
// 子层2：前馈网络 + 残差连接 + 层归一化
double[][] ffnOut = ffn.forward(norm1Out);
double[][] norm2Out = norm2.normalize(MatrixOps.add(norm1Out, ffnOut));
return norm2Out;
}
}
6. 主程序：序列转换演示
/**
* Transformer序列转换演示
* 任务：学习将输入序列映射到目标序列（如复制或移位）
*/
public class TransformerDemo {
public static void main(String[] args) {
// 模型超参数
int seqLen = 8; // 序列长度
int dModel = 32; // 模型维度
int numHeads = 4; // 注意力头数
int dFF = 128; // 前馈维度
int numLayers = 2; // 编码器层数
System.out.println("=== Transformer自注意力机制演示 ===");
System.out.println("模型配置:");
System.out.println(" 序列长度: " + seqLen);
System.out.println(" 模型维度: " + dModel);
System.out.println(" 注意力头数: " + numHeads);
System.out.println(" 前馈维度: " + dFF);
System.out.println(" 编码器层数: " + numLayers);
System.out.println();
// 创建位置编码
PositionalEncoding posEnc = new PositionalEncoding(seqLen, dModel);
// 创建多层编码器
TransformerEncoderLayer[] encoders = new TransformerEncoderLayer[numLayers];
for (int i = 0; i < numLayers; i++) {
encoders[i] = new TransformerEncoderLayer(dModel, numHeads, dFF, 42 + i);
}
// 构造输入序列：简单的数值模式
double[][] input = new double[seqLen][dModel];
for (int i = 0; i < seqLen; i++) {
input[i][0] = Math.sin(i * 0.5); // 第一个维度携带序列模式
input[i][1] = Math.cos(i * 0.5);
// 其余维度用小的随机值填充
for (int j = 2; j < dModel; j++) {
input[i][j] = (j % 2 == 0) ? 0.01 : -0.01;
}
}
System.out.println("输入序列特征（前2维）:");
for (int i = 0; i < seqLen; i++) {
System.out.printf(" 位置 %d: [%.4f, %.4f]%n", i, input[i][0], input[i][1]);
}
// 添加位置编码
double[][] encoded = posEnc.addPositionalEncoding(input);
System.out.println("\n位置编码已添加，模型现在能区分不同位置的token");
// 前向传播
double[][] output = encoded;
for (int layer = 0; layer < numLayers; layer++) {
output = encoders[layer].forward(output);
System.out.println("\n编码器层 " + (layer + 1) + " 输出特征分布:");
printFeatureStats(output, layer + 1);
}
// 分析注意力效果：计算输出序列的自相关性
System.out.println("\n=== 注意力效果分析 ===");
System.out.println("输出序列各位置与输入的相关性（余弦相似度）:");
for (int i = 0; i < seqLen; i++) {
double[] similarities = new double[seqLen];
for (int j = 0; j < seqLen; j++) {
similarities[j] = cosineSimilarity(input[i], output[j]);
}
System.out.printf(" 输入位置 %d 最关注输出位置: %d (相似度: %.4f)%n",
i, maxIndex(similarities), similarities[maxIndex(similarities)]);
}
System.out.println("\nTransformer通过自注意力机制，使每个位置都能");
System.out.println("直接关注序列中任意其他位置的信息，突破了RNN的串行限制。");
}
/** 计算两个向量的余弦相似度 */
private static double cosineSimilarity(double[] a, double[] b) {
double dot = 0, normA = 0, normB = 0;
for (int i = 0; i < a.length; i++) {
dot += a[i] * b[i];
normA += a[i] * a[i];
normB += b[i] * b[i];
}
return dot / (Math.sqrt(normA) * Math.sqrt(normB) + 1e-8);
}
/** 找到数组最大值的索引 */
private static int maxIndex(double[] arr) {
int maxIdx = 0;
for (int i = 1; i < arr.length; i++) {
if (arr[i] > arr[maxIdx]) maxIdx = i;
}
return maxIdx;
}
/** 打印特征统计信息 */
private static void printFeatureStats(double[][] matrix, int layer) {
double sum = 0, min = Double.POSITIVE_INFINITY, max = Double.NEGATIVE_INFINITY;
int count = 0;
for (double[] row : matrix) {
for (double v : row) {
sum += v;
min = Math.min(min, v);
max = Math.max(max, v);
count++;
}
}
System.out.printf(" 均值: %.4f, 最小: %.4f, 最大: %.4f%n", sum / count, min, max);
}
}
复杂度分析
| 组件 | 时间复杂度 | 空间复杂度 |
|---|---|---|
| 自注意力（单头） | O(seqLen² × dHead) | O(seqLen²) |
| 多头注意力 | O(seqLen² × dModel) | O(seqLen² × numHeads) |
| 前馈网络 | O(seqLen × dModel × dFF) | O(seqLen × dFF) |
| 完整编码器层 | O(seqLen² × dModel + seqLen × dModel × dFF) | O(seqLen × dModel + seqLen²) |
Transformer的主要优势是并行化能力（时间步之间无依赖），但自注意力的O(seqLen²)复杂度使其在长序列上开销较大。
总结
本文实现了Transformer的三个核心组件：
- 自注意力机制通过Query-Key-Value三元组计算序列内任意两个位置的相关性，彻底消除了RNN的串行依赖
- 多头注意力将注意力空间划分为多个子空间，使模型能够同时关注不同语义层面的信息
- 位置编码通过正弦函数为每个位置生成唯一标识，使模型获得序列顺序信息
Transformer的”Attention Is All You Need”理念彻底改变了自然语言处理领域。本文的Java实现虽为教学简化版，但完整保留了核心算法的数学结构，有助于深入理解大语言模型的工作机制。
