Extended Abstract Track
Sheaf Attention Networks
Editors: Sophia Sanborn, Christian Shewmake, Simone Azeglio, Arianna Di Bernardo, Nina Miolane
Abstract
Attention has become a central inductive bias for deep learning models irrespective of domain.
However, increasing theoretical and empirical evidence suggests that Graph Attention Networks (GATs) suffer from the same pathological issues affecting many other Graph
Neural Networks (GNNs).
First, GAT’s features tend to become progressively smoother as more layers are stacked, and second, the model performs poorly in heterophilic graphs.
Sheaf Neural Networks (SNNs), a new class of models inspired by algebraic topology and geometry, have shown much promise in tackling these two issues.
Building upon the recent success of SNNs and the wide adoption of attention-based architectures, we propose Sheaf
Attention Networks (SheafANs).
By making use of a novel and more expressive attention mechanism equipped with geometric inductive biases, we show that this type of construction generalizes popular attention-based GNN models to cellular sheaves.
We demonstrate that these models help tackle the oversmoothing and heterophily problems and show that, in practice, SheafANs consistently outperform GAT on synthetic and real-world benchmarks.
1. Introduction
Graph Neural Networks (GNNs) are Neural Networks which operate over graphs.
Their main competitive advantage being that they are able to exploit the graph’s local neighbourhood structure in the computations of the latent feature vectors.
Two problematic phenomena are associated to many GNN architectures: oversmoothing and the heterophily problem.
The former refers to the fact that building very deep neural networks often leads to poor performance.
The latter instead is related to the fact that GNNs are commonly built with the inductive bias that neighbouring nodes are likely to be similar.
One affected and popular GNN model is the Graph Attention Network (GAT).
It has been shown that the heterophily and oversmoothing problems are, from a topological perspective, intimately connected (Bodnar et al., 2022).
GNNs which operate over cellular sheaves, Sheaf Neural Networks (SNNs) (Hansen and Gebhart, 2020; Bodnar et al., 2022), offer a principled way of mitigating the aformentioned issues.
In this work, motivated by the great success of attention mechanisms, we propose the Sheaf Attention Network (SheafAN), a generalization of the popular Graph Attention Network to cellular sheaves.
We find that SheafAN is able mitigate the issues of oversmoothing and heterophily in GAT by leveraging the more complex geometry present in the sheaf.
2. Background
Definition 1 A cellular sheaf ( G, F ) on an undirected graph G =
( V, E ) consists of: (1) a vector space F ( v ) for each v ∈ V , (2) a vector space F ( e ) for each e ∈ E and, (3) a linear map F v ⊴ e : F ( v ) → F ( e ) for each incident node-edge pair v ⊴ e .
The vector spaces of the node and edges are called stalks , while the linear maps are called restriction maps .
The sheaf F “bundles” the stalks and restriction maps together.
The underlying graph specifies the nodes and edges, while the sheaf specifies a network of linear transformations over the nodes and edges.
We define the space of 0-cochains C
( G, F ) as the direct sum over the node stalks C 0 ( G, F ) :
= (cid:76) v ∈ V F ( v ).
Similarly, the space of 1-cochains
C 1 ( G, F ) as the direct sum over the edge stalks C 1 ( G, F ) :
= (cid:76) e ∈ E F ( e ).
It is then natural to construct a co-boundary operator δ : C 0 ( G, F ) →
C 1 ( G, F ) as δ ( x ) e = F v ⊴ e x v −F u ⊴ e x u .
Through the co-boundary operator, we can construct a sheaf Laplacian.
The sheaf Laplacian of a sheaf is a map L F : C 0 ( G, F ) → C 0 ( G, F ) defined as L F = δ ⊤ δ .
The normalized sheaf Laplacian ∆ F is defined as ∆ F = D − 12 L F D − 12 where D is the block-diagonal of
L F .
The sheaf Laplacian is a symmetric positive semi-definite (by construction) block matrix.
The diagonal blocks are L F v,v = (cid:80) v ⊴ e F ⊤ v ⊴ e F v ⊴ e , while the off-diagonal blocks are
L F v,u = −F ⊤ v ⊴ e F u ⊴ e .
With this machinery, one is able to build a Sheaf Neural Network layer (Hansen and
Gebhart, 2020; Bodnar et al., 2022), a GNN layer which operates over cellular sheaves.
This arises from the Euler discretization of a sheaf heat diffusion PDE (Hansen and Ghrist, 2021).
Consider a graph G with n nodes and d -dimensional node feature vectors x v ∈ F ( v ).
Their concatenation x is a 0-cochain such that x ∈ C
( G, F ).
Allowing for f feature channels, we can construct a ( nd ) ×
f feature matrix X .
The Neural Sheaf Diffusion (Bodnar et al., 2022) model takes the form:
[Equation]
Extended Abstract Track where σ is a non-linearity, I n is a n × n identity matrix, and W t 1 and W t 2 are learnable weight matrices.
The sheaf Laplacian at time t , ∆ F ( t ) , is typically learnt through gradient-based approaches by learning the individual restriction maps with a MLP Φ, i.e. F v ⊴ e :=( v,u ) =
Φ( x v , x u ).
More recently, Barbero et al.
(2022) have shown that the sheaf can also be learnt using tools inspired by Riemannian geometry.
3. Sheaf Attention Networks
The sheaf-attention mechanism we employ relies on an attention matrix Λ, which is a n × n row-stochastic matrix.
Each entry Λ ij is computed with an attention function a ( x i , x j ):
[Equation]
where a and W are learnable weights of appropriate dimension, and || denotes vector concatenation.
We employ multi-head attention which empirically tends to outperform single-head mechanisms (Velickovic et al., 2017; Vaswani et al., 2017).
Note that this is the same attention mechanism utilised in GAT.
To generalise Λ to d -dimensional sheaves, we construct the sheaf-attention matrix ˆΛ such that ˆΛ = Λ ⊗ 1 d where 1 d is a d × d -dimensional matrix with every element 1 and ⊗ is the Kronecker product.
We then consider the following attentive sheaf diffusion PDE, which evolves the node features X over time:
[Equation]
Here ⊙ indicates element-wise multiplication, and ˆ A F is the sheaf adjacency matrix with added self-loops.
That is ˆ A F ( i, j ) =
F ⊤ i ⊴ e F j ⊴ e =: P ij .
The Euler discretisation with unit time-step ( τ = 1) of Equation 3 has the form:
[Equation]
We can equip Equation 4 with weight matrices W 1 t ∈ d × d and W 2 t ∈ f t × f t +1 and a nonlinearity σ to derive our new Sheaf Attention Network (SheafAN) layer (see Figure 1):
[Equation]
where f t and f t +1 are the input and output channels sizes respectively, and ˆΛ ( X t ) is the attention matrix for layer t .
It is natural to call this a Sheaf Attention Network, since the original GAT model is a particular instantiation of the aforementioned model for which restriction maps are 1-dimensional ( d = 1) and equal to 1.
We also consider a different parametrization of Equation 4:
[Equation]
which we call Res-SheafAN.
The Res-SheafAN model is different to the SheafAN model as it uses a type of residual parametrization, where the feature updates are of the form Extended Abstract Track
h t +1 =
h t
+ f ( h t , θ t ).
This is motivated from the spectral perspective of being able to act as both a high-pass and low-pass filter (Di Giovanni et al., 2022).
In particular, without a residual connection the model may be dominated by the low frequency signals over the graph as more layers are added, contributing to oversmoothing.
We make the simplification that all restriction maps are orthogonal, i.e. F v ⊴ e ∈ O ( d ), effectively learning an O ( d )-vector bundle.
O ( d ) restriction maps are advantageous because they have d ( d − 1) / 2 free parameters compared to the d 2 free parameters in general invertible linear maps, which can be understood as a form of regularization.
On top of this, orthogonal maps have eigenvalues ± 1 which help with the stability of the training process.
Orthogonal maps have been shown to be theoretically advantageous as they are able to more efficiently separate nodes using the same stalk width d when compared to, for instance, diagonal linear maps (Bodnar et al., 2022).
Encoding signed messages (using rotation angles of orthogonal maps for example) is crucial in settings of low homophily
(Yan et al., 2021).
From an opinion dynamics perspective (Hansen and Ghrist, 2021), a negatively signed message may be thought of as modelling a node’s opinion which contradicts another node’s opinion.
As one can imagine, this behaviour is related to low homophily, where two nodes are more likely nodes to “disagree” on their class.
We therefore expect this higher-dimensional disagreement in the form of higher-dimensional orthogonal restriction maps to be very advantageous in heterophilic settings.
Layers Best [Table] Evaluation
We evaluate the performance of SheafAN and Res-SheafAN in different settings of heterophily and with varying number of layers.
We follow the analysis structure proposed by Yan et al.
(2021) by evaluating our sheaf attention models on Cora, Citeseer,
Cornell and Chameleon with layers increasing in powers of 2 from 2 to 64.
The accuracy and standard deviation is computed with respect to 10 fixed splits provided by (Pei et al., 2020), with each split containing 48% / 32% / 20% of nodes for training/calibration/testing
Extended Abstract Track
respectively.
We run our experiments on a machine equipped with an NVIDIA TITAN X
GPU (12 GB) and an Intel(R) Core(TM) i7-6700 CPU @ 3.40 GhZ.
We consider models of different classes.
1) Classical GNN models: GCN (Kipf and
Welling, 2016), and GAT (Velickovic et al., 2017).
2) Models designed to address oversmoothing: GCNII (Chen et al., 2020), and PairNorm (Zhao and Akoglu, 2019).
3) Models designed for heterophilic settings: GGCN (Yan et al., 2021), Geom-GCN (Pei et al., 2020),
H2GCN (Zhu et al., 2020), and GPRGNN (Chien et al., 2020).
Table 1 shows the oversmoothing results.
We see how models which are not designed for oversmoothing exhibit various issues, ranging from accuracy dropping rapidly (GCN and Geom-GCN), to memory issues (H2GCN), to numerical instability (GAT).
We find that our sheaf attention mechanism models perform very competitively when compared to state of the art techniques such as GCNII.
Furthermore, SheafAN consistenly outperforms GAT.
We attribute this to the more complex geometry of the underlying sheaf.
The model is able to leverage the additional stalk width to tamper the effect of oversmoothing, in accordance with the analysis of Bodnar et al.
SheafAN performs particularly well on the Cornell and Chameleon datasets, when compared to the other models, which are particularly heterophilic datasets.
While other models are generally designed to combat either oversmoothing or heterophily, SheafAN seems to be particularly well-equipped for both at once.
Conclusion
We proposed a generalization of the popular Graph Attention Network to cellular sheaves.
SheafAN demonstrates competitive performance with state-of-the-art models on tasks related to oversmoothing and heterophily.
We believe to have found an exciting connection between attention mechanisms and cellular sheaves and hope to encourage the further application of Sheaf Neural Networks to broader areas of machine learning in which attention mechanisms shine.