Recorrido de la Red de Atención de Grafos (GAT) con Implementación Visualizada

Recorrido de la Red de Atención de Grafos (GAT) - Implementación Visualizada

Visualización de alto nivel del material de recorrido

Comprender las Redes Neuronales Gráficas (GNNs) es cada vez más relevante a medida que los transformadores siguen abordando problemas de gráficos como los del Open Graph Benchmark. Incluso si el lenguaje natural es todo lo que necesita un gráfico, las GNNs siguen siendo una fuente fructífera de inspiración para futuros métodos.

En esta publicación, voy a explicar la implementación de una capa GNN básica y luego mostrar las modificaciones para una capa de atención gráfica básica, tal como se describe en el artículo de ICLR titulado Redes de Atención Gráfica.

Matriz de adyacencia de 5 x 5 para el gráfico de documentos

Inicialmente, imagina que tenemos un gráfico de documentos de texto representado como un grafo dirigido acíclico (DAG). El Documento 0 tiene una conexión con los documentos 1, 2 y 3, por lo que hay unos en la fila 0 para esas columnas.

Para realizar la implementación visualizada, voy a utilizar Graphbook, una herramienta de modelado de IA visual. Consulta mi otra publicación para obtener más información sobre cómo entender las representaciones visuales en Graphbook.

También tenemos algunas características de nodo para cada documento. He introducido cada documento en BERT como un único arreglo unidimensional de textos [5] para producir una forma de incrustaciones [5, 768] en la salida del Pooler.

Con fines instructivos, tomaré solo las primeras 8 dimensiones de la salida de BERT como las características de nodo para que podamos seguir más fácilmente las formas de los datos. Ahora, tenemos nuestra matriz de adyacencia y nuestras características de nodo.

Capa GNN

La fórmula general para una Capa GNN es que, para cada nodo, tomamos todos los vecinos de cada nodo y sumamos las características multiplicadas por una matriz de pesos, y luego pasamos por una función de activación. He creado un bloque en blanco con esta fórmula como título y he pasado la matriz de adyacencia y las características del nodo, y voy a implementar esa fórmula dentro del bloque.

Cuando implementamos esta fórmula, no queremos ejecutar realmente un bucle. Si podemos vectorizar completamente esto, entonces cualquier entrenamiento e inferencia con GPUs será mucho más rápido porque la multiplicación puede ser un solo paso de cálculo. Entonces, en su lugar, repetimos (es decir, transmitimos) las características del nodo en una forma tridimensional, por lo que teníamos una forma [5, 8] de características del nodo, ahora tendremos una forma [5, 5, 8] donde cada celda en la dimensión 0 es una repetición de las características del nodo. Ahora podemos pensar en la última dimensión como características de “vecino”. Cada nodo tiene un conjunto de 5 vecinos posibles.

No podemos transmitir directamente las características del nodo de [5, 8] a una forma [5, 5, 8]. En cambio, primero tenemos que transmitir a [25, 8] porque al transmitir, cada dimensión en la forma debe ser mayor o igual que la dimensión original. Por eso obtenemos las partes 5 y 8 de la forma (get_sub_arrays) y luego multiplicamos la primera para obtener 25, luego las concatenamos todas juntas. Finalmente, remodelamos el resultado [25, 8] resultante a [5, 5, 8], y de hecho podemos verificar en Graphbook que cada conjunto de características de nodo en las últimas 2 dimensiones es idéntico.

A continuación, también queremos transmitir la matriz de adyacencia al mismo tamaño. Esto significa que para cada 1 en la matriz de adyacencia en la fila i y columna j, hay una fila de 1.0s de num_feats en la dimensión [i, j]. Entonces, en esta matriz de adyacencia, la fila 0 tiene un 1 en las columnas 1, 2 y 3, por lo que hay una fila de num_feats 1.0s en las filas 1, 2 y 3 en la celda 0 (es decir, [0, 1:3, :]).

La implementación es muy simple aquí, simplemente analizamos la matriz de adyacencia a decimal y transmitimos de una forma [5, 5] a una forma [5, 5, 8]. Ahora podemos multiplicar elemento a elemento esta máscara de adyacencia por nuestras características de vecinos de nodos en mosaico.

