Graph Machine Learning — Features, Node2Vec, GraphSAGE & GAT

Why Graphs Need Special Treatment

Graph data is non-Euclidean: nodes and edges have no natural ordering, and the same graph can be represented by many equivalent adjacency matrices. Standard tabular models treat rows as independent, but in a graph a node's label often depends on its neighbors, its neighbors' neighbors, and the local structure around it. Graph ML builds representations that respect this relational inductive bias.

Graph Basics & Features

A graph G = (V, E) has a set of nodes V and edges E. It can be represented by an adjacency matrix A where A[i, j] = 1 if an edge exists between node i and node j. For undirected graphs A is symmetric.

Node-level structural features

• Degree — number of edges incident to a node. The simplest measure of importance.
• Degree centrality — normalized degree, deg(v) / (|V| - 1).
• Betweenness centrality — how often a node lies on shortest paths between other nodes. High values indicate gatekeepers.
• Clustering coefficient — fraction of a node's neighbors that are connected to each other. Measures local cliqueness.
• Eccentricity / closeness — distance-based measures of how central a node is in the graph.

Graph-level features

• Number of nodes, edges, and density — basic size and sparsity.
• Average degree, degree distribution — captures connectivity patterns.
• Number of triangles — a fundamental motif for social and biological graphs.
• Graph diameter / average shortest path — global reachability.
• Weisfeiler-Lehman features — iterative neighborhood hashing that captures multi-hop structural identity.

import networkx as nx

G = nx.karate_club_graph()

features = {
    "degree": dict(G.degree()),
    "degree_centrality": nx.degree_centrality(G),
    "betweenness": nx.betweenness_centrality(G),
    "clustering": nx.clustering(G),
    "triangles": nx.triangles(G),
    "pagerank": nx.pagerank(G),
}

# Average shortest path for connected graphs
print("Diameter:", nx.diameter(G))

Adjacency matrix pitfalls

• Sparsity — most real-world graphs have very few edges; dense matrices waste memory and compute.
• Permutation variance — reordering nodes changes the matrix but not the graph; models must be invariant to node ordering.
• Feature leakage — using the test graph at training time (e.g., computing PageRank on the whole graph) leaks labels and edges.

Graph Embeddings

Graph embeddings map nodes (or whole graphs) into low-dimensional vectors so that downstream models can operate on them. The goal is to place similar nodes close together in the embedding space, where similarity is defined by structural role, connectivity, or node attributes.

Classical approaches

• Adjacency/Spectral embeddings — use the top-k eigenvectors of the Laplacian as node features. Fast and deterministic, but limited expressiveness.
• Matrix factorization — factorize the adjacency matrix or a higher-order proximity matrix (e.g., GraRep) to learn embeddings.
• Random-walk methods — treat graph walks as sentences and learn node embeddings with Skip-gram (DeepWalk, Node2Vec).

Node2Vec

Node2Vec is a random-walk embedding method. It generates short walks starting from each node, then feeds those walks into a Skip-gram model (Word2Vec) to learn node embeddings. The key novelty is a biased random walk that interpolates between two exploration strategies.

Return and in-out parameters

• Return parameter p — controls the likelihood of returning to the previous node. Low p encourages backtracking, capturing local communities.
• In-out parameter q — controls the likelihood of moving outward from the start node. Low q encourages DFS-like exploration, capturing structural roles.

A walk of length L is generated by starting at node u and repeatedly sampling the next node with probability proportional to 1/p for returning to the previous node, 1 for staying close, and 1/q for moving outward.

Node2Vec in Python

import networkx as nx
from node2vec import Node2Vec

G = nx.karate_club_graph()

# p=1, q=1 recovers DeepWalk-like behavior
n2v = Node2Vec(G, dimensions=64, walk_length=30, num_walks=200,
               p=1.0, q=1.0, workers=4)

model = n2v.fit(window=10, min_count=1, batch_words=4)

# Get embedding for a single node
embedding = model.wv["0"]  # nodes are stringified in node2vec

