PyG中LSTM聚合为何需排序edge_index？相关疑问求解

我使用PyG库的GraphSAGE进行节点嵌入，选择LSTM作为聚合函数。输入包括归一化的2D节点特征（V×2）、格式为(2, |E|)的edge_index。使用LSTM聚合器时必须排序edge_index，否则会报错：

ValueError: Can not perform aggregation since the 'index' tensor. is not sorted.

执行sorted_edge_index, _ = torch.sort(edge_index, dim=1)后，edge_index变为全自环形式如(0,0)、(1,1)等，原边关系丢失。现存在两个疑问：

1. 为何LSTM聚合需要排序edge_index？
2. 如此排序后模型如何识别节点连接，是否存在弊端？

问题解答 #

1. LSTM聚合要求排序edge_index的原因 #

GraphSAGE的LSTM聚合器是把每个节点的邻居特征作为序列输入LSTM进行建模，而PyG中该聚合器的实现逻辑依赖同一目标节点的所有邻居边必须连续排列。edge_index的第二行是目标节点的索引，只有当这一行的索引是有序的，聚合器才能准确将同一个节点的所有邻居归为一组，按序列喂给LSTM处理。如果不排序，同一节点的邻居边会分散在edge_index中，聚合器无法正确分组，因此抛出报错。

2. 错误排序的问题与正确处理方式 #

你当前使用的torch.sort(edge_index, dim=1)是错误的——这个操作会对edge_index的两行分别独立排序，直接破坏了源节点与目标节点的对应关系，才会出现全自环的异常情况。

正确的排序方式是依据edge_index的第二行（目标节点索引）对整个edge_index进行排序，这样既能保证同一目标节点的邻居边连续排列，又能完整保留原有的边关系。正确代码如下：

# 获取按目标节点索引排序后的位置索引
sorted_indices = torch.argsort(edge_index[1])
# 用该索引重新排列整个edge_index
sorted_edge_index = edge_index[:, sorted_indices]

采用正确排序方式后，模型可以正常识别节点连接：同一目标节点的所有邻居边会连续排列，LSTM能按序列处理这些邻居特征，同时源节点与目标节点的对应关系完全保留，不存在任何弊端，这正是LSTM聚合器要求的标准输入格式。

内容的提问来源于stack exchange，提问作者omegabuz