También queremos incluir un auto-bucle en nuestra matriz de adyacencia, para que cuando sumemos las características de los vecinos, también incluyamos las características propias de ese nodo.

Después de hacer una multiplicación elemento a elemento (e incluyendo el auto-bucle), obtenemos las características de los vecinos para cada nodo y ceros para los nodos que no están conectados por una arista (no son vecinos). Para el nodo 0, esto incluye características para los nodos 0 al 3. Para el nodo 3, esto incluye los nodos 3 y 4.

A continuación, vamos a remodelar a [25, 8] para que cada característica del vecino sea su propia fila, y pasar eso a través de una capa lineal parametrizada con el tamaño oculto deseado. Aquí elegí 32 y lo guardé como una constante global para que pueda ser reutilizado. La salida de la capa lineal será [25, hidden_size]. Simplemente remodela esa salida, crea una forma [5, 5, hidden_size] y ¡ahora finalmente estamos listos para la parte de suma de la fórmula!

Sumamos sobre la dimensión intermedia (índice de dimensión 1) para que sumemos las características de los vecinos para cada nodo. El resultado es un conjunto de embeddings de nodo [5, hidden_size] que han pasado por 1 capa. Simplemente encadena estas capas juntas y tendrás una red GNN, y sigue las guías de https://www.youtube.com/@Graphbook para saber cómo entrenar.

Capa de Atención de Grafo

Según el artículo, el ingrediente secreto detrás de la capa de atención de grafo es el coeficiente de atención, dado en la fórmula anterior. En esencia, estamos concatenando los embeddings de nodo que están en una arista y pasándolos a través de otra capa lineal, antes de aplicar softmax.

Estos coeficientes de atención se utilizan luego para calcular una combinación lineal de las características correspondientes a las características originales del nodo.

Lo que tenemos que hacer es duplicar las características de cada nodo para cada vecino, y luego concatenar eso con las características de los vecinos del nodo.

El ingrediente secreto es obtener las características del nodo en mosaico para cada vecino. Para hacer eso, intercambiamos las dimensiones 0 y 1 de las características del nodo en mosaico antes de aplicar una máscara.

El resultado sigue siendo una matriz de forma [5, 5, 8], pero ahora cada fila en [i, :, :] es la misma y corresponde a la característica del nodo i. Ahora, podemos usar la multiplicación de elementos para crear las características del nodo repitiéndolas solo cuando contienen un vecino. Finalmente, concatenamos eso con las características del vecino tal como las creamos para la GNN y producimos las características concatenadas.

¡Ya casi estamos allí! Ahora que tenemos las características concatenadas, podemos pasarlas por una capa lineal. Necesitamos volver a dar forma a [5, 5, hidden_size] para poder aplicar softmax sobre la dimensión central y producir nuestros coeficientes de atención.

Ahora que tenemos nuestros coeficientes de atención con forma [5, 5, hidden_size], que es esencialmente una incrustación por cada borde del grafo para nuestro grafo de n nodos. El artículo dice que estos deben ser transpuestos (dimensiones intercambiadas), así que he vuelto y hecho eso antes de ReLU, y ahora aplico softmax sobre la última dimensión para que estén normalizados por índice de dimensión a lo largo de la dimensión de tamaño oculto. Vamos a multiplicar estos coeficientes por las incrustaciones originales de los nodos. Recuerda, las incrustaciones originales de los nodos tenían forma [5, 5, 8], donde 8 proviene arbitrariamente de cortar las primeras 8 características de las codificaciones de BERT de nuestros documentos de texto.

Multiplicar una forma [5, hidden_size, 5] por una forma [5, 5, 8] produce una forma [5, hidden_size, 8]. Luego sumamos sobre la dimensión del tamaño oculto para finalmente obtener una salida con forma [5, 8], que coincide con nuestra forma de entrada. También podríamos pasar esto por una no linealidad como otra ReLU y luego encadenar esta capa varias veces.

Conclusión

Hasta ahora, hemos repasado una implementación visual de capas GNN individuales y capas GAT. Puedes encontrar el proyecto en este repositorio de GitHub. En el artículo, también explican cómo extienden el método para la atención de múltiples cabezas. Hazme saber en los comentarios si también te gustaría que cubra esta parte o si hay algo más que te gustaría que cubra utilizando Graphbook.