Why don't Graph Transformers attend to the same token?

AI summary

Overview

Research area: Graph Transformers, attention mechanisms, and Geometric Deep Learning — specifically the study of "attention sinks" (tokens that absorb a disproportionate share of attention) in graph-structured models versus language models.

Technical level: Advanced. The paper combines empirical benchmarking of graph transformers with formal proofs bounding how much attention a single node can receive.

Scope: A single-author study comparing the state-of-the-art graph transformer GRIT against dot-product Graph Transformers on molecular benchmarks, isolating which architectural component (the graph, the positional encodings, message passing, or the attention scoring function) determines whether an attention sink forms.

What This Paper Is About

Large language models concentrate much of their attention on a single semantically meaningless token such as <BOS>, a phenomenon called an "attention sink" that is thought to protect deep models from over-mixing information. Graph Transformers face a comparable pressure from over-smoothing, so it seems natural that they would also form sinks — yet a graph has no "first token" and permutation equivariance forbids arbitrarily privileging one node.

The paper asks where a sink can form on a graph at all, and which architectural ingredient decides whether one actually does. It answers by measuring sink behavior in GRIT and several dot-product graph transformers on molecules, proving bounds on how much attention any node can receive, and then swapping individual components to find the cause.

Key Contributions

1. A measurement of sink behaviour in graph transformers. The paper defines a sink score for a node (the attention it receives averaged over all queries) and a sink rate (the largest fraction of heads sharing one sink node at the peak-sinking layer), then measures both for GRIT, GraphGPS variants, a plain Transformer with Laplacian encodings, and small language models under the same estimator.

2. Three provable bounds on received attention (Proposition 3.1). A symmetry/orbit bound, a locality bound based on how much attention a query sends beyond a radius, and a logit-spread bound. The authors show each can be computed head by head to test whether it actually prevents a sink, and derive the minimum logit spread a sink requires.

3. A component-by-component ablation isolating the cause. Changing only GRIT's logit from a gated score to a dot product raises its sink rate sharply, while removing message passing from GraphGPS does not remove its sinks — evidence that the scoring function, not message passing or the graph itself, drives sink formation.

4. Evidence that sinks in dot-product models are near-inert. Ablating attention to the most-attended node barely changes test error for dot-product models but roughly doubles it for GRIT, suggesting GRIT's most-attended nodes carry real predictive information.

Main Findings

• GRIT rarely forms sinks on molecules. On ZINC, GRIT's sink rate is 0.090 ± 0.028, versus 0.415 ± 0.194 for GraphGPS and 0.95–1.00 for small language models at the same context length. GRIT stays below GraphGPS at every threshold from 0.1 to 0.8 (closest at ε = 0.8: 0.02 versus 0.05).

• The size trend runs backwards relative to LLMs. Language models trained on longer contexts form stronger sinks; GRIT's sink rate falls from 0.238 on the smallest Peptides-func quartile to 0.045 on the largest.

• GRIT does grow large-norm nodes without attending to them. Its largest squared node norm grows across layers from 2.4 to 4.8 times the squared mean norm, yet sinks stay rare.

• Neither graph-side bound is tight on real molecules. Atoms rarely have symmetric counterparts, attention is nonlocal (median λ₁ ≥ 0.91), and the bounds rule out a sink in at most 5.1% of heads in any run.

• Symmetry does bind on genuinely symmetric inputs. On cycles C₈ to C₂₄ of identical atoms and bonds, every equivariant-positional-encoding checkpoint spreads attention evenly, with received attention uniform to 7.8×10⁻⁷ in every equivariant-PE checkpoint except GRIT-dotbias and GRIT-dotlogit. Laplacian encodings, whose eigenvector signs and bases are arbitrary, break equivariance, and a dot-product Transformer then puts almost all attention (0.999) on one node.

• The scoring function is the main cause. Replacing only GRIT's logit with a dot product (GRIT-dotlogit) raises the sink rate from 0.090 ± 0.028 to 0.680 ± 0.098 on identical graphs, above GraphGPS's. Removing message passing from GraphGPS does not lower its sink rate either (0.507 ± 0.159).

