GraphStorm 0.3: Scalable graph multi-task studying with a user-friendly API

GraphStorm is a low-code enterprise graph machine studying (GML) framework for constructing, coaching, and deploying graph ML options on advanced enterprise-scale graphs in days as an alternative of months. With GraphStorm, you possibly can construct options that straight take into account the construction of relationships or interactions between billions of entities which can be inherently embedded in most real-world knowledge, together with fraud detection situations, suggestions, neighborhood detection, search/retrieval issues, and extra.

In the present day, we launched GraphStorm 0.3, including native assist for multi-task studying on graphs. Particularly, GraphStorm 0.3 means that you can outline a number of coaching targets for various nodes and edges inside a single coaching loop. As well as, GraphStorm 0.3 provides new APIs for customizing GraphStorm pipelines, which now require solely 12 traces of code to implement a customized node classification coaching loop. That can assist you get began with the brand new APIs, we now have revealed two instance Jupyter notebooks, one for node classification and one for hyperlink prediction duties. We additionally launched a complete research on joint coaching of language fashions (LMs) and graph neural networks (GNNs) on large-scale graphs with wealthy textual content options utilizing the Microsoft Educational Graph (MAG) dataset. KDD 2024 PaperOn this research, we current the efficiency and scalability of GraphStorm on text-rich graphs, in addition to finest practices for configuring GML coaching loops to enhance efficiency and effectivity.

Native assist for multi-task studying on graphs

Many enterprise purposes have graph knowledge related to a number of duties on completely different nodes and edges. For instance, retail organizations wish to carry out fraud detection for each sellers and patrons. Scientific publishers wish to discover extra associated analysis to quote of their papers and want to pick out the proper topic to make their publications discoverable. To higher mannequin such purposes, our prospects have requested us to assist multi-task studying on graphs.

GraphStorm 0.3 helps graph multi-task studying for the six commonest duties: node classification, node regression, edge classification, edge regression, hyperlink prediction, and node characteristic reconstruction. You possibly can specify the coaching targets by a YAML configuration file. For instance, a scientific writer can concurrently outline a paper topic classification activity with the next YAML configuration: paper Node and Hyperlink Prediction Duties paper-citing-paper Edge use circumstances for scientific publishers:

model: 1.0
    gsf:
        primary: # primary settings of the spine GNN mannequin
            ...
        ...
        multi_task_learning:
            - node_classification:         # outline a node classification activity for paper topic prediction.
                target_ntype: "paper"      # the paper nodes are the coaching targets.
                label_field: "label_class" # the node characteristic "label_class" accommodates the coaching labels.
				mask_fields:
                    - "train_mask_class"   # prepare masks is called as train_mask_class.
                    - "val_mask_class"     # validation masks is called as val_mask_class.
                    - "test_mask_class"    # take a look at masks is called as test_mask_class.
                num_classes: 10            # There are complete 10 completely different lessons (topic) to foretell.
                task_weight: 1.0           # The duty weight is 1.0.
                
            - link_prediction:                # outline a hyperlink prediction paper quotation suggestion.
                num_negative_edges: 4         # Pattern 4 damaging edges for every optimistic edge throughout coaching
                num_negative_edges_eval: 100  # Pattern 100 damaging edges for every optimistic edge throughout analysis
                train_negative_sampler: joint # Share the damaging edges between optimistic edges (to speedup coaching)
                train_etype:
                    - "paper,citing,paper"    # The goal edge sort for hyperlink prediction coaching is "paper, citing, paper"
                mask_fields:
                    - "train_mask_lp"         # prepare masks is called as train_mask_lp.
                    - "val_mask_lp"           # validation masks is called as val_mask_lp.
                    - "test_mask_lp"          # take a look at masks is called as test_mask_lp.
                task_weight: 0.5              # The duty weight is 0.5.

For extra info on the right way to carry out graph multi-task studying in GraphStorm, see: Multi-Task Learning in GraphStorm It is acknowledged in our documentation.

New API for customizing GraphStorm pipelines and parts

Since GraphStorm was launched in early 2023, prospects have primarily used the command-line interface (CLI), which abstracts the complexities of graph ML pipelines and means that you can shortly construct, prepare, and deploy fashions utilizing widespread recipes. Nonetheless, prospects have informed us they need an interface that lets them extra simply customise GraphStorm coaching and inference pipelines to their particular necessities. Based mostly on buyer suggestions on the experimental API we launched in GraphStorm 0.2, GraphStorm 0.3 introduces a refactored graph ML pipeline API. With the brand new API, solely 12 traces of code are wanted to outline a customized node classification coaching pipeline, as proven within the following instance:

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)

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

coach = gs.coach.GSgnnNodePredictionTrainer(mannequin)
coach.setup_evaluator(evaluator)

coach.match(train_dataloader, val_dataloader, test_dataloader, num_epochs=5)

That can assist you get began with the brand new API, New Jupyter notebook examples Amongst Us Documentation and tutorials page.

A Complete Research of LM+GNN for Giant Graphs with Wealthy Textual content Options

