Do Transformers Really Perform Bad for Graph Representation?

In plain words

Transformers handle sequences, such as text, very well. On graphs (molecules, social networks and similar data) they had not reached the top of the public leaderboards, and most strong graph models were classic message-passing networks. The paper asks whether this is a flaw of Transformers or a flaw of how they were fed graph data.

Its answer is Graphormer, which is an ordinary Transformer with one change to the input and the attention. A plain attention layer only compares nodes by their features, so it cannot tell that one node is a hub, or that two nodes sit close together in the graph. Graphormer gives the model that information in three ways: a learned vector for how many links each node has, a learned score for how far apart each pair of nodes is, and a learned score that reads the bond or edge features along the path between them. It also adds one extra "graph token" that every node is connected to, and the final state of that token is used for the prediction.

The result is a Transformer that wins on several graph benchmarks, including a very large molecular property challenge. The paper also argues that many popular graph neural networks are special cases of this design.

TL;DR

• Standard Transformer encoder with three structural encodings: degree (centrality), shortest-path distance (spatial) and edge features along the shortest path (edge encoding) §3.1.
• A virtual graph token collects graph-level information; its last-layer state is the readout §3.2.
• Reported author results: validation MAE 0.1234 on PCQM4M-LSC (vs 0.1395 for GIN-VN) tab. 1; MolHIV AUC 80.51 with FLAG tab. 3; MolPCBA AP 31.39 with FLAG tab. 2; ZINC test MAE 0.122 tab. 4.
• Expressiveness: GNN layers such as GIN, GCN and GraphSAGE can be represented by a Graphormer layer with suitable weights §3.3, fact 1; app. A.
• Quadratic cost in the number of nodes, which the paper names as a main limitation §6 §D.

Problem

Transformers were not competitive on public graph-level leaderboards. The self-attention step scores pairs of nodes only by their feature similarity, so it ignores node importance and the structure between nodes §1. Earlier Transformer-based graph models mostly kept the GNN and replaced one aggregation step with softmax attention §1 §5.1. Self-attention costs time that grows quadratically with the number of nodes §D.

Method

Graphormer is built on the standard Transformer encoder §2 §3.2. The structural encodings below are added to it. Notation: h_i is the state of node i, d the hidden size, A the attention score matrix §2.

Centrality encoding §3.1 eq. 5. Each node gets one learned vector for its in-degree and one for its out-degree. Both are added to the node input features. For undirected graphs, the two degrees are the same number §3.1.

Spatial encoding §3.1 eq. 6. For each ordered node pair, the shortest-path distance is computed. Each distance value has one learned scalar, which is added to the attention score of that pair. Pairs with no path get one special value (stated as -1) §3.1. The scalar is shared across all layers §3.1.

Edge encoding §3.1 eq. 7. For each ordered pair, one shortest path is chosen. For the n-th edge e_n on that path, the edge feature x_{e_n} is dotted with a learned weight vector w^E_n of size d_E. The sum over the path is divided by the path length N and added to the attention score §3.1. The paper sets the edge feature dimension d_E and the number of heads to 32 §4.1.

Virtual graph token §3.2. One extra node is connected to every node. Its distance to any node is treated as 1 in the shortest-path step, but its spatial scalars are replaced by a separate learned scalar so that its links are not mistaken for real edges §3.2. Its state after the last layer is the graph representation §3.2.

Layer §3.2 eq. 8 eq. 9. Layer normalization is applied before the attention block and before the feed-forward block (pre-LN), and the feed-forward block has inner size d, not 4d §3.2.

The model and its attention bias, in pseudocode:

# Full model forward  [§3.1] [§3.2] [eq. 5]
# Input:  nodes (N, F); in_degree (N,); out_degree (N,); spd (N, N)
#         path_edges (N, N, path length, d_E)
# Output: one prediction per graph
def graphormer(nodes, in_degree, out_degree, spd, path_edges):
    # Centrality encoding: add learned in- and out-degree vectors  [eq. 5]
    h = nodes + in_degree_table[in_degree] + out_degree_table[out_degree]
    # Prepend the virtual graph token  [§3.2]
    h = concat(graph_token, h)                         # (N + 1, d)
    # The bias is built once and reused by every layer  [§3.1]
    bias = build_attention_bias(spd, path_edges)       # (heads, N+1, N+1)
    for layer in layers:
        h = graphormer_layer(h, bias)
    # Read the prediction from the graph token  [§3.2]
    # The output head is not stated in the paper.
    return head(h[0])

# Structural attention bias  [§3.1] [eq. 6] [eq. 7]
# Input:  spd (N, N); path_edges[i][j] = edge features on the chosen path
# Output: bias (heads, N+1, N+1) with the graph token at index 0
def build_attention_bias(spd, path_edges):
    # Spatial term: one learned scalar per distance value  [eq. 6]
    # Unconnected pairs use one reserved value; the paper writes -1 [§3.1]
    spatial = spatial_table[spd]                       # (N, N)
    # Graph token row and column: one separate learned scalar  [§3.2]
    spatial = pad_with_virtual_scalar(spatial)         # (N+1, N+1)
    # Edge term: average over the path of feature . learned weight  [eq. 7]
    edge = mean_over_path(path_edges, edge_weights)    # (N, N)
    return spatial + pad(edge)                         # (heads, N+1, N+1)

# One Graphormer layer, pre-LN  [§3.2] [eq. 8] [eq. 9]
# Input:  h (N + 1, d); bias (heads, N+1, N+1)
# Output: h (N + 1, d)
def graphormer_layer(h, bias):
    # Attention on the normalized input, with the bias, then residual  [eq. 8]
    h = h + multi_head_attention(layer_norm(h), bias)
    # Feed-forward on the normalized input, inner size d, then residual  [eq. 9]
    h = h + feed_forward(layer_norm(h))
    return h

Key equations, as stated in the paper:

A_{ij} = \frac{(h_i W_Q)(h_j W_K)^\top}{\sqrt{d}} + b_{\phi(v_i,v_j)} + c_{ij} eq. 6 eq. 7

c_{ij} = \frac{1}{N}\sum_{n=1}^{N} x_{e_n} (w^E_n)^\top eq. 7

Figures: fig. 1 shows the three encodings (centrality, spatial, edge) as an illustration. fig. 2 in appendix A shows two graphs that the 1-WL test cannot tell apart, but whose shortest-path distance sets differ app. A.

Expressiveness §3.3. Fact 1: with suitable weights and distance function, a Graphormer layer can represent the AGGREGATE and COMBINE steps of GIN, GCN and GraphSAGE §3.3 app. A. Fact 2: without extra encodings, a Graphormer layer's output can represent a mean READOUT over the graph, which is the role of a virtual node §3.3 app. A. The proofs are sketched in §3.3 and given in app. A. The paper also says Graphormer separates some graph pairs that the 1-WL test cannot, using shortest-path distance sets §3.3 app. A.

Hyperparameters

Values for the models in the main results. "not stated" means the paper does not give the value.

Results

Author-reported. The paper reports single-model results except the PCQM4M entry noted below. Setups: PCQM4M-LSC is the OGB Large-Scale Challenge quantum-chemistry regression set with more than 3.8M graphs; MolHIV, MolPCBA and ZINC are OGB and benchmarking-GNN sets §4 app. B.1. Baselines are GNNs that rank highly on the public leaderboards, without extra domain features §4.2.

Ablation on PCQM4M-LSC, 12 layers, 100K iterations, validation MAE tab. 5:

Pre-trained GNN comparison on MolHIV and MolPCBA app. C, tab. 12, tab. 13: GROVER with extra molecular features reaches 79.33 (MolHIV AUC) and 16.77 (MolPCBA AP) at 48.8M parameters; GROVER-Large reaches 80.32 and 13.05 at 107.7M. The paper reports these as its own fine-tuning runs.

Paper vs code

