gnn-augment

Boost a graph neural network for semi-supervised node classification by augmenting the graph with node-similarity representatives.

A graph neural network (GNN) propagates node features through the graph's adjacency matrix, so at layer l only l-hop neighbours influence a node. gnn-augment replaces the adjacency with a node-similarity representative — Katz, Rooted PageRank, or Graph Gravity (a gravitational link-prediction score: degree as mass, shortest path as distance) — letting distant-but-similar nodes influence each other even in a shallow network, which improves node-classification accuracy.

The trick is architecture-agnostic — it changes which graph the network sees — so it works across four GNN architectures: the Kipf–Welling GCN, GraphSAGE, GAT (graph attention), and GIN.

This is a packaged, reusable implementation of the augmentation method from:

A. Wahid-Ul-Ashraf, M. Budka, K. Musial. Simulation and Augmentation of Social Networks for Building Deep Learning Models. arXiv:1905.09087

The Graph Gravity representative is the same gravitational method whose link-prediction reference implementation is akanda-method.

Installation

pip install gnn-augment                   # release
pip install "gnn-augment[semantic]"       # + sentence-transformers, for the semantic (LLM-embedding) view
pip install git+https://github.com/AkandaAshraf/gnn-augment   # latest main

Requires Python ≥ 3.10; depends on torch, networkx, numpy, scipy and scikit-learn. A CUDA build of PyTorch is used automatically when a GPU is available (device="auto") — install the torch wheel you want (CPU or a specific CUDA) first if pip's default is not it. From a clone: pip install -e ".[dev,semantic]" (see Development and tests below).

Quick start

from gnn_augment import GNNClassifier, load_cora, make_semi_supervised_split

G, X, y = load_cora()                                    # networkx graph, features, labels
y_semi, train, val, test = make_semi_supervised_split(y, seed=0)  # -1 = unlabeled

clf = GNNClassifier(model="gcn", representative="rpr", threshold="auto")
clf.fit(G, X, y_semi, train_mask=train, val_mask=val)
print("test accuracy:", clf.score(G, X, y, mask=test))

Pick any architecture with model= — "gcn", "sage" (GraphSAGE), "gat" (graph attention), or "gin". The augmentation (representative, threshold) is independent of the architecture, so any combination works.

The estimator follows the scikit-learn fit / predict / predict_proba / score convention. Semi-supervised learning is expressed the scikit-learn way — set -1 for unlabeled nodes in y — or pass explicit train_mask / val_mask arrays. The graph may be a networkx graph, a dense adjacency matrix, or a SciPy sparse matrix.

The augmented representatives

Similarity matrices are L2 row-normalised and thresholded (paper §9.1): pass threshold=(t1, t2) (entries ≤ t1 → 0, > t2 → 1), "auto" (binarise at the mean), None (keep weights), or "select" to pick the best threshold on the validation set automatically. weighted_features=True adds the learnable per-feature weight vector (Eq. 25).

Build a representative on its own for use in another model:

from gnn_augment import build_representative
G_matrix = build_representative(A, kind="gravity", threshold="auto")  # normalised Laplacian

Multi-view training: a different augmented graph every epoch

The paper's conclusions suggest using several augmented matrices of the same graph as a regulariser. views= turns that on: each epoch the network trains on a different view of the graph (graph data augmentation), and at inference the class probabilities can be averaged over the views (test-time augmentation):

clf = GNNClassifier(
    model="gin",
    views=["adjacency", ("rpr", (0.0, 0.5)), ("gravity", (0.0, 1.0))],
    view_schedule="random",   # "random" | "cycle" | "mix" (Dirichlet mixture)
    predict_views="mean",     # "mean" (TTA) | "first"
)

On Cora (3 splits; benchmarks/multiview-cora.md), multi-view training with TTA beats the raw adjacency on every architecture without any view selection, and on the GCN it is the best configuration of all (0.819 vs 0.816 for the best single view). When one view clearly dominates (RPR on GraphSAGE/GIN), committing to it still wins — multi-view is the robust "don't know the best view" strategy. Most of the gain comes from averaging over views at inference, so predict_views="mean" is the recommended mode.

Is it a real regulariser? Tested properly — high-capacity models (up to 256 hidden × 4 layers), no dropout, no weight decay, no early stopping, so they memorise the training set (benchmarks/overfitting-cora.md): multi-view training beats the raw adjacency at 3 of 4 capacities on the GCN (up to +2.3 points at 256 × 4) and the single RPR view at all 4 — a genuine but modest effect that complements early stopping (the dominant regulariser) rather than replacing it. A fixed augmented view, by contrast, does not regularise: its gains depend on early stopping. One limitation: deep unregularised GIN fails to train under multi-view (its un-normalised sum aggregation cannot handle the message-scale jump between views) — use a single view, "cycle", or scaled views for deep GIN.

The semantic view: LLM embeddings of node text as a graph (Graph + LLM)

Nodes of citation graphs are papers with real text. "semantic" builds a graph representative from text embeddings: cosine-similarity kNN over sentence-transformer (or any LLM) embeddings of each node's text, with the original edges kept — a text-similarity view next to the topological ones. Embeddings are computed locally (MiniLM, 22M parameters; no API):

