Predicción de vínculos en gráficos con redes neuronales gráficas y DGL.ai

Los gráficos son una forma eficaz de representar datos, ya que capturan de forma única las relaciones entre entidades en una amplia variedad de aplicaciones. Ya sea que esté modelando redes sociales, interacciones de proteínas, sistemas de transporte o motores de recomendación, los gráficos representan y analizan de forma natural estas interdependencias complejas. En el mundo actual, impulsado por los datos, comprender las relaciones entre entidades suele ser tan importante como comprender las entidades en sí mismas: aquí es donde los gráficos realmente brillan.

La predicción de enlaces es una de las tareas fundamentales en el análisis de grafos, que implica la predicción de conexiones (o enlaces) entre nodos (las entidades representadas en un grafo). Imagine recomendar nuevos amigos en una red social, predecir posibles colaboraciones en un gráfico de citas académicas o pronosticar futuras interacciones entre usuarios y productos en un entorno de comercio electrónico: todos estos son ejemplos de predicción de enlaces en acción. Esta tarea ayuda a expandir redes, inferir información faltante y detectar anomalías. Para aplicaciones que van desde mejorar la experiencia del usuario hasta mejorar la detección de fraudes, la predicción de enlaces es un ingrediente clave para el éxito.

Para ilustrar la predicción de enlaces, utilizaremos el conjunto de datos de la red social Twitch del Stanford Network Analysis Project (SNAP). Este conjunto de datos captura las conexiones sociales entre los usuarios de la plataforma de streaming Twitch, donde los nodos representan a los usuarios de Twitch y los bordes representan las amistades entre ellos. El conjunto de datos está bien estructurado, lo que facilita su preprocesamiento y el trabajo con él.

Si sigue este tutorial, aprenderá a configurar un proyecto, preprocesar datos, crear un modelo y evaluarlo para la predicción de enlaces en un conjunto de datos del mundo real.

Redes neuronales gráficas y biblioteca de gráficos profundos

Trabajar con datos estructurados en gráficos presenta desafíos únicos, y aquí es donde entran en juego las redes neuronales de gráficos (GNN) . Las GNN son un tipo de red neuronal diseñada específicamente para trabajar con datos de gráficos. A diferencia de las redes neuronales tradicionales que funcionan con una entrada de tamaño fijo, las GNN pueden manejar tamaños de gráficos arbitrarios y aprovechar los patrones de conectividad dentro de los datos. Al agregar información de los vecinos de un nodo, las GNN aprenden representaciones que capturan tanto los atributos del nodo como la estructura del gráfico, lo que las hace muy efectivas para tareas como la clasificación de nodos, la predicción de enlaces y la clasificación de gráficos.

La biblioteca Deep Graph Library ( DGL.ai ) es un potente conjunto de herramientas para crear redes neuronales de gran alcance con facilidad y eficiencia. Con DGL, los desarrolladores pueden aprovechar las arquitecturas de redes neuronales de gran alcance de última generación para abordar una variedad de tareas, incluida la predicción de enlaces. DGL proporciona una variedad de utilidades para trabajar con gráficos homogéneos y heterogéneos, lo que lo convierte en una herramienta versátil tanto para investigadores como para profesionales. Al simplificar la implementación de las redes neuronales de gran alcance, DGL le permite centrarse más en el desarrollo de soluciones innovadoras en lugar de empantanarse en las complejidades técnicas subyacentes.

Con esta base en mente, profundicemos en la construcción de un modelo de predicción de enlaces utilizando GNN y DGL.ai.

Configuración del proyecto

El primer paso es configurar nuestro proyecto importando las bibliotecas necesarias:

 import json import numpy as np import pandas as pd import dgl from dgl.data import DGLDataset from dgl.nn import SAGEConv import torch import torch.nn as nn from torch.nn.functional import binary_cross_entropy_with_logits, relu, dropout from torch.nn.utils import clip_grad_norm_ from torch.optim.lr_scheduler import ReduceLROnPlateau import itertools import scipy.sparse as sp from sklearn.metrics import roc_auc_score

Preprocesamiento de datos y división de entrenamiento y prueba

Para preparar los datos para el entrenamiento, primero cargaremos el conjunto de datos de Twitch, lo representaremos como un gráfico y luego lo dividiremos en conjuntos de entrenamiento y prueba. Crearemos una clase personalizada que herede de DGLDataset , que ayudará a estructurar el proceso de carga de datos y agilizará las operaciones relacionadas con los gráficos.

1. Carga de los datos del gráfico : en la función de proceso, comenzamos leyendo la lista de aristas del conjunto de datos. Dado que las aristas del conjunto de datos original no tienen dirección (lo que representa amistades mutuas), agregamos aristas inversas para garantizar que las conexiones sean bidireccionales en nuestro gráfico.
2. Carga de características de nodos : a continuación, cargamos las características de los nodos desde un archivo JSON. Cada nodo tiene una lista de características que puede variar en longitud, por lo que rellenamos las listas de características más cortas con ceros para mantener una longitud constante en todos los nodos.

Aquí está el código para crear y preprocesar el conjunto de datos:

 # create a dataset that inherits DGLDataset class SocialNetworkDataset(DGLDataset): def __init__(self): super().__init__(name='social_network') def process(self): # load edges edges_df = pd.read_csv('./twitch/ENGB/musae_ENGB_edges.csv') # ensure edges are bidirectional edges_df_rev = edges_df.copy() edges_df_rev.columns = ['to', 'from'] edges_df_rev = edges_df_rev[['from', 'to']] edges_df = pd.concat([edges_df, edges_df_rev], ignore_index=True) edges_df.drop_duplicates(inplace=True) # create a graph using DGL max_node_id = max(edges_df['from'].max(), edges_df['to'].max()) edges_src = torch.from_numpy(edges_df['from'].to_numpy()) edges_dst = torch.from_numpy(edges_df['to'].to_numpy()) self.graph = dgl.graph( (edges_src, edges_dst), num_nodes=max_node_id + 1, ) # load and node features with open('./twitch/ENGB/musae_ENGB_features.json') as f: node_features_dict = json.load(f) # feature lists have various lengths, pad them with zeros max_feature_list_len = max([len(l) for l in node_features_dict.values()]) for k in node_features_dict: if len(node_features_dict[k]) < max_feature_list_len: node_features_dict[k] += [0] * (max_feature_list_len - len(node_features_dict[k])) # set node features in graph node_features_df = pd.DataFrame.from_dict(node_features_dict).T.astype('float64') node_features_np = node_features_df.to_numpy() self.graph.ndata['feat'] = torch.from_numpy(node_features_np).float() def __len__(self): return 1 # only the whole graph is returned def __getitem__(self, idx): return self.graph

Ahora inicializamos nuestro conjunto de datos para cargar los datos del gráfico.

 # init the dataset dataset = SocialNetworkDataset() g = dataset[0]

El siguiente paso es crear los conjuntos de entrenamiento y prueba . Dividiremos los bordes en una proporción 80/20 para el entrenamiento y la prueba. Generamos muestras positivas (bordes existentes) y negativas (bordes inexistentes) para ambos conjuntos. Para gráficos más grandes, las utilidades dgl.sampling de DGL pueden ser útiles, pero aquí, todo el gráfico cabe en la memoria. Aquí está el código para crear conjuntos de entrenamiento y prueba:

 # pick edges for train and test sets (80/20 split) # (for larger graphs, we can use dgl.sampling.negative etc) u, v = g.edges() edge_ids = np.random.permutation(g.num_edges()) test_set_size = int(len(edge_ids) * 0.2) train_set_size = len(edge_ids) - test_set_size # positive samples: existing edges test_positive_u, test_positive_v = u[edge_ids[:test_set_size]], v[edge_ids[:test_set_size]] train_positive_u, train_positive_v = u[edge_ids[test_set_size:]], v[edge_ids[test_set_size:]] # negative samples: nonexistent edges adj = sp.coo_matrix((np.ones(len(u)), (u.numpy(), v.numpy()))) adj_negative = 1 - adj.todense() - np.eye(g.num_nodes()) negative_u, negative_v = np.where(adj_negative != 0) negative_edge_ids = np.random.choice(len(negative_u), g.num_edges()) test_negative_u, test_negative_v = ( negative_u[negative_edge_ids[:test_set_size]], negative_v[negative_edge_ids[:test_set_size]], ) train_negative_u, train_negative_v = ( negative_u[negative_edge_ids[test_set_size:]], negative_v[negative_edge_ids[test_set_size:]], ) # create a training graph by copying the original graph and removing test edges train_g = dgl.remove_edges(g, edge_ids[:test_set_size]) # define positive and negative graphs for train and test sets train_positive_g = dgl.graph((train_positive_u, train_positive_v), num_nodes=g.num_nodes()) train_negative_g = dgl.graph((train_negative_u, train_negative_v), num_nodes=g.num_nodes()) test_positive_g = dgl.graph((test_positive_u, test_positive_v), num_nodes=g.num_nodes()) test_negative_g = dgl.graph((test_negative_u, test_negative_v), num_nodes=g.num_nodes())

Modelo: Capas convolucionales GraphSAGE y predictor MLP

Utilizaremos una red neuronal convolucional Graph Sample and Aggregate (GraphSAGE) para aprender representaciones de nodos, también conocidas como incrustaciones , que capturan tanto la estructura como las características de cada nodo dentro del gráfico. GraphSAGE funciona agregando información de características de los vecinos de cada nodo para crear una representación significativa para cada nodo. Este proceso, conocido como agregación de vecinos , permite que el modelo aprenda patrones ricos y localizados en el gráfico.

En cada capa de GraphSAGE, el modelo aplica una función de agregación (en este caso, la función "media") para recopilar información de los nodos vecinos, que luego se combina con las características propias del nodo. Apilar múltiples capas convolucionales permite que el modelo capture información de nodos cada vez más distantes , lo que expande de manera efectiva la vista de cada nodo dentro del gráfico.

Para mejorar el rendimiento del modelo y reducir el sobreajuste, aplicaremos la eliminación después de cada capa.

Construyamos ahora el modelo GraphSAGE con tres capas convolucionales , junto con una función forward para definir cómo fluyen los datos a través de él:

 # define the GraphSAGE model with 3 convolutional layers class GraphSAGE(nn.Module): def __init__( self, input_features, hidden_features, output_features, dropout_probability=0.3, ): super(GraphSAGE, self).__init__() self.conv_layer_1 = SAGEConv(input_features, hidden_features, "mean") self.conv_layer_2 = SAGEConv(hidden_features, hidden_features, "mean") self.conv_layer_3 = SAGEConv(hidden_features, output_features, "mean") self.dropout_probability = dropout_probability def forward(self, graph, input_features): # first layer with ReLU activation and dropout h = relu(self.conv_layer_1(graph, input_features)) h = dropout(h, p=self.dropout_probability) # second layer with ReLU activation and dropout h = relu(self.conv_layer_2(graph, h)) h = dropout(h, p=self.dropout_probability) # third layer without dropout h = self.conv_layer_3(graph, h) return h

La salida después de la tercera capa ( h ) contiene las incrustaciones de los nodos. Para predecir la probabilidad de que exista una arista (o vínculo) entre dos nodos cualesquiera, utilizaremos un predictor de perceptrón multicapa (MLP) . Este MLP toma las incrustaciones de dos nodos como entrada y calcula una puntuación que indica la probabilidad de que exista una arista entre ellos.

 # define the MLP predictor class MLPPredictor(nn.Module): def __init__(self, hidden_features): super().__init__() # first linear layer to combine node embeddings self.W1 = nn.Linear(hidden_features * 2, hidden_features) # second linear layer to produce a single score output self.W2 = nn.Linear(hidden_features, 1) def apply_edges(self, edges): # concatenate source and destination node embeddings h = torch.cat([edges.src["h"], edges.dst["h"]], dim=1) # pass through MLP layers to get the edge score score = self.W2(relu(self.W1(h))).squeeze(1) return {'score': score} def forward(self, g, h): with g.local_scope(): g.ndata["h"] = h g.apply_edges(self.apply_edges) return g.edata["score"]

El predictor MLP funciona de la siguiente manera:

• Tamaño de entrada: el tamaño de entrada del predictor se basa en las incrustaciones de nodos aprendidas. Dado que cada borde conecta dos nodos, concatenamos sus incrustaciones (cada una de tamaño hidden_features), lo que da como resultado un tamaño de entrada de hidden_features * 2.
• Capa 1 (W1): esta capa procesa las incrustaciones concatenadas y reduce la dimensión de salida a hidden_features, preservando información relacional importante entre los nodos.
• Capa 2 (W2): esta capa final produce una única puntuación escalar para cada par de nodos, que representa la probabilidad de que exista un borde entre ellos.

Este enfoque en capas permite al predictor capturar relaciones complejas entre pares de nodos y calcular una puntuación de borde que puede interpretarse como la probabilidad de existencia de un borde.

Función de pérdida y AUC

Para entrenar nuestro modelo de manera efectiva, necesitamos una función de pérdida que pueda cuantificar el desempeño del modelo en la predicción de enlaces. Dado que esta tarea es un problema de clasificación binaria (en el que cada enlace existe o no), utilizamos la entropía cruzada binaria (BCE) como nuestra función de pérdida. La entropía cruzada binaria mide la discrepancia entre las puntuaciones predichas del modelo y las etiquetas reales (1 para un enlace existente, 0 para ningún enlace). Utilizamos la versión _with_logits porque nuestro modelo genera puntuaciones brutas (logits) en lugar de probabilidades. Esta versión de BCE es más estable cuando se trabaja con logits, ya que combina la función sigmoidea y la entropía cruzada en un solo paso.

Aquí está el código que calcula la pérdida:

 def compute_loss(positive_logits, negative_logits): # concatenate positive and negative scores y_predicted = torch.cat([positive_logits, negative_logits]) # create true labels (1 for existing links, 0 for nonexistent) y_true = torch.cat([torch.ones(positive_logits.shape[0]), torch.zeros(negative_logits.shape[0])]) return binary_cross_entropy_with_logits(y_predicted, y_true)

Para evaluar el modelo, utilizamos la métrica Área bajo la curva ROC (AUC) . La AUC es adecuada para la predicción de enlaces porque maneja datos desequilibrados de manera eficaz, donde las muestras negativas (bordes inexistentes) son mucho más comunes que las muestras positivas. La puntuación AUC nos da una idea de qué tan bien el modelo clasifica los enlaces existentes por encima de los inexistentes.

Aquí está el código para calcular el AUC:

 def compute_auc(positive_logits, negative_logits): y_predicted = torch.cat([positive_logits, negative_logits]).detach().numpy() y_true = torch.cat([torch.ones(positive_logits.shape[0]), torch.zeros(negative_logits.shape[0])]).detach().numpy() return roc_auc_score(y_true, y_predicted)

Nota: utilizamos detach() para eliminar los tensores del gráfico de cálculo, lo que nos permite calcular el AUC sin afectar los gradientes.

Entrenando el modelo

Ahora estamos listos para entrenar el modelo. Para comenzar, crearemos una instancia del modelo, el predictor y el optimizador, y definiremos el ciclo de entrenamiento . También especificaremos la tasa de aprendizaje, el tamaño de la capa oculta y la tasa de abandono, entre otros hiperparámetros. Si bien no cubriremos la optimización de hiperparámetros aquí, usaremos un programador de tasa de aprendizaje para ajustar la tasa de aprendizaje si la pérdida se estabiliza, lo que significa que deja de disminuir durante una cantidad determinada de épocas (en este caso, 25). El programador reduce a la mitad la tasa de aprendizaje cuando esto sucede, lo que puede ayudar a que el modelo converja de manera más efectiva.

Aquí está el código para inicializar el modelo y configurar el entrenamiento:

 # init the model num_hidden_features = 32 model = GraphSAGE( train_g.ndata['feat'].shape[1], num_hidden_features, num_hidden_features, ) predictor = MLPPredictor(num_hidden_features) # create an optimizer and a learning rate scheduler learning_rate = 0.01 optimizer = torch.optim.Adam( itertools.chain(model.parameters(), predictor.parameters()), lr=learning_rate, ) lr_scheduler = ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=25) # train the model num_epochs = 1000 for epoch in range(num_epochs + 1): # forward h = model(train_g, train_g.ndata['feat']) positive_logits = predictor(train_positive_g, h) negative_logits = predictor(train_negative_g, h) loss = compute_loss(positive_logits, negative_logits) # backward optimizer.zero_grad() loss.backward() # clip gradients clip_grad_norm_(model.parameters(), 1.0) optimizer.step() # adjust learning rate based on the loss lr_scheduler.step(loss) # print loss, current learning rate, and AUC if epoch % 100 == 0: last_lr = lr_scheduler.get_last_lr()[0] train_auc = compute_auc(positive_logits, negative_logits) print(f'Epoch: {epoch}, learning rate: {last_lr}, loss: {loss}, AUC: {train_auc}')

Ejecutemos el código y examinemos la salida:

 Epoch: 0, learning rate: 0.01, loss: 262.4156188964844, AUC: 0.4934097124994463 Epoch: 100, learning rate: 0.01, loss: 0.5642552375793457, AUC: 0.7473735298706314 Epoch: 200, learning rate: 0.01, loss: 0.4622882306575775, AUC: 0.8431058751115716 Epoch: 300, learning rate: 0.01, loss: 0.40566185116767883, AUC: 0.8777374138645864 Epoch: 400, learning rate: 0.01, loss: 0.38118976354599, AUC: 0.8944719038039551 Epoch: 500, learning rate: 0.01, loss: 0.3690297603607178, AUC: 0.9039401673234729 Epoch: 600, learning rate: 0.005, loss: 0.3579995930194855, AUC: 0.9112366798940639 Epoch: 700, learning rate: 0.005, loss: 0.3557407557964325, AUC: 0.9128097572016495 Epoch: 800, learning rate: 0.005, loss: 0.3510144352912903, AUC: 0.9152937255697913 Epoch: 900, learning rate: 0.00125, loss: 0.3425179123878479, AUC: 0.9202487786553115 Epoch: 1000, learning rate: 0.00015625, loss: 0.3432360589504242, AUC: 0.9198250134354529

Como podemos ver, el modelo alcanza un AUC de alrededor de 0,92, lo que indica un sólido rendimiento predictivo . Observe que la tasa de aprendizaje se reduce entre las épocas 500 y 600 cuando la pérdida se estabiliza. Este ajuste permite un ajuste más fino, lo que conduce a una pequeña disminución de la pérdida. Después de cierto punto, la pérdida y el AUC se estabilizan, lo que indica que el modelo ha convergido .

Evaluemos el modelo en los datos de prueba (que no se usaron durante el entrenamiento) y veamos si se generaliza bien:

 # evaluate the model on the test data with torch.no_grad(): test_positive_scores = predictor(test_positive_g, h) test_negative_scores = predictor(test_negative_g, h) test_loss = compute_loss(test_positive_scores, test_negative_scores) test_auc = compute_auc(test_positive_scores, test_negative_scores) print(f'Test loss: {test_loss}, Test AUC: {test_auc}')

El resultado es:

 Test loss: 0.675215482711792, Test AUC: 0.866213400711374

El AUC de prueba es ligeramente inferior al AUC de entrenamiento, lo que indica un sobreajuste menor. Sin embargo, el AUC de 0,866 muestra que el modelo todavía funciona bien con datos no vistos . Un ajuste adicional de los hiperparámetros podría mejorar la generalización, especialmente si el sobreajuste es una preocupación.

Utilizando el modelo entrenado para predecir nuevos enlaces

Con nuestro modelo entrenado, ahora podemos predecir la probabilidad de vínculos entre nodos en el gráfico. Generaremos predicciones para todos los pares de nodos posibles , lo que nos permitirá identificar posibles nuevas conexiones.

1. Generación de pares de nodos candidatos : primero generamos todos los pares de nodos posibles, excluyendo los bucles propios (conexiones de un nodo consigo mismo). Esto nos da los pares de nodos candidatos para los que queremos predecir las probabilidades de enlace.
2. Creación de un gráfico candidato: luego creamos un gráfico DGL con estos bordes candidatos y usamos el modelo con incrustaciones de nodos del gráfico original. Estas incrustaciones, almacenadas en el atributo ndata['h'] del gráfico candidato, servirán como entradas para la predicción de enlaces.

Aquí está el código para estos pasos:

 # build node pairs, avoid self-loops (with_replacement=False) node_pairs = torch.combinations(torch.arange(g.num_nodes()), r=2, with_replacement=False) candidate_u = node_pairs[:, 0] candidate_v = node_pairs[:, 1] # build a graph with all node pairs candidate_graph = dgl.graph((candidate_u, candidate_v)) candidate_graph_node_embeddings = model(g, g.ndata['feat']) # we use embeddings from the original graph candidate_graph.ndata['h'] = candidate_graph_node_embeddings # use the predictor to predict the existence of links between nodes predicted_scores = predictor(candidate_graph, candidate_graph_node_embeddings)

Ahora que tenemos predicciones para todos los pares de candidatos, podemos verificar la probabilidad de enlace entre cualquier nodo específico . Por ejemplo, examinemos la puntuación y la probabilidad de un enlace entre los nodos 1773 y 7005, que no están conectados directamente en el conjunto de datos inicial:

 # find the index of the node pair (1773, 7005) pair_index = torch.where((candidate_u == 1773) & (candidate_v == 7005))[0] print(f'Pair index: {pair_index}') # get the logit score for this pair and compute probability of link existence pair_link_score = predicted_scores[pair_index].item() # logit score print(f'Pair link score: {pair_link_score}') link_probability = torch.sigmoid(torch.tensor(pair_link_score)).item() # apply sigmoid to convert score into probability print(f'Link probability: {link_probability * 100}%')

Este es el resultado:

 Pair index: tensor([11066978]) Pair link score: 0.7675977945327759 Link probability: 68.30010414123535%

Según nuestro modelo, existe un 68,3% de probabilidad de que exista un vínculo entre el usuario 1773 y el 7005 .

Conclusión

En esta publicación, construimos con éxito un modelo para predecir nuevos enlaces en un gráfico social, demostrando el uso de redes neuronales de gráficos y DGL para la predicción de enlaces. El uso de un conjunto de datos relativamente pequeño nos permitió trabajar de manera eficiente en una máquina local. Sin embargo, a medida que los gráficos se amplían a millones o miles de millones de nodos y aristas, su manejo requiere soluciones más avanzadas, como el entrenamiento distribuido en clústeres de GPU.

Como siguiente paso, exploraremos enfoques para manejar gráficos a gran escala e implementar la predicción de enlaces en entornos de nube, lo que le permitirá aplicar estas técnicas a conjuntos de datos de nivel de producción.