Romain Menegaux
email: romain.menegaux@inria.fr

Emmanuel Jehanno
email: emmanuel.jehanno@inria.fr

Margot Selosse
email: margot.selosse@gmail.com

Julien Mairal
email: julien.mairal@inria.fr

Self-Attention in Colors: Another Take on Encoding Graph Structure in Transformers

We introduce a novel self-attention mechanism, which we call CSA (Chromatic Self-Attention), which extends the notion of attention scores to attention filters, independently modulating the feature channels. We showcase CSA in a fully-attentional graph Transformer CGT (Chromatic Graph Transformer) which integrates both graph structural information and edge features, completely bypassing the need for local message-passing components. Our method flexibly encodes graph structure through node-node interactions, by enriching the original edge features with a relative positional encoding scheme. We propose a new scheme based on random walks that encodes both structural and positional information, and show how to incorporate higher-order topological information, such as rings in molecular graphs. Our approach achieves state-of-the-art results on the ZINC benchmark dataset, while providing a flexible framework for encoding graph structure and incorporating higher-order topology.

Introduction

The field of graph representation learning and graph-structured data has seen significant growth and maturity in recent years, with successful applications in a wide range of domains such as pharmaceutics, drug discovery [START_REF] Gaudelet | Utilizing graph machine learning within drug discovery and development[END_REF], recommender systems [START_REF] Ying | Graph convolutional neural networks for web-scale recommender systems[END_REF], or navigation systems [START_REF] Derrow-Pinion | ETA Prediction with Graph Neural Networks in Google Maps[END_REF]. Key to these advances has been the success of graph neural networks (GNNs), which up until 2021 took mostly the form of local message-passing neural networks (MPNNs). In addition, the field has benefited from the development of a flurry of high-quality benchmark datasets [START_REF] Prakash Dwivedi | Benchmarking graph neural networks[END_REF]Hu et al., 2020;2021;Dwivedi et al., 2022b).

Local MPNNs have been known to suffer from limitations such as over-smoothing [START_REF] Oono | Graph neural networks exponentially lose expressive power for node classification[END_REF], bottlenecks in information propagation [START_REF] Alon | On the bottleneck of graph neural networks and its practical implications[END_REF] and limited expressiveness [START_REF] Xu | How powerful are graph neural networks? 7th International Conference on Learning Representations[END_REF][START_REF] Morris | Weisfeiler and leman go machine learning: The story so far[END_REF]. This issue is alleviated by (i) adding positional or structural encodings to node features [START_REF] Bouritsas | Improving graph neural network expressivity via subgraph isomorphism counting[END_REF][START_REF] Abboud | The Surprising Power of Graph Neural Networks with Random Node Initialization[END_REF]Dwivedi et al., 2022a;[START_REF] Eliasof | Graph positional encoding via random feature propagation[END_REF] and (ii) softly decoupling the message passing flow from the graph structure, by enhancing the graph with virtual nodes or higher order

The output of the standard Transformer self-attention mechanism [START_REF] Vaswani | Attention is all you need[END_REF] is invariant to permutations. This issue is mostly solved for sequences by concatenating positional encodings to the element features. The natural equivalent for graphs are positional encodings based on the eigenvectors of the graph Laplacian [START_REF] Prakash Dwivedi | Benchmarking graph neural networks[END_REF][START_REF] Kreuzer | Rethinking graph transformers with spectral attention[END_REF]. Albeit theoretically justified, in practice they are outperformed by other structural encodings specially tailored for graphs (Dwivedi et al., 2022a). Furthermore, even state-of-the-art node-wise positional encodings alone are not enough to bridge the performance gap between GTs and MPNNs [START_REF] Rampášek | Recipe for a General, Powerful , Scalable Graph Transformer[END_REF]. Furthermore edge features -which describe semantic relationships between nodes, such as the type of chemical bond in molecules -can be essential for valid graph representations, yet there is no standard or clear-cut way to incorporate them in GTs. In this work we give an answer the question:

How to expressively incorporate both graph structural information and edge features in graph Transformers?

Some recent GTs delegate this task to a local MPNN module [START_REF] Wu | Representing Long-Range Context for Graph Neural Networks with Global Attention[END_REF][START_REF] Chen | Structure-aware transformer for graph representation learning[END_REF][START_REF] Rampášek | Recipe for a General, Powerful , Scalable Graph Transformer[END_REF], but otherwise all standalone GTs directly bias the attention matrix with relative positional encodings [START_REF] Shaw | Self-attention with relative position representations[END_REF] of different natures, engineering in both edge features and the adjacency matrix. We propose a unifying framework for the self-attention layer that flexibly subsumes most of the GT-flavors aforementioned.

Contributions :

