FDUDSDE / MAGIC

Codes and data for USENIX Security 24 paper "MAGIC: Detecting Advanced Persistent Threats via Masked Graph Representation Learning"
MIT License
64 stars 10 forks source link

About the GATConv class defined #7

Closed kamelferrahi closed 6 months ago

kamelferrahi commented 6 months ago

I would Like to ask what's the difference between the defined class here and the standard GraphConv dgl class ? Have you modified the class to handle different edge types when computing attention heads?

Jimmyokok commented 6 months ago

Yes. The standard dgl version of the GATConv class implements an additive attention mechanism between the source and destination nodes:

el = (feat_src * self.attn_l).sum(dim=-1).unsqueeze(-1)
er = (feat_dst * self.attn_r).sum(dim=-1).unsqueeze(-1)
graph.srcdata.update({'ft': feat_src, 'el': el})
graph.dstdata.update({'er': er})
graph.apply_edges(fn.u_add_v('el', 'er', 'e'))
e = self.leaky_relu(graph.edata.pop('e'))
graph.edata['a'] = self.attn_drop(edge_softmax(graph, e))

where the attention weight $a=Dropout(Softmax(LeakyReLU(Linear(feat{src} | feat{dst})))$ Our version includes edge features into this attention mechanism:

eh = (feat_src * self.attn_h).sum(-1).unsqueeze(-1)
et = (feat_dst * self.attn_t).sum(-1).unsqueeze(-1)
graph.srcdata.update({'hs': feat_src, 'eh': eh})
graph.dstdata.update({'et': et})
ee = (feat_edge * self.attn_e).sum(-1).unsqueeze(-1)
graph.edata.update({'ee': ee})
graph.apply_edges(fn.u_add_e('eh', 'ee', 'ee'))
graph.apply_edges(fn.e_add_v('ee', 'et', 'e'))
e = self.leaky_relu(graph.edata.pop('e'))
graph.edata['a'] = self.attn_drop(edge_softmax(graph, e))

where the attention weight $a=Dropout(Softmax(LeakyReLU(Linear(feat{src} | feat{edge} | feat_{dst})))$

kamelferrahi commented 6 months ago

Thank you for the clarification