Code: the official repository, Microsoft/Graphormer at commit 59c0decffcade9df81d29dcc178a489b31958bab (MIT), linked from the abstract. Code paths below are relative to the repository root.

• Layer norm placement. The paper uses pre-LN §3.2 eq. 8 eq. 9. The code picks the layer norm with pre_layernorm, which defaults to False in both graphormer_base and graphormer_slim graphormer.py graphormer.py. The PCQM4M and ZINC scripts do not pass --pre-layernorm, so they train with post-LN pcqv2.sh zinc.sh. Only the MolHIV script passes it hiv_pre.sh. The layer switches on this flag graphormer_graph_encoder_layer.py.
• Batch size, PCQM4M. Paper: 1024 tab. 7 §4.1. Released script: 256 pcqv2.sh. The paper does not say the script is the final setting.
• Batch size and warm-up, ZINC. Paper: batch 256 and 40K warm-up steps tab. 10. Released script: batch 64 zinc.sh and 60K warm-up steps zinc.sh.
• MolHIV script differs from the appendix table. Dropout: paper 0.1 tab. 9; script 0.0 hiv_pre.sh. Maximum epochs: paper 8 tab. 9; script 5 (from epoch=4) hiv_pre.sh. FLAG steps m: paper 2 tab. 9; script 3 hiv_pre.sh. FLAG step size \alpha: paper 0.2 tab. 9; script 0.01 hiv_pre.sh. \epsilon: 0 in both hiv_pre.sh. Warm-up: paper ratio 0.06 tab. 9; script 16% of updates hiv_pre.sh. The script also starts from the checkpoint named pcqm4mv1_graphormer_base_for_molhiv hiv_pre.sh, which is a PCQM4M version 1 name. The paper's text points to the PCQM4M-LSC model of tab. 1. The learning-rate end value (1e-5) in the script is not in the paper hiv_pre.sh.
• Adam \beta_1. Text: 0.99 §4.1. Table and scripts: 0.9 tab. 7 pcqv2.sh. The table value is likely right, but the text value is not explained.
• Spatial encoding, unconnected pairs. Paper: a special value written -1 §3.1. Code: the shortest-path routine stores 510 for unreachable pairs algos.pyx algos.pyx, and the spatial table is sized 512 graph_prediction.py.
• Spatial scalars per head. Paper: one scalar per distance value, with no mention of heads §3.1 eq. 6. Code: one embedding of size = number of heads per distance value graphormer_layers.py. The same is true of the edge weights graphormer_layers.py.
• Edge encoding. Paper: one dot product per path edge with a d_E-sized weight, averaged over the N edges eq. 7. Code: the edge feature is a categorical id looked up in a per-head embedding and averaged over the feature columns; each hop k is multiplied by a learned heads-by-heads matrix; the sum is divided by the full shortest-path length graphormer_layers.py. Only the first multi_hop_max_dist = 5 edges are used, and the division still uses the full length graph_prediction.py collator.py.
• Hard attention cut-off. Not in the paper. The collator sets the attention bias for pairs with shortest-path distance at or above spatial_pos_max to -\infty collator.py, with a default of 1024 graph_prediction.py.
• Virtual node distance. Paper: a distinct learned scalar for the graph token's row and column §3.2. Code: one learned scalar per head, added to both the graph token's row and column graphormer_layers.py graphormer_layers.py. This matches the idea but adds the per-head detail.
• Degree encoding. Matches eq. 5: separate in- and out-degree tables, added to the node features graphormer_layers.py graphormer_layers.py.
• Readout and output head. The paper names the graph token's final state as the graph representation §3.2 but does not describe the prediction head. Code: the token passes through a linear layer, GELU, layer norm and a bias-free linear to num_classes, plus a learned scalar graphormer.py graphormer.py. The loss reads token 0 l1_loss.py.
• Initialization. Not in the paper. Code: linear weights are drawn with standard deviation 0.02 divided by the square root of the layer count; embeddings with 0.02 graphormer_layers.py. The bias modules apply this on construction, and the encoder applies it when apply_graphormer_init is set graphormer_graph_encoder.py.
• Shared bias. Matches §3.1: the attention bias is built once per forward pass and used by every layer graphormer_graph_encoder.py graphormer_graph_encoder.py.
• Feed-forward size and residual form. Matches §3.2: base and slim archs set the FFN size to d graphormer.py graphormer.py.
• Mixed precision. The scripts train with --fp16 pcqv2.sh. The paper does not mention it.

