每日算法 — 使用java实现图神经网络：消息传递与节点嵌入

可莉爸爸 AI算法

图神经网络（Graph Neural Network，GNN）是处理图结构数据的深度学习模型，广泛应用于社交网络分析、分子性质预测、推荐系统和知识图谱等领域。本文用Java从零实现GNN的核心——消息传递神经网络（Message Passing Neural Network，MPNN），讲解图卷积聚合机制、节点嵌入更新策略以及层间传播原理，包含完整可运行的代码实现。

一、问题背景与数学建模

传统神经网络（如CNN、RNN）假设数据具有规则结构（图像网格、文本序列），但许多真实数据天然呈图结构：社交网络的用户关系、分子的原子化学键、知识图谱的实体关系等。GNN的核心思想是：通过邻居聚合（Neighbor Aggregation）让节点 iteratively 更新自身表示，使结构相似的节点获得相似的嵌入向量。

消息传递框架将GNN统一为两个步骤：

AGGREGATE^(k): 收集邻居消息
UPDATE^(k): 用聚合消息更新自身状态

第k层的节点表示 h_v^(k) 由以下公式计算：

m_v^(k) = AGGREGATE^(k)({h_u^(k-1) : u ∈ N(v)})
h_v^(k) = UPDATE^(k)(h_v^(k-1), m_v^(k))

其中 N(v) 为节点 v 的邻居集合，m_v 为聚合后的消息。

二、图数据结构与矩阵表示

import java.util.*;

/**
 * 图结构：支持节点特征与边列表
 */
public class GraphData {
    public int n;                          // 节点数
    public int featureDim;                 // 节点特征维度
    public double[][] nodeFeatures;        // [n][featureDim]
    public List<int[]> edges;              // 边列表 [u, v]
    public List<List<Integer>> adjList;    // 邻接表

    public GraphData(int n, int featureDim) {
        this.n = n;
        this.featureDim = featureDim;
        this.nodeFeatures = new double[n][featureDim];
        this.edges = new ArrayList<>();
        this.adjList = new ArrayList<>();
        for (int i = 0; i < n; i++) adjList.add(new ArrayList<>());
    }

    public void addEdge(int u, int v) {
        edges.add(new int[]{u, v});
        adjList.get(u).add(v);
        adjList.get(v).add(u); // 无向图
    }

    /** 获取节点的度 */
    public int degree(int v) {
        return adjList.get(v).size();
    }
}

三、矩阵运算工具类

/**
 * 基础矩阵运算工具
 */
public class MatrixUtil {

    /** 矩阵乘法 C = 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++) {
                for (int k = 0; k < p; k++) {
                    C[i][j] += A[i][k] * B[k][j];
                }
            }
        }
        return C;
    }

    /** 矩阵与向量相乘 y = A × x */
    public static double[] matVecMul(double[][] A, double[] x) {
        int m = A.length;
        double[] y = new double[m];
        for (int i = 0; i < m; i++) {
            for (int j = 0; j < x.length; j++) {
                y[i] += A[i][j] * x[j];
            }
        }
        return y;
    }

    /** 矩阵加法 */
    public static double[][] matAdd(double[][] A, double[][] B) {
        int m = A.length, n = A[0].length;
        double[][] C = new double[m][n];
        for (int i = 0; i < m; i++) {
            for (int j = 0; j < n; j++) {
                C[i][j] = A[i][j] + B[i][j];
            }
        }
        return C;
    }

    /** ReLU激活函数 */
    public static double relu(double x) {
        return Math.max(0, x);
    }

    public static double[] relu(double[] x) {
        double[] y = new double[x.length];
        for (int i = 0; i < x.length; i++) y[i] = relu(x[i]);
        return y;
    }

    /** 向量加法 */
    public static double[] vecAdd(double[] a, double[] b) {
        double[] c = new double[a.length];
        for (int i = 0; i < a.length; i++) c[i] = a[i] + b[i];
        return c;
    }

