Let the Heads Talk: Beyond Diagonal Graph Attention

2026-10-01

AI Summary

This study addresses the limited inter-head communication and diagonal mapping constraints in standard multi-head attention that restrict the expressive power of graph neural networks. To overcome these limitations, this work proposes a topological attention mechanism that reveals the inherent diagonal nature of multi-head attention. By introducing edge-dependent non-diagonal routing, it constructs cross-head communication primitives and leverages quiver representation theory to realize matrix-valued transport mappings, thereby supporting heterogeneous graph learning and algorithmic reasoning. Experimental results demonstrate that the proposed method achieves significant performance improvements on relational reasoning and out-of-distribution generalization tasks. These findings validate the effectiveness of cross-head transport as an independent computational primitive for enhancing graph representation learning.

Abstract

Sheaf Neural Networks generalize scalar-weighted message passing by replacing scalar edge weights with linear transport maps between local feature spaces. Yet the role of this matrix-valued transport is entangled with the broader sheaf-diffusion construction. We isolate the transport primitive through quiver representations and establish a direct connection with multi-head attention. Treating attention heads as coordinates of a local transport space reveals that standard multi-head attention implements diagonal edge maps: along each directed interaction, a source head can contribute only to the corresponding receiver head. Allowing off-diagonal entries instead enables edge-conditioned communication across heads before neighborhood aggregation. We show that this operation cannot, in general, be absorbed into a single shared linear map applied after aggregation. Building on this characterization, we introduce Topological Attention (Top-A), a multi-head attention that learns edge-dependent off-diagonal routes while preserving the original same-head paths and exactly recovering vanilla attention when the additional routing vanishes. We evaluate Top-A on relational reasoning, heterogeneous graph learning, and algorithmic reasoning, including out-of-distribution generalization, with heterophilic node classification as a contrast setting. The results show that cross-head transport is most useful when the task benefits from interaction-dependent transformations, while heterophily alone provides no systematic advantage. These findings identify edge-conditioned cross-head communication as a distinct computational primitive of matrix-valued transport.

Problem

• multi-head attention
• sheaf neural networks
• cross-head communication
• quiver representations
• diagonal edge maps

Innovation

• Topological Attention
• Multi-head Attention
• Sheaf Neural Networks
• Quiver Representations
• Cross-head Transport