SwapGT: Tokenized Graph Transformer

• SwapGT is a tokenized graph Transformer model that enriches node representations by using probabilistic token swapping to sample higher-order neighborhoods.
• The model employs multi-view similarity-based token sequence generation combined with a dedicated Transformer encoder to effectively integrate attribute and topology information.
• A center alignment loss regularizes multi-sequence embeddings, enhancing robustness and generalizability in node classification tasks.

SwapGT is a tokenized graph Transformer model for node classification, introducing a token swapping operation to enrich the diversity and informativeness of token sequences used as Transformer input. Unlike standard tokenized GTs, which primarily generate token sequences by considering only first-order neighbor similarity, SwapGT expands the sampled neighborhoods to include higher-order structure, generating more varied token sequences and improving node representation learning. The method comprises three central innovations: multi-view similarity-based token sequence generation, a probabilistic token swapping operation, and a center alignment loss for multiple sequence embeddings. SwapGT utilizes a Transformer backbone for node representation aggregation and is trained end-to-end with cross-entropy and center alignment objectives (Chen et al., 12 Feb 2025).

1. Token Sequence Generation with Multi-View Similarities

Given an attributed graph $G=(V,E,X)$ where $X\in\mathbb{R}^{n\times d}$, SwapGT constructs two distinct "views" (Editor's term) for each node $v_i$:

• Attribute view: The raw node feature vector $x_i$.
• Topology view: A propagated feature vector $x_i'$ derived via K-step Approximate Personalized PageRank (APPNP) propagation:

$$X' = \varphi(\hat{A}, X, K),\quad \text{where}~ \hat{A} = (D+I)^{-1/2}(A+I)(D+I)^{-1/2}$$

and $\varphi$ applies $K$ iterations:

$H \leftarrow (1-\alpha_p)\hat{A} H + \alpha_p X$

with teleport probability $\alpha_p$.

For each view, pairwise cosine similarities are computed (attribute: $X\in\mathbb{R}^{n\times d}$0, topology: $X\in\mathbb{R}^{n\times d}$1). Each node forms an initial token set by selecting the top-$X\in\mathbb{R}^{n\times d}$2 most similar nodes:

$X\in\mathbb{R}^{n\times d}$3

These constitute $X\in\mathbb{R}^{n\times d}$4-NN neighborhoods per view, forming the basis for subsequent token sequence construction.

2. Probabilistic Token Swapping to Enlarge Neighborhoods

To move beyond the limitations of fixed 1-hop neighborhoods and foster sequence diversity, SwapGT introduces the token swapping operation. For each view and each node $X\in\mathbb{R}^{n\times d}$5:

• Algorithm 1 (Token Swapping for a Single View):

1. Initialize the token set $X\in\mathbb{R}^{n\times d}$6.
2. For $X\in\mathbb{R}^{n\times d}$7 iterations:
   • For each $X\in\mathbb{R}^{n\times d}$8:
     • With probability $X\in\mathbb{R}^{n\times d}$9, keep $v_i$0.
     • Otherwise, substitute $v_i$1 with a uniformly random member of $v_i$2 (i.e., $v_i$3).
   • Update $v_i$4 to the swapped token set.
3. Output $v_i$5 as an augmented token set.

By running this procedure $v_i$6 times, $v_i$7 token sets (including the original) per view are produced:

$v_i$8

For each token set, a token sequence is constructed:

$v_i$9

Collectively, these yield $x_i$0.

3. Transformer-based Representation Learning

Each view is processed independently by a dedicated Transformer encoder:

• Projection: Each $x_i$1 is transformed via a learned linear layer $x_i$2 to model dimension $x_i$3.
• Stacked Transformer Layers: $x_i$4 standard Transformer layers are applied, each consisting of:
  • Multi-head self-attention:

  $x_i$5 - Position-wise FFN layers with GELU activation and LayerNorm.

After $x_i$6 layers, the first token embedding (corresponding to $x_i$7 itself) is extracted from each of the $x_i$8 sequences, resulting in $x_i$9. Per-view readout is computed by concatenating the first sequence's embedding with the mean of the remaining ones:

$x_i'$0

Attribute and topology view outputs, $x_i'$1 and $x_i'$2, are fused as a convex combination:

$x_i'$3

Prediction is made via a standard MLP head:

$x_i'$4

Classification loss is calculated via cross-entropy on labeled nodes.

4. Center Alignment Loss for Consistent Multi-Sequence Embeddings

SwapGT introduces the center alignment loss to encourage agreement among embeddings from multiple swapped token sequences per node:

• Define the per-node center embedding as:

$x_i'$5

• The per-node cosine alignment penalty per view is:

$x_i'$6

• The total batch center alignment loss is:

$x_i'$7

• The combined objective becomes:

$x_i'$8

with $x_i'$9 balancing classification and alignment.

5. Training Workflow and Hyperparameterization

Training consists of the following pipeline:

Step	Input/Operation	Output/Data Shape
Preprocessing	Graph, adjacency, APPNP, similarity matrices	Token sets per view per node
Token Swapping	$$X' = \varphi(\hat{A}, X, K),\quad \text{where}~ \hat{A} = (D+I)^{-1/2}(A+I)(D+I)^{-1/2}$$0 token sets via Algorithm 1	Sequences $$X' = \varphi(\hat{A}, X, K),\quad \text{where}~ \hat{A} = (D+I)^{-1/2}(A+I)(D+I)^{-1/2}$$1 and $$X' = \varphi(\hat{A}, X, K),\quad \text{where}~ \hat{A} = (D+I)^{-1/2}(A+I)(D+I)^{-1/2}$$2
Transformer Encode	Stack $$X' = \varphi(\hat{A}, X, K),\quad \text{where}~ \hat{A} = (D+I)^{-1/2}(A+I)(D+I)^{-1/2}$$3 layers per view	Embeddings $$X' = \varphi(\hat{A}, X, K),\quad \text{where}~ \hat{A} = (D+I)^{-1/2}(A+I)(D+I)^{-1/2}$$4, $$X' = \varphi(\hat{A}, X, K),\quad \text{where}~ \hat{A} = (D+I)^{-1/2}(A+I)(D+I)^{-1/2}$$5
Readout & Fusion	Per-view pooling, convex combination	$$X' = \varphi(\hat{A}, X, K),\quad \text{where}~ \hat{A} = (D+I)^{-1/2}(A+I)(D+I)^{-1/2}$$6, $$X' = \varphi(\hat{A}, X, K),\quad \text{where}~ \hat{A} = (D+I)^{-1/2}(A+I)(D+I)^{-1/2}$$7, $$X' = \varphi(\hat{A}, X, K),\quad \text{where}~ \hat{A} = (D+I)^{-1/2}(A+I)(D+I)^{-1/2}$$8
Classification Head	MLP prediction, loss computation	Scalar losses, predictions
Optimization	AdamW, backprop, early stopping	Trained model parameters

Critical hyperparameters encompass:

• $$X' = \varphi(\hat{A}, X, K),\quad \text{where}~ \hat{A} = (D+I)^{-1/2}(A+I)(D+I)^{-1/2}$$9 (token set size), $\varphi$0 (swapped sets), $\varphi$1 (swap rounds), $\varphi$2 (swap probability)
• $\varphi$3 (feature dimension), $\varphi$4 (attention heads), $\varphi$5 (Transformer layers)
• $\varphi$6, $\varphi$7, learning rate $\varphi$8, dropout $\varphi$9
• AdamW optimizer with $K$0, $K$1, $K$2

6. Empirical Performance and Practical Significance

Empirical evaluation on diverse node classification benchmarks demonstrates that SwapGT achieves superior classification performance compared to existing tokenized graph Transformers and other baselines. The key advantages derive from:

• Enhanced diversity and informativeness of token sequences, resulting in more expressive node representations.
• The probabilistic token swapping operation, which enables sampling from beyond first-order neighborhoods.
• Center alignment loss, which regularizes multi-token sequence embeddings to avoid representation collapse.

A plausible implication is that token-swapping and center alignment jointly contribute to both robustness and generalization, particularly in settings where node feature heterogeneity and higher-order structural information are significant.

7. Implementation Summary

The end-to-end pipeline, including preprocessing, token swapping, Transformer forward pass, and loss computation, is specified by the pseudocode in the primary source. All operations—from spectral preprocessing to Transformer encoding and optimization—are implemented in standard deep learning frameworks, involving no custom CUDA kernels. Early stopping is driven by validation set performance. The approach is applicable to any attributed graph with arbitrary node features, with hyperparameters amenable to grid search (Chen et al., 12 Feb 2025).