Una Graph Attention Network (GAT) es un tipo de red neuronal de grafos (GNN) que incorpora mecanismos de atención en el proceso de paso de mensajes. En las GNN estándar, cada nodo actualiza su representación agregando información de sus vecinos, a menudo con pesos iguales o predefinidos. Las GAT, en cambio, calculan coeficientes de atención que asignan diferente importancia a cada vecino, permitiendo que el modelo se centre en las partes más relevantes del grafo para una tarea dada. Este enfoque, introducido por Petar Veličković y sus colegas en 2018, se ha convertido en una arquitectura fundamental en el aprendizaje profundo geométrico y se utiliza ampliamente en aplicaciones que van desde el análisis de redes sociales hasta la predicción de propiedades moleculares.
La idea central de las GAT es aplicar el mecanismo de atención, originalmente popularizado en los modelos Transformer (architecture) para procesamiento del lenguaje natural, a datos estructurados como grafos. En un transformer, los pesos de atención se calculan entre todos los pares de tokens en una secuencia. En una GAT, la atención se calcula solo entre un nodo y sus vecinos inmediatos, lo que la convierte en una operación localizada y equivariante a permutaciones. Este diseño preserva la propiedad clave de las GNN: la salida es invariante al orden de los nodos, lo cual es crucial porque los grafos no tienen un orden canónico de nodos.
Antecedentes: Redes Neuronales de Grafos
Las redes neuronales de grafos son una clase de redes neuronales artificiales diseñadas para tareas donde las entradas son grafos, como moléculas, redes sociales o redes de citas. A diferencia de las imágenes o el texto, los grafos carecen de una cuadrícula fija o estructura secuencial, y los nodos pueden tener un número variable de conexiones. Las GNN abordan esto mediante capas equivariantes a permutaciones que actualizan las representaciones de los nodos a través de un paso de mensajes por pares. Cada nodo agrega mensajes de sus vecinos, y después de múltiples capas, el campo receptivo se expande para incluir nodos más distantes.
Una limitación clave de las primeras arquitecturas GNN, como las Redes Neuronales Convolucionales de Grafos (GCN), es que tratan a todos los vecinos por igual al agregar información. Por ejemplo, en un grafo molecular, un átomo de carbono podría estar conectado tanto a un átomo de hidrógeno como a uno de oxígeno, pero el átomo de oxígeno puede ser más significativo químicamente para predecir toxicidad. Las GCN asignarían el mismo peso a ambos vecinos, mientras que las GAT pueden aprender a asignar mayor atención al átomo de oxígeno.
Mecanismo de Atención en las GAT
El mecanismo de atención en las GAT opera de la siguiente manera. Para un nodo \(u\) con vector de características \(\mathbf{x}_u\), y su vecino \(v\), el modelo calcula un coeficiente de atención \(e_{uv}\) usando una transformación lineal compartida y un vector de pesos aprendible. Este coeficiente se normaliza típicamente entre todos los vecinos usando una función softmax, asegurando que los pesos de atención sumen uno. Los coeficientes normalizados se usan luego para calcular una suma ponderada de las características transformadas de los vecinos, que se convierte en la representación actualizada del nodo \(u\).
Formalmente, el coeficiente de atención se calcula como:
\[ e_{uv} = \text{LeakyReLU}(\mathbf{a}^T [\mathbf{W}\mathbf{x}_u \| \mathbf{W}\mathbf{x}_v]) \]
donde \(\mathbf{W}\) es una matriz de pesos compartida, \(\mathbf{a}\) es un vector aprendible, y \(\|\\) denota concatenación. Los coeficientes se normalizan usando softmax sobre todos los vecinos \(v \in N_u\). La representación actualizada del nodo es entonces:
\[ \mathbf{h}_u = \sigma\left(\sum_{v \in N_u} \alpha_{uv} \mathbf{W}\mathbf{x}_v\right) \]
donde \(\alpha_{uv}\) son los coeficientes de atención normalizados y \(\sigma\) es una no linealidad.
Este mecanismo es análogo a la atención de múltiples cabezas en los transformers, donde se usan múltiples cabezas de atención independientes para capturar diferentes tipos de relaciones. En las GAT, la atención de múltiples cabezas se puede aplicar calculando varias agregaciones ponderadas por atención en paralelo y concatenando o promediando sus salidas. Esto aumenta el poder expresivo del modelo y estabiliza el entrenamiento.
Variantes Arquitectónicas
Se han propuesto varias variantes de las GAT desde el artículo original de 2018. Una variante notable es GATv2, introducida en 2021, que aborda una limitación de la GAT original donde los coeficientes de atención se calculan usando una operación lineal después de la concatenación. GATv2 usa un mecanismo de atención más expresivo que permite al modelo calcular puntuaciones de atención más sensibles a las características de entrada, mejorando el rendimiento en tareas que requieren discriminación fina.
Otra variante es la Graph Attention Network con características de aristas, que incorpora características de aristas en el cálculo de la atención. En grafos moleculares, las características de aristas podrían representar tipos de enlace (simple, doble, aromático), e incorporarlas permite al modelo ponderar a los vecinos de manera diferente según la naturaleza de su conexión. Esto es particularmente útil en aplicaciones de química y biología.
Además, las GAT se pueden combinar con otros componentes de GNN como conexiones residuales y normalización de capas para mejorar la estabilidad del entrenamiento y el rendimiento. Estas mejoras son comunes en las arquitecturas GNN modernas.
Aplicaciones
Las GAT se han aplicado a una amplia gama de dominios. En biología molecular y descubrimiento de fármacos, las GAT se usan para predecir propiedades moleculares, como solubilidad, toxicidad o eficacia contra bacterias específicas como E. coli. Las moléculas se representan como grafos con átomos como nodos y enlaces como aristas, y las GAT pueden aprender a centrarse en grupos funcionales que son críticos para la actividad biológica.
En el análisis de redes sociales, las GAT se usan para clasificación de nodos y predicción de enlaces. Por ejemplo, en una red de citas, las GAT pueden clasificar artículos en temas de investigación atendiendo a los artículos citados más influyentes. En sistemas de recomendación, las GAT pueden modelar interacciones usuario-ítem como un grafo bipartito, donde la atención ayuda a identificar los ítems más relevantes para un usuario.
Las GAT también se usan en visión por computadora para tareas como detección de objetos y generación de grafos de escenas, donde las imágenes se representan como grafos de objetos y sus relaciones. En física, las GAT se han aplicado a la reconstrucción de trayectorias de partículas y al modelado de sistemas dinámicos.
Relación con los Transformers
Existe una relación estrecha entre las GAT y los modelos Transformer (architecture). Como se señaló en un artículo de posición de 2022 sobre aprendizaje profundo geométrico, una capa de transformer se puede interpretar como una GNN aplicada a un grafo completo donde cada token está conectado con todos los demás. En esta visión, el mecanismo de autoatención en los transformers es una forma de paso de mensajes con pesos de atención. A la inversa, las GAT se pueden ver como transformers adaptados a estructuras de grafos arbitrarias, donde la atención se restringe a aristas existentes en lugar de todos los pares.
Esta conexión ha llevado a una polinización cruzada entre los dos campos. Técnicas desarrolladas para transformers, como codificaciones posicionales y atención cruzada, se han adaptado para GNN. Por ejemplo, las codificaciones posicionales basadas en autovectores del laplaciano del grafo pueden proporcionar información estructural que complementa el paso de mensajes basado en atención.
Implementaciones y Bibliotecas
Varias bibliotecas de código abierto implementan GAT y otras arquitecturas GNN. PyTorch Geometric (PyG) es una biblioteca popular construida sobre PyTorch que proporciona implementaciones eficientes de capas GAT, junto con utilidades para manejar datos de grafos. TensorFlow GNN ofrece funcionalidad similar para el ecosistema TensorFlow. La Deep Graph Library (DGL) es una biblioteca independiente del marco que admite tanto backends de PyTorch como de TensorFlow. Para usuarios de JAX, la biblioteca jraph proporciona implementaciones de GNN, y para usuarios de Julia, están disponibles GraphNeuralNetworks.jl y GeometricFlux.jl.
Estas bibliotecas suelen incluir capas GAT preconstruidas que se pueden integrar fácilmente en modelos personalizados. También proporcionan conjuntos de datos y puntos de referencia para evaluar GNN, como redes de citas (Cora, CiteSeer) y conjuntos de datos de predicción de propiedades moleculares.
Limitaciones y Extensiones
Como todas las GNN de paso de mensajes, las GAT están sujetas a los límites de poder expresivo de la prueba de isomorfismo de grafos de Weisfeiler-Lehman. Esto significa que existen pares de grafos no isomorfos que las GAT no pueden distinguir, lo que puede ser una limitación para tareas que requieren discriminación estructural fina. Para superar esto, los investigadores han propuesto GNN de orden superior que operan en complejos simpliciales o usan transformers de grafos con atención global, aunque a menudo con un costo computacional mayor.
Otra limitación es la escalabilidad. Calcular atención sobre todos los vecinos puede ser costoso para grafos con nodos de alto grado, aunque esto es generalmente menos severo que el costo cuadrático de los transformers completos. Técnicas como el muestreo de vecinos o la atención dispersa pueden mitigar este problema.
A partir de 2024, las GAT siguen siendo una arquitectura ampliamente utilizada y activamente investigada. Su capacidad para ponderar adaptativamente la importancia de los vecinos las ha convertido en una herramienta estándar en el arsenal de las GNN, y continúan inspirando nuevas variantes y aplicaciones en aprendizaje automático y inteligencia artificial.