from gnn_augment import GNNClassifier, embed_texts, load_cora, load_cora_text

G, X, y = load_cora()
E = embed_texts(load_cora_text())          # n x 384, local sentence-transformer

clf = GNNClassifier(representative="semantic", threshold=None, semantic_k=10)
clf.fit(G, X, y_semi, train_mask=train, val_mask=val, node_embeddings=E)

# or as one view among several in multi-view training:
clf = GNNClassifier(views=["adjacency", ("rpr", (0.0, 0.5)), ("semantic", None)])

Raw texts: load_cora_text() (aligned with load_cora) and load_pubmed_raw() (graph + features + labels + texts), both reading the TAPE raw-text releases (data/raw_text/{cora_orig,PubMed_orig}.zip, from the Google-Drive links in the TAPE README). embed_texts(model_name=...) also accepts a local model directory — handy offline or behind TLS-intercepting proxies: download the sentence-transformers/all-MiniLM-L6-v2 files once (e.g. with curl) into data/models/all-MiniLM-L6-v2 and the example uses it automatically. cache_representatives=True caches augmented adjacencies in-process so repeated fits on the same graph (seeds, architectures, views) reuse the expensive Katz/RPR/gravity builds.

On Cora the semantic view, added to multi-view training, gives the best accuracy yet on every architecture (3 splits; benchmarks/semantic.md):

Text complements topology rather than replacing it: the semantic graph alone is weaker than the citation graph, but as an extra view it adds what the topological views lack — and it is the first augmentation that moves GAT.

At 7× the scale (PubMed, 19,717 nodes, run in sparse mode, 2 splits) the semantic view transfers: semantic kNN + adjacency is the most consistent single view (+2.4 / +3.3 / +3.4 points on GCN / GraphSAGE / GIN; the best single view on GCN and GraphSAGE), and multi-view training that includes it is best or tied-best on GraphSAGE (+3.4) and GIN (+4.0, with gravity as a fourth view) — every architecture gains at least 2.4 points where the topological views alone gave at most +2.3 at this scale.

The augmentation helps every architecture

The paper applied augmentation only to the GCN, but the trick generalises. On Cora, the RPR representative improves all four architectures (mean test accuracy over 3 splits; see benchmarks/architectures-cora.md):

python examples/cora_demo.py            # vanilla vs augmented GCN
python examples/architectures_demo.py   # augmentation across GCN/SAGE/GAT/GIN
python examples/pubmed_demo.py          # scale up to ~20k nodes (GCN/SAGE/GIN)

Datasets: load_cora(), load_planetoid("cora" | "citeseer" | "pubmed"), plus load_cora_text() / load_pubmed_raw() for node texts. Dense views fit to roughly 20k nodes on an 8 GB GPU (see benchmarks/pubmed.md); GAT's O(N²·heads) attention caps it at ~8–10k nodes; beyond that, sparse mode (next paragraph).

Large graphs — sparse mode. GNNClassifier(sparse="auto") (the default) switches above 8k nodes to chunked, sparse views: every representative is built in row chunks keeping only its view_top_k strongest entries per row (RPR via a sparse LU solve, gravity via chunked BFS, Katz via chunked powers, semantic via chunked kNN) and stored as a SciPy CSR matrix — O(n·k) memory instead of O(n²) — with sparse propagation on the GPU for GCN / GraphSAGE / GIN. A 20k-node view drops from ~1.5 GB to ~10 MB, and much larger graphs fit (the PubMed table in benchmarks/semantic.md ran this way in a few GB of host RAM). Force it with sparse=True; GAT needs the dense path.

The best representative is dataset-dependent (as the paper notes): RPR is the most robust across architectures on Cora, while Graph Gravity suits friendship/social networks rather than citation graphs. At PubMed scale the purely topological gains shrink and become architecture-specific (GIN benefits most; benchmarks/pubmed.md), whereas the semantic view keeps paying off there (+2.4 to +4.0 points; benchmarks/semantic.md).

Development and tests

git clone https://github.com/AkandaAshraf/gnn-augment && cd gnn-augment
pip install -e ".[dev,semantic]"
pytest                      # unit + integration (synthetic graphs, offline, ~1 min on CPU)
pytest -m unit              # fast, isolated tests of each function / class
pytest -m integration       # end-to-end through the public API, dense and sparse paths
pytest -m "not slow"        # skip the real-data tests (Cora / local MiniLM; auto-skipped when absent)
python -m build && twine check dist/*

Layout: src/gnn_augment/ (classifier, representatives, sparse, models, data), tests/unit/, tests/integration/, examples/, and benchmarks/ (every result table in this README, with the script that produced it). CI (ci.yml) runs the unit and integration suites on Python 3.10 and 3.12 with the CPU build of torch, then builds the sdist/wheel, installs the wheel into a clean venv and runs the integration tests against it; tagged releases are published to PyPI by release.yml (see RELEASING.md).

Relation to other work

The "gravity" representative is the node-topological gravitational score (centrality as mass, path length as distance). It is unrelated to gravity-inspired graph autoencoders, which apply a gravity term to learned embeddings — see the akanda-method README for that distinction.

Citing

Please cite the arXiv paper above. See CITATION.cff.

License

MIT — see LICENSE.