• Accuracy and sink formation are separable. Without any graph structure, GRIT's score and a dot product reach the same accuracy (test MAE 0.542 and 0.543) yet sink at rates of 0.214 and 0.532. GraphGPS with a virtual node matches GRIT's accuracy (0.066) yet sinks six times as often (0.559). Adding node RRWP and a bounded pair bias to GRIT-dot recovers over half the lost accuracy (test MAE 0.268) but leaves its sink rate unchanged (0.521).

• GRIT's score saturates. In layers 2 to 10, at least 96% of its coordinates sit at 0 or 5, so the logit is a sum of readout weights selected by which coordinates fire, not by input magnitude. The spread cap 5‖a‖₁ ranges from 2.7 to 11.7 across layers — above the spread of 2.2 a sink needs at ZINC size — so the cap alone does not explain the low sink rate. Removing the clip (GRIT-noclamp) still sinks rarely (0.113).

• Scaling one node's hidden state confirms the mechanism. Doubling a node's state makes it a sink in up to 38.8% of heads in GRIT-dotlogit but at most 0.2% in GRIT; even a 32-fold increase reaches only 1.3% in GRIT, though it reaches 54.0% in GRIT-noclamp.

• A virtual node is not enough to make GRIT sink. A virtual node leaves GRIT's sink rate low (0.095 ± 0.019), even with a logit penalty toward it (0.079 ± 0.027). On the symmetric cycles, GraphGPS sends all its attention to the virtual node while GRIT sends only about half to two-thirds on average, inconsistently across seeds.

• Sinks in dot-product models are close to no-ops; GRIT's most-attended nodes are not. Removing all attention to the most-attended node leaves dot-product models' test MAE close to that for a random node (1.00 to 1.22 times), but GRIT's is twice as high (2.06). With renormalised rows, GRIT's ratio falls from 2.06 to 1.35, and every gated arm (1.25 to 1.61) stays above every dot-product arm (1.00 to 1.23) except GraphGPS-none, whose ratio varies widely across seeds (2.03).

• A threshold-free measure separates the two families. The peak share n·maxⱼ sⱼ separates gated arms (2.5 to 3.6) from dot-product arms (6.0 to 13.1) with no overlap. On Peptides-func, GRIT's largest sink score falls with graph size (0.135 to 0.066) while its peak share rises (7.4 to 16.1).

• Over-mixing is not observed at this depth. From first to last layer the Dirichlet energy relative to the squared feature norm grows in every arm (on average 1.13- to 5.45-fold) except GRIT + virtual node (PE rec.), where it falls to 0.71–0.92 of its input value, and the distance from consensus rises in every arm: at 10 layers there is little over-smoothing for a sink to prevent.

• Rewarding concentrated attention does not reproduce the effect. An attention-entropy penalty lifts GRIT's sink rate only to 0.131 ± 0.027 (strength 1) and 0.127 ± 0.031 (strength 2).

Methodology in Plain English

The authors take a state-of-the-art graph transformer, GRIT, which mixes information between nodes through attention alone (no message passing), and measure how much attention each atom receives on average. They set a threshold — a node counts as a sink in a head when its sink score exceeds ε = 0.3, roughly seven times the uniform share on a 23-node graph — and track the largest fraction of heads sharing one sink node.

They compare against GraphGPS, a hybrid of message passing and dot-product attention, on two molecular datasets: ZINC-12k (regression, graphs of roughly 23 nodes) and Peptides-func (classification, larger graphs). Peptides-func is also split by node count into four quartiles with mean sizes 56 to 245, each trained on 2,023 graphs at batch size 16 to remove the dataset-size confound. Results are means with a 95% CI on the full test set over three seeds (five for GraphGPS and the quartiles).

To find the cause, they change one component at a time — swapping GRIT's gated logit for a dot product while keeping RRWP, the pair representation and the values; removing the clip; removing graph structure entirely; adding a virtual node (the graph analogue of <BOS>); adding a bounded pair bias. They also run controls at matched accuracy, so that "the model sinks because it is worse" can be ruled out.

