GraphStorm 0.3: Scalable, multi-task studying on graphs with user-friendly APIs

GraphStorm is a low-code enterprise graph machine studying (GML) framework to construct, practice, and deploy graph ML options on complicated enterprise-scale graphs in days as a substitute of months. With GraphStorm, you may construct options that instantly keep in mind the construction of relationships or interactions between billions of entities, that are inherently embedded in most real-world information, together with fraud detection situations, suggestions, group detection, and search/retrieval issues.

Immediately, we’re launching GraphStorm 0.3, including native help for multi-task studying on graphs. Particularly, GraphStorm 0.3 permits you to outline a number of coaching targets on totally different nodes and edges inside a single coaching loop. As well as, GraphStorm 0.3 provides new APIs to customise GraphStorm pipelines: you now solely want 12 strains of code to implement a customized node classification coaching loop. That will help you get began with the brand new API, now we have printed two Jupyter pocket book examples: one for node classification, and one for a hyperlink prediction activity. We additionally launched a complete research of co-training language fashions (LM) and graph neural networks (GNN) for giant graphs with wealthy textual content options utilizing the Microsoft Tutorial Graph (MAG) dataset from our KDD 2024 paper. The research showcases the efficiency and scalability of GraphStorm on textual content wealthy graphs and the perfect practices of configuring GML coaching loops for higher efficiency and effectivity.

Native help for multi-task studying on graphs

Many enterprise functions have graph information related to a number of duties on totally different nodes and edges. For instance, retail organizations wish to conduct fraud detection on each sellers and consumers. Scientific publishers wish to discover extra associated works to quote of their papers and wish to pick the proper topic for his or her publication to be discoverable. To higher mannequin such functions, clients have requested us to help multi-task studying on graphs.

GraphStorm 0.3 helps multi-task studying on graphs with six commonest duties: node classification, node regression, edge classification, edge regression, hyperlink prediction, and node function reconstruction. You possibly can specify the coaching targets by means of a YAML configuration file. For instance, a scientific writer can use the next YAML configuration to concurrently outline a paper topic classification activity on paper nodes and a hyperlink prediction activity on paper-citing-paper edges for the scientific writer use case:

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 function "label_class" comprises the coaching labels.
				mask_fields:
                    - "train_mask_class"   # practice masks is known as as train_mask_class.
                    - "val_mask_class"     # validation masks is known as as val_mask_class.
                    - "test_mask_class"    # check masks is known as as test_mask_class.
                num_classes: 10            # There are whole 10 totally different lessons (topic) to foretell.
                task_weight: 1.0           # The duty weight is 1.0.
                
            - link_prediction:                # outline a hyperlink prediction paper quotation advice.
                num_negative_edges: 4         # Pattern 4 adverse edges for every constructive edge throughout coaching
                num_negative_edges_eval: 100  # Pattern 100 adverse edges for every constructive edge throughout analysis
                train_negative_sampler: joint # Share the adverse edges between constructive 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"         # practice masks is known as as train_mask_lp.
                    - "val_mask_lp"           # validation masks is known as as val_mask_lp.
                    - "test_mask_lp"          # check masks is known as as test_mask_lp.
                task_weight: 0.5              # The duty weight is 0.5.

For extra particulars about learn how to run graph multi-task studying with GraphStorm, check with Multi-task Studying in GraphStorm in our documentation.

New APIs to customise GraphStorm pipelines and elements

Since GraphStorm’s launch in early 2023, clients have primarily used its command line interface (CLI), which abstracts away the complexity of the graph ML pipeline so that you can rapidly construct, practice, and deploy fashions utilizing frequent recipes. Nevertheless, clients are telling us that they need an interface that enables them to customise the coaching and inference pipeline of GraphStorm to their particular necessities extra simply. Primarily based on buyer suggestions for the experimental APIs we launched in GraphStorm 0.2, GraphStorm 0.3 introduces refactored graph ML pipeline APIs. With the brand new APIs, you solely want 12 strains of code to outline a customized node classification coaching pipeline, as illustrated by the next 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 will help you get began with the brand new APIs, we even have launched new Jupyter pocket book examples in our Documentation and Tutorials web page.

Complete research of LM+GNN for giant graphs with wealthy textual content options