• We enrich this framework by introducing Chromatic (or Channel-wise) Self-Attention (CSA): selectively modulating message channels, damping or highlighting them based on the relative positions of two nodes.

• A novel relative positional encoding scheme for graphs, based on random walks. We inject them in CSA and reach state-of-the-art results on several graph benchmarks.

• We showcase the generality of our framework by seamlessly integrating higher-order graph topological features such as rings. This augmented version CSA-rings greatly improves the state-of-the-art on the ZINC molecular dataset.

We show in rigorous ablation study on a set of graph benchmarks the added benefits of our proposed methods.

Setting and Related Work

Graph neural networks

Regular graph neural networks (GNNs) consider a graph as a set of node features X ∈ R N nodes ×d , an adjacency matrix A ∈ {0; 1} N nodes ×N nodes and (optionally) a set of edge features E ∈ R N edges ×d . The information contained in the adjacency matrix A describes the connections between nodes and is henceforth referred to as the graph structure.

GNNs node representations h (initialized by h (0) = X) are iteratively updated in layers, by receiving and aggregating messages M from other nodes. The edge features can also be updated, which we do not write here. h

(ℓ+1) i ← AGG j∈N (i) M h (ℓ) j , e ij . ( 1
)
On the one hand, in local message passing networks (local-MPNNs) the aggregation spans only the neighbors j of node i. The types of local-MPNNs - [START_REF] Kipf | Semi-supervised classification with graph convolutional networks[END_REF][START_REF] Bresson | Residual Gated Graph ConvNets[END_REF][START_REF] Corso | Principal neighbourhood aggregation for graph nets[END_REF] to cite just a few of the most popular variants -mostly differ on how the aggregation is performed. On the other hand, Graph Transformers (GTs), which can be interpreted as dense MPNNs, aggregate messages from all nodes together.

Transformer Architecture

In vanilla transformers, messages are linear transformations of node representations -the node values V -, and the aggregation is a simple weighted sum:

h i ← h i + N nodes j=1 ã(i, j) V j . ( 2
)
The coefficients in this weighted sum, the attention scores ã(i, j) ∈ [0, 1], are obtained by a softmax normalization such that ã(i, j) = a(i, j)/ k a(i, k) with

a(i, j) = exp 1 √ d Q i • K j . ( 3
)
From now on, we omit the normalization term 1

√ d . Q = W Q h; K = W K h; V = W V h
are commonly named the query, key and value matrices, respectively. In the multi-head setting, N heads separated attention scores are computed, modulating the values V by blocks of size d/N heads .

In this basic formulation, the update rule is invariant to permutations in the input nodes, and in fact oblivious to the graph structure and edge features E ij . It is therefore crucial to include these, either by encoding and injecting them in the node features, or by tweaking the update rule.

Encoding graph structure

In this section, we give a brief overview of existing methods to encode graph structure and edge features in GTs.

Local-MPNN hybrids

One natural way to include the graph structure is simply to mix in local-MPNN layers, by either interleaving [START_REF] Rampášek | Recipe for a General, Powerful , Scalable Graph Transformer[END_REF] or stacking [START_REF] Wu | Representing Long-Range Context for Graph Neural Networks with Global Attention[END_REF][START_REF] Chen | Structure-aware transformer for graph representation learning[END_REF] them with dense Transformer layers. The graph structural information, as well as the edge features, are handled by the local components, conveniently leveraging the rich literature on local MPNNs. In this work however we restrict ourselves to purely attentional networks, for which appropriate positional encodings (PE) are needed. We briefly go over the different flavors of graph PEs in the literature, see [START_REF] Rampášek | Recipe for a General, Powerful , Scalable Graph Transformer[END_REF] for a more systematic overview and a thorough categorization. We adopt their distinction of structural and positional encodings.

Node structural encodings (SE) Most successful node encodings for graph rely on local structural information, such as centrality measures [START_REF] Ying | Do transformers really perform badly for graph representation?[END_REF] or local substructure counts [START_REF] Bouritsas | Improving graph neural network expressivity via subgraph isomorphism counting[END_REF]. One such encoding, related to the one we will introduce in section 3 -node-RWSE (Dwivedi et al., 2022a)is based on random walks. Instrumental to the success of GTs, they have also shown to improve performance and expressivity in local MPNNs. However they cannot fully convey the relative spatial information in a Transformer.

For the sake of intuition, take the key and query matrices to be identity, and suppose we have concatenated a positional encoding p to the features h ← h ∥ p. The attention score can then be written as the product of a semantic part and a positional one:

a(i, j) = exp (h i • h j ) exp (p i • p j ) .
The original cosine positional embeddings p designed for sequences have the desirable property that p i • p j is a function of the relative distance between i and j: |i -j|. For most of the successful graph node encodings presented, these inner products can at most represent structural similarity i.e. how the local neighborhoods of nodes i and j compare to each other. They are independent of the distance between i and j within the graph, whereas the model could benefit from this information, for instance by attenuating signals coming from far away nodes.

