GraphStorm 0.3: aprendizaje escalable y multitarea en gráficos con API fáciles de usar

GraphStorm es un marco de trabajo de aprendizaje automático de gráficos (GML) empresarial de bajo código para crear, entrenar e implementar soluciones de aprendizaje automático de gráficos en gráficos complejos a escala empresarial en cuestión de días en lugar de meses. Con GraphStorm, puede crear soluciones que tengan en cuenta directamente la estructura de las relaciones o interacciones entre miles de millones de entidades, que están inherentemente integradas en la mayoría de los datos del mundo real, incluidos los escenarios de detección de fraudes, las recomendaciones, la detección de la comunidad y los problemas de búsqueda y recuperación.

Hoy lanzamos GraphStorm 0.3, que agrega compatibilidad nativa para el aprendizaje multitarea en gráficos. Específicamente, GraphStorm 0.3 le permite definir múltiples objetivos de entrenamiento en diferentes nodos y bordes dentro de un solo ciclo de entrenamiento. Además, GraphStorm 0.3 agrega nuevas API para personalizar los pipelines de GraphStorm: ahora solo necesita 12 líneas de código para implementar un ciclo de entrenamiento de clasificación de nodos personalizado. Para ayudarlo a comenzar con la nueva API, hemos publicado dos ejemplos de cuadernos Jupyter: uno para la clasificación de nodos y otro para una tarea de predicción de enlaces. También publicamos un estudio integral de coentrenamiento de modelos de lenguaje (LM) y redes neuronales de gráficos (GNN) para gráficos grandes con características de texto enriquecido utilizando el conjunto de datos Microsoft Academic Graph (MAG) de nuestro Documento KDD 2024El estudio muestra el rendimiento y la escalabilidad de GraphStorm en gráficos ricos en texto y las mejores prácticas para configurar bucles de entrenamiento GML para lograr un mejor rendimiento y eficiencia.

Soporte nativo para el aprendizaje multitarea en gráficos

Muchas aplicaciones empresariales tienen datos de gráficos asociados con múltiples tareas en diferentes nodos y bordes. Por ejemplo, las organizaciones minoristas quieren realizar una detección de fraudes tanto en vendedores como en compradores. Los editores científicos quieren encontrar más trabajos relacionados para citar en sus artículos y necesitan seleccionar el tema correcto para que su publicación sea detectable. Para modelar mejor dichas aplicaciones, los clientes nos han pedido que admitamos el aprendizaje multitarea en gráficos.

GraphStorm 0.3 admite el aprendizaje multitarea en gráficos con seis tareas más comunes: clasificación de nodos, regresión de nodos, clasificación de bordes, regresión de bordes, predicción de enlaces y reconstrucción de características de nodos. Puede especificar los objetivos de entrenamiento a través de un archivo de configuración YAML. Por ejemplo, un editor científico puede usar la siguiente configuración YAML para definir simultáneamente una tarea de clasificación de temas de artículos en paper nodos y una tarea de predicción de enlaces en paper-citing-paper Aristas para el caso de uso de editorial científica:

version: 1.0
    gsf:
        basic: # basic settings of the backbone GNN model
            ...
        ...
        multi_task_learning:
            - node_classification:         # define a node classification task for paper subject prediction.
                target_ntype: "paper"      # the paper nodes are the training targets.
                label_field: "label_class" # the node feature "label_class" contains the training labels.
				mask_fields:
                    - "train_mask_class"   # train mask is named as train_mask_class.
                    - "val_mask_class"     # validation mask is named as val_mask_class.
                    - "test_mask_class"    # test mask is named as test_mask_class.
                num_classes: 10            # There are total 10 different classes (subject) to predict.
                task_weight: 1.0           # The task weight is 1.0.
                
            - link_prediction:                # define a link prediction paper citation recommendation.
                num_negative_edges: 4         # Sample 4 negative edges for each positive edge during training
                num_negative_edges_eval: 100  # Sample 100 negative edges for each positive edge during evaluation
                train_negative_sampler: joint # Share the negative edges between positive edges (to speedup training)
                train_etype:
                    - "paper,citing,paper"    # The target edge type for link prediction training is "paper, citing, paper"
                mask_fields:
                    - "train_mask_lp"         # train mask is named as train_mask_lp.
                    - "val_mask_lp"           # validation mask is named as val_mask_lp.
                    - "test_mask_lp"          # test mask is named as test_mask_lp.
                task_weight: 0.5              # The task weight is 0.5.

