NodeFormer: A Scalable Graph Structure Learning Transformer for Node Classification

In plain words

Most graph networks only let a node hear from the neighbours it is linked to. NodeFormer lets every node listen to every other node in each layer. Doing that naively is far too expensive on a big graph, so the model uses a mathematical shortcut that never builds the full table of node pairs. A randomised choice decides which partners matter, so each layer ends up with a sparse set of links, and the links differ from layer to layer.

When the data comes with a graph, the model is nudged to keep the links that are already there. When it does not, the model learns its own links. The aim is node labelling on large networks, and also on data with no graph at all, such as image or text features.

TL;DR

NodeFormer is a graph Transformer for node classification in which each layer can pass messages between any pair of nodes. A kernelized Gumbel-Softmax operator, built on random features, keeps the cost linear in the number of nodes and learns a sparse latent graph per layer §3.1. Optional terms use the input graph as a relational bias and as a likelihood regulariser §3.3 §3.4. Author-reported results are competitive on small benchmarks and reach graphs with 2M nodes, where prior structure learners failed §4 §1.

Problem

Propagating messages only along the observed edges fails in known ways: over-squashing, heterophily, long-range dependence, incomplete edges, and no graph at all §1. Learning a new adjacency for each layer can address this, but full all-pair attention costs O(N²) in the number of nodes N §3, eq. 2. The softmax denominator also shrinks attention weights toward zero on large graphs, which the paper calls over-normalizing §3.1. Earlier structure learners cut the cost with neighbour sampling, anchors or hashing, which the paper says can lose precision §1.

The paper targets node-level prediction on one large graph, not graph-level tasks on many small graphs §2.

Method

Layer rule. Each layer first estimates a latent graph from the current node representations, then propagates features over that graph §3, eq. 1. The full version of this is a softmax attention over all node pairs §3.1, eq. 2.

Linear-time attention. The dot-then-exponentiate kernel is written as an inner product of feature maps §3.1, eq. 3. The paper uses positive random features, with m random directions drawn from a standard Gaussian §3.1, eq. 4. Once the feature maps are applied, the sums over all nodes can be computed once and reused for every query node. That gives O(N) cost for the whole layer and avoids the N×N similarity matrix §3.1, eq. 5.

Sparse latent structure. Sampling a neighbour directly is not differentiable. The paper adds Gumbel noise to the attention logits, scaled by a temperature τ, which gives a continuous relaxation of sampling one neighbour per node §3.1, eq. 6. Kernelizing that form gives the operator used in practice: each key feature is multiplied by e^(g/τ) before the sums are taken §3.1, eq. 7. One draw per node is taken for each sample, and the output is averaged over K draws §3.1. Algorithm 1 gives the same computation in matrix form, including the averaging and the relational bias step alg. 1.

# Kernelized Gumbel-Softmax message passing, one head  [alg. 1] [eq. 7]
# queries, keys, values: (N, d) from linear maps of the layer input
def latent_message_passing(queries, keys, values, tau, num_draws):
    # Divide by sqrt(tau) before the random feature map  [eq. 7]
    queries = queries / sqrt(tau)
    keys = keys / sqrt(tau)
    # Random feature map phi with m features  [eq. 4]
    query_features = random_features(queries)   # (N, m)
    key_features = random_features(keys)        # (N, m)
    outputs = zeros(N, d)
    for draw in range(num_draws):
        # One Gumbel(0, 1) noise value per node, per draw  [alg. 1, line 4]
        gumbel = sample_gumbel(N)               # (N,)
        # Multiply key features by e^(g / tau)  [eq. 7]
        noisy_keys = exp(gumbel / tau)[:, None] * key_features   # (N, m)
        # Sums over all nodes are computed once and shared by all queries
        # [eq. 5] [alg. 1, lines 6-7]
        key_value_sum = noisy_keys.T @ values   # (m, d)
        key_sum = noisy_keys.sum(axis=0)        # (m,)
        numerator = query_features @ key_value_sum   # (N, d)
        denominator = query_features @ key_sum       # (N,)
        outputs += numerator / denominator[:, None]
    # Average over the draws  [alg. 1, line 8]
    return outputs / num_draws