Node positional encodings (PE)

This distance-encoding role could be fulfilled by spectral positional encodings based on the eigenvectors of the graph Laplacian. These have been adapted into a node-PE [START_REF] Prakash Dwivedi | Benchmarking graph neural networks[END_REF][START_REF] Kreuzer | Rethinking graph transformers with spectral attention[END_REF], but suffer from implementation difficulties such as their invariance to sign flip of the eigenvectors, in practice hindering their performance (Dwivedi et al., 2022a). SignNet [START_REF] Lim | Sign and Basis Invariant Networks for Spectral Graph Representation Learning[END_REF] propose to learn these encodings via a sign-invariant GNN.

Regardless of their particular type, current node-wise positional encodings are a helpful addition but do not suffice on their own to make global attention even remotely competitive with sparse message-passing [START_REF] Rampášek | Recipe for a General, Powerful , Scalable Graph Transformer[END_REF].

Relative positional encodings (RPE)

Each flavor of graph transformer comes with its own way of further incorporating the graph structure. One of the first graph transformers, SAN [START_REF] Kreuzer | Rethinking graph transformers with spectral attention[END_REF] learns two different sets of attention parameters: one for adjacent pairs, the other for non-connected nodes. Graphormer [START_REF] Ying | Do transformers really perform badly for graph representation?[END_REF] adds a learnable bias b ϕ(i,j) to the attention score a(i, j) = exp

Q i • K j + b ϕ(i,j)
, indexed on the shortest path length ϕ(i, j) between i and j. GraphiT [START_REF] Mialon | GraphiT: Encoding Graph Structure in Transformers[END_REF] bias the attention matrix with a graph kernel K r encoding the structural similarity:

a(i, j) = exp (Q i • K j ) K r (i, j).
The choice of the kernel K r correspond to imposing different inductive biases. For instance K r (i, j) = exp (p i • p j ) -with p one of the structural encodings mentioned above -would encode structural similarity. On the other hand choosing the heat diffusion or the PageRank kernel as K r would encode a notion of walking distance between nodes. If K r is chosen as the adjacency matrix or the 1-step random walk then GraphiT becomes a local-MPNN, very similar to GAT [START_REF] Veličković | Graph Attention Networks[END_REF].

What we propose is an extension of GraphiT, and a generalization of Graphormer, where the attention mask K r is learned based on structural edge encodings.

Method

Unifying framework

Most of the self-attention modules of the aforementioned GTs can be rewritten as in equation 4, with e ij a scalar encoding the relationship between nodes i and j.

a(i, j) = exp (Q i • K j + e ij ) ∈ R, (4)
By choosing e ij = b ϕ(i,j) we recover Graphormer (the expression for e ij is slightly more complex in the presence of bond features, we do not write it here), and with e ij = log(K r (i, j)) we replicate GraphiT. In another recent work -GRPE [START_REF] Park | Grpe: Relative positional encoding for graph transformer[END_REF] -re-use the query and key vector to compute the RPE bias

e ij = Q i • E Q ij + K j • E K ij .
Note that if e ij can be factorized as an inner product p i • p j (or more exactly as W Q p i • W K p j ), we recover the standard Transformer with positional encodings p i .

To the best of our knowledge, this framework encompasses all previously proposed GTs, with the exception of EGT (Hussain et al., 2022), who scale the attention score after the ℓ 1 -normalization ã(i, j) with a lognormalized gating term on edge features σ(E gate ij ) log(1 + k σ(E gate ik )). Another recent GT [START_REF] Kim | Pure transformers are powerful graph learners[END_REF]) falls out of the scope of this framework entirely, by treating both nodes and edges as separate tokens.

σ(Q ⋅ K + E att ) Graph G h (l+1) h (l) ∑ A ij ⊙ (V j + E val ij ) V E val E att

Monochrome attention

In the coarsest form of our proposed model, CSA-mono, e ij is simply a scalar projection of the initial edge features e ij = W E • E (0) ij . This is the "black and white", or mono-chromatic version of CSA.

Chromatic Self-Attention

Attention filters In the self-attention formulation given above by (4 and 2), the entirety of the node-node interaction -both their structural and semantic relationship -is contracted into a single scalar a(i, j) (resp. N heads scalars in multi-headed attention) which then modulates the node message V j .

We refine this scheme by making the attention scores multi-dimensional, allowing a separate attention score a c (i, j) per output channel c ∈ (1, d), generalizing the scalar attention score into an attention filter

a(i, j) ∈ R d . a(i, j) = exp (Q i • K j + E ij ) ∈ R d (5) ℓ 1 -normalization is then performed across each channel ã(i, j) = a(i, j)/ k a(i, k)
These filters are applied to the messages V via an element-wise product.