# Train a classifier on embeddings
import numpy as np
from sklearn.linear_model import LogisticRegression

X = np.array([model.wv[str(node)] for node in G.nodes()])
# y would be the node labels, split with a graph-aware scheme
clf = LogisticRegression(max_iter=1000).fit(X_train, y_train)

When to use Node2Vec

• Static graphs where you can afford to retrain periodically.
• When interpretability is less important than a strong, fast baseline.
• When you need embeddings for unrelated downstream tasks (clustering, similarity search).

Graph Neural Networks

Graph Neural Networks (GNNs) learn representations by aggregating information from a node's neighborhood. A single GNN layer computes a node's new representation as a function of its own features and the aggregated features of its neighbors. Stacking layers expands the receptive field to multi-hop neighborhoods.

Here h_v^{(l)} is the embedding of node v at layer l, N(v) is its neighbor set, AGG is a permutation-invariant aggregation function, and σ is a non-linearity.

Common aggregation choices

• Mean — simple and stable; works well for continuous features and dense neighborhoods.
• Max / Sum — better for sparse or discrete features, but less stable.
• Attention — learn a weight for each neighbor, allowing the model to focus on relevant nodes.

GraphSAGE

GraphSAGE (Sample and Aggregate) introduced inductive learning to GNNs. Instead of learning a fixed embedding for each node, it learns an aggregation function that can be applied to unseen nodes. This makes it possible to train on one graph and apply the model to a new graph with new nodes and edges.

The GraphSAGE forward pass

For each layer, GraphSAGE first samples a fixed number of neighbors (to control computational cost), aggregates their embeddings, and then concatenates them with the node's own embedding before applying a linear transformation.

S(v) is a sampled subset of neighbors. The final node representation is then L2-normalized.

GraphSAGE variants by aggregator

• Mean aggregator — average neighbor embeddings. Equivalent to a simplified GCN.
• LSTM aggregator — apply an LSTM to a random permutation of neighbors. More expressive but not permutation-invariant unless permutations are averaged.
• Pooling aggregator — pass each neighbor through an MLP, then max-pool. Captures non-linear interactions.

GraphSAGE with PyTorch Geometric

import torch
import torch.nn.functional as F
from torch_geometric.nn import SAGEConv
from torch_geometric.datasets import Planetoid

dataset = Planetoid(root="/tmp/Cora", name="Cora")
data = dataset[0]

class GraphSAGE(torch.nn.Module):
    def __init__(self, in_channels, hidden_channels, out_channels):
        super().__init__()
        self.conv1 = SAGEConv(in_channels, hidden_channels, aggr="mean")
        self.conv2 = SAGEConv(hidden_channels, out_channels, aggr="mean")

    def forward(self, x, edge_index):
        x = self.conv1(x, edge_index)
        x = F.relu(x)
        x = F.dropout(x, p=0.5, training=self.training)
        x = self.conv2(x, edge_index)
        return x

model = GraphSAGE(dataset.num_features, 64, dataset.num_classes)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

# Train on the training mask; validate on val; test on test
for epoch in range(200):
    model.train()
    optimizer.zero_grad()
    out = model(data.x, data.edge_index)
    loss = F.cross_entropy(out[data.train_mask], data.y[data.train_mask])
    loss.backward()
    optimizer.step()

model.eval()
with torch.no_grad():
    pred = model(data.x, data.edge_index).argmax(dim=1)
    acc = (pred[data.test_mask] == data.y[data.test_mask]).sum().item() / data.test_mask.sum().item()
    print(f"Test accuracy: {acc:.4f}")

Graph Attention Network (GAT)

GAT replaces the fixed aggregation weights with learned attention coefficients. For each edge, the model computes a score that measures how important the neighbor is to the target node. The neighbor's features are then weighted by these scores and summed.

Attention mechanism

For edge (i, j), GAT first computes an attention score using a learned attention vector a and a weight matrix W:

The scores are normalized across neighbors with a softmax, and the update becomes a weighted sum:

where α_{ij} = softmax_j(e_{ij}).