Many enterprise purposes have graphs with textual content options. For instance, in a retail search software, buying log knowledge supplies insights into how text-rich product descriptions, search queries, and buyer habits are associated. The underlying large-scale language fashions (LLMs) alone should not appropriate to mannequin such knowledge as a result of the distribution and relationships within the underlying knowledge don’t match what the LLMs have discovered from the pre-training knowledge corpus. Alternatively, GML is good for modeling associated knowledge (graphs), however till now GML practitioners have needed to manually mix GML fashions with LLMs to mannequin textual content options and get the most effective efficiency for his or her use circumstances. This handbook effort is tough and time-consuming, particularly when the underlying graph dataset is giant.

In GraphStorm 0.2, we launched built-in methods for effectively coaching Language Fashions (LM) and GNN fashions at scale on giant text-rich graphs. Since then, prospects have been asking for steering on the right way to use GraphStorm’s LM+GNN methods to optimize efficiency. To handle this, in GraphStorm 0.3, we launched LM+GNN benchmarks on two customary graph ML duties (node ​​classification and hyperlink prediction) utilizing Microsoft Educational Graph (MAG), a large-scale graph dataset. The graph dataset is a heterogeneous graph, containing a whole lot of hundreds of thousands of nodes and billions of edges, with most nodes assigned wealthy textual content options. Detailed statistics of the dataset are supplied within the following desk.

GraphStorm benchmarks two main LM-GNN strategies: pre-trained BERT+GNN, a broadly adopted baseline methodology, and fine-tuned BERT+GNN, launched by GraphStorm builders in 2022. Within the pre-trained BERT+GNN methodology, we first use a pre-trained BERT mannequin to compute embeddings for node textual content options, after which prepare a GNN mannequin for prediction. Within the fine-tuned BERT+GNN methodology, we first fine-tune a BERT mannequin on graph knowledge, use the ensuing fine-tuned BERT mannequin to compute embeddings, after which use it to coach a GNN mannequin for prediction. GraphStorm presents other ways to fine-tune a BERT mannequin relying on the kind of activity. For node classification, we fine-tune a BERT mannequin on the coaching set utilizing the node classification activity. For hyperlink prediction, we fine-tune a BERT mannequin utilizing the hyperlink prediction activity. In our experiments, we use eight r5.24xlarge situations for knowledge processing and 4 g5.48xlarge situations for mannequin coaching and inference. Our fine-tuned BERT+GNN strategy achieves as much as 40% higher efficiency (hyperlink prediction on MAG) in comparison with pre-trained BERT+GNN.

The next desk exhibits the mannequin efficiency of the 2 strategies and the general computation time of your complete pipeline ranging from knowledge processing and graph building. NC means node classification, LP means hyperlink prediction. LM time value means the time spent on computing BERT embeddings and fine-tuning the BERT mannequin for pre-trained BERT+GNN and fine-tuned BERT+GNN, respectively.

We additionally carried out benchmarking on giant artificial graphs to exhibit the scalability of GraphStorm. We generate three artificial graphs with 1 billion, 10 billion, and 100 billion edges. The corresponding coaching set sizes are 8 million, 80 million, and 800 million, respectively. The next desk exhibits the computation time for graph preprocessing, graph partitioning, and mannequin coaching. Total, GraphStorm permits graph building and mannequin coaching on 100 billion-scale graphs in just a few hours.

For particulars and outcomes of the benchmark, please see KDD 2024 Paper.

Conclusion

GraphStorm 0.3 is launched below the Apache-2.0 license that will help you sort out large-scale graph ML challenges, offering native assist for multi-task studying and new APIs for customizing pipelines and different parts of GraphStorm. GraphStorm GitHub repository and documentation let’s begin.

In regards to the Writer

Shomatsu He’s a Senior Utilized Scientist with AWS AI Analysis and Schooling (AIRE) the place he develops deep studying frameworks together with GraphStorm, DGL, and DGL-KE. He led the event of Amazon Neptune ML, a brand new characteristic for Neptune that makes use of graph neural networks on graphs saved in graph databases. He at the moment leads the event of GraphStorm, an open supply graph machine studying framework for enterprise use circumstances. He acquired his PhD in Laptop Methods and Structure from Fudan College, Shanghai in 2014.

Jean Jean is a senior utilized scientist who has helped prospects resolve numerous issues comparable to fraud detection, adorned picture era, and so forth. utilizing machine studying methods. He has efficiently developed options in graph-based machine studying, particularly graph neural networks, for patrons in China, the US, and Singapore. As an evangelist for AWS graph capabilities, Zhang has given many public displays on GNNs, Deep Graph Library (DGL), Amazon Neptune, and different AWS providers.

Florian Sope He’s a Principal Technical Product Supervisor for AWS AI/ML Analysis, supporting scientific groups such because the Graph Machine Studying group and the ML Methods workforce engaged on large-scale distributed coaching, inference, and fault tolerance. Previous to becoming a member of AWS, Florian led Technical Product Administration for Autonomous Driving at Bosch, was a Technique Guide at McKinsey & Firm, and labored as a Management Methods/Robotics Scientist, the place he holds a PhD.