h i ← h i + N nodes j=1 ã(i, j) ⊙ V j (6)
This allows to selectively attend to, or dampen certain channels in the incoming messages. Note that the two formulations (4) and ( 5) are equivalent if we impose E ij to be a constant vector, or a piecewise-constant vector in the multi-head setting -constant on the channels of each head.

We emphasize that (5) comes as an enhancement over standard multi-headed attention, and is notably different from using a multi-head version of (4) with n heads = n channels . This is explained in detail in Appendix B.

Enhancing messages Most MPNNs that handle edge features combine them with the node features to form the message M (h j , E ij ), with M being a simple function such as concatenation, addition or elementwise product (Hu* et al., 2020;[START_REF] Corso | Principal neighbourhood aggregation for graph nets[END_REF]. Similarly, we also enrich the individual node messages with relational information, by adding an edge value vector E val ij to the node value vector V j , the edge values being a linear transformation of the edge features

E val ij = W EV E (0) ij : h i ← h i + N nodes j=1 ã(i, j) ⊙ (V j + E val ij ) (7)
Note that the update rule 7 opens up other possibilities for the attention filters, such as additive attention [START_REF] Bahdanau | Neural machine translation by jointly learning to align and translate[END_REF]. a(i, j) = exp (Q i + K j + E ij ) instead, completely foregoing the need for separate attention heads.

Design of the edge features

Fundamental to the success and expressive power of this scheme is the choice of these initial edge features.

Semantic edge features

After transforming the original or semantic edge features into suitable vector form E bond ∈ R N edges ×d , we extend them to node-pairs that were not originally connected in A, by introducing two learned embeddings ∈ R d : one for self-connections E self (if not already given), the other for non-connected nodes E n-c .

E bond ij =      E bond ij if i and j are connected E self if i = j E n-c otherwise (8)
To encode both semantic and structural information in the final edge features E we concatenate them with a relative positional encoding E RPE :

E (0) ij = E bond ij ∥ E RPE ij . ( 9
)
Here we choose concatenation but any simple function such as addition or multiplication should yield similar results.

This system is general and flexible enough to leverage all possible pairwise relations in graphs that can be encoded onto vectors.

Choice of the relative positional encoding

For the RPE E RPE ij , we draw upon the family of functions of the successive powers of the random walk matrix RW = D -1 A (where D is the diagonal degree matrix):

E RPE ij = Φ RW 1 ij , . . . , RW p ij ∈ R p , ( 10
)
With an appropriate choice of Φ, these constitute powerful descriptors of the graph relations [START_REF] Li | Distance encoding: Design provably more powerful neural networks for graph representation learning[END_REF]. Indeed, one can recover positional information -the index of the first non-zero component is the shortest path distance (SPD) -but also structural information. For instance RW 1 ij is nonzero iff i and j are neighbors and, in the absence of self-connections, RW 3

ii is nonzero iff i is part of a triangle. We give two possible RPEs constructed from these features. One of them -SPDE -is present in some form in other GTs [START_REF] Ying | Do transformers really perform badly for graph representation?[END_REF][START_REF] Park | Grpe: Relative positional encoding for graph transformer[END_REF][START_REF] Shamim Hussain | Global Self-Attention as a Replacement for Graph Convolution[END_REF]. The other -RWSE -is commonly used as a node-wise SE but has to our knowledge never been used before as an RPE.

Shortest Path Distance Encoding (SPDE)

The shortest path distance encoding is obtained by choosing Φ as a d-dimensional embedding of the index of the first nonzero component: with W SPDE ∈ R d×p the embedding vectors and e ϕ(i,j) ∈ R p the one-hot encoding of the shortest path length ϕ(i, j). Similarly to the semantic features above, we complete the feature matrix with a special embedding for nodes unreachable in less than p hops, and another one for self connections.

E SPDE ij = W SPDE e ϕ(i,j) (11)

Random Walk Structural Encoding (RWSE)

In our proposed RPE, Φ is simply a matrix multiplication with learnable weights W RWSE E :

E RWSE ij = W RWSE E • RW 1 ij , . . . , RW p ij ( 12
)
If the weights matrix W RWSE E is constrained to be non-negative, these features can be considered as re-weighted diffusion kernels between nodes, hence constituting a learnable and multi-dimensional extension of GraphiT. We argue that RWSE is a richer RPE than SPDE as it also preserves some local structural information. A simple illustrative example is shown in Figure 2.

Adding substructures

Incorporating higher-order structure information can be beneficial to performance. In molecular graphs for example, atoms in a ring (a chordless cycle) share covalent electrons, which influence their distribution in space and their chemical properties. Hierarchical message-passing schemes [START_REF] Fey | Hierarchical inter-message passing for learning on molecular graphs[END_REF][START_REF] Bodnar | Weisfeiler and Lehman go topological: Message passing simplicial networks[END_REF]a) were designed both to propagate information effectively along meaningful substructures and to suitably modulate the messages. We emulate this form of hierarchical message-passing in CSA by injecting specific topological information directly into the edge features (and hence in both the attention matrix and the messages).

