每日算法 — 使用java实现Transformer：自注意力机制与多头注意力计算

引言

Transformer架构自2017年提出以来，彻底改变了自然语言处理领域。与RNN和LSTM不同，Transformer完全基于注意力机制，摒弃了循环结构，实现了高效的并行计算。本文将用Java从零实现Transformer的核心组件——自注意力机制（Self-Attention）与多头注意力（Multi-Head Attention），帮助读者深入理解其内部计算流程。

核心原理

自注意力机制的核心思想是：对于序列中的每个位置，计算其与序列中所有其他位置的关联权重，然后基于这些权重对全序列信息进行加权聚合。这种”全局视野”使模型能够捕获长距离依赖关系。

多头注意力则将注意力计算拆分为多个并行的”头”，每个头学习不同的特征子空间，最终拼接所有头的输出，增强模型的表达能力。

数学表达

给定查询矩阵 Q、键矩阵 K、值矩阵 V，注意力输出为：

Attention(Q, K, V) = softmax(QK^T / √d_k) · V

其中 d_k 是键向量的维度，除以 √d_k 是为了防止点积结果过大导致softmax梯度消失。

Java实现

项目结构

src/
├── Matrix.java          // 矩阵运算工具类
├── SelfAttention.java   // 单头自注意力实现
├── MultiHeadAttention.java // 多头注意力实现
└── TransformerDemo.java // 演示主类

矩阵运算工具类

import java.util.Random;

/**
 * 矩阵运算工具类，支持基本的矩阵操作
 */
public class Matrix {
    private final double[][] data;
    public final int rows;
    public final int cols;

    public Matrix(int rows, int cols) {
        this.rows = rows;
        this.cols = cols;
        this.data = new double[rows][cols];
    }

    public Matrix(double[][] data) {
        this.rows = data.length;
        this.cols = data[0].length;
        this.data = new double[rows][cols];
        for (int i = 0; i < rows; i++) {
            System.arraycopy(data[i], 0, this.data[i], 0, cols);
        }
    }

    public void set(int row, int col, double value) {
        data[row][col] = value;
    }

    public double get(int row, int col) {
        return data[row][col];
    }

    /**
     * 使用Xavier初始化填充矩阵
     */
    public void xavierInit(Random rand) {
        double scale = Math.sqrt(2.0 / (rows + cols));
        for (int i = 0; i < rows; i++) {
            for (int j = 0; j < cols; j++) {
                data[i][j] = rand.nextGaussian() * scale;
            }
        }
    }

    /**
     * 矩阵乘法: this * other
     */
    public Matrix multiply(Matrix other) {
        if (this.cols != other.rows) {
            throw new IllegalArgumentException("矩阵维度不匹配");
        }
        Matrix result = new Matrix(this.rows, other.cols);
        for (int i = 0; i < this.rows; i++) {
            for (int j = 0; j < other.cols; j++) {
                double sum = 0;
                for (int k = 0; k < this.cols; k++) {
                    sum += this.data[i][k] * other.data[k][j];
                }
                result.data[i][j] = sum;
            }
        }
        return result;
    }

    /**
     * 矩阵转置
     */
    public Matrix transpose() {
        Matrix result = new Matrix(cols, rows);
        for (int i = 0; i < rows; i++) {
            for (int j = 0; j < cols; j++) {
                result.data[j][i] = data[i][j];
            }
        }
        return result;
    }

    /**
     * 矩阵加法
     */
    public Matrix add(Matrix other) {
        if (this.rows != other.rows || this.cols != other.cols) {
            throw new IllegalArgumentException("矩阵维度不匹配");
        }
        Matrix result = new Matrix(rows, cols);
        for (int i = 0; i < rows; i++) {
            for (int j = 0; j < cols; j++) {
                result.data[i][j] = this.data[i][j] + other.data[i][j];
            }
        }
        return result;
    }

    /**
     * 按行应用softmax
     */
    public Matrix softmax() {
        Matrix result = new Matrix(rows, cols);
        for (int i = 0; i < rows; i++) {
            double max = Double.NEGATIVE_INFINITY;
            for (int j = 0; j < cols; j++) {
                max = Math.max(max, data[i][j]);
            }
            double sum = 0;
            for (int j = 0; j < cols; j++) {
                result.data[i][j] = Math.exp(data[i][j] - max);
                sum += result.data[i][j];
            }
            for (int j = 0; j < cols; j++) {
                result.data[i][j] /= sum;
            }
        }
        return result;
    }

    /**
     * 标量除法
     */
    public Matrix divide(double scalar) {
        Matrix result = new Matrix(rows, cols);
        for (int i = 0; i < rows; i++) {
            for (int j = 0; j < cols; j++) {
                result.data[i][j] = data[i][j] / scalar;
            }
        }
        return result;
    }

    @Override
    public String toString() {
        StringBuilder sb = new StringBuilder();
        for (int i = 0; i < Math.min(rows, 5); i++) {
            for (int j = 0; j < Math.min(cols, 5); j++) {
                sb.append(String.format("%.4f ", data[i][j]));
            }
            if (cols > 5) sb.append("...");
            sb.append("\n");
        }
        if (rows > 5) sb.append("...\n");
        return sb.toString();
    }
}

多头注意力实现

import java.util.Random;

