Inputs for SelfAttention mechanism

google-deepmind / graph_nets

Build Graph Nets in Tensorflow

https://arxiv.org/abs/1806.01261

Apache License 2.0

5.34k stars 783 forks source link

Inputs for SelfAttention mechanism #85

Closed dxcl closed 5 years ago

dxcl commented 5 years ago

Apologies if this is inappropriate for an issue or relates too closely to #77. I've had a look at the documentation for SelfAttention in modules.py and its unit test, but I'm still unsure of the following: if (for example), I've got a graph with nodes A, B, and C, feature vectors Av, Bv and Cv for all 3 nodes, and sender and receiver information for those nodes, what are my "keys", "values", and "queries"? Are queries the relationships a given node has with other nodes, and keys the indices of the nodes?

alvarosg commented 5 years ago

You would need a tensor of keys, a tensor of queries (same feature_size as keys) and a tensor of values for each node in the batch.

Each of these tensors will have shape: [total_num_nodes, feature_size]

There are many ways you can get these three tensors, for example:

Take the nodes that are output by a graph network, and split the feature size in 3 (keys, queries, values = tf.split(output_graph.nodes, [key_size, key_size, value_size], axis=-1))
Take the nodes that are output by a graph network, and use 3 different MLPs/LinearLayer with consistent output sizes to create the three tensors. (keys = snt.Linear(output_size=query_size)(output_graph.nodes) and so on for queries and values).

Hope this helps!

dxcl commented 5 years ago

Thanks very much for the prompt response! So just to clarify, the flow would be:

Input graph --> Graph Network --> Output graph 1 --> Generate keys, queries etc. via splitting / layers --> feed into SelfAttention module --> Output graph 2

where the 2nd output graph has the aggregated attended value for each of the nodes?

alvarosg commented 5 years ago

Not necessarily, that was just an example of how to get some tensors nodes. In that flow you will have communication both via self attention and via message passing in the first GraphNetwork.

If you want to have communication only via self-attention, instead of using a modules.GraphNetwork you could use a modules.GraphIndependent (providing a node_model_fn to encode the nodes independently), or a blocks.NodeBlock (with both use_*_edges options disabled) if you want the nodes to also be conditioned by some global vector in the graph.

dxcl commented 5 years ago

Ok great - thanks so much! I'll try and get a flow with pure communication going ASAP.