Controllable Graph Diffusion for Community Search (CGD)

Reference implementation of CGD, a discrete graph-diffusion model for community search (CS) and attributed community search (ACS). The package covers the full method: the query-conditioned graph transformer denoiser, the Zero-GCN copilot, the attribute encoder, the discrete diffusion schedule and posterior, training, and stepwise-denoising inference.

Installation

pip install -e .# or, without installing the package:
pip install -r requirements.txt

Python >= 3.10 and PyTorch >= 2.5 are required.

Datasets

This implementation reads graphs from local files only; it does not download any dataset. Supply each dataset as a single NumPy .npz archive at <root>/<name>.npz or <root>/<name>/graph.npz (with <root> set by dataset.root and <name> by dataset.name in the config). The archive must contain:

array	dtype	shape	meaning
x	float32	[num_nodes, num_features]	node features
edge_index	int64	[2, num_edges]	COO edges, both directions for an undirected graph
labels	int64	[num_nodes]	node class / community id
attribute_names	str (optional)	[num_features]	feature names
communities	object (optional)	list of int64 arrays	ground-truth community memberships

Node ids in edge_index index [0, num_nodes); self-loops are dropped on load. The schema is described in src/cgd/data/datasets.py.

Usage

Train and evaluate a community-search model:

python scripts/train_cs.py --config configs/best/cs_cora_stable.yaml
python scripts/eval_cs.py --config configs/best/cs_cora_stable.yaml \
--checkpoint runs/best_cs_cora_stable/cs_model.pt

Attributed community search fine-tunes on top of a CS checkpoint:

python scripts/train_acs.py --config configs/best/acs_cora_example.yaml
python scripts/eval_acs.py --config configs/best/acs_cora_example.yaml \
--checkpoint runs/acs_cora_example/acs_model.pt

configs/best/ ships the community-only diffusion configurations for Cora and CiteSeer. With the standard Planetoid Cora/CiteSeer splits supplied in the format above, these configurations reach roughly 0.9x F1 on the 4-shot community-search task.

Tests

pytest

The unit tests run on toy tensors and require no dataset.