GraphFormers: GNN-nested Transformers for Representation Learning on Textual Graph

In plain words

Many datasets link items that each carry text, such as papers that cite each other or products that get viewed together. The usual recipe runs in two stages: a language model turns each item's text into a vector, then a graph network mixes those vectors across the links. This paper argues the two stages should be fused. GraphFormers places a small graph-mixing step inside every layer of the language model, so each text is read together with its neighbours' context at every layer.

Training uses two extra ideas. First, the model practises on copies where some words are hidden, so it cannot rely only on the item's own text. Then it trains on the clean data. Second, a simplification lets each neighbour's encoding be computed once and reused. On three large link-prediction tasks, the authors report better accuracy than the two-stage recipe at about the same running cost.

TL;DR

GraphFormers nests a graph-attention step inside each layer of a BERT-style encoder (UniLM-base), so a node's text and its neighbours' text are processed together rather than in two separate stages. It is trained to predict links with in-batch negatives, first on span-masked inputs and then on clean inputs. The unidirectional variant lets neighbour encodings be cached. The authors report the best P@1, NDCG and MRR of all baselines tested on three textual graphs at near-equal running time; these results are author-reported, and the official code differs from the paper in several details (see Paper vs code).

Problem

Textual-graph representation learning needs each node's text and its graph neighbourhood. The common "cascaded" design encodes every node's text on its own first, then aggregates the vectors with a GNN. Because no information passes between nodes during text encoding, the text model cannot use the neighbourhood to interpret a node's words §1 fig. 1.

A fused design is more expressive but naively costly: every neighbour depends on the others, so each query re-encodes all of them §3.2. Training has its own failure mode. When the centre node's text is informative enough to predict its links alone, the model can ignore the GNN part and leave it undertrained §3.3.

The experiments assume link prediction as the target, a text per node, and at most five sampled neighbours per node §4.1.

Method

Each node x has a text, and G_x denotes x together with its neighbour set N_x. The model learns an embedding h_x from x's text and N_x, and it is trained so that the embedding similarity predicts whether two nodes are linked §3.

1. Tokenise and embed. Each node's text is tokenised with uncased WordPiece. A [CLS] token is padded in front, and the input embedding sums word and position embeddings §3.1 §4.1.
2. Layer 0 is an ordinary Transformer layer. It encodes every node independently §3.1 alg. 1.
3. Graph aggregation, for layers 1 to L-1. Each node's [CLS] state is its node-level embedding. A multi-head attention over all nodes of the graph, with a learnable position bias B, produces one graph vector per node §3.1 eq. 1.
4. Dispatch. The graph vector is concatenated to the node's token states, giving graph-augmented tokens §3.1.
5. Asymmetric Transformer. Queries come from the plain tokens of the previous layer, while keys and values come from the augmented tokens. The output keeps the input length §3.1 eq. 4.
6. Unidirectional simplification. Only the centre node queries the neighbourhood. Neighbours are encoded independently, so their per-layer states can be cached and reused §3.2 eq. 5.
7. Readout. The centre node's [CLS] state after the last layer is its representation §3.1.

The GNN parameters are shared across layers by default §4.4.

Full workflow as pseudocode alg. 1:

# GraphFormers encoding of a centre node and its neighbours  [alg. 1]
# Input:  node_tokens  (N, L) token ids; row 0 is the centre node  [§3.1]
# Output: embedding of the centre node  (d,)

def encode_graph(node_tokens):
    # Layer 0 encodes every node alone, with no graph information  [alg. 1]
    hidden = [layers[0](tokens) for tokens in node_tokens]    # N x (L, d)
    for layer_index in range(1, num_layers):
        # Node-level embedding is the [CLS] state of each node  [§3.1]
        node_vectors = stack([tokens[0] for tokens in hidden])  # (N, d)
        # Layerwise GNN; only the centre node queries  [§3.2] [eq. 5]
        graph_vector = graph_attention(node_vectors, position_bias)
        # Centre: queries from plain tokens, keys/values from augmented  [eq. 4]
        center_next = asymmetric_layer(
            queries=hidden[0],
            keys_values=prepend(graph_vector, hidden[0]),
        )
        # Neighbours skip the graph vector and can be cached  [§3.2]
        neighbours_next = [
            layers[layer_index](tokens) for tokens in hidden[1:]
        ]
        hidden = [center_next] + neighbours_next
    # Output is the centre's [CLS] state after the last layer  [§3.1]
    return hidden[0][0]

