Let the Heads Talk: Beyond Diagonal Graph Attention

Motivation and gap

• Classical sheaf neural networks entangle local transport with restriction-map factorization and sheaf-Laplacian diffusion, leaving the contribution of richer transport unclear.The paper therefore studies the additional message-passing operation enabled by matrix-valued transport rather than asking only when sheaf networks outperform conventional GNNs.Source: §1 · p. 1

Approach

• Quiver representations isolate directed edge maps from sheaf factorization and diffusion, connecting matrix-valued transport directly to multi-head attention.With heads as coordinates, standard attention is diagonal: each source head contributes only to its corresponding receiver head.Source: §1 · p. 2

Contribution

• Top-A adds edge-conditioned off-diagonal routes before aggregation, preserves same-head paths, and exactly recovers vanilla attention when routing vanishes.The paper also establishes that this edge-specific mixing cannot generally be absorbed into a shared post-aggregation projection.Source: §1 · p. 2

Evaluation scope

• The evaluation spans algorithmic reasoning, relational reasoning, heterogeneous graph learning, and heterophilic node classification as a contrast setting.These benchmarks probe edge-specific computation across relational composition, directed heterogeneous interactions, algorithmic size generalization, and whether heterophily alone makes cross-head transport beneficial.Source: §1 · p. 2

Message Passing Neural Networks

• Message-passing GNNs aggregate information from neighboring node representations, commonly using scalar coefficients to modulate each neighbor's contribution.The figure contrasts fixed scalar coefficients in convolutional layers with data-dependent coefficients in attention.Source: Figure 1 · p. 2, §2 · p. 3

Sheaf Neural Networks

• Sheaf neural networks assign local vector spaces to nodes and edges and use incidence restriction maps to induce linear transformations between adjacent node representations.Neural Sheaf Diffusion builds propagation dynamics from the resulting sheaf Laplacian.Source: §2 · p. 3

Graph Quivers

• A graph-quiver representation assigns a vector space to each vertex and a linear map to each directed interaction, representing transport without retaining its sheaf factorization.The paper adopts this abstraction to study the effective directed map independently of the restriction-map factorization and diffusion operator.Source: §2 · p. 3

Graph Attention as Message Passing

• Attention makes a sender's contribution depend on its interaction with the receiver, while its scoring rule and attended neighborhood vary by architecture.The joint multi-head representation is left implicit in the standard formulation, motivating an explicit head-space view.Source: §2 · pp. 3–4

Multi-head attention as diagonal transport

• The quiver view takes the attention heads as coordinates of the local transport space, with each head's value representation propagated in parallel.The interaction's diagonal map acts on the head coordinates while leaving within-head value coordinates unchanged.Source: §3 · p. 4
• Standard multi-head attention induces a diagonal map on head space, so each source head contributes only to the corresponding receiver head on a directed interaction.Different heads may assign different weights to the same interaction, but the source-to-receiver head correspondence remains fixed.Source: §3 · p. 4

Beyond diagonal head transport

• Off-diagonal entries in an interaction's transport map let a source head contribute to a different receiver head before neighborhood aggregation.Because the map is attached to a directed edge, different senders to the same receiver can induce different transformations.Source: §3 · pp. 4–5, Figure 2 · p. 5

Relation to sheaf transport

• Sheaf-induced transport has the same cross-coordinate routing role, but undirected sheaf maps are coupled by transpose reciprocity while the studied directed maps need not be.The paper also notes that symmetric multi-head attention can be realized as diagonal cellular-sheaf transport.Source: §3 · p. 5

4 TOPOLOGICAL ATTENTION

• Top-A extends diagonal attention with edge-dependent communication across heads before incoming messages are aggregated, while retaining the original same-head paths.When the additional routing vanishes, the operator reduces exactly to standard multi-head attention.Source: §4 · p. 5, §4.1 · p. 6

4.1 EDGE-CONDITIONED CROSS-HEAD TRANSPORT

• For each directed edge, a learned routing matrix adds only off-diagonal head routes, while the original attention coefficients remain on the diagonal.The routing is computed from an interaction-specific representation; different directed edges can induce different patterns.Source: §4.1 · pp. 5–6

4.2 ROUTING BEFORE NEIGHBORHOOD AGGREGATION

• A shared post-aggregation map works for every sender assignment only if it reproduces the routed transformation on every incoming edge.Because Top-A routing may vary by edge, that condition is not generally satisfied; aggregation has already erased the identities of individual incoming messages.Source: §4.2 · pp. 6–7

5 EXPERIMENTS

