How AWS uses graph neural networks to meet customer needs

Information extraction, drug discovery, and software analysis are just a few applications of this versatile tool.

Graphs are an information-rich way to represent data. A graph consists of nodes — typically represented by circles — and edges — typically represented as line segments between nodes. In a knowledge graph, for instance, the nodes represent entities, and the edges represent relationships between them. In a social graph, the nodes represent people, and an edge indicates that two of those people know each other.

At Amazon Web Services, the use of machine learning (ML) to make the information encoded in graphs more useful to our customers has been a major research focus. In this post, we’ll showcase a variety of graph ML applications that customers have developed in collaboration with AWS scientists, from malicious-account detection and automated document processing to knowledge-graph-assisted drug discovery and protein property prediction.

Introduction to graph learning

Graphs can be homogenous, meaning the nodes represent a single type of entity (say, airports), and the edges represent a single type of relationship (say, scheduled flights). Or they can be heterogeneous, meaning they integrate multiple types of relationships among different entities, such as a graph of customers and products connected by both purchase histories and interests, or a knowledge graph of drugs, diseases, genes, and biological pathways connected by relationships such as indication and regulation. Nodes are often associated with data features, such as a product’s price or text description.

Heterogeneous knowledge graph
In a heterogenous knowledge graph, nodes can represent different classes of objects.

Graph neural networks

In the past 10 years, deep learning has revolutionized a host of AI applications, from natural-language processing to speech synthesis to computer vision.

Graph neural networks (GNNs) extend the performance benefits of deep learning to graph data. Like other popular neural networks, a GNN model has a series of layers, which progress toward higher levels of abstraction.

For instance, the first layer of a GNN computes a representation — or embedding — of the data represented by each node in the graph, while the second layer computes a representation of each node based on the prior embedding and the embeddings of the node’s nearest neighbors. In this way, every layer expands the scope of a node’s embedding, from one-hop neighbors, to two-hop neighbors, and for some applications, even further.

Graph neural network
A demonstration of how graph neural networks use recursive embedding to condense all the information in a two-hop graph into a single vector. Relationships between entities — such as "produce" and "write" in a movie database (red and yellow arrows, respectively) — are encoded in the level-0 embeddings of the entities themselves (red and orange blocks).
Stacy Reilly

GNN tasks

The individual node embeddings can then be used for node-level tasks, such as predicting properties of a node. The embeddings can also be used for higher-level inferences. For instance, using representations across a pair of nodes or across all nodes from the graph, GNNs can perform link-level or graph-level tasks, respectively.

Related content
Amazon’s George Karypis will give a keynote address on graph neural networks, a field in which “there is some fundamental theoretical stuff that we still need to understand.”

In this section, we demonstrate the versatility of GNNs across all three levels of tasks and examine how our customers are using GNNs to tackle a variety of problems.

Node-level tasks

Using GNNs, we can infer the behavior of an individual node in the graph based on the relationships it has to other nodes. One common task is node classification, where the objective is to infer nodes’ missing labels by looking at their neighbors’ labels and features. This method is used in applications such as financial-fraud detection, publication categorization, and disease classification.

In AWS, we have successfully used Amazon Neptune and Deep Graph Library (DGL) to apply GNN node representation learning to customers’ fraud detection use cases. For a large e-commerce sports gadgets customer, for instance, scientists in the Amazon Machine Learning Solutions Lab successfully used GNN models implemented in DGL to detect malicious accounts among billions of registered accounts.

Fraud graph.png
An example of how a graph representation can be used to detect fraud.

These malicious accounts were created in large quantities to abuse usage of promotional codes and block general public access to the vendor’s best-selling items. Using data from e-commerce sites, we built a massive heterogenous graph in which the nodes represented accounts and other entities, such as products purchased, and the edges connected nodes based on usage histories. To identify malicious accounts, we trained a GNN model to propagate labels from accounts that were known to be malicious to unlabeled accounts.

With this method, we were able to detect 10 times as many malicious accounts as a previous rule-based detection method could. Such performance improvements could not be achieved by traditional methods for doing machine learning on tabular datasets, such as CatBoost, which take only account features as inputs, without considering the relationships between accounts captured by the graph.