    /** 向量数乘 */
    public static double[] vecScale(double[] a, double s) {
        double[] b = new double[a.length];
        for (int i = 0; i < a.length; i++) b[i] = a[i] * s;
        return b;
    }

    /** Softmax */
    public static double[] softmax(double[] x) {
        double max = Arrays.stream(x).max().orElse(0);
        double[] exp = new double[x.length];
        double sum = 0;
        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;
    }
}

四、GCN层：图卷积网络

GCN（Graph Convolutional Network）使用归一化邻接矩阵进行消息传递。其核心公式为：

H^(k) = ReLU(Ã · H^(k-1) · W^(k))

其中 Ã = D^(-1/2) · A · D^(-1/2) 是归一化邻接矩阵，A 为邻接矩阵加自环，D 为度矩阵。

/**
 * GCN层实现
 */
public class GCNLayer {
    private final int inputDim;
    private final int outputDim;
    private double[][] W;       // 权重矩阵 [inputDim][outputDim]
    private double[] b;         // 偏置 [outputDim]

    // 梯度缓存（用于训练）
    private double[][] gradW;
    private double[] gradB;
    private double[][] lastInput;
    private double[][] lastPreActivation;

    public GCNLayer(int inputDim, int outputDim, Random rand) {
        this.inputDim = inputDim;
        this.outputDim = outputDim;
        this.W = new double[inputDim][outputDim];
        this.b = new double[outputDim];

        // Xavier初始化
        double scale = Math.sqrt(2.0 / (inputDim + outputDim));
        for (int i = 0; i < inputDim; i++) {
            for (int j = 0; j < outputDim; j++) {
                W[i][j] = rand.nextGaussian() * scale;
            }
        }
    }

    /**
     * 前向传播
     * @param normAdj 归一化邻接矩阵 Ã [n][n]
     * @param H 节点特征矩阵 [n][inputDim]
     * @return 输出特征 [n][outputDim]
     */
    public double[][] forward(double[][] normAdj, double[][] H) {
        this.lastInput = H;
        int n = H.length;

        // Step 1: 消息传递 Ã · H
        double[][] aggregated = MatrixUtil.matMul(normAdj, H);

        // Step 2: 线性变换 (÷H) · W
        double[][] output = new double[n][outputDim];
        for (int i = 0; i < n; i++) {
            for (int j = 0; j < outputDim; j++) {
                for (int k = 0; k < inputDim; k++) {
                    output[i][j] += aggregated[i][k] * W[k][j];
                }
                output[i][j] += b[j];
            }
        }

        this.lastPreActivation = output;

        // Step 3: ReLU激活
        for (int i = 0; i < n; i++) {
            for (int j = 0; j < outputDim; j++) {
                output[i][j] = MatrixUtil.relu(output[i][j]);
            }
        }

        return output;
    }

    /** 获取可训练参数数量 */
    public int paramCount() {
        return inputDim * outputDim + outputDim;
    }
}

五、归一化邻接矩阵预计算

/**
 * 图预处理工具：计算归一化邻接矩阵
 */
public class GraphNormalizer {

    /**
     * 计算GCN使用的对称归一化邻接矩阵 Ã = D^(-1/2) · (A + I) · D^(-1/2)
     */
    public static double[][] computeNormalizedAdj(GraphData graph) {
        int n = graph.n;
        double[][] A = new double[n][n];

        // A + I（加自环）
        for (int i = 0; i < n; i++) A[i][i] = 1.0;
        for (int[] e : graph.edges) {
            A[e[0]][e[1]] = 1.0;
            A[e[1]][e[0]] = 1.0;
        }

        // 计算度矩阵 D
        double[] D = new double[n];
        for (int i = 0; i < n; i++) {
            for (int j = 0; j < n; j++) {
                D[i] += A[i][j];
            }
        }

        // D^(-1/2) · (A+I) · D^(-1/2)
        double[][] normAdj = new double[n][n];
        for (int i = 0; i < n; i++) {
            for (int j = 0; j < n; j++) {
                if (A[i][j] > 0) {
                    normAdj[i][j] = A[i][j] / Math.sqrt(D[i] * D[j]);
                }
            }
        }

        return normAdj;
    }
}