Para obtener más detalles sobre cómo ejecutar el aprendizaje multitarea de gráficos con GraphStorm, consulte Aprendizaje multitarea en GraphStorm en nuestra documentación.

Nuevas API para personalizar componentes y pipelines de GraphStorm

Desde el lanzamiento de GraphStorm a principios de 2023, los clientes han utilizado principalmente su interfaz de línea de comandos (CLI), que abstrae la complejidad de la secuencia de comandos de ML de gráficos para que usted pueda crear, entrenar e implementar rápidamente modelos utilizando recetas comunes. Sin embargo, los clientes nos dicen que quieren una interfaz que les permita personalizar la secuencia de comandos de entrenamiento e inferencia de GraphStorm según sus requisitos específicos con mayor facilidad. Según los comentarios de los clientes sobre las API experimentales que lanzamos en GraphStorm 0.2, GraphStorm 0.3 presenta API de secuencia de comandos de ML de gráficos refactorizadas. Con las nuevas API, solo necesita 12 líneas de código para definir una secuencia de comandos de entrenamiento de clasificación de nodos personalizada, como se ilustra en el siguiente ejemplo:

import graphstorm as gs
gs.initialize()

acm_data = gs.dataloading.GSgnnData(part_config='./acm_gs_1p/acm.json')

train_dataloader = gs.dataloading.GSgnnNodeDataLoader(dataset=acm_data, target_idx=acm_data.get_node_train_set(ntypes=['paper']), fanout=[20, 20], batch_size=64)
val_dataloader = gs.dataloading.GSgnnNodeDataLoader(dataset=acm_data, target_idx=acm_data.get_node_val_set(ntypes=['paper']), fanout=[100, 100], batch_size=256, train_task=False)
test_dataloader = gs.dataloading.GSgnnNodeDataLoader(dataset=acm_data, target_idx=acm_data.get_node_test_set(ntypes=['paper']), fanout=[100, 100], batch_size=256, train_task=False)

model = RgcnNCModel(g=acm_data.g, num_hid_layers=2, hid_size=128, num_classes=14)
evaluator = gs.eval.GSgnnClassificationEvaluator(eval_frequency=100)

trainer = gs.trainer.GSgnnNodePredictionTrainer(model)
trainer.setup_evaluator(evaluator)

trainer.fit(train_dataloader, val_dataloader, test_dataloader, num_epochs=5)

Para ayudarlo a comenzar con las nuevas API, también hemos lanzado Nuevos ejemplos de cuadernos Jupyter en nuestro Página de documentación y tutoriales.

Estudio exhaustivo de LM+GNN para gráficos grandes con funciones de texto enriquecidas

Muchas aplicaciones empresariales tienen gráficos con características de texto. En las aplicaciones de búsqueda minorista, por ejemplo, los datos de registro de compras brindan información sobre cómo se relacionan las descripciones de productos ricas en texto, las consultas de búsqueda y el comportamiento del cliente. Los modelos de lenguaje grande (LLM) básicos por sí solos no son adecuados para modelar dichos datos porque las distribuciones y relaciones de datos subyacentes no se corresponden con lo que los LLM aprenden de sus corpus de datos de preentrenamiento. GML, por otro lado, es excelente para modelar datos relacionados (gráficos), pero hasta ahora, los profesionales de GML tenían que combinar manualmente sus modelos GML con LLM para modelar las características del texto y obtener el mejor rendimiento para sus casos de uso. Especialmente cuando el conjunto de datos del gráfico subyacente era grande, este trabajo manual era desafiante y consumía mucho tiempo.

En GraphStorm 0.2, GraphStorm introdujo técnicas integradas para entrenar modelos de lenguaje (LM) y modelos GNN de manera eficiente y a escala en gráficos masivos con gran cantidad de texto. Desde entonces, los clientes nos han pedido orientación sobre cómo se deben emplear las técnicas LM+GNN de GraphStorm para optimizar el rendimiento. Para abordar esto, con GraphStorm 0.3, lanzamos un punto de referencia LM+GNN utilizando el conjunto de datos de gráficos grandes, Microsoft Academic Graph (MAG), en dos tareas estándar de ML de gráficos: clasificación de nodos y predicción de enlaces. El conjunto de datos de gráficos es un gráfico heterogéneo, contiene cientos de millones de nodos y miles de millones de aristas, y la mayoría de los nodos están atribuidos con características de texto enriquecido. Las estadísticas detalladas de los conjuntos de datos se muestran en la siguiente tabla.

