Graph Neural Networks: обучение на графовых данных

opensourceaimlgnndeep-learninggraphsit
← Back to Blog

Введение: Графы в машинном обучении

Почему графы?

Большинство ML работает с:
  - Таблицы (CSV, SQL)
  - Изображения (сетка пикселей)
  - Последовательности (текст, время)

Реальность - графы:
  - Социальные сети (люди + связи)
  - Молекулы (атомы + связи)
  - Рекомендации (пользователи + товары)
  - Дороги (перекрёстки + дороги)
  - Знания (сущности + отношения)

Факт: 70% данных в enterprise - это графовые данные.


Основы графов

Определение графа

Graph G = (V, E)

V = {v1, v2, ..., vn}  # Вершины (nodes)
E = {(v1,v2), (v2,v3), ...}  # Рёбра (edges)

Пример: социальная сеть
  V = {Alice, Bob, Charlie, Diana}
  E = {(Alice,Bob), (Bob,Charlie), (Charlie,Diana), (Alice,Diana)}

Свойства:
  - Направленный / ненаправленный
  - Взвешенный / невзвешенный
  - Атрибуты вершин и рёбер

Матричное представление

import torch
import torch_geometric

# Adjacency matrix
#     A  B  C  D
A = [[0, 1, 0, 1],  # Alice
     [1, 0, 1, 0],  # Bob
     [0, 1, 0, 1],  # Charlie
     [1, 0, 1, 0]]  # Diana

# Feature matrix
#     feature1  feature2
X = [[1, 0],       # Alice
     [0, 1],       # Bob
     [1, 1],       # Charlie
     [0, 0]]       # Diana

# In PyTorch Geometric:
from torch_geometric.data import Data

edge_index = torch.tensor([
    [0, 1, 1, 2, 2, 3, 3, 0],  # source
    [1, 0, 2, 1, 3, 2, 0, 3]   # target
])

x = torch.tensor([[1., 0.],
                  [0., 1.],
                  [1., 1.],
                  [0., 0.]], dtype=torch.float)

data = Data(x=x, edge_index=edge_index)

Graph Neural Networks

Message Passing

Core GNN operation: Message Passing

For each edge (j, i):
  1. Message: m_ji = φ(h_j, h_i, e_ji)
  2. Aggregate: m_i = Σ_j m_ji
  3. Update: h_i' = ψ(h_i, m_i)

Where:
  h = node features
  e = edge features
  φ, ψ = learnable functions
class MessagePassing(torch.nn.Module):
    def __init__(self, hidden_dim):
        super().__init__()
        self.message_fn = torch.nn.Sequential(
            torch.nn.Linear(hidden_dim * 2 + edge_dim, hidden_dim),
            torch.nn.ReLU(),
        )
        self.update_fn = torch.nn.GRUCell(hidden_dim, hidden_dim)
    
    def forward(self, x, edge_index, edge_attr):
        """
        x: (num_nodes, hidden_dim)
        edge_index: (2, num_edges)
        edge_attr: (num_edges, edge_dim)
        """
        src, dst = edge_index  # source and target nodes
        
        # Create messages for each edge
        x_src = x[src]  # (num_edges, hidden_dim)
        x_dst = x[dst]  # (num_edges, hidden_dim)
        
        # Combine features
        combined = torch.cat([x_src, x_dst, edge_attr], dim=-1)
        messages = self.message_fn(combined)  # (num_edges, hidden_dim)
        
        # Aggregate messages (sum)
        out = torch.zeros_like(x)
        out.index_add_(0, dst, messages)
        
        # Update node states
        h_new = self.update_fn(out, x)
        
        return h_new

Graph Convolution (GCN)

GCN: Spectral graph convolution (approximation)

Update rule:
  H' = σ(D^(-1/2) A D^(-1/2) H W)

Where:
  A = adjacency matrix (with self-loops)
  D = degree matrix
  H = node features
  W = learnable weights
  σ = activation function

Simplified:
  h_i' = σ(Σ_j∈N[i] (1/c_ij) h_j W)
  
  N[i] = neighbors of i (including self)
  c_ij = normalization constant
