Fundamentos y Motivación

La evaluación de la similitud entre estructuras de grafos es una tarea fundamental en áreas como la quimioinformática, el análisis de redes sociales y la visión por computadora. Métricas tradicionales como la Distancia de Edición de Grafos (GED) o el Subrgafo Común Máximo (MCS) proporcionan medidas rigurosas, pero su complejidad computacional es NP-completa, lo que las hace inviables para conjuntos de datos a gran escala.

Para abordar este cuello de botella, la arquitectura SimGNN (propuesta en WSDM 2018) introduce un enfoque basado en Redes Neuronales de Grafos (GNN). El objetivo principal es aprender una función de mapeo que proyecte un par de grafos directamente a un valor de similitud escalar, reduciendo drásticamente el tiempo de infreencia.

Estrategias de Diseño del Modelo

La arquitectura de SimGNN integra dos mecanismos complementarios para capturar tento la estructura global como las diferencias locales:

• Resumen Global con Atención: Se utiliza una función de incrustación (embedding) aprendible para mapear cada grafo a un vector de representación global. Se incorpora un mecanismo de atención que pondera la importancia de los nodos específicos según la métrica de similitud objetivo.
• Comparación Detallada de Nodos: Para evitar la pérdida de información de alto nivel que ocurre al agrupar (pooling) los nodos, se diseña un módulo de comparación por pares. Este módulo calcula las similitudes entre todos los nodos de ambos grafos y genera un histograma de frecuencias que captura la distribución de las similitudes locales.

Propiedades Clave

El modelo está diseñado para cumplir con tres propiedades esenciales:

1. Invariancia a permutaciones: El orden de los nodos en la matriz de adyacencia no debe alterar la puntuación de similitud resultante.
2. Capacidad inductiva: El modelo debe generalizar y calcular similitudes para pares de grafos no vistos durante la fase de entrenamiento.
3. Adaptabilidad: Los parámetros de la red se optimizan para ajustarse a cualquier métrica de similitud subyacente proporcionada en las etiquetas de verdad fundamental (ground truth).

Implementación en PyTorch

A continuación, se presenta la reestructuración del código fuente para la arquitectura completa, incluyendo el módulo principal, el pooling por atención y la red de tensores para la interacción de características.

import torch
import torch.nn.functional as F
from torch_geometric.nn import GCNConv, GINConv
from torch_geometric.utils import to_dense_batch
from torch_scatter import scatter_mean, scatter_add