Besides applications for inherently relational, graph-structured data, such as social-network and citation-network data, there have been extensions of GNNs for data normally presented in Euclidean space, such as images and texts. By transforming data in Euclidean space to graphs based on spatial proximity, GNNs can solve problems that are typically solved by convolutional neural networks (CNNs) and recurrent neural networks (RNNs), which were designed to handle visual data and sequential data.

Related content
New method enables two- to 14-fold speedups over best-performing predecessors.

For example, researchers have explored GNN models to improve the accuracy of information extraction, a task typically handled by RNNs. GNNs turn out to be better at incorporating the nonlocal and nonsequential relationships captured by graph representations of word dependencies.

In a recent collaboration, the Amazon Machine Learning Solutions Lab and United Airlines developed a customized GNN model (DocGCN) to improve the accuracy of automatic information extraction from self-uploaded passenger documents, including travel documents, COVID-19 test results, and vaccine cards. The team built a graph for each scanned travel document that connected textual units based on their spatial proximities and orientations in the document.

Then, the DocGCN model reasoned over the relationships among textual units (nodes of the graph) to improve the identification of relevant textual information. DocGCN also generalized to complex forms with different formats by leveraging graphs to capture relationships between texts in tables, key-value pairs, and paragraphs. This improvement expedited the automation of international travel readiness verification.

Link-level tasks

Another important learning task in graphs is link prediction, which is central to applications such as product or ad recommendation and friendship suggestion. Given two nodes and a relation, the goal is to determine whether the nodes are connected by the relation.

Typically, the prediction is provided by a decoder that consumes the embeddings of the source and destination nodes, as in the work on knowledge graph embedding at scale that members of our team presented at SIGIR 2020. The decoder is trained to correctly predict existing edges in the graph.

DRKG.png
The high-level structure of DRKG. Numerals indicate the number of different types of relationships between classes of entities; terms between parentheses are examples of those relationships.
Credit: Glynis Condon

An exciting opportunity area in this context is drug discovery. AWS has recently provided a drug-repurposing knowledge graph (DRKG) that employs link prediction to identify new targets for existing drugs. Built by scientists at AWS, DRKG is a comprehensive biological knowledge graph that relates human genes, chemical compounds, biological processes, drug side effects, diseases, and symptoms. By performing link prediction around COVID-19 in DRKG, researchers were able to identify 41 drugs that were potentially effective against COVID-19 — 11 of which were already in clinical trials.

AWS also publicly released this solution, built by leveraging DRKG, as the COVID-19 Knowledge Graph (CKG). CKG organizes and represents the information in the COVID-19 Open Research Dataset (CORD-19), enabling fast discovery and prioritization of drug candidates. It can also be employed to identify papers relevant to COVID-19, thereby reducing the scale of human effort required to study, summarize, and interpret findings relevant to the pandemic.

Graph-level tasks

Graph-level tasks involve the analysis of large collections of small and independent graphs. A chemical library of organic compounds is a common example of a graph-level application, where each organic compound is represented as a graph of atoms connected by chemical bonds. Graph-level analyses of chemical libraries are often vital for drug development and discovery use cases; applications include predicting organic compounds’ chemical properties and predicting biological activities such as binding affinity to protein targets.

Code graph.png
An example of a program dependence graph.

Another example of data that can benefit from graph-level representation is code snippets in programming languages. A piece of code can be represented by a program dependence graph (PDG), where variables, operators, and statements are nodes connected by their dependencies (links).

At PAKDD 2021, we presented a new method for using GNNs to represent code snippets. Recently, we have been using that method to identify similar code snippets, to find opportunities to make code more modular and easier to maintain.

GNNs can also be used to encode global properties of the underlying systems and incorporate them into graph embeddings, in a way that is difficult with other deep-learning methods. We recently worked with scientists from Janssen Biopharmaceuticals to predict the function of proteins from their 3-D structure, which is useful for research and development in the pharmaceutical and biotech industries.