Theory. Theorem 1 bounds the random-feature error for the softmax kernel with probability at least 1 − ε, given norm bounds r on queries and keys. The bound depends on m and τ, not on N §3.2. Theorem 2 states that as τ → 0 and m is large enough, the kernelized Gumbel-Softmax selects each neighbour with the original categorical probability §3.2. The paper's remark: a large τ pushes every edge weight toward 1/N, so the layer becomes mean pooling, while a small τ needs a large m for an accurate approximation §3.2. Proofs are in appendix B.

Input graph as bias. If an input graph exists, an observed edge (u, v) adds σ(b) to the attention weight, where b is a learnable scalar per layer. The update adds σ(b) times the sum of the observed neighbours' values §3.3, eq. 8. This costs O(N + E) §3.3. Higher-order adjacency is allowed, as in MixHop §3.3.

Edge likelihood. The model also estimates π_uv, the probability of each observed edge under the latent attention. The objective adds an edge term that rewards high π on observed edges, weighted by 1/d_u §3.4, eq. 10 eq. 11 eq. 12. The total loss is the supervised cross-entropy plus λ times the edge term §3.4, eq. 9. Section 5 reads this objective as an evidence lower bound whose optimum sets the latent graph close to the posterior §5, eq. 13.

# Training objective  [§3.4] [eq. 9] [eq. 10] [eq. 11]
# edge_probs: per layer, pi_uv for each observed edge  [eq. 12]
# edge_weights: 1 / d_u for each observed edge  [eq. 11]
def training_loss(logits, labels, train_nodes, edge_probs, edge_weights,
                  lam):
    # Supervised classification loss on labelled nodes  [eq. 9]
    supervised = cross_entropy(logits[train_nodes], labels[train_nodes])
    # Edge term: weighted log-likelihood of observed edges, every layer
    # [eq. 10]
    edge_term = 0
    for probs in edge_probs:
        edge_term += sum(edge_weights * log(probs))
    edge_term = -edge_term / (num_nodes * num_layers)
    return supervised + lam * edge_term

Architecture. The overview figure shows the three parts: all-pair message passing (always on), a relational bias and an edge loss (both used only when an input graph exists) fig. 1. Each layer's feed-forward pass is Algorithm 1 alg. 1, and the model ends in a one-layer MLP prediction head §4. For transductive runs the paper uses two layers of width 32, an input MLP with ELU, and no activation in the propagation layers app. C.1. Graph-enhanced runs use the same design, with the width grid-searched app. C.3. For the two large graphs it uses three layers of width 64 with one head, ELU in every layer, an identity σ, and a first-order adjacency app. C.2.

Hyperparameters

Values the paper states:

The paper defers the NodeFormer settings to its public code app. C.1. The values below are what the released scripts set for each dataset. None of them are stated in the paper itself.

The no-graph runs set the relational bias order to 0 run.sh. The kNN graph uses k = 5 in these scripts, which is the default parse.py; Table 4 also reports k = 10, 15, 20.

Results

Author-reported. Setup: five runs, random 50/25/25 splits, test accuracy at the best validation epoch §4 app. C.1.

Transductive node classification §4.1. Figure 2 plots results with error bars on Cora, Citeseer, Deezer and Actor against GCN, GAT, MixHop, JKNet, DropEdge, LDS and IDGL fig. 2. The baseline bars have no printed values, so the figure cannot be tabulated here. The full model's numbers are in the ablation table:

Large graphs §4.2. Mini-batch training, with 10,000 nodes per batch for OGB-Proteins and 100,000 for Amazon2M §4.2.

The two ablation rows in each table replace the Gumbel operator with an ordinary softmax, at τ = 1.0 or τ = 0.25 §4.4.

