GraphTokenLM

GraphTokenLM is a Graph-Language Model (GLM) based on GraphToken (Perozzi et al., 2024). A GNN encoder maps an input graph into a small number of soft-prompt vectors ("graph tokens"), which are prepended to the input embeddings of a frozen pre-trained LLM. This checkpoint was trained on the MotifQA dataset in a multitask setting, and is used in our study that proposed a method for explaining GLM predictions via edge importance (see Citation).

Source code, training and evaluation scripts: N-Shimoda/GLMExplainer.

Architecture

Component	Setting
Pre-trained LLM	Qwen/Qwen3-4B-Base (hidden size 2560, 36 layers, 32 heads)
LLM parameters	Frozen (freeze_llm = true), no LoRA (enable_lora = false)
GNN encoder	GIN, 3 layers, hidden / output dim 64
Graph pooling	mean
Projection layers	2-layer MLP, 64-dim GNN output → 2560-dim graph tokens
Number of graph tokens	4
Node features	Laplacian positional encoding (lpe_dim = 8), no degree embedding
Learned node positional embedding	8-dim (pos_emb_dim = 8)
Max nodes per graph	20
dtype	float32

Only the GNN encoder and the projection layers are trained; the LLM weights are identical to Qwen/Qwen3-4B-Base.

Training

Item	Setting
Dataset	MotifQA (multitask over ba_shapes, ba_two_motifs, tree_cycle, tree_grid_v2)
Epochs	32
Optimizer	AdamW
Learning rate	5e-3
Weight decay	1e-2
LR scheduler	cosine, warmup ratio 0.05
Per-device train batch size	2
Gradient accumulation steps	4
GPUs	2 (torchrun --nproc_per_node=2)
Seed	42

The equivalent training command from the source repository:

torchrun --nproc_per_node=2 train.py \
   --dataset MotifQA \
   --subset ba_shapes ba_two_motifs tree_cycle tree_grid_v2 \
   --lpe-dim 8 --pos-emb-dim 8 \
   --gnn-type GIN \
   --gnn-hidden-dim 64 --gnn-out-dim 64 \
   --num-gnn-layers 3 --graph-pooling mean \
   --num-proj-layers 2 --num-graph-tokens 4 \
   --epochs 32 \
   --optim adamw --lr 5e-3 --weight-decay 1e-2 \
   --lr-scheduler-type cosine --warmup-ratio 0.05

Usage

The model ships with custom code (glm.py), so trust_remote_code=True is required. torch-geometric must be installed in addition to transformers.

from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained(
    "naos-ku/GraphTokenLM",
    trust_remote_code=True,
    load_llm_weights=False,  # LLM weights are already included in this repo
)
tokenizer = AutoTokenizer.from_pretrained("naos-ku/GraphTokenLM", trust_remote_code=True)

For end-to-end evaluation and explanation, use eval.py and explain.py in the source repository:

torchrun --nproc_per_node=2 eval.py \
   --dataset MotifQA --subset ba_shapes tree_cycle \
   --model-path "naos-ku/GraphTokenLM" \
   --num-trials 5 --per-device-batch-size 5

Citation

@article{shimoda2026glmexplainer,
  title={Identifying Important Subgraphs in Graph-Language Models via Representative Value Aggregation},
  author={Naoki Shimoda and Akihiro Yamamoto},
  journal={JSAI Technical Report, SIG-FPAI},
  volume={137},
  pages={36-43},
  year={2026},
  month=sep,
  doi={10.11517/jsaifpai.137.0_36}
}

References

• Perozzi et al. (2024). Let Your Graph Do the Talking: Encoding Structured Data for LLMs. (GraphToken)
• Fatemi et al. (2024). Talk like a Graph: Encoding Graphs for Large Language Models. (GraphQA)