For a chosen set of structures (graphs) S we define a binary relation S ∼ on nodes of a graph in definition 3.1.

S

∼ is symmetric but not necessarily reflexive or transitive.

Definition 3.1. Given a set of structures S, two nodes i and j of a graph G are S-bound -written i S ∼ j -if and only if there exists an element S ∈ S and a subgraph G of G containing both i and j such that G is

isomorphic to S. i S ∼ j ⇐⇒ ∃S ∈ S, G ∈ Sub(G) s.t. i, j ∈ G and S ∼ G (13)
We then design boolean features based on this relation e S ij = 1(i S ∼ j). In the more generic form of our topology encoding, which can be applied to any type of data, these features are embedded and then added to the edge features:

E ij ← E ij + Embed(e S ij )
In the cases where the original edge features E bond are categorical, i.e. belong to a discrete set C, one can give more degrees of freedom to the substructure encoding by forming the Cartesian product of the semantic and the topology features, effectively doubling the number of edge types.

E bond ij ← Embed E bond ij , e S ij
Note this could also be done on the SPDE positional encoding for instance. Although this choice of encoding is less easily applicable (due to the restrictive condition on the edge features) it is also more expressive, and yields slightly improved results empirically.

In practice we have implemented this topology encoding for rings in molecular graphs, by choosing S as the set of rings up to a certain length, but the same could be done for any type of substructure deemed relevant to the task at hand.

Discussion and implementation

Model architecture

The main components of our architecture are illustrated in Figure 1. As is standard in Transformer layers, our CSA layer is composed of the self-attention block followed by 2 linear layers with batch normalization. We stack L of these layers which yield final representations h (L) . These are fed to a prediction head that depends on the task (graph/node/edge classification/regression).

h (l+1) = CSA(h (l) , E (0) ) E (0) = EdgeEncoder(E bond , A) h (0) = NodeEncoder(X, A)
Regularization Aside from standard regularization techniques, we perform dropout directly on the attention tensor, silencing either whole nodes (node-ablation), connections (edge-ablation) or channels (feature dropout).

All of these are options in our model and our code.

Generality Semantic edge feature information are preserved (this is guaranteed by adding it to the messages directly), and structural information is also taken into account, modulating both the attention matrix and the messages.

Note that this framework is general enough to be applied to other types of data such as sequences, images or 3D point clouds. The graph structure is not essential to the CSA formulation, which can be transposed to other fields by designing an appropriate RPE -for instance embedding 3D distance and angles in 3D points clouds, or pixel distances in images.

Limitations There are several limitations with our proposed model, the main one being the scalability to large and sparse graphs. Indeed the computational complexity of our method scales as O(|N nodes | 2 d), compared to the O(|N edges | d)-complexity of local MPNNs. For small graphs however this is partially offset by the efficiency of GPU dense matrix computation. To scale our method to larger graphs one could explore the numerous attempts at a linear or sub-quadratic Transformers [START_REF] Zaheer | Big bird: Transformers for longer sequences[END_REF][START_REF] Jaegle | Perceiver IO: A general architecture for structured inputs & outputs[END_REF], transposing them to our framework. The pre-processing steps also incur a one-time cost of O(|N nodes | 3 p) for the computation of the RWSE or the SPDE base values.

Another limitation could be the reliance on hand-crafted features in the RPEs -although the same could be said of all GTs -but in fact there are very few hyper-parameters to tune, and fixed step sizes of RWSE-16, SPDE-8 work across a variety of datasets.

Implementation We integrate our proposed CSA layer in the general GraphGPS framework, which enables testing on several benchmark datasets, choice of a wide range of positional encodings as well as adding local MPNN layers in parallel (we do not use this feature in this paper though). We also add our implementations of the relative positional encodings Edge-RWSE and SPDE. The code of the model, and scripts to reproduce the experiments are freely available at https://github.com/inria-thoth/csa.

Experiments

Experimental setup

We evaluate the performance of our model on 3 medium-scale benchmark datasets from (Dwivedi et al., 2020): ZINC, PATTERN and CLUSTER and on the large molecular dataset PCQM4Mv2 [START_REF] Hu | OGB-LSC: A large-scale challenge for machine learning on graphs[END_REF]. We arrange the compared methods in 3 broad categories: MPNNs, extended (or hierarchical) MPNNs, and graph transformers. We chose to include the best performing ones (to our knowledge) in each category, with an emphasis on the more comparable graph transformers. We show SOTA results in all 4 datasets.