Graph-enhanced tasks §4.3. A k-NN graph is built from input features, or no graph is used at all tab. 4.

The paper says NodeFormer is best in seven of these eight kNN settings. At 20News with k = 5, LDS scores 66.15 ± 0.36 against NodeFormer's 66.01 ± 1.18 tab. 4.

Cost §4.4. Against LDS, the paper reports reductions of up to 93.1% in training time, 97.9% in inference time and 75.6% in memory. Against IDGL, with its anchor approximation, the reductions are up to 61.8%, 80.8% and 80.6% fig. 3 §4.4.

Ablations tab. 6 §4.4. Both the edge regulariser and the relational bias help on Cora and Citeseer. On Deezer and Actor, the differences are inside one standard deviation.

Paper vs code

Code: official repository github.com/qitianwu/NodeFormer, cloned at 5390b6908326905b2eb840f65a70b0f4082aef43. The repository is named in the abstract and owned by the first author. It has no LICENSE file, so this card cites the code by path and line only and quotes nothing from it.

• Random features. The paper draws i.i.d. standard Gaussian directions §3.1, eq. 4. The code uses orthogonal blocks from a QR decomposition, with rows rescaled by Gaussian norms, and redraws them on every forward pass with a seed taken from the sum of the queries nodeformer.py nodeformer.py.
• Kernel scaling. The paper's kernel is κ(q/√τ, k/√τ) §3.1, eq. 3, eq. 7. The code also multiplies queries and keys by d^(-1/4), where d is the per-head width. The effective exponent is q·k/(τ√d), not q·k/τ. Here the per-head width equals the hidden size, because the query projection outputs hidden × heads nodeformer.py nodeformer.py nodeformer.py nodeformer.py.
• Gumbel only in training. The paper describes the feed-forward pass with Gumbel noise §3.1 and does not say what happens at test time. The code draws Gumbel noise only when the model is in training mode and otherwise uses the deterministic kernel softmax at the same τ nodeformer.py.
• Relational bias. The paper adds σ(b) times the raw sum over observed neighbours, with one scalar per layer §3.3, eq. 8 alg. 1, line 9. The code normalises by the square root of both endpoint degrees, keeps one scalar per order and per head, initialises it to 0.1 before the sigmoid, and adds the bias before the output projection nodeformer.py nodeformer.py nodeformer.py nodeformer.py.
• Head combination. Appendix C.1 says heads are combined by averaging app. C.1. The code concatenates the heads and applies a learned linear map (Wo) nodeformer.py.
• Edge-loss scale. The paper normalises the edge term by 1/(N·L), summed over edges §3.4, eq. 10. The code takes the mean over edges, heads and layers, and then multiplies by λ nodeformer.py main.py. The code's term is therefore the paper's term times N/E for a single head, with E counting the directed edges in the code, self-loops included. λ values from the paper are not directly comparable to the code's.
• Edge set. The paper defines the observed edges E and never mentions self-loops or symmetrisation §3. The code adds self-loops to the edge set, which feeds both the relational bias and the edge term, and symmetrises the graph unless --directed is set main.py main.py main-batch.py main-batch.py.
• Architecture not in Algorithm 1. Algorithm 1 has no input MLP, no normalisation, no residual link and no output projection, and its readout uses all layers alg. 1. The code adds an input linear layer with LayerNorm, ELU and dropout nodeformer.py nodeformer.py, LayerNorm and a residual link after each layer nodeformer.py, and reads out the last layer unless --use_jk is set nodeformer.py. The flag --use_bn controls LayerNorm, not batch norm nodeformer.py nodeformer.py. --use_jk is set only for the large-graph scripts run.sh.
• Number of draws K. The text gives K = 5 as an example §3.1. The transductive and graph-enhanced scripts use K = 10 run.sh run.sh; the large-graph scripts use K = 5 run.sh.
• Sigma for large graphs. Section 4 says σ is sigmoid for all datasets §4. Appendix C.2 and the large-graph script use identity app. C.2 run.sh.
• Epochs for graph-enhanced runs. Appendix C.3 says training is the same as for transductive runs, which means 1000 epochs app. C.3. The scripts use 200 to 300 run.sh run.sh run.sh run.sh.
• Mini-batch sampling. Appendix C.2 says the mini-batches partition all nodes app. C.2. The code shuffles only the training nodes into batches main-batch.py and evaluates on the full graph on the CPU main-batch.py.
• Learning-rate schedule (OGB-Proteins only). A MultiStepLR schedule halves the rate at steps 100 and 200 and is stepped once per mini-batch, not per epoch main-batch.py main-batch.py. The paper does not mention it.
• Splits. The paper says five repeats with different initializations §4. The code draws a new random split for each run main.py main.py.
• Numbers. The Cora full-model result is 88.69 ± 0.46 in Table 6 tab. 6 but 88.80 ± 0.26 in the README. The Actor standard deviation is 1.29 in the paper and 0.89 in the README Readme.md (search "Key results"). Appendix D.1 gives Deezer as "nearly 20K nodes", but Table 5 gives 28,281 app. D.1 tab. 5.

