Potential bug in the CoMPT attention implementation

[qiyan98]

Hi,

Thanks for open sourcing your wonderful work, "Learning Attributed Graph Representations with Communicative Message Passing Transformer." I found a potential bug in your in_scores cross-attention computation. It seems problematic when the edge feature matrix is not symmetric.

Please see my reproduction at:
https://colab.research.google.com/drive/1LLMc9Ue47MoHrTDqJgM-rMMLUacgXM_R?usp=sharing

And by einsum definition, the Eq. (8) and your code seem inconsistent. I could understand the Eq. (8), where the broadcast of node/token embedding for M_i is along the column axis. But your code implementation is not the same as such.

SAIL/Repeat/CoMPT/model_node.py

Lines 181 to 182 in 878c59e

	out_scores = torch.einsum('b h m d, b h m n d -> b h m n', query, key) / math.sqrt(d)
	in_scores = torch.einsum('b h m d, b h m n d -> b h n m', query, key) / math.sqrt(d)

Any help is appreciated.

Many thank,
Qi

[qiyan98]

Your original code seems right except for the transpose that I can't understand.

https://github.com/jcchan23/CoMPT/blob/50bb4a83fc4ac843ee92bf56421cec0ba631a90f/transformer_graph.py#L326

out_scores = torch.einsum('bhmd,bhmnd->bhmn', query, key) / math.sqrt(d_e)
in_scores = torch.einsum('bhnd,bhmnd->bhnm', query, key) / math.sqrt(d_e)

Simply put, why to have 'bhnd,bhmnd->bhnm' instead of 'bhnd,bhmnd->bhmn'? Is there any reason to transpose the in_scores matrix?

[qiyan98]

Ok, maybe I got it. The transpose of in_scores is to make sure the matrix multiplication of matmul(M, V) is consistent with the O(N) node embeddings from V. Each row of M always means some edges attending to the same node.

But in that case, the element-wise multiplication of h(E) = M * K looks inconsistent.
As M = σ(M_o) + σ(M_i) - σ(M_d), h(E) = (σ(M_o) + σ(M_i) - σ(M_d)) * K is related to three M matrices.
M_o-related term is still good with each edge's feature multiplied by the attention coefficient of the same edge. M_i, however, is transposed and the element-wise product leads to inconsistent results. For example, M_i[i,j] refers to the similarity score between E[j, i] and X[i] (roughly speaking).

That being said, σ(M_i[i,j]) * K[i,j] is using a similarity score between E[j, i] and X[i] to update E[i, j]. This doesn't look like typical message passing where the edge update is strongly dependent on related nodes or its own features from previous layers.

Any explanation is appreciated. Thanks!