Many enterprise functions have graphs with textual content options. In retail search functions, for instance, buying log information gives insights on how text-rich product descriptions, search queries, and buyer conduct are associated. Foundational giant language fashions (LLMs) alone should not appropriate to mannequin such information as a result of the underlying information distributions and relationships don’t correspond to what LLMs study from their pre-training information corpuses. GML, alternatively, is nice for modeling associated information (graphs) however till now, GML practitioners needed to manually mix their GML fashions with LLMs to mannequin textual content options and get the perfect efficiency for his or her use circumstances. Particularly when the underlying graph dataset was giant, this handbook work was difficult and time-consuming.

In GraphStorm 0.2, GraphStorm launched built-in strategies to coach language fashions (LMs) and GNN fashions collectively effectively at scale on large text-rich graphs. Since then, clients have been asking us for steering on how GraphStorm’s LM+GNN strategies ought to be employed to optimize efficiency. To handle this, with GraphStorm 0.3, we launched a LM+GNN benchmark utilizing the big graph dataset, Microsoft Tutorial Graph (MAG), on two normal graph ML duties: node classification and hyperlink prediction. The graph dataset is a heterogeneous graph, comprises a whole bunch of thousands and thousands of nodes and billions of edges, and nearly all of nodes are attributed with wealthy textual content options. The detailed statistics of the datasets are proven within the following desk.

We benchmark two fundamental LM-GNN strategies in GraphStorm: pre-trained BERT+GNN, a baseline methodology that’s extensively adopted, and fine-tuned BERT+GNN, launched by GraphStorm builders in 2022. With the pre-trained BERT+GNN methodology, we first use a pre-trained BERT mannequin to compute embeddings for node textual content options after which practice a GNN mannequin for prediction. With the fine-tuned BERT+GNN methodology, we initially fine-tune the BERT fashions on the graph information and use the ensuing fine-tuned BERT mannequin to compute embeddings which might be then used to coach a GNN fashions for prediction. GraphStorm gives alternative ways to fine-tune the BERT fashions, relying on the duty varieties. For node classification, we fine-tune the BERT mannequin on the coaching set with the node classification duties; for hyperlink prediction, we fine-tune the BERT mannequin with the hyperlink prediction duties. Within the experiment, we use 8 r5.24xlarge situations for information processing and use 4 g5.48xlarge situations for mannequin coaching and inference. The fine-tuned BERT+GNN strategy has as much as 40% higher efficiency (hyperlink prediction on MAG) in comparison with pre-trained BERT+GNN.

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

We additionally benchmark GraphStorm on giant artificial graphs to showcase its scalability. 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 reveals the computation time of graph preprocessing, graph partition, and mannequin coaching. General, GraphStorm permits graph building and mannequin coaching on 100 billion scale graphs inside hours!

Extra benchmark particulars and outcomes can be found in our KDD 2024 paper.

Conclusion

GraphStorm 0.3 is printed below the Apache-2.0 license that can assist you deal with your large-scale graph ML challenges, and now gives native help for multi-task studying and new APIs to customise pipelines and different elements of GraphStorm. Check with the GraphStorm GitHub repository and documentation to get began.

Concerning the Creator

Xiang Track is a senior utilized scientist at AWS AI Analysis and Training (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 functionality of Neptune that makes use of graph neural networks for graphs saved in graph database. He’s now main the event of GraphStorm, an open-source graph machine studying framework for enterprise use circumstances. He obtained his Ph.D. in pc methods and structure on the Fudan College, Shanghai, in 2014.

Jian Zhang is a senior utilized scientist who has been utilizing machine studying strategies to assist clients clear up varied issues, comparable to fraud detection, ornament picture technology, and extra. He has efficiently developed graph-based machine studying, notably graph neural community, options for purchasers in China, USA, and Singapore. As an enlightener of AWS’s graph capabilities, Zhang has given many public displays in regards to the GNN, the Deep Graph Library (DGL), Amazon Neptune, and different AWS companies.

Florian Saupe is a Principal Technical Product Supervisor at AWS AI/ML analysis supporting science groups just like the graph machine studying group, and ML Methods groups engaged on giant scale distributed coaching, inference, and fault resilience. Earlier than becoming a member of AWS, Florian lead technical product administration for automated driving at Bosch, was a technique guide at McKinsey & Firm, and labored as a management methods/robotics scientist – a subject through which he holds a phd.