A protein is composed of a sequence of amino acids folded in a particular way. We developed a graph representation of proteins in which each node was an amino acid, and the interactions between amino acids in the folded protein structure determined whether two nodes were linked or not.

Protein graphs.png
Examples of graph representations of proteins.

This allowed us to encode fine-grained biological information, including the distance, angle, and direction of contact between neighboring amino acid residues. When we combined a GNN trained on these graph representations with a model trained to parse billions of protein sequences, we improved performance on various protein function prediction tasks of real-world importance.

Graph-level tasks for GNNs have different data-engineering requirements than the previous tasks. Node-level and link-level tasks usually operate on a single giant graph, whereas graph-level tasks operate on a large number of independent small graphs.

To help customers scale GNNs up for graph-level tasks, we developed a cloud-based architecture that leverages the highly performant open-source GNN library DGL, the ML resource orchestration tool SageMaker, and Amazon DocumentDB for managing graph data.

Getting started on your GNN journey

Related content
Approach that uses a hierarchical graph neural network improves F-score by 49% relative to predecessors.

In this article, we presented a few examples of GNN applications at all three levels of graph-related tasks to showcase the value of GNNs to various enterprise and research problems. AWS provides several options for customers looking to build and deploy GNN-powered ML solutions. Customers looking to get started quickly can use Amazon Neptune ML to build GNN models directly on graph data stored in Amazon Neptune without writing any code. Amazon Neptune ML can train models to tackle node-level and link-level tasks like those described above. Customers looking to get more hands-on can implement GNN models using DGL on Amazon SageMaker. In the meantime, we will continue to advance the science of GNNs to build more products and solutions to make GNNs more accessible to all our customers.

Acknowledgments: Guang Yang, Soji Adeshina, Jasleen Grewal, Miguel Romero Calvo, Suchitra Sathyanarayana

Research areas

Related content