class GCNLayer(torch.nn.Module):
    def __init__(self, in_features, out_features):
        super().__init__()
        self.linear = torch.nn.Linear(in_features, out_features)
        self.activation = torch.nn.ReLU()
    
    def forward(self, x, edge_index, edge_weight=None):
        """
        x: (num_nodes, in_features)
        edge_index: (2, num_edges)
        edge_weight: (num_edges,) optional
        """
        # Prepare adjacency
        row, col = edge_index
        num_nodes = x.shape[0]
        
        # Add self-loops
        self_loop = torch.arange(num_nodes).to(edge_index)
        edge_index_full = torch.cat([
            torch.stack([self_loop, self_loop], dim=0),
            edge_index
        ], dim=1)
        
        # Compute normalization
        deg = degree(edge_index_full[1], num_nodes)
        deg_inv_sqrt = deg.pow(-0.5)
        deg_inv_sqrt[deg_inv_sqrt == float('inf')] = 0
        
        # Normalized adjacency
        for i in range(edge_index_full.shape[1]):
            edge_weight_norm[i] = deg_inv_sqrt[edge_index_full[0, i]] * \
                                  deg_inv_sqrt[edge_index_full[1, i]]
        
        # Message passing
        out = self.linear(x)
        
        # Aggregate
        row, col = edge_index_full
        out_col = out[col]
        out_col = out_col * edge_weight_norm.unsqueeze(1)
        
        agg = torch.zeros_like(out)
        agg.index_add_(0, row, out_col)
        
        return self.activation(agg)

class GCN(torch.nn.Module):
    def __init__(self, in_channels, hidden_channels, out_channels):
        super().__init__()
        self.conv1 = GCNLayer(in_channels, hidden_channels)
        self.conv2 = GCNLayer(hidden_channels, out_channels)
    
    def forward(self, x, edge_index):
        x = self.conv1(x, edge_index)
        x = torch.nn.functional.relu(x)
        x = torch.nn.functional.dropout(x, training=self.training)
        x = self.conv2(x, edge_index)
        return x

GraphSAGE

GraphSAGE: Inductive representation learning

Key idea: Learn aggregation functions
  instead of fixed graph Laplacian

Aggregators:
  - Mean (original)
  - LSTM
  - Pooling

Works on unseen nodes!
class SAGELayer(torch.nn.Module):
    def __init__(self, in_features, out_features, aggregator='mean'):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.aggregator = aggregator
        
        # Self features transformation
        self.self_linear = torch.nn.Linear(in_features, out_features)
        
        # Aggregation
        if aggregator == 'mean':
            self.aggregator_linear = torch.nn.Linear(in_features, out_features)
        elif aggregator == 'pool':
            self.pool = SAGEPool(in_features, out_features)
        elif aggregator == 'lstm':
            self.lstm = nn.LSTM(in_features, out_features, batch_first=True)
        
        self.activation = torch.nn.ReLU()
    
    def forward(self, x, edge_index):
        row, col = edge_index
        num_nodes = x.shape[0]
        
        # Transform self
        h_self = self.self_linear(x)
        
        # Aggregate neighbors
        h_neighbors = self.aggregate(x, row, col)
        
        # Combine
        out = torch.cat([h_self, h_neighbors], dim=-1)
        out = self.activation(out)
        
        # L2 normalize (like original paper)
        out = torch.nn.functional.normalize(out, p=2, dim=1)
        
        return out
    
    def aggregate(self, x, row, col):
        if self.aggregator == 'mean':
            # Mean aggregator
            neighbor_features = x[col]
            deg = torch.bincount(row, minlength=x.shape[0])
            deg_inv = 1.0 / deg.clamp(min=1)
            
            out = torch.zeros_like(x)
            out.index_add_(0, row, neighbor_features * deg_inv[row].unsqueeze(1))
            return self.aggregator_linear(out)
        
        elif self.aggregator == 'pool':
            return self.pool(x, row, col)

GAT (Graph Attention)

GAT: Self-attention on graphs

For each node i:
  h_i' = Σ_j α_ij Wh_j
  
  α_ij = softmax_i(e_ij)
  e_ij = LeakyReLU(a^T [Wh_i || Wh_j])
  
