Node Classification

Node classification predicts labels for entities in a graph by leveraging their connections. Unlike tabular ML where each sample is independent, graph ML exploits relational structure - a user's friends, a paper's citations, a molecule's bonds. When relationships matter, graphs often outperform feature-only approaches.

Where It Is Used

Node classification powers network-based applications:

• Social networks: User profiling, bot detection, community detection
• E-commerce: Product categorization, fraud ring detection
• Biology: Protein function prediction, drug-target interaction
• Academic: Paper classification, author expertise
• Security: Malware detection, network intrusion
• Finance: Account classification, money laundering detection

Products: Pinterest PinSage, LinkedIn graph ML, Google Knowledge Graph.

Typical Data

What You Do in Practice

1. Construct the graph - Define nodes, edges, features
2. Analyze graph structure - Density, degree distribution, components
3. Handle label imbalance - Often severe in node classification
4. Choose transductive vs inductive - Fixed graph or new nodes?
5. Select GNN architecture - GCN, GAT, GraphSAGE
6. Prevent label leakage - Careful with message passing
7. Evaluate properly - Train/val/test node splits
8. Scale for production - Mini-batching, sampling

Baseline First

Start with simple approaches before GNNs:

import numpy as np
from sklearn.linear_model import LogisticRegression
from sklearn.preprocessing import StandardScaler
import networkx as nx

class NodeClassificationBaseline:
    """Simple baselines for node classification."""

    def __init__(self, graph, features, labels):
        self.graph = graph
        self.features = features
        self.labels = labels

    def feature_only_baseline(self, train_mask, test_mask):
        """Classify using node features only (no graph)."""
        scaler = StandardScaler()
        X_train = scaler.fit_transform(self.features[train_mask])
        X_test = scaler.transform(self.features[test_mask])

        model = LogisticRegression(max_iter=1000)
        model.fit(X_train, self.labels[train_mask])

        predictions = model.predict(X_test)
        return predictions

    def label_propagation(self, train_mask, max_iter=100, alpha=0.8):
        """Semi-supervised label propagation."""
        n_nodes = len(self.labels)
        n_classes = len(np.unique(self.labels[self.labels >= 0]))

        # Initialize label matrix
        Y = np.zeros((n_nodes, n_classes))
        for i in np.where(train_mask)[0]:
            Y[i, self.labels[i]] = 1

        # Normalized adjacency
        A = nx.adjacency_matrix(self.graph).toarray()
        D = np.diag(A.sum(axis=1) + 1e-10)
        D_inv_sqrt = np.diag(1.0 / np.sqrt(np.diag(D)))
        A_norm = D_inv_sqrt @ A @ D_inv_sqrt

        # Propagate
        Y_init = Y.copy()
        for _ in range(max_iter):
            Y = alpha * A_norm @ Y + (1 - alpha) * Y_init

        predictions = Y.argmax(axis=1)
        return predictions

    def node2vec_baseline(self, train_mask, test_mask, dimensions=64):
        """Node2Vec embeddings + classifier."""
        from node2vec import Node2Vec

        # Generate embeddings
        node2vec = Node2Vec(self.graph, dimensions=dimensions,
                            walk_length=30, num_walks=200, workers=4)
        model = node2vec.fit(window=10, min_count=1)

        # Get embeddings for all nodes
        embeddings = np.array([model.wv[str(i)] 
                              for i in range(len(self.labels))])

        # Combine with features
        X = np.hstack([self.features, embeddings])

        # Train classifier
        scaler = StandardScaler()
        X_train = scaler.fit_transform(X[train_mask])
        X_test = scaler.transform(X[test_mask])

        clf = LogisticRegression(max_iter=1000)
        clf.fit(X_train, self.labels[train_mask])

        return clf.predict(X_test)

Why baselines matter: - Feature-only shows value of graph structure - Label propagation is strong semi-supervised baseline - Node2Vec captures structural information - GNNs should beat these meaningfully

GNN Implementation

Graph Neural Networks for node classification:

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch_geometric.nn import GCNConv, GATConv, SAGEConv