class FastGraphSimilarityNetwork(torch.nn.Module):
    """
    Red neuronal para la estimación rápida de similitud entre grafos.
    Combina representaciones globales mediante atención y distribuciones 
    de similitud nodal a través de histogramas.
    """

    def __init__(self, config, num_node_features):
        super(FastGraphSimilarityNetwork, self).__init__()
        self.config = config
        self.num_node_features = num_node_features
        self._determine_bottleneck_dim()
        self._initialize_architecture()

    def _determine_bottleneck_dim(self):
        self.bottleneck_dim = self.config.tensor_neurons
        if self.config.use_histogram:
            self.bottleneck_dim += self.config.histogram_bins

    def _initialize_architecture(self):
        # Capas de convolución de grafos (GCN o GIN)
        if self.config.gnn_type == "gcn":
            self.gnn_layer_1 = GCNConv(self.num_node_features, self.config.filters_1)
            self.gnn_layer_2 = GCNConv(self.config.filters_1, self.config.filters_2)
            self.gnn_layer_3 = GCNConv(self.config.filters_2, self.config.filters_3)
        elif self.config.gnn_type == "gin":
            mlp_1 = torch.nn.Sequential(torch.nn.Linear(self.num_node_features, self.config.filters_1), torch.nn.ReLU(), torch.nn.BatchNorm1d(self.config.filters_1))
            mlp_2 = torch.nn.Sequential(torch.nn.Linear(self.config.filters_1, self.config.filters_2), torch.nn.ReLU(), torch.nn.BatchNorm1d(self.config.filters_2))
            mlp_3 = torch.nn.Sequential(torch.nn.Linear(self.config.filters_2, self.config.filters_3), torch.nn.ReLU(), torch.nn.BatchNorm1d(self.config.filters_3))
            self.gnn_layer_1 = GINConv(mlp_1, train_eps=True)
            self.gnn_layer_2 = GINConv(mlp_2, train_eps=True)
            self.gnn_layer_3 = GINConv(mlp_3, train_eps=True)
        else:
            raise ValueError("Tipo de GNN no soportado.")

        # Módulos de pooling y comparación
        self.node_attention_pooler = NodeAttentionPooling(self.config)
        self.tensor_interaction = TensorNetworkInteraction(self.config)
        
        # Capas fully connected finales
        self.fc_bottleneck = torch.nn.Linear(self.bottleneck_dim, self.config.bottle_neck_neurons)
        self.fc_score = torch.nn.Linear(self.config.bottle_neck_neurons, 1)

    def generate_similarity_histogram(self, node_feats_g1, node_feats_g2, batch_idx_g1, batch_idx_g2):
        """
        Calcula el histograma de las similitudes por pares entre nodos de ambos grafos.
        """
        dense_feats_1, valid_mask_1 = to_dense_batch(node_feats_g1, batch_idx_g1)
        dense_feats_2, valid_mask_2 = to_dense_batch(node_feats_g2, batch_idx_g2)

        max_nodes_g1 = valid_mask_1.sum(dim=1)
        max_nodes_g2 = valid_mask_2.sum(dim=1)
        max_nodes = torch.max(max_nodes_g1, max_nodes_g2)

        # Matriz de similitud por pares (desacoplada del grafo computacional)
        pairwise_sim = torch.matmul(dense_feats_1, dense_feats_2.transpose(1, 2)).detach()

        histograms = []
        for i in range(pairwise_sim.size(0)):
            valid_sim = pairwise_sim[i, :max_nodes[i], :max_nodes[i]]
            flattened_sim = torch.sigmoid(valid_sim).view(-1)
            
            hist = torch.histc(flattened_sim, bins=self.config.histogram_bins)
            normalized_hist = hist / hist.sum()
            histograms.append(normalized_hist.view(1, -1))

        return torch.cat(histograms, dim=0)

    def extract_node_embeddings(self, edge_index, node_features):
        """
        Propagación a través de las capas convolucionales del grafo.
        """
        x = self.gnn_layer_1(node_features, edge_index)
        x = F.dropout(F.relu(x), p=self.config.dropout, training=self.training)
        
        x = self.gnn_layer_2(x, edge_index)
        x = F.dropout(F.relu(x), p=self.config.dropout, training=self.training)
        
        return self.gnn_layer_3(x, edge_index)

    def forward(self, graph_pair_data):
        edge_idx_1, feats_1, batch_1 = graph_pair_data["g1"].edge_index, graph_pair_data["g1"].x, graph_pair_data["g1"].batch
        edge_idx_2, feats_2, batch_2 = graph_pair_data["g2"].edge_index, graph_pair_data["g2"].x, graph_pair_data["g2"].batch

        # 1. Extracción de características nodales
        node_emb_1 = self.extract_node_embeddings(edge_idx_1, feats_1)
        node_emb_2 = self.extract_node_embeddings(edge_idx_2, feats_2)

        # 2. Cálculo de histograma de similitud nodal
        if self.config.use_histogram:
            hist_features = self.generate_similarity_histogram(node_emb_1, node_emb_2, batch_1, batch_2)

        # 3. Pooling a nivel de grafo mediante atención
        graph_repr_1 = self.node_attention_pooler(node_emb_1, batch_1)
        graph_repr_2 = self.node_attention_pooler(node_emb_2, batch_2)

        # 4. Interacción de tensores para obtener vector de similitud
        interaction_scores = self.tensor_interaction(graph_repr_1, graph_repr_2)

        # 5. Concatenación y predicción final
        if self.config.use_histogram:
            interaction_scores = torch.cat([interaction_scores, hist_features], dim=1)

        hidden = F.relu(self.fc_bottleneck(interaction_scores))
        final_score = torch.sigmoid(self.fc_score(hidden)).view(-1)
        
        return final_score

class NodeAttentionPooling(torch.nn.Module):
    """
    Módulo de atención para generar una representación a nivel de grafo,
    ponderando la importancia de cada nodo basándose en el contexto global.
    """

    def __init__(self, config):
        super(NodeAttentionPooling, self).__init__()
        self.attention_weights = torch.nn.Parameter(torch.Tensor(config.filters_3, config.filters_3))
        torch.nn.init.xavier_uniform_(self.attention_weights)

    def forward(self, node_embeddings, graph_indices, target_size=None):
        num_graphs = graph_indices[-1].item() + 1 if target_size is None else target_size

        # Calcular la representación media global para cada grafo en el batch
        global_mean = scatter_mean(node_embeddings, graph_indices, dim=0, dim_size=num_graphs)

        # Transformar el contexto global mediante una proyección lineal y tangente hiperbólica
        context_vector = torch.tanh(torch.matmul(global_mean, self.attention_weights))

        # Expandir el contexto para que coincida con la dimensionalidad de los nodos
        expanded_context = context_vector[graph_indices]
        
        # Calcular coeficientes de atención (producto punto seguido de sigmoide)
        attention_scores = torch.sigmoid((node_embeddings * expanded_context).sum(dim=1))

        # Ponderar las incrustaciones de los nodos y agregarlas de vuelta al nivel de grafo
        weighted_nodes = attention_scores.unsqueeze(-1) * node_embeddings
        return scatter_add(weighted_nodes, graph_indices, dim=0, dim_size=num_graphs)

class TensorNetworkInteraction(torch.nn.Module):
    """
    Capa de red de tensores que modela interacciones de alto orden 
    entre las representaciones globales de los dos grafos.
    """

    def __init__(self, config):
        super(TensorNetworkInteraction, self).__init__()
        self.config = config
        
        # Tensor de interacción 3D y matriz de combinación lineal
        self.interaction_tensor = torch.nn.Parameter(torch.Tensor(config.filters_3, config.filters_3, config.tensor_neurons))
        self.weight_matrix_block = torch.nn.Parameter(torch.Tensor(config.tensor_neurons, 2 * config.filters_3))
        self.bias = torch.nn.Parameter(torch.Tensor(config.tensor_neurons, 1))
        
        torch.nn.init.xavier_uniform_(self.interaction_tensor)
        torch.nn.init.xavier_uniform_(self.weight_matrix_block)
        torch.nn.init.xavier_uniform_(self.bias)

    def forward(self, graph_repr_1, graph_repr_2):
        batch_sz = graph_repr_1.size(0)
        feat_dim = self.config.filters_3
        tensor_dim = self.config.tensor_neurons

        # Producto tensorial: proyectar g1 y multiplicar por g2
        reshaped_tensor = self.interaction_tensor.view(feat_dim, -1)
        projected_g1 = torch.matmul(graph_repr_1, reshaped_tensor)
        projected_g1 = projected_g1.view(batch_sz, tensor_dim, feat_dim)
        
        tensor_scores = torch.matmul(projected_g1, graph_repr_2.unsqueeze(-1)).squeeze(-1)

        # Combinación lineal de las representaciones concatenadas
        combined_repr = torch.cat([graph_repr_1, graph_repr_2], dim=1)
        linear_scores = torch.matmul(combined_repr, self.weight_matrix_block.transpose(0, 1))

        # Fusión no lineal de ambas puntuaciones
        final_scores = F.relu(tensor_scores + linear_scores + self.bias.squeeze())
        
        return final_scores