We conduct ablation studies to showcase the individual contributions of (i) chromatic attention, (ii) attending to distant nodes (iii) the choice of the RPE and (iv) incorporating topological sub-structures.

Model ZINC MAE ↓

MPNN

GCN [START_REF] Kipf | Semi-supervised classification with graph convolutional networks[END_REF] 0.367 ± 0.011 GatedGCN (Dwivedi et al., 2022a) 0.090 ± 0.001 GPS [START_REF] Rampášek | Recipe for a General, Powerful , Scalable Graph Transformer[END_REF] 0.070 ± 0.004

h-MPNN

CIN (Bodnar et al., 2021a) 0.079 ± 0.006 CRaWL [START_REF] Toenshoff | Graph Learning with 1D Convolutions on Random Walks[END_REF] 0.085 ± 0.004 GIN-AK+ [START_REF] Zhao | From stars to subgraphs: Uplifting any GNN with local structure awareness[END_REF] 0.080 ± 0.001

Transformers SAN (Kreuzer et al., 2021) 0.139 ± 0.006 Graphormer [START_REF] Ying | Do transformers really perform badly for graph representation?[END_REF] 0.122 ± 0.006 SAT [START_REF] Chen | Structure-aware transformer for graph representation learning[END_REF] 0.094 ± 0.008 EGT [START_REF] Shamim Hussain | Global Self-Attention as a Replacement for Graph Convolution[END_REF] 0.108 ± 0.009 GRPE [START_REF] Park | Grpe: Relative positional encoding for graph transformer[END_REF] 0.094 ± 0.002 CSA (ours) 0.070 ± 0.003 CSA-rings (ours) 0.056 ± 0.002

Table 1: Results on the ZINC dataset, ordered by the broad model class and performance. (GPS is placed among local MPNNs here, as their classification performance can be fully imputed to their local module).

Following [START_REF] Rampášek | Recipe for a General, Powerful , Scalable Graph Transformer[END_REF], the ablation studies are performed on ZINC and a subset of PCQM4Mv2 containing 10% of the whole training set (but the same validation set as the full version).

Little hyperparameter search was done, defaulting to the setup in GraphGPS or other comparable models whenever possible. We chose maximum ring-size as in (Bodnar et al., 2021a) k = 18 for ZINC, and k = 6 for OGB datasets. For ZINC we use the categorical ring encoding scheme described in section 3.4 and for OGB datasets we use the boolean encoding.

All models were trained on a single NVidia V100 GPU system, with 16GB or 32GB memory depending on the dataset. See Appendix A for the complete experimental details and hyperparameters. Larger scale datasets Finally, we assess the performance of our method on the large-scale benchmark PCQM4Mv2. The dataset is composed of 3.7M small undirected graphs (14.6 nodes on average). We test both a small and a deep version of our model in Table 3, which outperforms SOTA with less parameters. The difference in performance is however not as significant as for the ZINC dataset, we presume this is because the comparative advantages of our method diminishes with larger models and datasets.

Benchmarking

Model

Ablation study

Usefulness of the chromatic attention

In Figure 3 we visualize the different learned attention score matrices, for 3 output channels of a 16-channel model trained on ZINC. Each of these channels clearly focuses on structurally different nodes: the first channel on neighboring nodes, the second one on atoms in the same ring, and the third one focuses exclusively on faraway nodes. The figure shows both the complete node-node attention maps for each channel, and also some selected rows plotted on the graphs, for ease of visualization. These show the attention scores of the highlighted node, across 4 structurally diverse nodes.

The attention maps for all 16 channels are available in the appendix. To ensure that this is not a learned artifact and that the model actually benefits from these extra degrees of freedom, we empirically validate the increased expressiveness of CSA in Table 5 by successively ablating the chromatic nature of the attention scores, and the edge value vectors. On the two datasets we have tested, both additions make an impact.

RPE

Structural encoding

We confront our proposed positional and structural Random Walks Encodings with the Shortest Path Distance Encoding in Table 4 and confirm the intuition that RWSE outperforms SPDE, regardless of the topological encoding.

ZINC

Conclusion

We have proposed an extension of the Transformer self-attention, generalizing attention scores to attention filters. We adapt the resulting model to graph representation learning, further extending the self-attention mechanism by incorporating edge features in the message values. We combine our model with a novel relative positional encoding scheme -and a topology encoding one -and we show its added value on a diverse set of graph benchmarks. Table 7: Training times of CSA on a single A100 GPU, compared to GraphGPS [START_REF] Rampášek | Recipe for a General, Powerful , Scalable Graph Transformer[END_REF] Edge parameter sharing To reduce computation time and in some cases improve performance, we provide the option to share the edge representations E att and E val among all layers.