• The evaluation spans controlled relational reasoning, heterogeneous link prediction, and 30 algorithmic tasks, comparing GATv2 and Graph Transformer backbones with corresponding Top-A variants.The experiments also include generalization to longer STaR reasoning chains and larger CLRS instances.Source: §5 · p. 7, §5.1 · p. 7, §5.3 · p. 9

5.1 RELATIONAL REASONING ON STAR

• On STaR, Top-A improves mean accuracy for both attention backbones across RCC-8 and Allen’s Interval Algebra, including across a substantial range of unseen reasoning depths.The margin narrows on the most challenging configurations; the depth analysis holds the number of paths within its training range.Source: §5.1 · p. 8, Table 1 · p. 8

5.2 RELATION-CONDITIONED ROUTING ON HETEROGENEOUS GRAPHS

• On MovieLens-100K, Top-A improves mean AUROC and AUPR for both backbones; average gains are approximately 2.0 points for GATv2 and 10.5 points for GT.GT+Top-A performs comparably to GATv2+Top-A and reduces variability relative to the original GT backbone.Source: §5.2 · p. 8, Table 2 · p. 8

5.3 CLRS BENCHMARK

• Across all 30 CLRS tasks, Top-A improves GATv2 by 6 percentage points and GT by 8.5 points, with largest gains in sorting and searching and little or no improvement in greedy algorithms or dynamic programming.Table 3 reports official CLRS test scores, averaged equally across tasks within each seed; entries are means across 20 seeds.Source: §5.3 · p. 9, Table 3 · p. 9

6 CONCLUSION

• Off-diagonal edge maps let messages route information across heads, identifying edge-conditioned cross-head communication as a distinct transport primitive.Top-A retains standard attention as a special case, and the paper reports useful inductive bias across relational, heterogeneous, and algorithmic reasoning tasks.Source: §6 · p. 9

A EXTENDED RELATED WORK

• Top-A places transport between heads, whereas the cited copresheaf attention construction places it within each head’s channel space.Top-A acts on the head space while carrying within-head coordinates unchanged; the comparison is specific to Hajij et al.’s attention construction, not the broader copresheaf formalism.Source: A EXTENDED RELATED WORK · p. 13
• Talking-Heads Attention and DCMHA mix attention scores or weights, while Top-A routes value representations across heads for each directed interaction.Top-A retains same-head paths, adds interaction-dependent off-diagonal routes, and recovers standard multi-head attention when the added routing vanishes.Source: A EXTENDED RELATED WORK · p. 13

B.1 FROM CELLULAR SHEAVES TO SYMMETRIC ATTENTION

• Symmetric multi-head attention’s diagonal head-space transport admits an exact realization as cellular-sheaf transport.The appendix connects this diagonal attention operator to the local transport structure underlying Neural Sheaf Diffusion.Source: §B.1 · pp. 13–14

B.2 PROOF OF PROPOSITION 3.1

• Normalized multi-head attention induces diagonal transport on its native head space before any shared post-aggregation readout.Each head is scaled by its own attention coefficient, so a source head other than the receiver head cannot contribute to that receiver.Source: §B.2 · pp. 14–15
• The proof arranges sender values by head and applies the interaction’s attention coefficients as a diagonal matrix.The resulting headwise update scales each head independently.Source: §B.2 · p. 14

B.3 SHARED POST-AGGREGATION EQUIVALENCE

• The following section uses a theorem and a two-edge example to examine why edge-specific routing before aggregation cannot generally be replaced by one post-aggregation linear map.The passage introduces the proof and example without stating their detailed results.Source: §B.3 · p. 15

B.3.1 PROOF OF THEOREM 4.1

• A shared post-aggregation map reproduces Top-A’s routed aggregation only if it reproduces the routed transformation on every incoming edge; varying edge routes need not meet this condition.The theorem allows an arbitrary receiver-specific linear map, making the comparison more permissive than the attention backbone’s usual output projection.Source: §B.3.1 · pp. 15–16

B.3.2 A TWO-EDGE EXAMPLE

• In the two-edge example, distinct edge-specific routes produce different outputs from sender-value collections that standard attention collapses to the same aggregate.The example’s incompatible routing requirements and identical standard-attention aggregate show why one linear map after aggregation cannot reproduce both routed outputs.Source: §B.3.2 · p. 16

B.4 BACKBONE-SPECIFIC ROUTING INSTANTIATIONS

• Backbone-specific interaction representations supply the edge-dependent routing matrix; the appendix gives GATv2 and graph Transformer instantiations.For the graph Transformer, the routing input combines the sender key, receiver query, and edge- and graph-level embeddings.Source: §B.4 · pp. 16–17

