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

CA, ON, Toronto
Are you motivated to explore research in ambiguous spaces? Are you interested in conducting research that will improve associate, employee and manager experiences at Amazon? Do you want to work on an interdisciplinary team of scientists that collaborate rather than compete? Join us at PXT Central Science! The People eXperience and Technology Central Science Team (PXTCS) uses economics, behavioral science, statistics, and machine learning to proactively identify mechanisms and process improvements which simultaneously improve Amazon and the lives, wellbeing, and the value of work to Amazonians. We are an interdisciplinary team that combines the talents of science and engineering to develop and deliver solutions that measurably achieve this goal. Key job responsibilities As an Applied Scientist for People Experience and Technology (PXT) Central Science, you will be working with our science and engineering teams, specifically on re-imagining Generative AI Applications and Generative AI Infrastructure for HR. Applying Generative AI to HR has unique challenges such as privacy, fairness, and seamlessly integrating Enterprise Knowledge and World Knowledge and knowing which to use when. In addition, the team works on some of Amazon’s most strategic technical investments in the people space and support Amazon’s efforts to be Earth’s Best Employer. In this role you will have a significant impact on 1.5 million Amazonians and the communities Amazon serves and ample scope to demonstrate scientific thought leadership and scientific impact in addition to business impact. You will also play a critical role in the organization's business planning, work closely with senior leaders to develop goals and resource requirements, influence our long-term technical and business strategy, and help hire and develop science and engineering talent. You will also provide support to business partners, helping them use the best scientific methods and science-driven tools to solve current and upcoming challenges and deliver efficiency gains in a changing marke About the team The AI/ML team in PXTCS is working on building Generative AI solutions to reimagine Corp employee and Ops associate experience. Examples of state-of-the-art solutions are Coaching for Amazon employees (available on AZA) and reinventing Employee Recruiting and Employee Listening.
CA, ON, Toronto
Conversational AI ModEling and Learning (CAMEL) team is part of Amazon Devices organization where our mission is to build a best-in-class Conversational AI that is intuitive, intelligent, and responsive, by developing superior Large Language Models (LLM) solutions and services which increase the capabilities built into the model and which enable utilizing thousands of APIs and external knowledge sources to provide the best experience for each request across millions of customers and endpoints. We are looking for a passionate, talented, and resourceful Applied Scientist in the field of LLM, Artificial Intelligence (AI), Natural Language Processing (NLP), Recommender Systems and/or Information Retrieval, to invent and build scalable solutions for a state-of-the-art context-aware conversational AI. A successful candidate will have strong machine learning background and a desire to push the envelope in one or more of the above areas. The ideal candidate would also have hands-on experiences in building Generative AI solutions with LLMs, enjoy operating in dynamic environments, be self-motivated to take on challenging problems to deliver big customer impact, moving fast to ship solutions and then iterating on user feedback and interactions. Key job responsibilities As an Applied Scientist, you will leverage your technical expertise and experience to collaborate with other talented applied scientists and engineers to research and develop novel algorithms and modeling techniques to reduce friction and enable natural and contextual conversations. You will analyze, understand and improve user experiences by leveraging Amazon’s heterogeneous data sources and large-scale computing resources to accelerate advances in artificial intelligence. You will work on core LLM technologies, including Prompt Engineering and Optimization, Supervised Fine-Tuning, Learning from Human Feedback, Evaluation, Self-Learning, etc. Your work will directly impact our customers in the form of novel products and services.
CA, ON, Toronto
Conversational AI ModEling and Learning (CAMEL) team is part of Amazon Devices organization where our mission is to build a best-in-class Conversational AI that is intuitive, intelligent, and responsive, by developing superior Large Language Models (LLM) solutions and services which increase the capabilities built into the model and which enable utilizing thousands of APIs and external knowledge sources to provide the best experience for each request across millions of customers and endpoints. We are looking for a passionate, talented, and resourceful Applied Scientist in the field of LLM, Artificial Intelligence (AI), Natural Language Processing (NLP), Recommender Systems and/or Information Retrieval, to invent and build scalable solutions for a state-of-the-art context-aware conversational AI. A successful candidate will have strong machine learning background and a desire to push the envelope in one or more of the above areas. The ideal candidate would also have hands-on experiences in building Generative AI solutions with LLMs, enjoy operating in dynamic environments, be self-motivated to take on challenging problems to deliver big customer impact, moving fast to ship solutions and then iterating on user feedback and interactions. Key job responsibilities As an Applied Scientist, you will leverage your technical expertise and experience to collaborate with other talented applied scientists and engineers to research and develop novel algorithms and modeling techniques to reduce friction and enable natural and contextual conversations. You will analyze, understand and improve user experiences by leveraging Amazon’s heterogeneous data sources and large-scale computing resources to accelerate advances in artificial intelligence. You will work on core LLM technologies, including Prompt Engineering and Optimization, Supervised Fine-Tuning, Learning from Human Feedback, Evaluation, Self-Learning, etc. Your work will directly impact our customers in the form of novel products and services.
US, CA, San Diego
Do you want to join an innovative team of scientists who use machine learning and statistical techniques to help Amazon provide the best customer experience by preventing eCommerce fraud? Are you excited by the prospect of analyzing and modeling terabytes of data and creating state-of-the-art algorithms to solve real world problems? Do you like to own end-to-end business problems/metrics and directly impact the profitability of the company? Do you enjoy collaborating in a diverse team environment? If yes, then you may be a great fit to join the Amazon Buyer Risk Prevention (BRP) Machine Learning group. We are looking for a talented scientist who is passionate to build advanced algorithmic systems that help manage safety of millions of transactions every day. Key job responsibilities Use machine learning and statistical techniques to create scalable risk management systems Learning and understanding large amounts of Amazon’s historical business data for specific instances of risk or broader risk trends Design, development and evaluation of highly innovative models for risk management Working closely with software engineering teams to drive real-time model implementations and new feature creations Working closely with operations staff to optimize risk management operations, Establishing scalable, efficient, automated processes for large scale data analyses, model development, model validation and model implementation Tracking general business activity and providing clear, compelling management reporting on a regular basis Research and implement novel machine learning and statistical approaches
US, MA, Boston
The Artificial General Intelligence (AGI) team is looking for a passionate, talented, and inventive Applied Scientist with a strong deep learning background, to build industry-leading Generative Artificial Intelligence (GenAI) technology with Large Language Models (LLMs) and multimodal systems. Key job responsibilities As a Applied Scientist with the AGI team, you will work with talented peers to lead the development of novel algorithms and modeling techniques, to advance the state of the art with LLMs. Your work will directly impact our customers in the form of products and services that make use of speech and language technology. You will leverage Amazon’s heterogeneous data sources and large-scale computing resources to accelerate advances in spoken language understanding. About the team The AGI team has a mission to push the envelope in GenAI with LLMs and multimodal systems, in order to provide the best-possible experience for our customers.
US, WA, Seattle
The XCM (Cross Channel Cross-Category Marketing) team seeks an Applied Scientist to revolutionize our marketing strategies. XCM's mission is to build the most measurably effective, creatively impactful, and cross-channel campaigning capabilities possible, with the aim of growing "big-bet" programs, strengthening positive brand perceptions, and increasing long-term free cash flow. As a science team, we're tackling complex challenges in marketing incrementality measurement, optimization and audience segmentation. In this role, you'll collaborate with a diverse team of scientists and economists to build and enhance causal measurement, optimization and prediction models for Amazon's global multi-billion dollar fixed marketing budget. You'll also work closely with various teams to develop scientific roadmaps, drive innovation, and influence key resource allocation decisions. Key job responsibilities 1) Innovating scalable marketing methodologies using causal inference and machine learning. 2) Developing interpretable models that provide actionable business insights. 3) Collaborating with engineers to automate and scale scientific solutions. 4) Engaging with stakeholders to ensure effective adoption of scientific products. 5) Presenting findings to the Amazon Science community to promote excellence and knowledge-sharing.
US, WA, Seattle
Do you want to join an innovative team of scientists who use machine learning and statistical techniques to help Amazon provide the best customer experience by preventing eCommerce fraud? Are you excited by the prospect of analyzing and modeling terabytes of data and creating state-of-the-art algorithms to solve real world problems? Do you like to own end-to-end business problems/metrics and directly impact the profitability of the company? Do you enjoy collaborating in a diverse team environment? If yes, then you may be a great fit to join the Amazon Buyer Risk Prevention (BRP) Machine Learning group. We are looking for a talented scientist who is passionate to build advanced algorithmic systems that help manage safety of millions of transactions every day. Key job responsibilities Use machine learning and statistical techniques to create scalable risk management systems Learning and understanding large amounts of Amazon’s historical business data for specific instances of risk or broader risk trends Design, development and evaluation of highly innovative models for risk management Working closely with software engineering teams to drive real-time model implementations and new feature creations Working closely with operations staff to optimize risk management operations, Establishing scalable, efficient, automated processes for large scale data analyses, model development, model validation and model implementation Tracking general business activity and providing clear, compelling management reporting on a regular basis Research and implement novel machine learning and statistical approaches
US, WA, Seattle
The Global Cross-Channel and Cross- Category Marketing (XCM) org are seeking an experienced Economist to join our team. XCM’s mission is to be the most measurably effective and creatively breakthrough marketing organization in the world in order to strengthen the brand, grow the business, and reduce cost for Amazon overall. We achieve this through scaled campaigning in support of brands, categories, and audiences which aim to create the maximum incremental impact for Amazon as a whole by driving the Amazon flywheel. This is a high impact role with the opportunities to lead the development of state-of-the-art, scalable models to measure the efficacy and effectiveness of a new marketing channel. In this critical role, you will leverage your deep expertise in causal inference to design and implement robust measurement frameworks that provide actionable insights to drive strategic business decisions. Key Responsibilities: Develop advanced econometric and statistical models to rigorously evaluate the causal incremental impact of marketing campaigns on customer perception and customer behaviors. Collaborate cross-functionally with marketing, product, data science and engineering teams to define the measurement strategy and ensure alignment on objectives. Leverage large, complex datasets to uncover hidden patterns and trends, extracting meaningful insights that inform marketing optimization and investment decisions. Work with engineers, applied scientists and product managers to automate the model in production environment. Stay up-to-date with the latest research and methodological advancements in causal inference, causal ML and experiment design to continuously enhance the team's capabilities. Effectively communicate analysis findings, recommendations, and their business implications to key stakeholders, including senior leadership. Mentor and guide junior economists, fostering a culture of analytical excellence and innovation.
US, WA, Seattle
We are open to hiring candidates to work out of one of the following locations: Seattle, WA, USA Do you love using data to solve complex problems? Are you interested in innovating and developing world-class big data solutions? We have the career for you! EPP Analytics team is seeking an exceptional Data Scientist to recommend, design and deliver new advanced analytics and science innovations end-to-end partnering closely with our security/software engineers, and response investigators. Your work enables faster data-driven decision making for Preventive and Response teams by providing them with data management tools, actionable insights, and an easy-to-use reporting experience. The ideal candidate will be passionate about working with big data sets and have the expertise to utilize these data sets to derive insights, drive science roadmap and foster growth. Key job responsibilities - As a Data Scientist (DS) in EPP Analytics, you will do causal data science, build predictive models, conduct simulations, create visualizations, and influence data science practice across the organization. - Provide insights by analyzing historical data - Create experiments and prototype implementations of new learning algorithms and prediction techniques. - Research and build machine learning algorithms that improve Insider Threat risk A day in the life No two days are the same in Insider Risk teams - the nature of the work we do and constantly shifting threat landscape means sometimes you'll be working with an internal service team to find anomalous use of their data, other days you'll be working with IT teams to build improved controls. Some days you'll be busy writing detections, or mentoring or running design review meetings. The EPP Analytics team is made up of SDEs and Security Engineers who partner with Data Scientists to create big data solutions and continue to raise the bar for the EPP organization. As a member of the team you will have the opportunity to work on challenging data modeling solutions, new and innovative Quicksight based reporting, and data pipeline and process improvement projects. About the team Diverse Experiences Amazon Security values diverse experiences. Even if you do not meet all of the qualifications and skills listed in the job description, we encourage candidates to apply. If your career is just starting, hasn’t followed a traditional path, or includes alternative experiences, don’t let it stop you from applying. Why Amazon Security? At Amazon, security is central to maintaining customer trust and delivering delightful customer experiences. Our organization is responsible for creating and maintaining a high bar for security across all of Amazon’s products and services. We offer talented security professionals the chance to accelerate their careers with opportunities to build experience in a wide variety of areas including cloud, devices, retail, entertainment, healthcare, operations, and physical stores Inclusive Team Culture In Amazon Security, it’s in our nature to learn and be curious. Ongoing DEI events and learning experiences inspire us to continue learning and to embrace our uniqueness. Addressing the toughest security challenges requires that we seek out and celebrate a diversity of ideas, perspectives, and voices. Training & Career Growth We’re continuously raising our performance bar as we strive to become Earth’s Best Employer. That’s why you’ll find endless knowledge-sharing, training, and other career-advancing resources here to help you develop into a better-rounded professional. Work/Life Balance We value work-life harmony. Achieving success at work should never come at the expense of sacrifices at home, which is why flexible work hours and arrangements are part of our culture. When we feel supported in the workplace and at home, there’s nothing we can’t achieve.
IN, KA, Bengaluru
Do you want to join an innovative team of scientists who use machine learning and statistical techniques to create state-of-the-art solutions for providing better value to Amazon’s customers? Do you want to build and deploy advanced algorithmic systems that help optimize millions of transactions every day? Are you excited by the prospect of analyzing and modeling terabytes of data to solve real world problems? Do you like to own end-to-end business problems/metrics and directly impact the profitability of the company? Do you like to innovate and simplify? If yes, then you may be a great fit to join the Machine Learning and Data Sciences team for India Consumer Businesses. If you have an entrepreneurial spirit, know how to deliver, love to work with data, are deeply technical, highly innovative and long for the opportunity to build solutions to challenging problems that directly impact the company's bottom-line, we want to talk to you. Major responsibilities - Use machine learning and analytical techniques to create scalable solutions for business problems - Analyze and extract relevant information from large amounts of Amazon’s historical business data to help automate and optimize key processes - Design, development, evaluate and deploy innovative and highly scalable models for predictive learning - Research and implement novel machine learning and statistical approaches - Work closely with software engineering teams to drive real-time model implementations and new feature creations - Work closely with business owners and operations staff to optimize various business operations - Establish scalable, efficient, automated processes for large scale data analyses, model development, model validation and model implementation - Mentor other scientists and engineers in the use of ML techniques