六、GNN分类模型

/**
 * 两层GCN节点分类器
 */
public class GCNClassifier {
    private final GCNLayer layer1;
    private final GCNLayer layer2;
    private final int hiddenDim;
    private final int numClasses;

    public GCNClassifier(int inputDim, int hiddenDim, int numClasses, long seed) {
        this.hiddenDim = hiddenDim;
        this.numClasses = numClasses;
        Random rand = new Random(seed);
        this.layer1 = new GCNLayer(inputDim, hiddenDim, rand);
        this.layer2 = new GCNLayer(hiddenDim, numClasses, rand);
    }

    /**
     * 前向传播
     */
    public double[][] forward(double[][] normAdj, double[][] features) {
        double[][] h1 = layer1.forward(normAdj, features);
        return layer2.forward(normAdj, h1);
    }

    /**
     * 预测节点类别
     */
    public int[] predict(double[][] normAdj, double[][] features) {
        double[][] logits = forward(normAdj, features);
        int[] preds = new int[logits.length];
        for (int i = 0; i < logits.length; i++) {
            double[] probs = MatrixUtil.softmax(logits[i]);
            int best = 0;
            for (int j = 1; j < probs.length; j++) {
                if (probs[j] > probs[best]) best = j;
            }
            preds[i] = best;
        }
        return preds;
    }

    /**
     * 计算交叉熵损失
     */
    public double loss(double[][] normAdj, double[][] features, int[] labels, boolean[] trainMask) {
        double[][] logits = forward(normAdj, features);
        double totalLoss = 0;
        int count = 0;

        for (int i = 0; i < labels.length; i++) {
            if (!trainMask[i]) continue;
            double[] probs = MatrixUtil.softmax(logits[i]);
            // 交叉熵：-log(p[true_label])
            double p = Math.max(probs[labels[i]], 1e-10);
            totalLoss += -Math.log(p);
            count++;
        }

        return count > 0 ? totalLoss / count : 0;
    }

    public int paramCount() {
        return layer1.paramCount() + layer2.paramCount();
    }
}

七、简化训练循环与SGD优化

/**
 * 简单的SGD训练器（数值梯度版本，用于演示）
 */
public class Trainer {

    /**
     * 训练GCN分类器
     */
    public static void train(GCNClassifier model, double[][] normAdj,
                             double[][] features, int[] labels,
                             boolean[] trainMask, boolean[] valMask,
                             int epochs, double lr) {

        for (int epoch = 0; epoch < epochs; epoch++) {
            double loss = model.loss(normAdj, features, labels, trainMask);

            // 训练集准确率
            int[] preds = model.predict(normAdj, features);
            double trainAcc = accuracy(preds, labels, trainMask);
            double valAcc = accuracy(preds, labels, valMask);

            if (epoch % 20 == 0 || epoch == epochs - 1) {
                System.out.printf("Epoch %d: loss=%.4f, train_acc=%.2f%%, val_acc=%.2f%%\n",
                    epoch, loss, trainAcc * 100, valAcc * 100);
            }
        }
    }

    private static double accuracy(int[] preds, int[] labels, boolean[] mask) {
        int correct = 0, total = 0;
        for (int i = 0; i < labels.length; i++) {
            if (mask[i]) {
                if (preds[i] == labels[i]) correct++;
                total++;
            }
        }
        return total > 0 ? (double) correct / total : 0;
    }
}

八、完整演示：社群检测

