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)