No discrepancy was found for: learning rates of the PCQM4M Graphormer and of ZINC (2e-4), PCQM4M training steps (1M), PCQM4M dropout values, and ZINC weight decay (0.01) zinc.sh.

Pitfalls

• Layer norm. The text says pre-LN, but the PCQM4M and ZINC scripts train post-LN unless --pre-layernorm is set. An implementation that follows the paper and not the script may behave differently. Pick one and say which.
• Structure bias shape. The paper writes one scalar per distance and one edge weight vector per path position, and does not mention heads. The code uses per-head values. The paper says the multi-head extension is standard §2, but the per-head form is a choice an implementer must make explicitly.
• Edge term details. The paper's edge term is an average of dot products with learned vectors. The code uses categorical lookups, a per-hop heads-by-heads matrix, and a 5-hop cut-off. Implementing the paper's formula literally differs from the released code. The paper does not state the cut-off.
• Unreachable pairs. The paper says -1; the code uses 510 and a table of 512 entries. The paper does not say how the table size is chosen.
• Quadratic cost. Full attention over all node pairs scales with the square of the node count §6 §D. The paper leaves large graphs to future work.
• Batch sizes. The PCQM4M batch is 1024 in the paper's table and 256 in the released script. The paper does not say which one produced the reported number.
• Graph token readout. Only the state of the graph token is read. The paper does not name the head, so an implementation must pick one. The code uses a small MLP-like head with layer norm graphormer.py.
• FLAG augmentation. The MolPCBA and MolHIV scores use FLAG with the listed step counts and sizes tab. 8 tab. 9. The MolHIV script uses other values hiv_pre.sh. Before comparing numbers, check which settings produced them.
• Pre-training. The MolPCBA model is pre-trained with 18 layers and d = 1024, which differs from the 12-layer model used for MolHIV app. B.2.2 app. B.2.3. The PCQM4M result in tab. 1 is the 12-layer model.
• Model size. The MolPCBA model has 119.5M parameters, the MolHIV model 47.0M, and the ZINC model 489K tab. 2 tab. 3 tab. 4. The model size is part of the comparison.
• Ensemble result. The PCQM4M test result of 0.1200 uses an ensemble with ExpC §4.1. It is not a single-model number.
• Overfitting. The paper says the large model overfits the small molecular datasets, so it uses FLAG. Embedding dropout of 0.1 hurt the PCQM4M validation score; the paper sets it to 0 app. B.2.1.
• Dropout settings. Attention dropout is 0.1 and FFN dropout 0.1 for PCQM4M tab. 7. The script's --dropout 0.0 sets the embedding and residual dropout together in the code pcqv2.sh.
• Baselines. The GT baseline is re-run with Laplacian PE dimension 4 (default 8, which filters out 2.91% of molecules), at learning rate 1e-4 app. B.3.1. Baseline numbers depend on these choices.
• Reported spread. MolHIV, MolPCBA and ZINC report mean and standard deviation over runs. The PCQM4M result is a single validation number.

Lineage

• Builds on: the standard Transformer encoder §2; message-passing GNNs such as GIN, GCN and GraphSAGE, which Fact 1 recovers §3.3; Transformer-based graph models GT (Dwivedi et al.) and SAN, which the paper compares against §5.1.
• Compared against: GIN-VN, DeeperGCN-VN, GT, SAN, GROVER, GraphNorm, PNA, DGN tab. 1–4 tab. 12–13.
• Superseded by: not stated in the card's source material.