Presentamos ARGO: Un impulso para el entrenamiento de redes neuronales gráficas

Las Redes Neuronales de Grafos (GNNs) son un tipo de inteligencia artificial que funciona muy bien con datos estructurados como grafos. Los grafos están compuestos por nodos (puntos) y bordes (conexiones entre los puntos). Las GNNs ayudan en varias áreas, como predecir propiedades de moléculas, hacer recomendaciones en redes sociales y pronosticar el rendimiento en diferentes sistemas.

A medida que las GNNs se vuelven más populares, se han desarrollado varias bibliotecas para facilitar a los desarrolladores la creación de modelos de GNN. Dos de las bibliotecas más conocidas son Deep Graph Library (DGL) y PyTorch-Geometric (PyG). Estas bibliotecas ofrecen herramientas que simplifican la tarea de programar diferentes modelos de GNN e incluyen backends optimizados para realizar cálculos.

A pesar de su efectividad, estas bibliotecas tienen problemas para rendir bien al usar múltiples procesadores de núcleo. Esta limitación lleva a un uso insuficiente de los recursos disponibles y puede obstaculizar el rendimiento general. Entrenar GNNs requiere recursos significativos y a menudo implica acceder a los datos de manera irregular, lo que las bibliotecas existentes no manejan bien en términos de uso de memoria.

#La Necesidad de Mejorar el Rendimiento

El entrenamiento de GNN puede ser muy intensivo en recursos. Por ejemplo, al usar un modelo GraphSAGE de tres capas para entrenar con un gran conjunto de datos, escalar el programa a más de 16 núcleos muestra poca o ninguna mejora en velocidad. Este problema surge principalmente por la forma en que las bibliotecas actuales gestionan y programan la memoria, que no aprovechan al máximo las capacidades del hardware moderno.

Los modelos de GNN requieren tanto memoria como potencia de procesamiento. Sin embargo, las bibliotecas existentes a menudo programan tareas de una manera que lleva a un uso ineficiente de los recursos. Por ejemplo, al ejecutar una sola tarea de entrenamiento, las operaciones que requieren memoria y las que requieren cálculo no están bien equilibradas.

Para hacer un uso efectivo de los recursos, una opción es ejecutar múltiples tareas de entrenamiento de GNN al mismo tiempo. Al permitir que estas tareas se ejecuten en paralelo, el uso de la memoria puede distribuirse, lo que permite un mejor rendimiento general. Sin embargo, decidir cuántos procesos ejecutar en paralelo puede ser complicado. Ejecutar demasiados procesos puede aumentar innecesariamente la carga de trabajo, mientras que muy pocos procesos pueden no utilizar los recursos de manera efectiva.

#Presentando ARGO

Para abordar estos desafíos, se ha creado un nuevo sistema de ejecución llamado ARGO para entrenar GNNs. ARGO tiene como objetivo mejorar la escalabilidad de las GNNs haciendo un mejor uso de los procesadores de múltiples núcleos. Lo logra ejecutando múltiples procesos de entrenamiento al mismo tiempo, superponiendo efectivamente las tareas de cálculo y comunicación.

ARGO cuenta con un autoajustador que puede encontrar automáticamente la mejor configuración para ejecutar estos procesos. Este autoajustador no requiere intervención del usuario, lo que lo hace fácil de usar. Funciona en segundo plano y se adapta a diferentes hardware, modelos de GNN y conjuntos de datos, lo que le permite ofrecer un buen rendimiento en diversas situaciones.

Las pruebas preliminares de ARGO han demostrado que puede acelerar significativamente el entrenamiento de GNN en comparación con las bibliotecas existentes. En algunos casos, se ha demostrado que es hasta cinco veces más rápido en sistemas multinúcleo. Esta mejora no solo hace que el entrenamiento de GNN sea más eficiente, sino que también facilita a los desarrolladores integrar ARGO en sus bibliotecas de GNN existentes con solo unas pocas líneas de código.

#Cómo Funcionan las GNNs

Antes de hablar de ARGO en detalle, es útil entender cómo funcionan las GNNs. Las GNNs aprenden a crear representaciones de baja dimensión para los nodos en un grafo. Estas representaciones se conocen como embeddings de nodos. El proceso implica dos pasos clave: Agregación de Características y Actualización de Características.

Durante la etapa de Agregación de Características, las características de los nodos vecinos se combinan para crear un nuevo vector de características para un nodo objetivo. Esto se puede hacer utilizando operaciones como sumar, promediar o tomar el máximo de las características vecinas.

Después de la agregación, la etapa de Actualización de Características aplica una función, generalmente un perceptrón multicapa (MLP), a las características agregadas. Este paso transforma las características de entrada en un nuevo espacio latente, preparando el escenario para predicciones y otras aplicaciones posteriores.

#Entrenando GNNs en Lotes

Las GNNs pueden entrenarse ya sea en el grafo completo o en lotes más pequeños. El entrenamiento en grafos completos puede llevar a altos costos de memoria y puede requerir más ciclos de entrenamiento para alcanzar la convergencia. Por otro lado, el entrenamiento en mini-lotes puede muestrear una parte más pequeña del grafo cada vez, siendo una mejor opción para grafos a gran escala.

El entrenamiento en mini-lotes es generalmente más rápido y usa menos memoria. Muestrea subconjuntos del grafo, aplica el modelo de GNN y calcula gradientes de esta parte más pequeña. Existen varios algoritmos de muestreo para seleccionar qué partes del grafo usar durante el entrenamiento, como Muestreo de Vecinos y Muestreo de ShaDow.