Multi-head attention

GAT uses multiple independent attention heads to stabilize learning and capture different aspects of the neighborhood. The outputs are concatenated in intermediate layers and averaged in the final layer.

GAT with PyTorch Geometric

from torch_geometric.nn import GATConv

class GAT(torch.nn.Module):
    def __init__(self, in_channels, hidden_channels, out_channels, heads=8):
        super().__init__()
        self.conv1 = GATConv(in_channels, hidden_channels, heads=heads, dropout=0.6)
        # Output layer averages the heads
        self.conv2 = GATConv(hidden_channels * heads, out_channels, heads=1, concat=False, dropout=0.6)

    def forward(self, x, edge_index):
        x = F.dropout(x, p=0.6, training=self.training)
        x = F.elu(self.conv1(x, edge_index))
        x = F.dropout(x, p=0.6, training=self.training)
        x = self.conv2(x, edge_index)
        return x

model = GAT(dataset.num_features, 8, dataset.num_classes, heads=8)
optimizer = torch.optim.Adam(model.parameters(), lr=0.005, weight_decay=5e-4)

# Training loop is identical to GraphSAGE
for epoch in range(200):
    model.train()
    optimizer.zero_grad()
    out = model(data.x, data.edge_index)
    loss = F.cross_entropy(out[data.train_mask], data.y[data.train_mask])
    loss.backward()
    optimizer.step()

Graph Features vs Embeddings vs GNNs

Each approach has a sweet spot. Hand-crafted features are fast and interpretable. Random-walk embeddings give a strong, transferable vector for each node. GNNs learn end-to-end and can incorporate node attributes, but they are more expensive and harder to scale.

• Hand-crafted features — best when data is small, the graph is static, and you need interpretability.
• Node2Vec / DeepWalk — best for static, attribute-free graphs where you want a reusable embedding.
• GraphSAGE — best for large, evolving graphs with node attributes and when you need inductive generalization.
• GAT — best when different neighbors have very different importance and you have enough data to learn attention.

Evaluation & Validation

Graph data violates the i.i.d. assumption. A standard random train-test split can place connected nodes on opposite sides of the split, causing leakage. Use graph-aware splits.

• Transductive setting — train, validation, and test nodes all appear in the same graph. The model uses the full graph structure but only the training labels. Common in semi-supervised node classification.
• Inductive setting — train on one graph, test on a completely different graph. Required for production graphs that grow over time.
• Link prediction split — remove a fraction of edges for validation/test, train on the remaining graph, and evaluate whether the model can recover the removed edges. Negative edges must be sampled from non-edges.

from sklearn.model_selection import train_test_split

# For node classification with a graph-aware split
# Avoid putting strongly connected nodes in different splits
# Use graph-aware samplers such as GraphSAINT or Cluster-GCN when needed

# For link prediction: remove edges, then train on the reduced graph
edges = list(G.edges())
train_edges, test_edges = train_test_split(edges, test_size=0.2, random_state=42)
G_train = G.edge_subgraph(train_edges).copy()

Production Checklist

• Split before featurization — compute graph features only on the training subgraph to avoid leakage.
• Handle dynamic graphs — use inductive models (GraphSAGE, GAT with sampling) and re-embed only changed nodes.
• Scale with sampling — use neighbor sampling or graph clustering for training and inference on large graphs.
• Monitor graph drift — track degree distribution, connected-component sizes, and label distribution over time.
• Explain predictions — use attention weights or gradient-based explainers (GNNExplainer, Integrated Gradients) to show which neighbors mattered.
• Balance class labels — node labels are often imbalanced; use weighted loss or focal loss.

Summary

Graph ML starts with understanding structure: degrees, centralities, and triangles. Node2Vec turns random walks into embeddings but is transductive. GraphSAGE makes graph learning inductive through neighborhood sampling and aggregation. GAT adds attention so nodes decide how much to trust each neighbor. Choose the method based on graph size, dynamics, attributes, and the need for interpretability — and always validate with graph-aware splits.