Modelpublic

naos-ku/GraphTokenLM

sourceHugging Faceapache-2.0updated 5d agoView on Hugging Face

0likes376downloads

Model Card

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:

bash

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.

python

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:

bash

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

bibtex

@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)