class GCN(nn.Module):
    """Graph Convolutional Network."""

    def __init__(self, in_channels, hidden_channels, out_channels, 
                 num_layers=2, dropout=0.5):
        super().__init__()
        self.convs = nn.ModuleList()
        self.convs.append(GCNConv(in_channels, hidden_channels))

        for _ in range(num_layers - 2):
            self.convs.append(GCNConv(hidden_channels, hidden_channels))

        self.convs.append(GCNConv(hidden_channels, out_channels))
        self.dropout = dropout

    def forward(self, x, edge_index):
        for i, conv in enumerate(self.convs[:-1]):
            x = conv(x, edge_index)
            x = F.relu(x)
            x = F.dropout(x, p=self.dropout, training=self.training)

        x = self.convs-1
        return x

class GAT(nn.Module):
    """Graph Attention Network."""

    def __init__(self, in_channels, hidden_channels, out_channels,
                 heads=8, dropout=0.6):
        super().__init__()
        self.conv1 = GATConv(in_channels, hidden_channels, heads=heads,
                             dropout=dropout)
        self.conv2 = GATConv(hidden_channels * heads, out_channels, 
                             heads=1, concat=False, dropout=dropout)
        self.dropout = dropout

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

class GraphSAGE(nn.Module):
    """GraphSAGE for inductive learning."""

    def __init__(self, in_channels, hidden_channels, out_channels,
                 num_layers=2, dropout=0.5):
        super().__init__()
        self.convs = nn.ModuleList()
        self.convs.append(SAGEConv(in_channels, hidden_channels))

        for _ in range(num_layers - 2):
            self.convs.append(SAGEConv(hidden_channels, hidden_channels))

        self.convs.append(SAGEConv(hidden_channels, out_channels))
        self.dropout = dropout

    def forward(self, x, edge_index):
        for i, conv in enumerate(self.convs[:-1]):
            x = conv(x, edge_index)
            x = F.relu(x)
            x = F.dropout(x, p=self.dropout, training=self.training)

        x = self.convs-1
        return x

def train_gnn(model, data, optimizer, epochs=200):
    """Training loop for node classification."""
    model.train()

    for epoch in range(epochs):
        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()

        # Validation
        if epoch % 10 == 0:
            model.eval()
            with torch.no_grad():
                pred = model(data.x, data.edge_index).argmax(dim=1)
                val_acc = (pred[data.val_mask] == data.y[data.val_mask]).float().mean()
            model.train()
            print(f'Epoch {epoch}: Loss {loss:.4f}, Val Acc {val_acc:.4f}')

    return model

Handling Large Graphs

Mini-batching for scalability:

from torch_geometric.loader import NeighborLoader

class ScalableNodeClassification:
    """Handle large graphs with sampling."""

    def __init__(self, data, num_neighbors=[25, 10], batch_size=1024):
        self.data = data
        self.num_neighbors = num_neighbors
        self.batch_size = batch_size

    def get_loader(self, mask, shuffle=True):
        """Create neighbor sampling loader."""
        return NeighborLoader(
            self.data,
            num_neighbors=self.num_neighbors,
            batch_size=self.batch_size,
            input_nodes=mask,
            shuffle=shuffle
        )

    def train_epoch(self, model, loader, optimizer):
        """Train one epoch with mini-batching."""
        model.train()
        total_loss = 0

        for batch in loader:
            optimizer.zero_grad()

            out = model(batch.x, batch.edge_index)
            # Only compute loss on target nodes (not sampled neighbors)
            loss = F.cross_entropy(out[:batch.batch_size], 
                                   batch.y[:batch.batch_size])

            loss.backward()
            optimizer.step()

            total_loss += loss.item() * batch.batch_size

        return total_loss / self.data.train_mask.sum()

    @torch.no_grad()
    def inference(self, model, loader):
        """Full-batch inference with subgraph sampling."""
        model.eval()
        predictions = []

        for batch in loader:
            out = model(batch.x, batch.edge_index)
            pred = out[:batch.batch_size].argmax(dim=1)
            predictions.append(pred)

        return torch.cat(predictions)