Pitfalls

• Kernel and width. If you implement the paper's formula literally, you get a softmax with an effective temperature τ, not τ√d. Choose one convention, and say which one you used. The width d is the per-head width, which equals the hidden size in the released configs nodeformer.py.
• Random features and τ. The approximation error depends on m and τ, and the paper's remark says small τ needs large m §3.2. The paper uses τ = 0.25 with m not stated in its main text. Section 4.4 says Cora accuracy peaks near τ = 0.4 §4.4 fig. 6. Tune both together and report them.
• Numerical overflow. The Gumbel factor e^(g/τ) is applied directly, with no maximum subtraction nodeformer.py. In float32 it overflows once g/τ exceeds about 88, that is g above about 22. Such draws are rare per node, but the number of draws per step is large on the big graphs. The kernel path adds a 1e-6 stabiliser and subtracts maxima for stability nodeformer.py, which the paper does not describe.
• Train and test mismatch. Noise is drawn only in training. Test predictions use the deterministic softmax at the same τ nodeformer.py. Check that the evaluation path matches what you intend.
• Cost of sampling. Averaging over K draws multiplies the per-layer message-passing work by K §3.1. The O(E) terms (the edge likelihood and the relational bias) are not free: OGB-Proteins has 39,561,252 edges in the paper's dataset table tab. 5.
• Edge term sign and size. Subtracting λ times the averaged log-likelihood is correct, but the normalisation differs from the paper (see Paper vs code). Do not transfer λ values between the two without rescaling.
• Ablation strength. On Deezer the edge-regulariser difference is 0.02 with a standard deviation of 0.32. On Actor it is 0.16 with 1.32 tab. 6. The claim that both components help is clearly supported only for Cora and Citeseer.
• Mini-ImageNet features. The authors extracted 128-dimensional features with their own CNN, following a few-shot learning setup app. D.2. Results on these features do not transfer directly to raw images.
• Baselines. The transductive baselines have no printed numbers in the text fig. 2. Their hyperparameters were chosen by grid search on the validation set app. C.1.
• Licence. The repository has no LICENSE file. Implementers should not copy code from it. Use the paper's equations and pseudocode above.

Lineage

• Builds on: positive random features for the softmax kernel from Performer (ICLR 2021); the Gumbel-Softmax reparameterisation and the Concrete distribution (both ICLR 2017); and graph structure learning, in particular LDS (ICML 2019) and IDGL (NeurIPS 2020). Higher-order adjacency follows MixHop (ICML 2019). §3 §3.1 §3.3
• Compared against: GCN, GAT, MixHop, JKNet, DropEdge, SGC and GraphSAINT on the node benchmarks; LDS and IDGL on structure learning §4. Complexity and largest-demo sizes are compared in tab. 1.