US, CA, San Francisco
About Twitch Launched in 2011, Twitch is a global community that comes together each day to create multiplayer entertainment: unique, live, unpredictable experiences created by the interactions of millions. We bring the joy of co-op to everything, from casual gaming to world-class esports to anime marathons, music, and art streams. Twitch also hosts TwitchCon, where we bring everyone together to celebrate, learn, and grow their personal interests and passions. We’re always live at Twitch. Stay up to date on all things Twitch on Linkedin, Twitter and on our Blog. About the role: Twitch builds data-driven machine learning solutions across several rich problem spaces: Natural Language Processing (NLP), Recommendations, Semantic Search, Classification/Categorization, Anomaly Detection, Forecasting, Safety, and HCI/Social Computing/Computational Social Science. As an Intern, you will work with a dedicated Mentor and Manager on a project in one of these problem areas. You will also be supported by an Advisor and participate in cohort activities such as research teach backs and leadership talks. This position can also be located in San Francisco, CA or virtual. You Will: Solve large-scale data problems. Design solutions for Twitch's problem spaces Explore ML and data research
US, CA, San Francisco
About Twitch Launched in 2011, Twitch is a global community that comes together each day to create multiplayer entertainment: unique, live, unpredictable experiences created by the interactions of millions. We bring the joy of co-op to everything, from casual gaming to world-class esports to anime marathons, music, and art streams. Twitch also hosts TwitchCon, where we bring everyone together to celebrate, learn, and grow their personal interests and passions. We’re always live at Twitch. Stay up to date on all things Twitch on Linkedin, Twitter and on our Blog. About the role: Twitch builds data-driven machine learning solutions across several rich problem spaces: Natural Language Processing (NLP), Recommendations, Semantic Search, Classification/Categorization, Anomaly Detection, Forecasting, Safety, and HCI/Social Computing/Computational Social Science. As an Intern, you will work with a dedicated Mentor and Manager on a project in one of these problem areas. You will also be supported by an Advisor and participate in cohort activities such as research teach backs and leadership talks. This position can also be located in San Francisco, CA or virtual. You Will: Solve large-scale data problems. Design solutions for Twitch's problem spaces Explore ML and data research
LU, Luxembourg
Are you a talented and inventive scientist with a strong passion about modern data technologies and interested to improve business processes, extracting value from the data? Would you like to be a part of an organization that is aiming to use self-learning technology to process data in order to support the management of the procurement function? The Global Procurement Technology, as a part of Global Procurement Operations, is seeking a skilled Data Scientist to help build its future data intelligence in business ecosystem, working with large distributed systems of data and providing Machine Learning (ML) and Predictive Modeling expertise. You will be a member of the Data Engineering and ML Team, joining a fast-growing global organization, with a great vision to transform the Procurement field, and become the role model in the market. This team plays a strategic role supporting the core Procurement business domains as well as it is the cornerstone of any transformation and innovation initiative. Our mission is to provide a high-quality data environment to facilitate process optimization and business digitalization, on a global scale. We are supporting business initiatives, including but not limited to, strategic supplier sourcing (e.g. contracting, negotiation, spend analysis, market research, etc.), order management, supplier performance, etc. We are seeking an individual who can thrive in a fast-paced work environment, be collaborative and share knowledge and experience with his colleagues. You are expected to deliver results, but at the same time have fun with your teammates and enjoy working in the company. In Amazon, you will find all the resources required to learn new skills, grow your career, and become a better professional. You will connect with world leaders in your field and you will be tackling Data Science challenges to ensure business continuity, by taking the right decisions for your customers. As a Data Scientist in the team, you will: -be the subject matter expert to support team strategies that will take Global Procurement Operations towards world-class predictive maintenance practices and processes, driving more effective procurement functions, e.g. supplier segmentation, negotiations, shipping supplies volume forecast, spend management, etc. -have strong analytical skills and excel in the design, creation, management, and enterprise use of large data sets, combining raw data from different sources -provide technical expertise to support the development of ML models to facilitate intelligent digital services, such as Contract Lifecycle Management (CLM) and Negotiations platform -cooperate closely with different groups of stakeholders, e.g. data/software engineers, product/program managers, analysts, senior leadership, etc. to evaluate business needs and objectives to set up the best data management environment -create and share with audiences of varying levels technical papers and presentations -deal with ambiguity, prioritizing needs, and delivering results in a dynamic environment Basic qualifications -Master’s Degree in Computer Science/Engineering, Informatics, Mathematics, or a related technical discipline -3+ years of industry experience in data engineering/science, business intelligence or related field -3+ years experience in algorithm design, engineering and implementation for very-large scale applications to solve real problems -Very good knowledge of data modeling and evaluation -Very good understanding of regression modeling, forecasting techniques, time series analysis, machine-learning concepts such as supervised and unsupervised learning, classification, random forest, etc. -SQL and query performance tuning skills Preferred qualifications -2+ years of proficiency in using R, Python, Scala, Java or any modern language for data processing and statistical analysis -Experience with various RDBMS, such as PostgreSQL, MS SQL Server, MySQL, etc. -Experience architecting Big Data and ML solutions with AWS products (Redshift, DynamoDB, Lambda, S3, EMR, SageMaker, Lex, Kendra, Forecast etc.) -Experience articulating business questions and using quantitative techniques to arrive at a solution using available data -Experience with agile/scrum methodologies and its benefits of managing projects efficiently and delivering results iteratively -Excellent written and verbal communication skills including data visualization, especially in regards to quantitative topics discussed with non-technical colleagues
US, CA, San Francisco
About Twitch Launched in 2011, Twitch is a global community that comes together each day to create multiplayer entertainment: unique, live, unpredictable experiences created by the interactions of millions. We bring the joy of co-op to everything, from casual gaming to world-class esports to anime marathons, music, and art streams. Twitch also hosts TwitchCon, where we bring everyone together to celebrate, learn, and grow their personal interests and passions. We’re always live at Twitch. Stay up to date on all things Twitch on Linkedin, Twitter and on our Blog. About the role: Twitch builds data-driven machine learning solutions across several rich problem spaces: Natural Language Processing (NLP), Recommendations, Semantic Search, Classification/Categorization, Anomaly Detection, Forecasting, Safety, and HCI/Social Computing/Computational Social Science. As an Intern, you will work with a dedicated Mentor and Manager on a project in one of these problem areas. You will also be supported by an Advisor and participate in cohort activities such as research teach backs and leadership talks. This position can also be located in San Francisco, CA or virtual. You Will: Solve large-scale data problems. Design solutions for Twitch's problem spaces Explore ML and data research
US, CA, San Francisco
About Twitch Launched in 2011, Twitch is a global community that comes together each day to create multiplayer entertainment: unique, live, unpredictable experiences created by the interactions of millions. We bring the joy of co-op to everything, from casual gaming to world-class esports to anime marathons, music, and art streams. Twitch also hosts TwitchCon, where we bring everyone together to celebrate, learn, and grow their personal interests and passions. We’re always live at Twitch. Stay up to date on all things Twitch on Linkedin, Twitter and on our Blog. About the role: Twitch builds data-driven machine learning solutions across several rich problem spaces: Natural Language Processing (NLP), Recommendations, Semantic Search, Classification/Categorization, Anomaly Detection, Forecasting, Safety, and HCI/Social Computing/Computational Social Science. As an Intern, you will work with a dedicated Mentor and Manager on a project in one of these problem areas. You will also be supported by an Advisor and participate in cohort activities such as research teach backs and leadership talks. This position can also be located in San Francisco, CA or virtual. You Will: Solve large-scale data problems. Design solutions for Twitch's problem spaces Explore ML and data research
US, CA, San Francisco
About Twitch Launched in 2011, Twitch is a global community that comes together each day to create multiplayer entertainment: unique, live, unpredictable experiences created by the interactions of millions. We bring the joy of co-op to everything, from casual gaming to world-class esports to anime marathons, music, and art streams. Twitch also hosts TwitchCon, where we bring everyone together to celebrate, learn, and grow their personal interests and passions. We’re always live at Twitch. Stay up to date on all things Twitch on Linkedin, Twitter and on our Blog. About the role: Twitch builds data-driven machine learning solutions across several rich problem spaces: Natural Language Processing (NLP), Recommendations, Semantic Search, Classification/Categorization, Anomaly Detection, Forecasting, Safety, and HCI/Social Computing/Computational Social Science. As an Intern, you will work with a dedicated Mentor and Manager on a project in one of these problem areas. You will also be supported by an Advisor and participate in cohort activities such as research teach backs and leadership talks. This position can also be located in San Francisco, CA or virtual. You Will: Solve large-scale data problems. Design solutions for Twitch's problem spaces Explore ML and data research
US, CA, San Francisco
About Twitch Launched in 2011, Twitch is a global community that comes together each day to create multiplayer entertainment: unique, live, unpredictable experiences created by the interactions of millions. We bring the joy of co-op to everything, from casual gaming to world-class esports to anime marathons, music, and art streams. Twitch also hosts TwitchCon, where we bring everyone together to celebrate, learn, and grow their personal interests and passions. We’re always live at Twitch. Stay up to date on all things Twitch on Linkedin, Twitter and on our Blog. About the role: Twitch builds data-driven machine learning solutions across several rich problem spaces: Natural Language Processing (NLP), Recommendations, Semantic Search, Classification/Categorization, Anomaly Detection, Forecasting, Safety, and HCI/Social Computing/Computational Social Science. As an Intern, you will work with a dedicated Mentor and Manager on a project in one of these problem areas. You will also be supported by an Advisor and participate in cohort activities such as research teach backs and leadership talks. This position can also be located in San Francisco, CA or virtual. You Will: Solve large-scale data problems. Design solutions for Twitch's problem spaces Explore ML and data research
US, WA, Seattle
We are a team of doers working passionately to apply cutting-edge advances in deep learning in the life sciences to solve real-world problems. As a Senior Applied Science Manager you will participate in developing exciting products for customers. Our team rewards curiosity while maintaining a laser-focus in bringing products to market. Competitive candidates are responsive, flexible, and able to succeed within an open, collaborative, entrepreneurial, startup-like environment. At the leading edge of both academic and applied research in this product area, you have the opportunity to work together with a diverse and talented team of scientists, engineers, and product managers and collaborate with others teams. Location is in Seattle, US Embrace Diversity Here at Amazon, we embrace our differences. We are committed to furthering our culture of inclusion. We have ten employee-led affinity groups, reaching 40,000 employees in over 190 chapters globally. We have innovative benefit offerings, and host annual and ongoing learning experiences, including our Conversations on Race and Ethnicity (CORE) and AmazeCon (gender diversity) conferences. Amazon’s culture of inclusion is reinforced within our 14 Leadership Principles, which remind team members to seek diverse perspectives, learn and be curious, and earn trust Balance Work and Life Our team puts a high value on work-life balance. It isn’t about how many hours you spend at home or at work; it’s about the flow you establish that brings energy to both parts of your life. We believe striking the right balance between your personal and professional life is critical to life-long happiness and fulfillment. We offer flexibility in working hours and encourage you to find your own balance between your work and personal lives Mentor & Grow Careers Our team is dedicated to supporting new members. We have a broad mix of experience levels and tenures, and we’re building an environment that celebrates knowledge sharing and mentorship. Our senior members enjoy one-on-one mentoring and thorough, but kind, code reviews. We care about your career growth and strive to assign projects based on what will help each team member develop into a better-rounded engineer and enable them to take on more complex tasks in the future. Key job responsibilities • Manage high performing engineering and science teams • Hire and develop top-performing engineers, scientists, and other managers • Develop and execute on project plans and delivery commitments • Work with business, data science, software engineer, biological, and product leaders to help define product requirements and with managers, scientists, and engineers to execute on them • Build and maintain world-class customer experience and operational excellence for your deliverables
US, Virtual
The Amazon Economics Team is hiring Interns in Economics. We are looking for detail-oriented, organized, and responsible individuals who are eager to learn how to work with large and complicated data sets. Some knowledge of econometrics, as well as basic familiarity with Stata, R, or Python is necessary. Experience with SQL, UNIX, Sawtooth, and Spark would be a plus. These are full-time positions at 40 hours per week, with compensation being awarded on an hourly basis. You will learn how to build data sets and perform applied econometric analysis at Internet speed collaborating with economists, data scientists and MBAʼs. These skills will translate well into writing applied chapters in your dissertation and provide you with work experience that may help you with placement. Roughly 85% of interns from previous cohorts have converted to full time economics employment at Amazon. If you are interested, please send your CV to our mailing list at econ-internship@amazon.com.
US, WA, Seattle
Amazon internships are full-time (40 hours/week) for 12 consecutive weeks with start dates in May - July 2023. Our internship program provides hands-on learning and building experiences for students who are interested in a career in hardware engineering. This role will be based in Seattle, and candidates must be willing to work in-person. Corporate Projects (CPT) is a team that sits within the broader Corporate Development organization at Amazon. We seek to bring net-new, strategic projects to life by working together with customers and evolving projects from ZERO-to-ONE. To do so, we deploy our resources towards proofs-of-concept (POCs) and pilot programs and develop them from high-level ideas (the ZERO) to tangible short-term results that provide validating signal and a path to scale (the ONE). We work with our customers to develop and create net-new opportunities by relentlessly scouring all of Amazon and finding new and innovative ways to strengthen and/or accelerate the Amazon Flywheel. CPT seeks an Applied Science intern to work with a diverse, cross-functional team to build new, innovative customer experiences. Within CPT, you will apply both traditional and novel scientific approaches to solve and scale problems and solutions. We are a team where science meets application. A successful candidate will be a self-starter comfortable with ambiguity, strong attention to detail, and the ability to work in a fast-paced, ever-changing environment. As an Applied Science Intern, you will own the design and development of end-to-end systems. You’ll have the opportunity to create technical roadmaps, and drive production level projects that will support Amazon Science. You will work closely with Amazon scientists, and other science interns to develop solutions and deploy them into production. The ideal scientist must have the ability to work with diverse groups of people and cross-functional teams to solve complex business problems.