Where:
  || = concatenation
  a = learnable attention vector
class GATLayer(torch.nn.Module):
    def __init__(self, in_features, out_features, n_heads=1, dropout=0.6):
        super().__init__()
        self.n_heads = n_heads
        self.head_dim = out_features // n_heads
        
        # Linear projection
        self.W = torch.nn.Linear(in_features, out_features)
        
        # Attention mechanism
        self.a = torch.nn.Parameter(
            torch.randn(2 * self.head_dim, n_heads)
        )
        
        self.leakyrelu = torch.nn.LeakyReLU(0.2)
        self.softmax = torch.nn.Softmax(dim=-1)
        self.dropout = torch.nn.Dropout(dropout)
    
    def forward(self, x, edge_index):
        """
        x: (num_nodes, in_features)
        edge_index: (2, num_edges)
        """
        num_nodes = x.shape[0]
        
        # Linear transform
        H = self.W(x)  # (num_nodes, out_features)
        H = H.view(num_nodes, self.n_heads, self.head_dim)
        
        # Compute attention coefficients
        row, col = edge_index
        
        # Attention for source and target
        a_src = (H[row] * self.a[:self.head_dim]).sum(dim=-1)  # (num_edges, n_heads)
        a_dst = (H[col] * self.a[self.head_dim:]).sum(dim=-1)  # (num_edges, n_heads)
        
        e = self.leakyrelu(a_src + a_dst.T)  # (num_edges, n_heads)
        
        # Sparse attention
        alpha = torch.zeros(num_nodes, num_nodes, self.n_heads, device=x.device)
        alpha[row, col] = self.dropout(self.softmax(e))
        
        # Aggregate
        out = torch.zeros(num_nodes, self.n_heads * self.head_dim, device=x.device)
        
        for h in range(self.n_heads):
            H_h = H[:, h, :]  # (num_nodes, head_dim)
            alpha_h = alpha[:, :, h]  # (num_nodes, num_nodes)
            out[:, h*self.head_dim:(h+1)*self.head_dim] = (alpha_h @ H_h)
        
        return out

GNN Training

Node Classification

class GNNClassifier(torch.nn.Module):
    def __init__(self, in_channels, hidden_channels, num_classes):
        super().__init__()
        self.encoder = GCN(in_channels, hidden_channels, hidden_channels)
        self.classifier = torch.nn.Linear(hidden_channels, num_classes)
    
    def forward(self, x, edge_index):
        embeddings = self.encoder(x, edge_index)
        logits = self.classifier(embeddings)
        return logits

# Training loop
model = GNNClassifier(in_channels=128, hidden_channels=256, num_classes=7)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
criterion = torch.nn.CrossEntropyLoss()

for epoch in range(200):
    model.train()
    logits = model(data.x, data.edge_index)
    loss = criterion(logits[train_mask], data.y[train_mask])
    
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    
    # Evaluate
    model.eval()
    with torch.no_grad():
        logits = model(data.x, data.edge_index)
        pred = logits.argmax(dim=1)
        accuracy = (pred[test_mask] == data.y[test_mask]).float().mean()

Graph Classification

class GraphClassifier(torch.nn.Module):
    def __init__(self, in_channels, hidden_channels, num_classes):
        super().__init__()
        self.gnn = GCN(in_channels, hidden_channels, hidden_channels)
        
        # Readout functions
        self.readout = 'mean'  # mean, sum, max, attention
        
        self.classifier = torch.nn.Sequential(
            torch.nn.Linear(hidden_channels, hidden_channels),
            torch.nn.ReLU(),
            torch.nn.Dropout(0.5),
            torch.nn.Linear(hidden_channels, num_classes)
        )
    
    def forward(self, batch_data):
        """
        batch_data: batch of graphs
        """
        # GNN encoding
        x = torch.cat([data.x for data in batch_data])
        edge_index = torch.cat([data.edge_index for data in batch_data], dim=1)
        batch = torch.cat([torch.full((data.x.shape[0],), i) 
                          for i, data in enumerate(batch_data)])
        
        h = self.gnn(x, edge_index)
        
        # Readout (pooling)
        if self.readout == 'mean':
            graph_repr = scatter(h, batch, dim=0, reduce='mean')
        elif self.readout == 'sum':
            graph_repr = scatter(h, batch, dim=0, reduce='sum')
        elif self.readout == 'max':
            graph_repr = scatter(h, batch, dim=0, reduce='max')
        
        return self.classifier(graph_repr)