/**
 * 多头注意力机制实现
 * 
 * 参数说明:
 * - dModel: 模型维度(输入/输出维度)
 * - numHeads: 注意力头的数量
 * - dK: 每个头的键/查询维度 = dModel / numHeads
 */
public class MultiHeadAttention {
    private final int dModel;
    private final int numHeads;
    private final int dK;
    private final Random rand;

    // 权重矩阵: W_Q, W_K, W_V (每个头共享，也可每个头独立)
    private final Matrix wq;
    private final Matrix wk;
    private final Matrix wv;
    private final Matrix wo; // 输出投影矩阵

    public MultiHeadAttention(int dModel, int numHeads) {
        if (dModel % numHeads != 0) {
            throw new IllegalArgumentException("dModel必须能被numHeads整除");
        }
        this.dModel = dModel;
        this.numHeads = numHeads;
        this.dK = dModel / numHeads;
        this.rand = new Random(42);

        // 初始化权重矩阵
        this.wq = new Matrix(dModel, dModel);
        this.wk = new Matrix(dModel, dModel);
        this.wv = new Matrix(dModel, dModel);
        this.wo = new Matrix(dModel, dModel);

        wq.xavierInit(rand);
        wk.xavierInit(rand);
        wv.xavierInit(rand);
        wo.xavierInit(rand);
    }

    /**
     * 前向传播
     * @param input [seqLen, dModel] 输入序列
     * @return [seqLen, dModel] 注意力输出
     */
    public Matrix forward(Matrix input) {
        int seqLen = input.rows;

        // 1. 线性投影得到Q, K, V
        Matrix q = input.multiply(wq); // [seqLen, dModel]
        Matrix k = input.multiply(wk); // [seqLen, dModel]
        Matrix v = input.multiply(wv); // [seqLen, dModel]

        // 2. 拆分为多个头并计算注意力
        Matrix output = new Matrix(seqLen, dModel);

        for (int h = 0; h < numHeads; h++) {
            // 提取当前头的Q_h, K_h, V_h
            Matrix qh = extractHead(q, h);   // [seqLen, dK]
            Matrix kh = extractHead(k, h);   // [seqLen, dK]
            Matrix vh = extractHead(v, h);   // [seqLen, dK]

            // 3. 计算注意力: softmax(Q_h * K_h^T / sqrt(dK)) * V_h
            Matrix scores = qh.multiply(kh.transpose()); // [seqLen, seqLen]
            scores = scores.divide(Math.sqrt(dK));
            Matrix attnWeights = scores.softmax();       // [seqLen, seqLen]
            Matrix headOutput = attnWeights.multiply(vh); // [seqLen, dK]

            // 4. 将头的输出拼接回完整维度
            mergeHead(output, headOutput, h);
        }

        // 5. 最终线性投影
        return output.multiply(wo);
    }

    /**
     * 从拼接的Q/K/V中提取指定头的矩阵
     */
    private Matrix extractHead(Matrix m, int headIdx) {
        int seqLen = m.rows;
        Matrix head = new Matrix(seqLen, dK);
        int startCol = headIdx * dK;
        for (int i = 0; i < seqLen; i++) {
            for (int j = 0; j < dK; j++) {
                head.set(i, j, m.get(i, startCol + j));
            }
        }
        return head;
    }

    /**
     * 将单头输出合并到最终输出矩阵中
     */
    private void mergeHead(Matrix output, Matrix head, int headIdx) {
        int seqLen = output.rows;
        int startCol = headIdx * dK;
        for (int i = 0; i < seqLen; i++) {
            for (int j = 0; j < dK; j++) {
                double val = output.get(i, startCol + j) + head.get(i, j);
                output.set(i, startCol + j, val);
            }
        }
    }
}

演示主类

/**
 * Transformer注意力机制演示
 */
public class TransformerDemo {
    public static void main(String[] args) {
        // 模拟参数: 序列长度4, 模型维度8, 4个注意力头
        int seqLen = 4;
        int dModel = 8;
        int numHeads = 4;

        // 构建随机输入 [seqLen, dModel]
        Matrix input = new Matrix(seqLen, dModel);
        Random rand = new Random(123);
        for (int i = 0; i < seqLen; i++) {
            for (int j = 0; j < dModel; j++) {
                input.set(i, j, rand.nextGaussian() * 0.1);
            }
        }

        System.out.println("输入矩阵 (前5x5预览):");
        System.out.println(input);

        // 创建多头注意力层并前向传播
        MultiHeadAttention mha = new MultiHeadAttention(dModel, numHeads);
        Matrix output = mha.forward(input);

        System.out.println("多头注意力输出 (前5x5预览):");
        System.out.println(output);

        // 验证输出维度
        System.out.println("输入维度: " + input.rows + "x" + input.cols);
        System.out.println("输出维度: " + output.rows + "x" + output.cols);
        System.out.println("✓ 维度检查通过，注意力机制运行正常");
    }
}

复杂度分析

对于序列长度 n、模型维度 d、头数 h 的多头注意力：

其中 O(n²) 是注意力权重的计算开销，这是Transformer的主要瓶颈。对于长序列，可使用稀疏注意力、线性注意力或Flash Attention等优化方案。

总结

本文用Java完整实现了Transformer的多头注意力机制，包括矩阵运算、分头计算、softmax归一化与输出投影。读者可以基于此扩展实现完整Transformer（加入位置编码、前馈网络、层归一化与残差连接），进一步探索BERT、GPT等预训练模型的原理。