En GraphStorm, evaluamos dos métodos LM-GNN principales: BERT+GNN preentrenado, un método de referencia que se adopta ampliamente, y BERT+GNN perfeccionado, introducido por los desarrolladores de GraphStorm en 2022. Con el método BERT+GNN preentrenado, primero usamos un modelo BERT preentrenado para calcular incrustaciones para características de texto de nodos y luego entrenamos un modelo GNN para la predicción. Con el método BERT+GNN perfeccionado, inicialmente afinamos los modelos BERT en los datos del gráfico y usamos el modelo BERT perfeccionado resultante para calcular incrustaciones que luego se usan para entrenar modelos GNN para la predicción. GraphStorm ofrece diferentes formas de perfeccionar los modelos BERT, según los tipos de tareas. Para la clasificación de nodos, afinamos el modelo BERT en el conjunto de entrenamiento con las tareas de clasificación de nodos; para la predicción de enlaces, afinamos el modelo BERT con las tareas de predicción de enlaces. En el experimento, utilizamos 8 instancias r5.24xlarge para el procesamiento de datos y 4 instancias g5.48xlarge para el entrenamiento y la inferencia del modelo. El enfoque BERT+GNN ajustado tiene un rendimiento hasta un 40 % mejor (predicción de enlaces en MAG) en comparación con el BERT+GNN previamente entrenado.

La siguiente tabla muestra el rendimiento del modelo de los dos métodos y el tiempo total de cálculo de todo el proceso a partir del procesamiento de datos y la construcción de gráficos. NC significa clasificación de nodos y LP significa predicción de enlaces. El costo de tiempo de LM significa el tiempo empleado en calcular las incrustaciones de BERT y el tiempo empleado en ajustar los modelos BERT para BERT+GNN preentrenado y BERT+GNN ajustado, respectivamente.

También evaluamos GraphStorm en gráficos sintéticos de gran tamaño para demostrar su escalabilidad. Generamos tres gráficos sintéticos con mil millones, diez mil millones y cien mil millones de aristas. Los tamaños de conjunto de entrenamiento correspondientes son ocho millones, ochenta millones y ochocientos millones, respectivamente. La siguiente tabla muestra el tiempo de cálculo del preprocesamiento de gráficos, la partición de gráficos y el entrenamiento de modelos. En general, GraphStorm permite la construcción de gráficos y el entrenamiento de modelos en gráficos de escala de cien mil millones en cuestión de horas.

Más detalles y resultados de referencia están disponibles en nuestro Documento KDD 2024.

Conclusión

GraphStorm 0.3 se publica bajo la licencia Apache-2.0 para ayudarlo a enfrentar sus desafíos de aprendizaje automático de gráficos a gran escala y ahora ofrece soporte nativo para el aprendizaje multitarea y nuevas API para personalizar los pipelines y otros componentes de GraphStorm. Consulte la Repositorio de GraphStorm en GitHub y documentación Para empezar.

Sobre el Autor

Canción de Xiang es un científico aplicado sénior en AWS AI Research and Education (AIRE), donde desarrolla marcos de aprendizaje profundo, incluidos GraphStorm, DGL y DGL-KE. Lideró el desarrollo de Amazon Neptune ML, una nueva capacidad de Neptune que utiliza redes neuronales de gráficos para gráficos almacenados en bases de datos de gráficos. Ahora lidera el desarrollo de GraphStorm, un marco de aprendizaje automático de gráficos de código abierto para casos de uso empresarial. Obtuvo su doctorado en sistemas informáticos y arquitectura en la Universidad Fudan, Shanghái, en 2014.

Jian Zhang es un científico aplicado de alto nivel que ha estado utilizando técnicas de aprendizaje automático para ayudar a los clientes a resolver diversos problemas, como la detección de fraudes, la generación de imágenes de decoración y más. Ha desarrollado con éxito soluciones de aprendizaje automático basadas en gráficos, en particular redes neuronales gráficas, para clientes en China, EE. UU. y Singapur. Como ilustrador de las capacidades gráficas de AWS, Zhang ha realizado muchas presentaciones públicas sobre GNN, Deep Graph Library (DGL), Amazon Neptune y otros servicios de AWS.

Florian Saupe es gerente técnico principal de productos en AWS AI/ML Research y brinda apoyo a equipos científicos como el grupo de aprendizaje automático de gráficos y los equipos de sistemas de ML que trabajan en capacitación distribuida a gran escala, inferencia y resiliencia ante fallas. Antes de unirse a AWS, Florian dirigió la gestión técnica de productos para conducción automatizada en Bosch, fue consultor de estrategia en McKinsey & Company y trabajó como científico de sistemas de control/robótica, un campo en el que tiene un doctorado.