Transductive vs Inductive

class TransductiveVsInductive:
    """Handle both settings properly."""

    def transductive_split(self, data, train_ratio=0.6, val_ratio=0.2):
        """
        Transductive: All nodes in graph during training,
        but only some have labels.
        """
        n = data.num_nodes
        perm = torch.randperm(n)

        train_size = int(n * train_ratio)
        val_size = int(n * val_ratio)

        train_mask = torch.zeros(n, dtype=torch.bool)
        val_mask = torch.zeros(n, dtype=torch.bool)
        test_mask = torch.zeros(n, dtype=torch.bool)

        train_mask[perm[:train_size]] = True
        val_mask[perm[train_size:train_size+val_size]] = True
        test_mask[perm[train_size+val_size:]] = True

        return train_mask, val_mask, test_mask

    def inductive_split(self, data, train_ratio=0.6, val_ratio=0.2):
        """
        Inductive: Test nodes completely unseen during training.
        Must remove test nodes from training graph.
        """
        n = data.num_nodes
        perm = torch.randperm(n)

        train_size = int(n * train_ratio)
        val_size = int(n * val_ratio)

        train_nodes = perm[:train_size]
        val_nodes = perm[train_size:train_size+val_size]
        test_nodes = perm[train_size+val_size:]

        # Create training subgraph (no test nodes)
        train_val_nodes = torch.cat([train_nodes, val_nodes])
        train_graph = self.subgraph(data, train_val_nodes)

        return {
            'train_graph': train_graph,
            'train_nodes': train_nodes,
            'val_nodes': val_nodes,
            'test_nodes': test_nodes,
            'full_graph': data
        }

    def subgraph(self, data, node_idx):
        """Extract subgraph with selected nodes."""
        from torch_geometric.utils import subgraph

        edge_index, _ = subgraph(node_idx, data.edge_index, 
                                  relabel_nodes=True)

        return type(data)(
            x=data.x[node_idx],
            edge_index=edge_index,
            y=data.y[node_idx]
        )

Evaluation Metrics

Failure Modes and Debugging Checklist

• [ ] Over-smoothing: Too many layers, node representations converge
• [ ] Label leakage: Test labels visible through message passing
• [ ] Disconnected components: Some nodes isolated from training signal
• [ ] Class imbalance: Minority classes underrepresented
• [ ] Feature scaling: Node features not normalized
• [ ] Memory issues: Graph too large for GPU
• [ ] Heterophily: Assumption that neighbors are similar doesn't hold

Production Checklist

Before deployment: - [ ] Graph construction pipeline validated - [ ] Inductive learning if new nodes expected - [ ] Mini-batching for large graphs - [ ] Feature preprocessing standardized - [ ] Baseline comparison documented

After deployment: - [ ] New node handling tested - [ ] Graph update frequency defined - [ ] Prediction latency monitored - [ ] Edge case handling (isolated nodes) - [ ] Model retraining schedule

Tools and Stack

Example Mini-Case

Scenario: Bot detection in social network.

Graph construction: - Nodes: User accounts - Edges: Follows, mentions, retweets - Features: Account age, post frequency, text embeddings - Labels: Bot/human (small labeled set)

Approach: 1. Built follow graph with 10M nodes 2. Semi-supervised with 5% labeled 3. GraphSAGE for inductive learning (new accounts) 4. Neighbor sampling for scalability

Results: - 94% F1 on bot detection (vs 78% feature-only) - Catches coordinated bot networks - 50ms inference per new account - Retrains weekly on new labels

Related Pages

• Index: ML Field Guide
• Knowledge Graphs
• Fraud Detection
• Semantic Search

Further Reading

• Graph Representation Learning Book
• PyTorch Geometric Documentation
• Stanford CS224W: ML with Graphs

Try it live

Everything above runs in your browser — open Graph Attention Network: Node Classification and change the parameters while it is running. Nothing is installed, nothing is uploaded, the whole model lives in one tab.