Graph attention is standard scaled dot-product attention over the node-level vectors, with the learnable bias B added before the softmax eq. 1:

\hat{Z} = \mathrm{MHA}(Z), \qquad \mathrm{head}_j = \mathrm{softmax}\!\left(\frac{Q K^\top}{\sqrt{d}} + B\right) V eq. 1

The bias distinguishes three node relations: centre-to-centre, centre-to-neighbour and neighbour-to-neighbour §3.1. Eq. 1 writes this bias as B; the ablation with it removed is in §4.4.

Training objective: link prediction with in-batch negatives. The positive partner of one instance serves as a negative for every other instance in the same mini-batch §3.3:

\mathcal{L} = -\log \frac{\exp(\langle h_q, h_k\rangle)}{\exp(\langle h_q, h_k\rangle) + \sum_{r \in R} \exp(\langle h_q, h_r\rangle)} eq. 6

The paper removes pairs where one node is in the other's neighbour set §3.3, footnote.

Two-stage progressive learning §3.3:

1. Warm-up. Train with the loss of eq. 7 on inputs in which a subset of tokens is masked (15% of tokens, span masking, 80/10/10 app. A.1) until convergence.
2. Fine-tune. Continue training with eq. 6 on the unmodified data until convergence.

Encoding cost per layer is O(M² + M·P²) for M nodes of P tokens. The graph part is O(M²), the text part O(M·P²) §3.1.

Hyperparameters

Training hyperparameters per dataset tab. 6 app. B:

Other settings:

Results

Author-reported. Setup: test links scored with P@1, NDCG and MRR; backbone UniLM-base; baselines share the same text encoder and aggregator capacity §4.1 §4.2. Dataset statistics tab. 1:

Overall results tab. 2 §4.3:

The paper reports relative P@1 gains over the most competitive baseline, PLM+Max on all three datasets (underlined in the table), of 2.9% on Product, 4.8% on DBLP and 6.5% on Wiki §4.3. These match the table's values.

Neighbour-size sweep on DBLP tab. 3 §4.3:

Ablations tab. 4 §4.4. Each row removes one component from the default model:

Without progressive learning, the paper reports P@1 drops of 0.98, 1.71 and 1.18 points on Product, DBLP and Wiki §4.4. The paper writes these as percentages. They equal the absolute P@1 differences in tab. 4.

Time and memory per mini-batch with neighbour count #N from 3 to 200 tab. 5 §4.5. Batch of 32 instances, 16 tokens per node, Nvidia P100:

The paper reports GraphFormers' time overhead at #N = 200 as about 3.5% §4.5; the table gives 2801.67 vs 2706.35 ms.

Online A/B test on Bing Search ads retrieval §4.6 fig. 3: relative to the previous production system, the authors report RPM +1.87%, click yield +0.96% and CPC +0.91%, measured over an 11-day window of millions of impressions per day.

Paper vs code

Official code: https://github.com/microsoft/GraphFormers, MIT licence, at commit 06636cd7196cc2b3eecdb5a49fd2d405d2ea4e4b. Searched via the paper's abstract, which links it directly.

