Implementing graph co-attention in PyTorch Geometric

[davidbuterez]

Hi @rusty1s,

I noticed that currently there is no graph co-attention implementation available in PyTorch Geometric. I am interested in exploring this model and attempted to implement it. The goal is to model interactions between two graphs with an attention mechanism, leading to two "outer" messages as well as two "inner" (usual) messages.The tricky part to implement in PyG seems to be the outer message co-attention, which requires every node i of graph 1 to interact with every node j of graph 2 (as opposed to just every other node j of graph 1, which is what the existing PyG propagate mechanism seems to support).

I am not used to implementing graph layers in PyG and I am not sure if PyG supports the required functionality in a natural way. I started from the GAT implementation and I was thinking of something along these lines, using repeats for the interactions between the two graphs. This is a first attempt so probably does not compute the correct thing but I was wondering if it would be the correct approach in PyG and if you have any advice!

class CoAttnConv(MessagePassing):
    def __init__(
        self,
        in_channels_1: Union[int, Tuple[int, int]],
        in_channels_2: Union[int, Tuple[int, int]],
        out_channels: int,
        outer_message_dim: int,
        heads: int = 1,
        concat: bool = True,
        negative_slope: float = 0.2,
        dropout: float = 0.0,
        add_self_loops: bool = True,
        edge_dim_1: Optional[int] = None,
        edge_dim_2: Optional[int] = None,
        fill_value: Union[float, Tensor, str] = 'mean',
        bias: bool = True,
        **kwargs,
    ):
        kwargs.setdefault('aggr', 'add')
        super().__init__(node_dim=0, **kwargs)

        self.in_channels_1 = in_channels_1
        self.in_channels_2 = in_channels_2
        self.out_channels = out_channels
        self.heads = heads
        self.concat = concat
        self.negative_slope = negative_slope
        self.dropout = dropout
        self.add_self_loops = add_self_loops
        self.edge_dim_1 = edge_dim_1
        self.edge_dim_2 = edge_dim_2
        self.fill_value = fill_value
        # self.outer_message_dim = outer_message_dim

        self.lin_node_initial_1 = Linear(in_channels_1, heads * out_channels, bias=False, weight_initializer='glorot')
        self.lin_node_initial_2 = Linear(in_channels_2, heads * out_channels, bias=False, weight_initializer='glorot')

        self.lin_node_1 = Linear(heads * out_channels, heads * out_channels, bias=False, weight_initializer='glorot')
        self.lin_node_2 = Linear(heads * out_channels, heads * out_channels, bias=False, weight_initializer='glorot')

        self.att_1 = Parameter(torch.Tensor(1, heads, out_channels))
        self.att_2 = Parameter(torch.Tensor(1, heads, out_channels))

        self.lin_edge_1 = Linear(edge_dim_1, heads * out_channels, bias=False, weight_initializer='glorot')
        self.lin_edge_2 = Linear(edge_dim_2, heads * out_channels, bias=False, weight_initializer='glorot')

        self.lin_outer_message = Linear(out_channels * heads, heads * out_channels, bias=False, weight_initializer='glorot')
        self.lin_outer_final = Linear(out_channels * heads, heads * out_channels, bias=False, weight_initializer='glorot')

        self.ln = torch.nn.LayerNorm(heads * out_channels)

        if bias and concat:
            self.bias_1 = Parameter(torch.Tensor(heads * out_channels))
            self.bias_2 = Parameter(torch.Tensor(heads * out_channels))
        elif bias and not concat:
            self.bias_1 = Parameter(torch.Tensor(out_channels))
            self.bias_2 = Parameter(torch.Tensor(out_channels))
        else:
            self.register_parameter('bias_1', None)
            self.register_parameter('bias_2', None)

        self.reset_parameters()


    def reset_parameters(self):
        self.lin_node_1.reset_parameters()
        self.lin_node_2.reset_parameters()
        self.lin_edge_1.reset_parameters()
        self.lin_edge_2.reset_parameters()
        self.lin_outer_message.reset_parameters()
        self.lin_outer_final.reset_parameters()
        self.lin_node_initial_1.reset_parameters()
        self.lin_node_initial_2.reset_parameters()

        glorot(self.att_1)
        glorot(self.att_2)
        zeros(self.bias_1)
        zeros(self.bias_2)


    def forward(self, x_1, edge_index_1, edge_attr_1, x_2, edge_index_2, edge_attr_2, size, return_attention_weights=None):
        H, C = self.heads, self.out_channels

        x_1 = self.lin_node_initial_1(x_1).view(-1, H, C)
        x_2 = self.lin_node_initial_2(x_2).view(-1, H, C)

        if self.add_self_loops:
            num_nodes_1 = x_1.size(0)
            num_nodes_1 = min(size) if size is not None else num_nodes_1
            edge_index_1, edge_attr_1 = remove_self_loops(edge_index_1, edge_attr_1)
            edge_index_1, edge_attr_1 = add_self_loops(edge_index_1, edge_attr_1, fill_value=self.fill_value, num_nodes=num_nodes_1)

            num_nodes_2 = x_2.size(0)
            num_nodes_2 = min(size) if size is not None else num_nodes_2
            edge_index_2, edge_attr_2 = remove_self_loops(edge_index_2, edge_attr_2)
            edge_index_2, edge_attr_2 = add_self_loops(edge_index_2, edge_attr_2, fill_value=self.fill_value, num_nodes=num_nodes_2)

        alpha_1 = (x_1 * self.att_1)
        alpha_2 = (x_2 * self.att_2)

        N1 = alpha_1.shape[0]
        N2 = alpha_2.shape[0]

        def outer_message(alpha_1, alpha_2, x_1, x_2, direction='1-to-2'):
            # alpha_1 has shape (num_nodes_1, H, C)
            # alpha_2 has shape (num_nodes_2, H, C)

            # num_nodes_1 * num_nodes_2
            total_interactions = N1 * N2

            if direction == '1-to-2':
                # alpha_1_pair has shape (num_nodes_1 * num_nodes_2, H, C)
                # alpha_2_pair has shape (num_nodes_1 * num_nodes_2, H, C)

                # repeat_interleave repeats a tensor [0, 1, 2] like this [0, 0, 1, 1, 2, 2]
                # repeat repeats a tensor [0, 1, 2] like this [0, 1, 2, 0, 1, 2]
                
                alpha_1_pair = alpha_1.repeat_interleave(total_interactions // N1, dim=0)
                alpha_2_pair = alpha_2.repeat(total_interactions // N2, 1, 1)

                # alpha_dot has shape (num_nodes_1 * num_nodes_2, H)
                alpha_dot = (alpha_1_pair * alpha_2_pair).sum(dim=-1)

                # For each node i in graph 1, we want to compute the softmax across all nodes j in graph 2
                # After the first line, alpha_dot has shape (num_atoms_graph_1, num_atoms_graph_2, attn_heads)
                alpha_dot = alpha_dot.view(N1, -1, H)
                alpha_dot = F.softmax(alpha_dot, dim=1)

                # Project features of graph 2
                x_lin = self.lin_outer_message(x_2.view(N2, H * C))
                x_lin = x_lin.reshape(x_lin.shape[0], H, C)

            elif direction == '2-to-1':
                # alpha_1_pair has shape (num_nodes_1 * num_nodes_2, H, C)
                # alpha_2_pair has shape (num_nodes_1 * num_nodes_2, H, C)
                alpha_2_pair = alpha_2.repeat_interleave(total_interactions // N2, dim=0)
                alpha_1_pair = alpha_1.repeat(total_interactions // N1, 1, 1)

                # alpha_dot has shape (num_nodes_1 * num_nodes_2, H)
                alpha_dot = (alpha_2_pair * alpha_1_pair).sum(dim=-1)

                # For each node i in graph 2, we want to compute the softmax across all nodes j in graph 1
                # After the first line, alpha_dot has shape (num_atoms_graph_2, num_atoms_graph_1, attn_heads)
                alpha_dot = alpha_dot.view(N2, -1, H)
                alpha_dot = F.softmax(alpha_dot, dim=1)

                # Project features of graph 1
                x_lin = self.lin_outer_message(x_1.view(N1, H * C))
                x_lin = x_lin.reshape(x_lin.shape[0], H, C)

            # Outer message; has shape (num_nodes_1 or num_nodes_2, H, C)
            outer = torch.einsum('abc,bcd->acd', (alpha_dot, x_lin))

            # Simulate concatenation
            # After first line has shape (num_nodes_1 or num_nodes_2, H * C)
            outer = outer.reshape(outer.shape[0], outer.shape[1] * outer.shape[2])
            outer = self.lin_outer_final(outer)

            return outer      

        m_1 = self.propagate(x=x_1, edge_index=edge_index_1, edge_attr=edge_attr_1, graph_id=1)
        m_2 = self.propagate(x=x_2, edge_index=edge_index_2, edge_attr=edge_attr_2, graph_id=2)

        n_1 = outer_message(alpha_1, alpha_2, x_1, x_2, direction='1-to-2')
        n_2 = outer_message(alpha_1, alpha_2, x_1, x_2, direction='2-to-1')

        overall_message_1 = self.ln(x_1.view(N1, H * C) + m_1 + n_1)
        overall_message_2 = self.ln(x_2.view(N2, H * C) + m_2 + n_2)

        if self.concat:
            overall_message_1 = overall_message_1.view(-1, H * C)
            overall_message_2 = overall_message_2.view(-1, H * C)
        else:
            overall_message_1 = overall_message_1.view(-1, H, C)
            overall_message_2 = overall_message_2.view(-1, H, C)

            overall_message_1 = overall_message_1.mean(dim=1)
            overall_message_2 = overall_message_2.mean(dim=1)

        if self.bias_1 is not None and self.bias_2 is not None:
            overall_message_1 += self.bias_1
            overall_message_2 += self.bias_2

        return overall_message_1, overall_message_2


    def message(self, x_j: Tensor, edge_attr: Tensor, graph_id: int) -> Tensor:
        if graph_id == 1:
            x_j = x_j.view(x_j.shape[0], self.heads * self.out_channels)
            x_j = self.lin_node_1(x_j)
            edge_attr = self.lin_edge_1(edge_attr)
        elif graph_id == 2:
            x_j = x_j.view(x_j.shape[0], self.heads * self.out_channels)
            x_j = self.lin_node_2(x_j)
            edge_attr = self.lin_edge_2(edge_attr)

        return edge_attr * x_j

Thanks in advance!

[rusty1s]

In general, "co-attention" is applied for every node in one graph to another node in the other graph, right? In that case, I think it makes more sense to implement a "dense" version of it without the MessagePassing interface (similar to a Transformer block). For batch-wise mode, you can achieve this by utilizing to_dense_batch. If you still want to use the "sparse" mode via MessagePassing, utilizing repeat and repeat_interleave is correct. Nonetheless, it looks like you are missing the additional batch vector here as well (you generally do not want to perform co-attention between all nodes in all examples).

[davidbuterez]

Thanks for the quick reply! I have a few points:

1. The original implementation is presented as a variation of message passing so I still think it is the most natural. Would the "dense" alternative provide any advantages and are there any examples of this in PyG?
2. Is your point about batch mode referring to the fact that we usually get batched inputs (e.g. 100 graphs) and we don't want to compute the 100 x 100 pair-wise attention coefficients? So instead we should have an extra batch dimension and apply the computation only on the actual pairs.
3. Do we have a suitable (e.g. drug-drug interaction) dataset in PyG to test if the implementation works?