Why don't Graph Transformers attend to the same token?
Organizations: BOLD Lab University of Oxford
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.
Figures & tables
| Arm | Sink rate | Test MAE | Ablation ratio |
|---|---|---|---|
| GRIT (gated scoring, RRWP) | |||
| GRIT-dotlogit (dot-product logit only) | |||
| GRIT-blind (gated scoring, no structure) | |||
| GRIT-dot (dot-product scoring, no structure) | |||
| GRIT-dotbias (GRIT-dot + pair bias) | |||
| GraphGPS (GINE + dot attention) |
Appendix figures & tables3 assets
Supplementary material from the paper’s appendix.
Appendix
| Model | Dataset | Dim | Heads | Layers | MPNN | PE | LR | Epochs | Batch |
|---|---|---|---|---|---|---|---|---|---|
| GRIT | Peptides | 96 | 4 | 4 | – | RRWP-17 | 3e-4 | 150 | 8 |
| GRIT | ZINC | 64 | 8 | 10 | – | RRWP-21 | 1e-3 | 2000 | 32 |
| GraphGPS | ZINC | 64 | 4 | 10 | GINE | RWSE-28 | 1e-3 | 1500 | 32 |
| Arm | Sink rate | Test MAE | Abl. removed | Abl. renorm. | |
|---|---|---|---|---|---|
| GRIT | |||||
| GRIT-dotlogit | |||||
| GRIT-noclamp | |||||
| GRIT-official | |||||
| GRIT-blind | |||||
| GRIT-dot |
| GRIT | GRIT-noclamp | GRIT-official | GRIT-dotlogit | GRIT-dot | GraphGPS | |
|---|---|---|---|---|---|---|
Attention Sinks
Attention sinks are tokens that attract disproportionate Transformer attention, shaping information flow, long-context behavior, and model efficiency.
See all 55 papers