Con el Muestreo de Vecinos, los vecinos se seleccionan al azar para crear un subgrafo para cada nodo objetivo, mientras que el Muestreo de ShaDow construye primero un subgrafo localizado antes de muestrear de él. Este enfoque ayuda a evitar el problema de la Explosión de Vecinos, donde el número de vecinos crece demasiado con cada capa en el modelo de GNN.

#Resumen de los Componentes de ARGO

ARGO actúa como un puente entre el modelo de GNN del desarrollador y el hardware. Se compone de varios componentes clave:

1. Autoajustador: El autoajustador determina automáticamente las mejores configuraciones para ejecutar el entrenamiento de GNN. Explora una pequeña porción de todas las configuraciones posibles, permitiendo encontrar configuraciones efectivas sin necesidad de pruebas extensas.

2. Motor Multi-Proceso: Este componente gestiona la ejecución paralela de varios procesos de entrenamiento de GNN. Divide los datos de entrada equitativamente entre los procesos y asegura que compartan recursos de manera efectiva.

3. Vinculador de Núcleos: Esta función asigna núcleos de CPU a procesos de entrenamiento según las configuraciones encontradas por el autoajustador. Al gestionar cuidadosamente la asignación de núcleos, ARGO mejora el rendimiento a través de una mejor asignación de recursos.

Estos componentes trabajan juntos para asegurar que el entrenamiento de GNN pueda llevarse a cabo de manera eficiente, maximizando el uso del hardware disponible.

#Realizando Experimentos con ARGO

Se ha llevado a cabo una investigación para evaluar el rendimiento de ARGO probando sus capacidades en diferentes plataformas de hardware. Los experimentos se han centrado en usar dos tipos de procesadores multinúcleo: un sistema de cuatro zócalos con 112 núcleos y una máquina de dos zócalos con 64 núcleos.

En estos experimentos, se ha probado ARGO con dos modelos de GNN diferentes y una variedad de conjuntos de datos. Los resultados mostraron que ARGO puede acelerar efectivamente el entrenamiento de GNN, a menudo por más de cuatro veces en comparación con las bibliotecas existentes.

Por ejemplo, al usar ARGO, el tiempo de entrenamiento se redujo significativamente, particularmente en conjuntos de datos más grandes. Tanto el Muestreador de Vecinos como el Muestreador de ShaDow fueron examinados, y se notaron mejoras en todos los escenarios probados.

#Evaluación del Rendimiento

El rendimiento de ARGO se evaluó no solo en términos de velocidad, sino también en relación con la precisión de los modelos entrenados con él. El autoajustador pudo encontrar configuraciones que igualaban o superaban el rendimiento de las configuraciones manuales.

En pruebas donde la GNN fue entrenada usando tanto ARGO como bibliotecas existentes, los resultados mostraron que el entrenamiento alcanzó niveles similares de precisión. Esto indica que, aunque ARGO mejora la velocidad y la utilización de recursos, no sacrifica la calidad del modelo.

Además, las pruebas han demostrado que el autoajustador incurre en una sobrecarga mínima durante el entrenamiento, convirtiéndolo en una solución ligera que no obstaculiza el rendimiento general.

#Integrando ARGO en Frameworks Existentes

Una de las mayores ventajas de ARGO es lo fácilmente que puede agregarse a bibliotecas de GNN existentes como DGL y PyG. Los desarrolladores simplemente pueden incluir unas pocas líneas de código para aprovechar las capacidades de ARGO sin necesidad de reestructurar significativamente su base de código.

El proceso de integración está diseñado para ser sencillo. Un usuario puede especificar el algoritmo de muestreo, conjuntos de datos y configuraciones utilizando las APIs existentes proporcionadas por las bibliotecas. Esto permite a los usuarios mantener la estructura prevista de sus modelos de GNN, mientras obtienen los beneficios de una mejor escalabilidad y rendimiento a través de ARGO.

#Futuras Mejoras

Aunque ARGO ha demostrado mejoras notables en rendimiento, siempre hay áreas en las que trabajar. Los desarrollos futuros pueden centrarse en mejorar aún más la eficiencia en cómo se accede y se asignan los datos durante el entrenamiento.

Una área clave de exploración es el impacto de los patrones de acceso a la memoria en el rendimiento. Por ejemplo, acceder a datos a través de diferentes zócalos puede ser más lento que acceder a datos en el mismo zócalo. En el futuro, refinar ARGO para gestionar mejor estos patrones de acceso podría generar resultados aún mejores.

Además, con el rápido avance en los modelos y bibliotecas de GNN, las adaptaciones continuas a ARGO asegurarán que siga siendo relevante y efectivo para manejar los desafíos futuros en el entrenamiento de GNN.

#Conclusión

En resumen, ARGO representa un avance significativo en el panorama del entrenamiento de GNN. Al abordar los desafíos de las bibliotecas existentes e introducir un mecanismo de autoajuste efectivo, ARGO mejora tanto la velocidad como la eficiencia del entrenamiento de GNN en procesadores de múltiples núcleos.

Su integración en bibliotecas populares de GNN lo hace accesible para los desarrolladores, permitiéndoles aprovechar sus capacidades con un esfuerzo mínimo. A medida que la investigación en inteligencia artificial y GNNs sigue avanzando, ARGO se destaca como una herramienta prometedora para apoyar el progreso futuro en el campo.