Set2Set and Attention Readout: Order-Invariant Graph Summaries

Beyond Uniform Pooling

Mean and sum pooling treat all nodes identically. But for most tasks, nodes differ greatly in importance:

• In a molecule, the reactive functional group matters more than inert backbone atoms
• In a social network, hubs matter more than peripheral nodes
• In a citation graph, landmark papers matter more than derivative works

Attention readout learns these importance differences during training.

Attention Readout (Global Attention Pooling)

For each node \(v\), compute a scalar importance score:

Normalise the scores across the graph with a softmax:

Compute the graph embedding as a weighted sum:

This is a single-pass soft attention over all nodes. The model learns which nodes to weight highly for the specific prediction task. (The original gated readout in Gated Graph Sequence Neural Networks used an unnormalised sigmoid gate rather than a softmax; the softmax-normalised form above is the variant most libraries implement, and it makes \(h_G\) a convex combination, so its scale does not grow with graph size.)

Properties:

• Permutation-invariant: the softmax denominator is a sum over all nodes and the output is a sum of weighted terms, both unordered
• Differentiable: all operations are smooth
• Task-conditioned: \(\alpha_v\) depends on \(h_v\), which already encodes \(v\)’s local neighbourhood

Limitation: the logit \(a_v\) is computed from \(h_v\) alone. The softmax does couple the nodes, but only through a single global normaliser — it can rescale the weights, not change their relative order. So a node cannot be judged important because of what some other node contributes; the ranking of nodes is fixed before any comparison between them happens.

Set2Set (Vinyals et al., 2016)

Intuition first. Imagine reading a complex document by scanning it \(T\) times, each time looking for something different. On scan 1 you find the main claim; on scan 2 you look for supporting evidence; on scan 3 you check for caveats. Set2Set does the same for a graph: each LSTM step issues a different “query” that attends to a different subset of nodes, building a richer summary than any single pass could.

Set2Set produces a graph embedding using \(T\) steps of LSTM-driven attention. At each step \(t\) the LSTM emits a query vector \(q_t\), which is used to attend over all nodes:

Step \(t\):

The LSTM then consumes the read vector and produces the next query:

After \(T\) steps, the final graph embedding is the concatenation of the last query and the last attended message:

Why an LSTM Readout Is Still Order-Invariant

An LSTM is the archetypal order-sensitive module, so its presence in a readout looks like a contradiction. It is not, and the reason is worth being precise about: the LSTM does not consume the nodes. It is unrolled over \(T\) processing steps, a number fixed as a hyperparameter and completely independent of \(N\). The nodes enter only through \(e_v^t\), the softmax, and the weighted sum \(m_t\) — and every one of those three is a symmetric function of the node set. Permute the nodes and each \(\alpha_v^t\) follows its node, \(m_t\) is unchanged, so \(q_{t+1}\) is unchanged, and by induction \(h_G\) is unchanged.

This is exactly the point of the paper’s title, Order Matters: feeding a set to a sequence model in some arbitrary order makes the output depend on that order, which is wrong. Set2Set’s fix is to let the recurrence run over reads of the set rather than over its elements.

Worked Example: Set2Set on a 3-Node Graph

Consider a graph with 3 nodes and embeddings \(h_1 = [1, 0]\), \(h_2 = [0, 1]\), \(h_3 = [1, 1]\) (so \(d = 2\)). Run Set2Set with \(T = 2\) steps.

Step \(t=1\): initial query \(q_1 = [0.5, 0.5]\) (from the learned initial LSTM state)

• Scores: \(e_1^1 = q_1^{\top}h_1 = 0.5\), \(e_2^1 = 0.5\), \(e_3^1 = 1.0\)
• Softmax over \((0.5,\,0.5,\,1.0)\): \(\alpha^1 \approx [0.274,\; 0.274,\; 0.452]\) — node 3 wins, and nodes 1 and 2 are tied
• Attended message: \(m_1 = 0.274\,[1,0] + 0.274\,[0,1] + 0.452\,[1,1] \approx [0.726,\; 0.726]\)
• LSTM update: \((q_2, c_2) = \mathrm{LSTM}([q_1 ; m_1], c_1)\) — suppose it returns \(q_2 \approx [0.8, 0.2]\)

Step \(t=2\): the new query \(q_2 = [0.8, 0.2]\) emphasises the first dimension

• Scores: \(e_1^2 = 0.8\), \(e_2^2 = 0.2\), \(e_3^2 = 1.0\)
• Softmax over \((0.8,\,0.2,\,1.0)\): \(\alpha^2 \approx [0.361,\; 0.198,\; 0.441]\) — node 3 is still the largest, but node 1 has now overtaken node 2
• Attended message: \(m_2 = 0.361\,[1,0] + 0.198\,[0,1] + 0.441\,[1,1] \approx [0.802,\; 0.639]\)

Final embedding: \(h_G = [q_2 ; m_2] \approx [0.8,\; 0.2,\; 0.80,\; 0.64]\), of dimension \(2d = 4\).

Notice what the second step bought: step 1 could not separate nodes 1 and 2 at all — its query was symmetric in the two coordinates, so they received identical weights. Step 2’s query, conditioned on what step 1 read, breaks that tie. A single attention pass with a symmetric query would have left nodes 1 and 2 indistinguishable in the summary.

Set2Set vs Attention Readout vs Sum

When Set2Set Helps

Set2Set is particularly effective when:

• Graph-level prediction requires integrating information from multiple disjoint node subsets
• Different “aspects” of the graph matter for the prediction (Set2Set reads each in turn)
• The graph size varies widely across the dataset (attention readout adapts better than fixed pooling)

Set2Set is the readout used in the MPNN of Gilmer et al. (2017) for molecular property prediction on QM9, chosen there over sum or mean precisely because it is both order-invariant and able to attend to several parts of a molecule in turn.

Multi-head Attention Readout

A simpler extension of attention readout: compute \(K\) independent attention heads, each with its own gate MLP, and concatenate:

Each head learns to attend to a different subset of important nodes. This gives multi-aspect graph summarisation without the LSTM overhead of Set2Set.

Summary

For small graphs (molecules, proteins), Set2Set and multi-head attention are the readouts to reach for when a single weighted average is too blunt a summary. For large graphs, the \(O(TNd)\) cost of Set2Set and its inherently sequential \(T\) steps make single-pass attention readout the preferred choice.

References

• Vinyals, O., Bengio, S., & Kudlur, M. (2016). Order Matters: Sequence to Sequence for Sets. ICLR 2016 (Set2Set).
• Li, Y., Tarlow, D., Brockschmidt, M., & Zemel, R. (2016). Gated Graph Sequence Neural Networks. ICLR 2016 (gated global readout).
• Gilmer, J., Schoenholz, S. S., Riley, P. F., Vinyals, O., & Dahl, G. E. (2017). Neural Message Passing for Quantum Chemistry. ICML 2017 (uses Set2Set as the readout for molecular property prediction).