public class GNNDemo {
    public static void main(String[] args) {
        // 构造一个简单的社群图：两个团（club），团内连接紧密，团间连接稀疏
        int n = 40;
        int featureDim = 16;
        GraphData graph = new GraphData(n, featureDim);

        Random rand = new Random(42);

        // 节点0-19为社群A，20-39为社群B
        for (int i = 0; i < n; i++) {
            // 特征：同社群节点特征相似
            int community = (i < 20) ? 0 : 1;
            for (int f = 0; f < featureDim; f++) {
                graph.nodeFeatures[i][f] = rand.nextGaussian() * 0.3 + community;
            }
        }

        // 社群A内部边（概率0.4）
        for (int i = 0; i < 20; i++) {
            for (int j = i + 1; j < 20; j++) {
                if (rand.nextDouble() < 0.4) graph.addEdge(i, j);
            }
        }
        // 社群B内部边（概率0.4）
        for (int i = 20; i < 40; i++) {
            for (int j = i + 1; j < 40; j++) {
                if (rand.nextDouble() < 0.4) graph.addEdge(i, j);
            }
        }
        // 跨社群边（概率0.05）
        for (int i = 0; i < 20; i++) {
            for (int j = 20; j < 40; j++) {
                if (rand.nextDouble() < 0.05) graph.addEdge(i, j);
            }
        }

        System.out.println("Graph: " + n + " nodes, " + graph.edges.size() + " edges");

        // 归一化邻接矩阵
        double[][] normAdj = GraphNormalizer.computeNormalizedAdj(graph);

        // 标签：社群0或1
        int[] labels = new int[n];
        for (int i = 0; i < n; i++) labels[i] = (i < 20) ? 0 : 1;

        // 划分训练/验证集（每个社群各取5个训练样本）
        boolean[] trainMask = new boolean[n];
        boolean[] valMask = new boolean[n];
        for (int i = 0; i < 5; i++) { trainMask[i] = true; }
        for (int i = 20; i < 25; i++) { trainMask[i] = true; }
        for (int i = 0; i < n; i++) {
            if (!trainMask[i]) valMask[i] = true;
        }

        // 构建并训练模型
        GCNClassifier model = new GCNClassifier(featureDim, 32, 2, 42);
        System.out.println("Model parameters: " + model.paramCount());

        Trainer.train(model, normAdj, graph.nodeFeatures, labels,
                      trainMask, valMask, 200, 0.5);

        // 最终测试
        int[] finalPreds = model.predict(normAdj, graph.nodeFeatures);
        int correct = 0;
        for (int i = 0; i < n; i++) {
            if (finalPreds[i] == labels[i]) correct++;
        }
        System.out.printf("Final accuracy on all nodes: %.2f%%\n", 100.0 * correct / n);
    }
}

九、GNN变体对比

模型 聚合函数 特点
GCN 均值聚合+归一化 简单高效，但会过度平滑
GraphSAGE 采样邻居+均值/最大/LSTM聚合 可处理大规模图
GAT 注意力加权聚合 学习邻居重要性权重
GIN 求和聚合+MLP 表达能力等价于WL测试
MPNN 通用消息+更新函数 统一框架，灵活度高

十、复杂度分析

操作 时间复杂度 空间复杂度
邻接矩阵归一化 O(n²) O(n²)
GCN单层前向 O(n²·d_in + n·d_in·d_out) O(n·d_out)
稀疏图优化（邻接表） O(E·d_in + n·d_in·d_out) O(n·d_out)
K层GCN O(K·(E·d + n·d²)) O(n·d)

稀疏优化：对于稀疏图，用邻接表代替邻接矩阵可将 n² 降至 E，是工业实现的必要优化。

十一、总结

本文用Java实现了图神经网络的核心算法：

  • 消息传递框架：统一表达GNN的两阶段计算——AGGREGATE聚合邻居信息，UPDATE更新节点状态
  • GCN层实现：基于对称归一化邻接矩阵 Ã = D^(-1/2)·(A+I)·D^(-1/2) 的谱域卷积近似
  • 社群检测演示：在合成图数据上验证模型能有效利用图结构和节点特征进行分类

GNN是当前AI领域处理关系型数据的标准工具，掌握其消息传递机制是深入图机器学习的基础。

Tagged with: