Graph Neural Network (GNN)

Una red neuronal de grafos (GNN) es una red neuronal que opera directamente sobre grafos: conjuntos de nodos conectados por aristas, donde tanto los nodos como las aristas pueden incluir características. Las moléculas (átomos y enlaces), las redes sociales (usuarios y relaciones de seguimiento), las redes de carreteras (intersecciones y tramos de carretera), las mallas 3D y las articulaciones de un esqueleto humano son grafos. Un grafo no tiene una cuadrícula ni un orden fijos y sus nodos tienen distintos números de vecinos, por lo que no se pueden aplicar directamente los filtros deslizantes de una red neuronal convolucional (CNN) ni el procesamiento paso a paso de una red neuronal recurrente (RNN). Las GNN resuelven esto mediante el paso de mensajes: cada nodo actualiza repetidamente su vector de características a partir de los de sus vecinos, y así produce embeddings de nodos que codifican tanto los atributos propios de cada nodo como su posición en el grafo.

Cómo funciona el paso de mensajes#

Una capa GNN actualiza cada nodo en tres pasos. Cada vecino envía un mensaje, normalmente una transformación aprendida de sus características. El nodo agrega los mensajes recibidos con una función que ignora su orden, como una suma, una media o un máximo. Después combina el agregado con sus propias características para formar su nueva representación. Como la agregación ignora el orden, el resultado no depende de cómo estén numerados los nodos. En 2017, Justin Gilmer y sus colaboradores describieron muchas GNN como casos particulares de este marco de red neuronal de paso de mensajes. Tras la última capa, una cabeza específica de la tarea utiliza las incrustaciones de los nodos para etiquetarlos, predecir aristas ausentes o clasificar un grafo completo:

Cada capa amplía en un salto la perspectiva de un nodo, así que, tras k capas, la incrustación del nodo depende de todos los nodos que estén a un máximo de k aristas: su campo receptivo en el grafo. El fragmento de PyTorch que aparece a continuación crea un grafo de camino con seis nodos, aplica tres capas de convolución de grafos con la normalización de la GCN de Thomas Kipf y Max Welling y usa autograd para enumerar qué nodos de entrada afectan al nodo 0 después de cada capa:

import torch
from torch import nn

# A path graph 0-1-2-3-4-5 as an adjacency matrix with self-loops
A = torch.eye(6)
for i in range(5):
    A[i, i + 1] = A[i + 1, i] = 1
deg = A.sum(dim=1)
A_hat = A / torch.sqrt(deg[:, None] * deg[None, :])  # GCN normalization D^-1/2 (A + I) D^-1/2

torch.manual_seed(0)
x = torch.randn(6, 4, requires_grad=True)  # 4 input features per node
layers = [nn.Linear(4, 4) for _ in range(3)]

h = x
for k, layer in enumerate(layers, start=1):
    h = torch.tanh(A_hat @ layer(h))  # each node mixes in its neighbors' transformed features
    grad = torch.autograd.grad(h[0].sum(), x, retain_graph=True)[0]
    print(f"After layer {k}, node 0 depends on nodes {grad.abs().sum(dim=1).nonzero().flatten().tolist()}")

Imprime [0, 1], [0, 1, 2] y [0, 1, 2, 3]: el campo receptivo del nodo 0 crece exactamente un salto por capa.

Arquitecturas GNN habituales#

• GCN: La red convolucional de grafos calcula H′ = σ(D̃^(−1/2) à D̃^(−1/2) H W), donde à es la matriz de adyacencia con bucles propios y D̃ su matriz de grados: una media ponderada por grado de cada nodo y sus vecinos, seguida de una matriz de pesos W compartida por todos los nodos.
• GraphSAGE: Muestrea un número fijo de vecinos por nodo y aprende el agregador, por lo que se escala a grafos muy grandes y puede generar incrustaciones de nodos que aparecen después del entrenamiento.
• GAT: Las redes de atención de grafos asignan a cada vecino un peso mediante puntuaciones aprendidas de atención, en lugar de usar pesos fijos basados en el grado.
• GIN: La red de isomorfismo de grafos suma los vecinos y aplica un MLP, lo que, según demostraron Keyulu Xu y sus colaboradores, le permite distinguir grafos tan bien como la prueba de Weisfeiler-Lehman, el límite superior del paso de mensajes estándar.
• Transformadores de grafos: Aplican autoatención a todos los nodos y añaden codificaciones posicionales derivadas de la estructura del grafo. Un transformador estándar es en sí mismo una GNN sobre un grafo totalmente conectado de tokens.

Apilar muchas capas provoca un suavizado excesivo: las medias repetidas acercan las incrustaciones de los nodos hasta volverlas indistinguibles, por eso muchas GNN solo usan dos o tres capas, como la GCN original, que usaba dos. PyTorch Geometric y Deep Graph Library (DGL) implementan estas capas de forma eficiente sobre listas de aristas dispersas.

Aplicaciones en el mundo real#

• Predicción meteorológica: GraphCast de Google DeepMind, publicado en Science en 2023, ejecuta el paso de mensajes sobre una malla multiescala que cubre el planeta y genera predicciones a 10 días con una resolución de 0,25° en menos de un minuto; supera a los sistemas deterministas operativos más precisos en el 90 % de 1.380 objetivos de verificación.
• Descubrimiento de fármacos: En 2020, investigadores del MIT entrenaron una red de paso de mensajes dirigidos para predecir la actividad antibacteriana a partir de grafos moleculares y la usaron para identificar la halicina, un compuesto estructuralmente distinto de los antibióticos conocidos que eliminó muchas bacterias resistentes a los fármacos en pruebas de laboratorio.
• Sistemas de recomendación: PinSage de Pinterest aprendió incrustaciones de pines a partir de un grafo con 3.000 millones de pines y tableros y 18.000 millones de aristas, entrenado con 7.500 millones de ejemplos.
• Reconocimiento de acciones basado en esqueletos: ST-GCN trata las articulaciones del cuerpo como nodos, las conecta mediante huesos dentro de cada fotograma y con la misma articulación en los fotogramas vecinos, y clasifica la acción a partir del grafo espaciotemporal resultante. Los modelos de estimación de pose de YOLO26 devuelven 17 puntos clave COCO por persona en cada fotograma, las coordenadas articulares que sirven para construir ese grafo esquelético.

GNN frente a conceptos relacionados#

• Grafo de conocimiento: Estructura de datos que almacena hechos como tripletas entidad-relación-entidad. Una GNN es un modelo que puede aprender a partir de un grafo de conocimiento, por ejemplo, para predecir enlaces ausentes.
• Aprendizaje profundo geométrico: Campo más amplio que engloba redes basadas en las simetrías de sus datos, incluidas las CNN sobre cuadrículas, las GNN sobre grafos y los modelos sobre mallas y variedades.
• CNN: Una CNN aplicada a una imagen es un caso particular de paso de mensajes sobre un grafo en cuadrícula. Cada píxel tiene el mismo número de vecinos en posiciones relativas fijas, así que una CNN puede aprender un peso distinto para cada posición vecina, algo que una GNN sobre un grafo irregular no puede hacer.