Using PyG for Factor Graphs

❓ Questions & Help

Hi there, I was wondering if we could use PyG for Factor Graphs? I’m trying to integrate some form of belief propagation in Factor Graphs together with GNNs, and would like to use the same interface for the belief propagation. I’ve structured the data such that it looks like a bipartite graph from networkx, but can’t really figure out how to solve the MessagePassing protocol.

Here’s what it looks like:

The red nodes are the variable nodes, with a single-valued feature on them. Blue nodes are factor nodes which have the following form: m_{fac_i -> var_j} = -min[2] m_{var_k -> fac_i, where k \in N(j)}} This means that my factor basically returns the 2nd smallest value in the neighborhood, where the neighborhood does not include the message coming from the variable I’m sending to (variable j in this case).

I understand that at the moment, PyG only has mean, max, sum. Hence, min in this case can be turned to max pretty easily. But the problem I see is that I can’t return the 2nd highest max value, is there a solution to this? E.g. creating a custom aggregation function?

Secondly, I’m having a bit of issues trying to execute the variable-to-factor and factor-to-variable messages. At the moment, I’ve represented the data point as a type of BipartiteData, where x_s are my variable features, and x_t is just the identity values (since my factors don’t have features, I only want the messages). m_{var_i -> fac_j} = sum_{k \in N(i) \ j} m_{fac_k -> var_i}

Here’s my implementation:

class BipartiteData(pyg.data.Data):
    def __init__(self, edge_index, x_s, x_t):
        super(BipartiteData, self).__init__()
        self.edge_index = edge_index
        self.edge_index_T = torch.stack([edge_index[1], edge_index[0]], dim=0)
        self.x_s = x_s
        self.x_t = x_t

So, I’m thinking of defining 2 message passing classes, a Var2Fac and Fac2Var class to pass the messages accordingly. But I’m still a bit lost on how I can do this.

[rusty1s]

Hi, and thanks for your questions.

1. You can create a custom aggregation function by overwriting:

def aggregate(self, inputs, index):
    ...

but from a standpoint of efficiency, it is quite hard to retrieve the 2nd highest max value. This would require us to support some kind of scatter_top_k implementation, which is not supported at the moment. Your best bet is, in case efficiency is not a problem, to manually implement it via pure Python/PyTorch, e.g.:

def aggregate(self, inputs, index, dim_size):
    out = torch.zeros(dim_size, inputs.size(1))
    for i in range(dim_size):
        out[i] = inputs[index == i].sort(dim=0)[0][1]
    return out

1. You only need to define two message passing classes if messages between different node types should be treated differently. Otherwise, something like that will work:

    conv1 = SAGEConv((num_var_features, num_fac_features), 256)
    conv2 = SAGEConv((num_fac_fatures, num_var_features), 256)
    new_x_t = conv1((x_s, x_t), edge_index)
    new_x_s = conv2((x_t, x_s), edge_index_T) 

I hope that this is useful to you.

[CarloLucibello]

This paper https://arxiv.org/pdf/1705.08415.pdf defines an auxiliary graph where nodes are oriented edges of the original graph and the adjacency matrix is the non-backtracking matrix