B.5 INITIALIZATION AND COMPUTATIONAL COST

• Zero-initializing the routing generator makes Top-A match its attention baseline at initialization, while the routing branch can still receive first-order gradients when λ > 0.The appendix also states that dense H × H routing adds generation and application costs and, when materialized, O(|E|H^2) memory.Source: §B.5 · p. 17

C ADDITIONAL EXPERIMENTAL DETAILS, RESULTS, AND ANALYSIS

• The supplementary experiments were run on NVIDIA A100 SXM4 GPUs with 80 GB of memory, and the code was slated for release upon acceptance.Source: C ADDITIONAL EXPERIMENTAL DETAILS,… · p. 17

• STaR includes RCC-8 and Allen’s Interval Algebra, each with 57,600 training instances and 153,600 test instances.The domains contain eight qualitative spatial relations and thirteen temporal relations, respectively.Source: §C.1 · p. 17
• STaR asks models to infer a relation between queried nodes, with edge labels determining how relational states transform before propagation.Instances vary in path length k and number of paths b, requiring composition along paths and integration of potentially disjunctive information from multiple paths.Source: §C.1 · p. 17
• Models train on k ∈ {2, 3, 4} and b ∈ {1, 2, 3}, then evaluate through k = 15 while keeping b within the training range.Within each backbone family, baseline and Top-A share the input representation, architecture, prediction head, optimization, and tuning budget; Top-A starts from the baseline function.Source: §C.1 · pp. 17–18

C.2 MOVIELENS-100K: EXPERIMENTAL SETUP AND ANALYSIS

• MovieLens-100K is modeled as a bipartite link-prediction task with directed rates and rated-by edges between 943 users and 1,682 movies.The models use four attention heads, two message-passing layers, and a dot-product decoder; baseline and Top-A share their initialization for shared parameters, with routing zero-initialized.Source: §C.2 · p. 18
• MovieLens configurations are selected by validation AUROC, with a shared learning-rate search and a transport-scale search added for Top-A variants.The selected configuration is the grid point with the highest mean validation AUROC over three tuning seeds.Source: §C.2 · p. 18, Table 5 · p. 19, Table 4 · p. 19
• Top-A’s off-diagonal transport differs between rates and rated-by directions, while standard attention has only diagonal paths.The directional comparison uses edges joining the same user–movie pair with sender and receiver roles exchanged.Source: §C.2 · pp. 18–19, Figure 4 · p. 20
• Directional asymmetry compares the mean effective transport for rates and rated-by; it is zero for matching routes and one for one-sided or opposite-sign routes.The normalization is independent of a common rescaling of the two transports; the baseline is recovered when Ωe = 0.Source: §C.2 · p. 19
• Movie genre is associated with per-edge cross-head transport, most clearly for rated-by edges and more weakly for rates edges.The genre structure appears in a PCA projection of the twelve off-diagonal transport coefficients.Source: §C.2 · pp. 19–20, Figure 5 · p. 20

C.3 CLRS-30

• CLRS-30 tests 30 algorithmic tasks, using test problem size n = 64 after training at nominal size n = 16 to assess out-of-distribution generalization.The dataset provides graph-oriented encodings and intermediate execution states called hints; models are trained separately for each algorithm.Source: §C.3 · pp. 20–21
• GATv2 + Top-A raises the CLRS-30 macro-average from 53.22% to 59.26% (+6.04 percentage points), while GT + Top-A raises it from 49.81% to 58.36% (+8.56 percentage points).The macro-average weights all 30 tasks equally; the comparisons cover 20 paired seeds.Source: §C.3 · pp. 21–22, Table 7 · p. 22
• CLRS-30 gains are uneven: both architectures improve notably on binary search and sorting tasks, while heapsort declines for both and optimal BST has GT’s largest mean decline (-15.93 percentage points).GT also gains on shortest-path tasks, including Dijkstra (+21.50 percentage points), Bellman–Ford (+13.77), Floyd–Warshall (+12.23), and DAG shortest paths (+10.90).Source: §C.3 · p. 22

C.4 HETEROPHILIC NODE CLASSIFICATION

• Across 14 heterophilic dataset–backbone comparisons, Top-A does not consistently outperform its corresponding attention baseline.The clearest exception is GT on tolokers, where ROC-AUC rises from 77.89 to 84.76; most other changes are small and vary in sign.Source: §C.4 · p. 23, Table 8 · p. 23
• The heterophilic results suggest that cross-head transport alone is insufficient for consistent gains; its usefulness appears clearest when interactions call for different transformations before aggregation.Source: §C.4 · p. 23