• Attention in layers 1 and later. The paper's asymmetric form computes queries from the plain tokens H and keys/values from the augmented tokens Ĥ §3.1 eq. 4. The code uses the standard self-attention block on the augmented sequence, so queries also come from augmented tokens modeling_graphformers.py modeling.py. The paper's listing passes asymmetric = True app. A.2, but the repo's modeling code has no such flag.
• Masked warm-up covers only the centre node. The paper masks tokens of each input node §3.3. The code masks only index 0 of each sample, the centre node data_handler.py.
• Span length. The paper clips the geometric span length at 10 app. A.1. The code redraws until the length is at most 10 data_handler.py.
• Warm-up labels are unused. The code computes MLM labels, but the model's forward accepts and ignores them modeling_graphformers.py. The warm-up stage is therefore the contrastive loss on masked inputs, as the paper describes.
• Two-stage schedule is not orchestrated. The paper trains the warm-up stage to convergence, then fine-tunes §3.3. The code's --mlm flag defaults to off parameters.py. The shipped main.py never sets it and hard-codes cont = False main.py. Running both stages requires two manual runs with checkpoint loading run.py.
• Neighbour sampling. The paper samples five neighbours uniformly without replacement §4.1. The code keeps the first five neighbours in the input file and does no sampling data_handler.py. The data files are not in the repo, so any sampling done offline cannot be checked.
• Evaluation negatives. The paper ranks each query against one positive and 299 randomly sampled negatives §4.1. The code ranks each query against the other 299 keys in its evaluation batch of 300 pairs, with the diagonal as positive parameters.py modeling_graphformers.py. Whether the test files are built to match is not verifiable.
• Position bias. The paper's bias distinguishes three relations §3.1. Only centre queries occur, so the code's bias depends only on whether the key is the centre or a neighbour modeling_graphformers.py. The same linear layer also adds a token-level relative-position bias when rel_pos_bins > 0 modeling_graphformers.py modeling_graphformers.py. The paper does not mention this. The value comes from an external config, so whether it is active is not verifiable.
• Scaling. The paper writes √d without defining d eq. 1. The code divides by the square root of the per-head size modeling.py.
• Per-dataset settings. The paper gives learning rate, batch size and sequence length per dataset tab. 6. The code defaults to lr 1e-5, batch 30 and 32 tokens parameters.py parameters.py parameters.py. main.py sets neither per dataset nor the 8-GPU count main.py.
• Initialisation. The paper's backbone is UniLM-base §4.2. The code's default starting weights are a file named graphformers-dblp.pt, loaded with strict=False parameters.py run.py. Whether that file is a fine-tuned checkpoint is not stated.
• Matching the paper. Inner-product scoring with in-batch negatives eq. 6 modeling_graphformers.py; early stopping on validation P@1 with patience 2 run.py; Adam with the paper's betas and epsilon (PyTorch defaults) run.py; seed 42 parameters.py run.py; 12 nodes per instance data_handler.py.
• Not in the paper. The code also reports AUC modeling_graphformers.py and enables fp16 by default parameters.py.

Pitfalls

• Masking scope differs. The paper masks tokens of every input node §3.3; the reference code masks only the centre node data_handler.py. Choose one and record it.
• Attention form changes the model. Asymmetric attention (queries from plain tokens) is the paper's stated design eq. 4. Standard self-attention over augmented tokens changes what the centre's queries see.
• The footnoted pair filter is not in the code. Pairs where one node is in the other's neighbourhood are kept in the code's data loader data_handler.py. Whether the data is pre-filtered is not stated.
• In-batch negatives are batch-dependent. Negatives are the other instances' positives. The paper's batch sizes are 240, 240 and 160 tab. 6, so the loss and gradients depend on batch composition. Duplicate or near-duplicate pairs in one batch become false negatives. The paper does not address this.
• Masking hyperparameters. Span lengths are geometric with p = 0.2 and a cap of 10 app. A.1. The code uses a rejection loop, not a clip, so the length distribution is slightly different. The mask fraction of 15% is applied per node in the paper, but per centre node in the code.
• Additive masks use −10000, not −∞. Attention masks are added as −10000 modeling_graphformers.py. Half-precision training (on by default) is sensitive to large negative masks, so verify numerics if you change the masks.
• Neighbour count and memory. Memory grows with neighbours; at #N = 200 GraphFormers uses 7.33 GiB against 6.92 GiB for PLM+Max tab. 5, about 6% more, though the text reports only time overhead.
• Checkpoint key mismatch in the code. train() saves a raw state_dict run.py run.py, but test() reads checkpoint['model_state_dict'] run.py, and load_bert() reads the same key run.py. Expect a KeyError when you pass a checkpoint written by train() to test(). Also load_bert() uses strict=False, so missing weights are silently skipped run.py.
• Dataset files are external. The data and the pretrained checkpoint are on an external drive, not in the repo README.md. The exact splits, neighbour sampling and warm-up data cannot be reproduced from the repo alone.
• Stage switch is manual. The paper says the warm-up stage runs until convergence §3.3. The code has no automatic switch between stages, and its early stopping is applied per run, not per stage.
• Cost of the two stages. The paper reports 58.8 to 151.2 hours per dataset on 8 V100s app. B. The paper's efficiency comparison used a P100 §4.5.

Lineage

• Builds on: the cascaded PLM+GNN recipe, using GraphSAGE-style pooling and GAT aggregators §2 §4.2; UniLMv2 as the backbone §4.2; span masking after SpanBERT app. A.1; in-batch negatives after dense retrieval work §3.3; curriculum learning as motivation for progressive training §1.
• Compared against: PLM only, TNVE, IFTN, PLM+GAT, PLM+Max, PLM+Mean and PLM+Att §4.2.
• Superseded by: not stated in the wiki's content.