On the theory side, they prove three bounds on how much attention any single node can receive — from symmetry (equivariance forces the sink score to be constant on each orbit), from locality (a node's attention from far-away queries is limited by how much those queries leak beyond a radius), and from logit spread (softmax can only concentrate when logits differ enough). They compute each bound head by head on real checkpoints to see whether any of them actually bites. A separate probe pushes cycles of identical atoms and bonds through the trained checkpoints, doubling as a test of whether equivariance actually holds in the implementation.

Finally, they run two interventions: ablating all attention onto the most-attended node (on 500 test graphs) versus ablating a random node, and multiplying one node's hidden state by a factor ρ at layers 3, 5 or 8 (on 200 test graphs) to see whether large norms attract attention.

Why This Matters

Impact on research. The paper reframes the attention sink as a property of dot-product attention rather than of Transformers as such. Since GRIT's gated score barely responds to a node's norm, and dot-product scoring sinks with or without a designated token, the results suggest the sink is an artefact of one particular scoring function. This is consistent with reported sink disappearance without softmax normalisation or under an output gate. It also gives graph researchers a concrete, head-by-head diagnostic: compute the three bounds and the peak share to see whether a given architecture is even capable of sinking on a given input.

Real-world applications:

• Molecular property prediction and drug discovery. ZINC and Peptides-func are standard chemistry and peptide benchmarks; understanding which nodes a model actually relies on affects how much trust to place in graph transformers for screening compounds.
• Peptide and protein function classification. Because GRIT's most-attended atoms matter to its predictions, interpreting them may yield more chemically meaningful explanations than interpreting dot-product models, whose sinks are near-inert.
• Architecture selection for graph learning pipelines. Teams choosing between gated-score and dot-product graph transformers now have evidence that the choice changes attention structure independently of accuracy.
• Model debugging and interpretability tooling. The sink-rate estimator and ablation protocol transfer directly to auditing deployed graph models.

Industry relevance. Any group applying graph transformers to molecular or materials data — pharmaceutical machine learning, chemical informatics, peptide therapeutics — is affected by whether attention maps reflect genuine chemical signal or an architectural artefact. The finding that accuracy and sink formation decouple matters for anyone who assumes a well-performing model has interpretable attention. The paper's headline comparison — GRIT at 0.090 ± 0.028 versus GraphGPS at 0.415 ± 0.194 on identical ZINC graphs — is a concrete point in favour of gated scoring for applications where attention is read as explanation. The compute for this study came from the Isambard-AI National AI Research Resource.

Future Directions

• Extend beyond one gated architecture. The authors explicitly note the study covers one gated-score architecture against dot-product baselines; EGT, Graphormer, and SAN are named as next targets.
• Move beyond molecules. All experiments are on molecular graphs; non-molecular graphs remain untested.
• Resolve whether sinks protect graph transformers against over-mixing. The 10-layer models show little over-smoothing to prevent, so the protective role sinks play in LLMs is left open for graph models at greater depth.
• Separate reception from mixing. The paper measures how much attention a node receives, not how much a layer mixes; a perfect sink can also be a complete mixer, and the authors flag this distinction as a direction for further work. Several GRIT virtual-node quantities also remain seed-dependent and need firmer characterisation.

Target Audience

Researchers and graduate students working on Graph Transformers, Geometric Deep Learning, or attention-mechanism interpretability will get the most out of this paper, since it assumes familiarity with permutation equivariance, softmax attention, and standard graph benchmarks such as ZINC and Peptides-func. It is also relevant to practitioners in molecular machine learning deciding between gated and dot-product attention, and to LLM researchers interested in the origins of attention sinks, because the graph setting provides a clean control that language modelling cannot: a domain with no designated first token. Readers looking for a beginner-level introduction to attention sinks should start elsewhere; the proofs and ablation tables here are aimed at readers comfortable with the underlying formalism.

Authors’ abstract

In large language models (LLMs), attention concentrates heavily on the first token. This 'attention sink' is believed to protect deep models against over-mixing. Even though graph transformers face a related pressure, we find that a SoTA graph transformer, GRIT, rarely forms sinks on molecules, and that it sinks less when trained on larger graphs, the opposite of the trend in LLMs. In principle, a graph can prevent sinks: we prove and show empirically that its symmetry and locality limit the attention any node can receive. On molecules, however, with distinguishable atoms and global attention, neither limit prevents sinks; what stops them lies in the model. We find that the attention scoring function is the main cause. Replacing GRIT's gated logit with a dot product raises its sink rate from 0.09 to 0.68 on identical graphs. In Graph Transformers using dot-product attention, removing the sink gives about the same error as removing a random node, while GRIT's most-attended nodes matter to its predictions. Attention sinks are therefore a property of dot-product attention rather than of Transformers. In our experiments, they form when dot-product scoring has a distinguishable node to concentrate on.