Applications

Molecular Property Prediction

Molecule = Graph
  Nodes = Atoms (element, charge, hybridization)
  Edges = Bonds (type, length, angle)

Tasks:
  - Drug discovery (toxicity, binding affinity)
  - Material science (conductivity, stability)
  - Chemistry (reaction prediction)
class MoleculeGNN(torch.nn.Module):
    def __init__(self, num_atom_types, num_bond_types, hidden_dim):
        super().__init__()
        
        # Atom embeddings
        self.atom_embedding = torch.nn.Embedding(num_atom_types, hidden_dim)
        
        # Bond embeddings
        self.bond_embedding = torch.nn.Embedding(num_bond_types, hidden_dim)
        
        # GNN layers
        self.gnn_layers = torch.nn.ModuleList([
            GATLayer(hidden_dim, hidden_dim)
            for _ in range(3)
        ])
        
        # Property prediction heads
        self.property_heads = torch.nn.ModuleDict({
            'solubility': torch.nn.Linear(hidden_dim, 1),
            'toxicity': torch.nn.Linear(hidden_dim, 1),
            'binding': torch.nn.Linear(hidden_dim, 1),
        })
    
    def forward(self, molecule_data):
        # Node features
        atom_types = molecule_data.x.long()
        atom_features = self.atom_embedding(atom_types)
        
        # Edge features
        bond_types = molecule_data.edge_attr.long()
        bond_features = self.bond_embedding(bond_types)
        
        # Message passing
        h = atom_features
        for gnn_layer in self.gnn_layers:
            h = gnn_layer(h, molecule_data.edge_index)
            h = torch.nn.functional.relu(h)
        
        # Graph-level readout
        batch = molecule_data.batch
        graph_repr = scatter(h, batch, dim=0, reduce='mean')
        
        # Predict properties
        properties = {
            name: head(graph_repr)
            for name, head in self.property_heads.items()
        }
        
        return properties

Recommendation Systems

User-Item Graph:
  Users: age, gender, location
  Items: category, price, features
  Edges: purchases, clicks, ratings

GNN approach:
  1. Encode user and item features
  2. Message passing on bipartite graph
  3. Predict edge probability (will user click item?)

GNN Frameworks

# PyTorch Geometric
import torch_geometric.nn as pyg_nn

# Common layers
gcnn = pyg_nn.GCNConv(in_channels, out_channels)
sage = pyg_nn.SAGEConv(in_channels, out_channels)
gat = pyg_nn.GATv2Conv(in_channels, out_channels)
gin = pyg_nn.GINConv(torch.nn.Sequential(
    torch.nn.Linear(in_channels, out_channels),
    torch.nn.ReLU(),
    torch.nn.Linear(out_channels, out_channels)
))

# Pooling layers
mean_pool = pyg_nn.global_mean_pool
sum_pool = pyg_nn.global_add_pool
max_pool = pyg_nn.global_max_pool
attention_pool = pyg_nn.GlobalAttention(torch.nn.LSTM(1, 1))

# Datasets
from torch_geometric.datasets import ZINC, OGBNaMolecules, Planetoid
from torch_geometric.loader import DataLoader

dataset = Planetoid(root='data/Cora', name='Cora')
loader = DataLoader(dataset, batch_size=32, shuffle=True)

Заключение

GNNs provide:

  • Relational inductive bias: Learn from structure
  • Message passing: Information flows through graph
  • Flexibility: Works on any graph-structured data

Key architectures:

  1. GCN - spectral convolution
  2. GraphSAGE - inductive aggregation
  3. GAT - attention-based
  4. GIN - theoretically powerful

Applications:

  • Drug discovery
  • Recommendation systems
  • Social network analysis
  • Knowledge graphs
  • Materials science

Ресурсы