Graph transformers avoid focusing on the same node unlike language models
Why don't Graph Transformers attend to the same token?
Machine Learning
Summary
Attention mechanisms in language models often focus too much on the first word, creating an 'attention sink.' The authors studied graph transformers, especially the GRIT model on molecules, and found they rarely do this, even less on larger graphs. They show this difference is due to the model’s attention scoring method rather than the graph structure itself. Changing GRIT’s scoring to a dot product made attention sinks more common, suggesting such sinks mainly occur with dot-product attention.
What this means in practice
- •For machine learning engineers: Design graph transformer models for molecular data with attention mechanisms that avoid unhelpful focus on single nodes, improving interpretability and performance.
- •For drug discovery teams: Use graph transformers like GRIT with specialized attention scoring to better analyze molecular interactions without misleading focus on specific atoms.
Authors
Alexandre Bismuth
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.