Let the Heads Talk: Beyond Diagonal Graph Attention

The formalization of multi-head attention as diagonal transport is elegant and exact. Proposition 3.1 shows that $D_{i \leftarrow j} = \mathrm{Diag}(\alpha^{(1)}_{i \leftarrow j}, \dots, \alpha^{(H)}_{i \leftarrow j})$ acts as a head-space linear map with no off-diagonal terms, so ``Standard multi-head attention therefore occupies the diagonal subfamily of linear maps on the head space.'' The proof that this is a strict restriction---Theorem 4.1 and the subsequent two-edge example in Appendix B.3.2---firmly grounds the architecture: ``A single post-aggregation map therefore exists only if all incoming edges induce the same effective transformation.'' Because Top-A parameterizes edges by $T_e = D_e (I_H + \lambda \Omega_e)$ with $\mathrm{diag}(\Omega_e)=0$, it preserves the baseline pathway and recovers vanilla attention exactly when $\Omega_e = 0$, which is a strong and attractive design.

Several experimental and conceptual limitations temper the claims. First, the CLRS-30 improvements are highly task-dependent: sorting and searching gain substantially, but as reported in Appendix C.3, ``heapsort deteriorates for both architectures, and optimal BST exhibits the largest mean decline for GT (-15.93 percentage points).'' Second, the heterophilic-node-classification contrast in Appendix C.4 largely contradicts the original motivation for sheaf transport: ``cross-head transport provides no consistent advantage across these benchmarks.'' The paper frames this positively as specificity, but it also raises the question of why full sheaf models have been reported to help on these datasets, a tension the discussion does not fully resolve. Third, the learned routing is a dense $H \times H$ matrix per edge; while the authors note an $O(|E|H^2(p+C))$ cost, they do not report empirical runtimes or memory footprints, which could become prohibitive at large $H$.

The evidence supports the central claim that cross-head transport is useful when interactions demand distinct pre-aggregation transformations. On STaR, Top-A improves both backbones, and on MovieLens it recovers meaningful directional asymmetry in the learned transport. The paper also fairly distinguishes its value-routing mechanism from Talking-Heads and copresheaf attention, noting that the latter act on attention scores or within-head channels rather than mixing value heads before aggregation. The comparison to classical sheaf transport is nuanced---``classical sheaf transport is therefore an undirected instance of matrix-valued head-space transport''---though the heterophilic benchmark results complicate the narrative that this transport primitive alone explains prior sheaf-model gains. The omitted systematic comparison to full sheaf models on the same tasks makes it difficult to assess whether Top-A captures the useful subset of sheaf behavior or something orthogonal.

The experimental protocol is otherwise exemplary: zero-initialized routing generators, paired seeds, shared hyperparameter searches, and detailed appendices. As stated in Appendix B.5, ``With matched backbone parameters, Top-A therefore implements the same function as the baseline at initialization.'' However, the code is promised only ``upon acceptance'' (Appendix C), so independent reproduction is currently blocked. Minor capacity differences exist---e.g., GATv2+Top-A adds roughly 2k parameters on MovieLens-100K---but these are small relative to the total model size. Hyperparameter selection on CLRS uses a reduced eight-task tuning budget before confirmation, which is pragmatic but adds procedural complexity that should be documented precisely in any released repository.