图神经网络三元组（Triplet）特征编码全解析

一、核心总纲

所有代码的核心目标：在图神经网络中，对「三个节点 + 三条边 + 整张图」构成的“三元组（triplet）”进行特征编码，最终输出指定维度的特征向量。

二、输入参数统一说明

所有函数的输入参数完全一致，含义如下：

参数名	含义
z	节点特征张量
edge_fts	边特征张量
graph_fts	整张图的全局特征张量
nb_triplet_fts	中间特征投影的目标维度
out_size	最终输出特征的维度

三、通用执行流程（所有模型的底层逻辑）

所有函数均遵循以下4个核心步骤，仅在细节上有差异：

步骤1：特征维度统一投影

通过 hk.Linear（线性层）或 hk.nets.MLP（多层感知机），将节点/边/图的原始特征映射到 nb_triplet_fts 维度，确保不同类型特征的维度一致。

步骤2：维度扩展（核心操作）

通过 jnp.expand_dims(tensor, axis=perm) 为特征张量添加新维度，核心目的是：

• 让三个节点、三条边的特征在不同维度上对齐，支持广播相加
• 实现节点/边/图特征的“三元组交互”
• 常用扩展维度：(2, 3)、(1, 3)、(1, 2)、(1, 2, 3) 等

步骤3：特征聚合拼接

通过 functools.reduce(lambda x, y: x + y, [...]) 将所有扩展后的特征张量直接相加，形成包含所有三元组交互信息的大张量。

步骤4：最终聚合 + 输出映射

对拼接后的张量在指定维度（通常是axis=1）做聚合操作，再通过线性层/MLP映射到 out_size 维度：

• 聚合方式：max/mean/sum/max-min
• 输出映射：hk.Linear 或 hk.nets.MLP

四、各模型核心差异详解

1. QUADNODEMINMAX（get_triplet_msgs_quad）

特征使用	编码方式	聚合方式	输出处理
仅节点特征（4组）	hk.Linear 投影	max - min（最大值减最小值）	无额外线性层，直接返回聚合结果
无边缘/图特征	-	-	-

2. CONCATREP（get_triplet_msgs_concatrep）

特征使用	编码方式	聚合方式	输出处理
节点+边+图特征	自定义rep_fn：拼接 MLP(proj) + MLP(ff)	jnp.max（axis=1）	hk.Linear 映射到out_size
各特征各3组	-	-	-

3. TANHEXPANDTRIPLETS（get_triplet_msgs_tanhexplandtriplets）

特征使用	编码方式	聚合方式	输出处理
节点+边+图特征	节点/边：Linear；图：MLP	jnp.max（axis=1）	tanh激活后，hk.Linear映射到out_size
图特征扩展3种维度后加入	-	-	-

4. DIV2MEAN（get_triplet_msgs_div2mean）

特征使用	编码方式	聚合方式	输出处理
节点+边特征	节点特征先除以2，再Linear投影	jnp.mean（axis=1）	hk.Linear 映射到out_size
无图特征	-	-	-

5. MAXMEAN（get_triplet_msgs_maxmean）

特征使用	编码方式	聚合方式	输出处理
节点+边+图特征	Linear投影	先将张量限制≥-100，再jnp.mean（axis=1）	hk.Linear 映射到out_size
-	-	-	-

6. 种子模型v1~v9（CLRSSeedModels）

模型	核心差异点
v1	标准基线：节点+边+图（Linear）+ max聚合 + Linear输出
v2	输出层替换为MLP（[out_size, out_size]）
v3	聚合方式从max改为sum
v4	节点特征用MLP编码（Linear→MLP）
v5	边特征用MLP编码（Linear→MLP）
v6	图特征用MLP编码（Linear→MLP）
v7	节点用MLP + 移除图特征
v8	极简版：仅对z做MLP + tile扩维（无三元组交互）
v9	特征融合方式改为乘法（proj*ff），而非拼接

五、关键总结

1. 所有模型的核心逻辑一致：投影→扩维→拼接→聚合→输出，仅在“特征使用范围、编码方式、聚合方式”三个维度做差异化调整；
2. 维度扩展（expand_dims）是实现三元组交互的核心，不同的axis参数对应特征在张量中的不同交互位置；
3. 种子模型v1~v9是对基础三元组编码的“消融实验”，逐一替换编码层/聚合方式/特征类型，验证各组件的作用。