Self attention and feature fusion over graphs

username: kristianeschenburg
cooked: <p>Hi</p>
<p>This is a repost from <a href="https://discuss.dgl.ai/t/self-attention-and-feature-fusion-over-graphs/2455" rel="noopener nofollow ugc">here</a>, as I’m hoping someone in this forum might have some additional expertise in feature fusion.</p>
<p>Is it possible to perform feature fusion using a self-attention mechanism?  I have data distributed over a graph, and multiple feature modalities sampled per node that I would like to “optimally” combine for a node classification problem.  Let’s say that I have a graph $G=(V,E)$ with $N$ nodes.  For each node, we sample $M$ different modalities, such that each node $v$ is characterized by $M$ feature vectors $\in \mathbb{R}^{m_{1}}, \mathbb{R}^{m_{2}}…\mathbb{R}^{m_{M}}$.  For simplicity, assume $m_{1} = m_{2} = … = m_{M} = D$.  I’m interested in “fusing” these features in a  more intelligent way than simply concatenating them and passing them through a linear layer or MLP.</p>
<p>Assuming each node $v$ is characterzed by a feature matrix $x_{i} \in \mathbb{R}^{M \times D}$ (where $M$ is the number of modalities, and $D$ the input dimension), a transformer approach over modalities would yield something like this:</p>
<p>$$<br>
\begin{align}<br>
Q = x_{i}W_{q} \<br>
K = x_{i}W_{k} \<br>
V = x_{i}W_{v}<br>
\end{align}<br>
$$</p>
<p>where $W_{q,k,v} \in \mathbb{R}^{D \times p}.  We then have that</p>
<p>$$<br>
\begin{align}<br>
Z = softmax(\frac{QK^{T}}{\sqrt{D}}, axis=1)V<br>
\end{align}<br>
$$</p>
<p>where $Z \in \mathbb{R}^{M \times D} and $Z_{i}$ is the fusion of modalities with respect to modality $i$ (I think – correct me if I’m wrong, please).  I could then sum over the rows of $Z$ to compute the final combination.  Does this seem correct?  I havn’t found many papers on using self-attention for feature fusion, so any help is much appreciated.</p>
<p>k</p>

username: tom
cooked: <p>So I used the link for the better maths rendering, but you have two different uses of the index <em>i</em> there: in the input, it is an index over the nodes, in the output, you use it as an index over the modality in your question. But in self-attention, one would typically expect the inputs and the outputs to match in shape, so if your inputs are N x M x D in your notation, so would your output be of the same.</p>
<p>That said, attention seems like a thing to try in a situation like yours. If you want to go to N x D outputs, there would seem to several options</p>
<ul>
<li>pool over the features after the attention,</li>
<li>have a “node token” additional modality that works similar to the class token in BERT / vision transformers, in particular if you have several attention blocks stacked as in transformers,</li>
<li>go the “external attention” way of separating Q from KV and make Q independent of the modality, this would make your output independent of the modality. However, then the node must “know” which modalities to attend to.</li>
</ul>
<p>I am afraid I cannot help you with the literature bit.</p>
<p>Best regards</p>
<p>Thomas</p>

username: kristianeschenburg
cooked: <p>Hi <a class="mention" href="/u/tom">@tom</a></p>
<p>Thank you for the insights – to your first point, I did incorrectly use the $i$ index twice, my mistake</p>
<p>I think a permutation invariant function for aggregating attentions as suggested in your first point would likely be the easiest approach.</p>
<p>I’m not sure I understand the “external attention” approach, could you elaborate a bit here?</p>
<p>k</p>

username: tom
cooked: <p>So to me, the defining bit for <em>self</em>-attention is that QKV all come from the same input.<br>
Using the image from the famous <a href="https://arxiv.org/abs/1706.03762" rel="noopener nofollow ugc">Vasvani et al: Attention is all you need</a> paper: In the decoder, you have an attention where the Q comes from the decoder while KV is on the encoder (so the “decoder looks at the encoder” - the red circle):</p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/5/c/5c03aab2018f1cb1a8635d9a96cbb022362814f6.png" alt="attn" data-base62-sha1="d7ZMz4HcAZak5M8ZYsLSsI38j9c" width="292" height="431"></p>
<p>But in your case it may be natural to take the Q just from the node, so the linear for Q takes N x D and the linear for KV take N X M x D. This may or may not work better than just pooling over the modalities after the attention layer.</p>
<p>Best regards</p>
<p>Thomas</p>