To get a frame of reference for the execution times, we report the training times of GraphGPS [START_REF] Rampášek | Recipe for a General, Powerful , Scalable Graph Transformer[END_REF] in table 7.

In figure 4 we show all 16 channels of the first layer of a single attention head CSAmodel trained on ZINC. 3 of these were chosen in figure 3 of the main text. One can clearly see that channels are learning different attention patterns, with some being identifiable as immediate neighbors, nodes in the same ring, and also complementary channels concentrating on faraway nodes.

Figure 1 :

 1 Figure 1: Chromatic Graph Transformer. (a) Graph features are preprocessed into node features h (top-left) and two edge feature matrices E att and E val ∈ R N nodes ×N nodes ×d (bottom). (b) These are input to the CSA layers, which iteratively update the node representations h. (c) These representations are fed to a classification or regression head (not shown here).

Figure 2 :

 2 Figure2: Shortest Path Distance Encoding cannot distinguish the relationship i ↔ j in the two above graphlets, as they are both neighbors with ϕ(i, j) = 1. On the other hand the three first coordinates of the RWSE ij vectors are markedly different: [0.33, 0.33, 0.26] and [0.5, 0.0, 0.63] for the graphlets on the left and right, respectively.

Figure 3 :

 3 Figure 3: Attention scores for multiple nodes and channels at the last layer of a 16-channel model trained on ZINC. Scores are shown for the molecule ZINC44549294. Each column corresponds to a different channel, which we hand-picked based on their diversity and interpretability. The top row shows the attention channel scores for all node pairs, and each subsequent row corresponds to a different receiving node, highlighted in red.

Figure 4 :

 4 Figure 4: Attention scores for all 16 channels at the first layer of a single-head model trained on ZINC. Scores are shown for the molecule ZINC44549294.

  Xu et al. (ICLR'2019). How powerful are graph neural networks? Alon, U., & Yahav, E. (2020). On the Bottleneck of Graph Neural Networks and its Practical Implications

CSA ZINC and Benchmarking Graph Neural Networks ZINC

  is a popular molecular graph regression benchmark, for which Graph Transformers have historically lagged behind their sparse counterparts. The dataset splits are fixed as in[START_REF] Prakash Dwivedi | Benchmarking graph neural networks[END_REF], with 10K graphs in the training set, 1K in the both the validation and test sets. The graphs contain on average ∼ 20-30 nodes.

	Model	PATTERN Accuracy ↑	CLUSTER Accuracy ↑
	SAN	86.581 ± 0.037	76.691 ± 0.65
	Graphormer	-	-
	SAT	86.848 ± 0.037	77.856 ± 0.104
	EGT	86.821 ± 0.020	79.232 ± 0.348
	GPS	86.685 ± 0.059	78.016 ± 0.180
	CSA (ours)	87.011 ± 0.036 79.175 ± 0.126

Table 2 :

 2 Results on the node classification datasets PATTERN and CLUSTER. In both cases, graphs were generated using a Stochastic Block Model (SBM). In PATTERN, one must detect which nodes were generated with different SBM parameters than the rest of the graph, and in CLUSTER one must assign each node to one of the 6 clusters (or blocks) it belongs to. These datasets contain 14K and 12K graphs, each with 100 nodes and 2K edges on average (much denser than molecule graphs). As shown in Table2, we reach state-of-the-art performances on both datasets.

	Our results are presented in Table 1, in which the
	ring-augmented version of our model CSA-ring im-
	proves the state of the art by a significant margin.
	Adding ring information greatly improves the per-
	formance in ZINC, which is to be expected as the
	artificial objective partially depends on the counts
	of these patterns.
	However we note that CSA achieves state-of-the-art
	results even without this substructure information,
	far surpassing other standalone GTs.
	PATTERN and CLUSTER PATTERN and
	CLUSTER are synthetic datasets of inductive node-
	level classification.

Table 3 :

 3 Results on the PCQM4M-v2 dataset, ordered by the broad model class and performance. As the test set was kept private by challenge organizers, evaluation is done using the validation set.

			PCQM4Mv2
			Validation MAE ↓ # Param.
	MPNN	GCN GCN-virtual GIN GIN-virtual	0.1379 0.1153 0.1195 0.1083	2.0M 4.9M 3.8M 6.7M
	Transformers	Graphormer EGT GRPE GPS-small GPS-medium	0.0864 0.0869 0.0890 0.0938 0.0858	48.3M 89.3M 46.2M 6.2M 19.4M
		CSA-small (ours)	0.0898	2.8M
		CSA-deep (ours)	0.0853	8.3M

Table 4 :

 4 CSA performance with different relative positional and structural encodings (RPE).

		ZINC	PCQM4Mv2
		MAE ↓	MAE ↓
	None	0.112 ± 0.002 0.1192 ± 0.0001
	SPDE	0.068 ± 0.002 0.1123 ± 0.0006
	RWSE	0.066 ± 0.001 0.1109 ± 0.0002

Table 5 :

 5 Added value of including multidimensional attention filters ("Color" column) and injecting edge information in the value messages ("Edge Value" column)

	PCQM4Mv2
	subset

Table 6 :

 6 Hyperparameters used for all datasets. Ablation studies on the PCQM4Mv2-subset are all done based on the CSA-small configuration.

	Hyperparameter	ZINC	PATTERN CLUSTER	PCQM4Mv2 CSA-small CSA-deep
	# CSA Layers	10		6	16	5	16
	Hidden Dim	64		64	48	256	256
	# Heads	4		4	8		16	16
	Dropout	0		0	0.1	0	0
	Attention Dropout	0.5		0.5	0.5	0.5	0.5
	Node Positional Encoding	RWSE-20		RWSE-16	RWSE-16	None	None
	PE dim	28		16	16	-	-
	Relative Positional Encoding RWSE-20		RWSE-16	RWSE-16	RWSE-16 SPDE-8
	Edge parameter sharing	✓		✗	✗		✓	✓
	Rings	✓		✗	✗		✗	✓
	Max rings size	18		-	-		-	6
	Encoding	categorical	-	-		-	additive
	Batch Size	32		32	16	256	256
	Learning rate	0.001		0.0005	0.0005	0.0005	0.0001
	# Epochs	2000		100	100	300	150
	# Warmup epochs	50		5	5		10	10
	Weight decay	1e-5		1e-5	1e-5	0	0
	# Parameters	350k		259k	382k	2.8M	8.2M
	Time (per epoch)	21s		40s	85s	822s	1900s
	Hyperparameter	ZINC PATTERN CLUSTER	PCQM4Mv2 small model medium model
	GraphGPS (Time/epoch)	21s	32s	86s		619s	1124s
	CSA (Time/epoch)	21s	40s	85s		822s	1520s

Acknowledgments

This project was supported by ANR 3IA MIAI@Grenoble Alpes (ANR-19-P3IA-0003) and by ERC grant number 101087696 (APHELEIA project). The bulk of the computations were performed using HPC resources from GENCI-IDRIS (Project-[A0131012608]).

A Experimental Details

Implementation As previously mentioned, our code is integrated in the GraphGPS framework introduced by [START_REF] Rampášek | Recipe for a General, Powerful , Scalable Graph Transformer[END_REF]. In particular, the data-loaders and splitting strategies are kept the same to ensure a fair comparison with other methods. Results in the benchmark section were all obtained with 10 seeds each, except for PCQM4Mv2-full which was run on a single seed. Results in the ablation study are aggregated over 4 seeds. We report the mean loss along with its standard deviation over runs.

Hyper-parameters

The table 6 recaps the chosen hyperparameters for each dataset, as well as training execution time. In all our experiments we use the AdamW [START_REF] Loshchilov | Decoupled weight decay regularization[END_REF] optimizer with default values for β 1 , β 2 and ϵ, and a cosine learning rate scheduler with warmup.

For the ZINC ablation studies, the number of epochs was reduced from 2000 to 1000, and rings were not taken into account.

Color Heads

B Difference between CSA and multi-head attention

In short, our method enhances the standard multi-head attention (which we call CSA-mono in the paper).

Indeed, note d the number of channels and h the number of heads. For a given channel c ∈ [1, d], we denote d h = d/h the implicit dimension of the attention heads, and c h = ⌊c/d h ⌋ the index of the attention head which channel c belongs to. The c-th channel of the unnormalized attention score between two nodes q and k of dimension d, linked by the edge e is then given by:

• Standard multi-head attention (called CSA-mono in the paper) with d = h:

The edge bias e c h here is constant across all channels of the same attention head. (In fact in most other GTs this bias is constant across all attention heads: e c h = e 1 , ∀c).

• Colored multi-head attention (CSA-color with 1 < h < d), e ∈ R d :

The difference here is that the edge bias e c varies per channel instead of per head.

In multi-headed dot-product attention, choosing the hyperparameter h = n heads amounts to striking a balance between the selectivity of the per-head cosine similarities Q h • K h (of implicit dimension d/h), and the potential diversity of those similarities across attention heads (h different scores).

In CSA we maintain this balancing option intact, but also keep a separate edge bias e c per channel, without loss of expressivity.

We compare CSA with its monochrome version in table 8, varying the number of attention heads.


Paolo Perna

Laurence Méchin

Marie-Pierre Chauvat

Pierre Ruterana

Ch Simon

U Scotti

Di Uccio

    High Curie temperature for La0.7Sr0.3MnO3 thin films deposited on CeO2/YSZ-based buffered silicon substrates