Amazon’s open-source tools make embedding knowledge graphs much more efficient

Tools include optimizations for multicore, multiple-GPU, and distributed-training settings.

Knowledge graphs are data structures consisting of entities (the nodes of the graph) and relationships between them (the edges, usually depicted as line segments connecting nodes). The entities “Nile River” and “Africa”, for instance, might be connected in a knowledge graph by the relationship “is in”.

At Amazon, we use knowledge graphs to represent relationships between products and to encode information for Alexa’s question-answering service, among other things. Recently, we also built a knowledge graph that represented medical and biological information to help find treatments for COVID-19.

Today, many applications of knowledge graphs involve knowledge graph embedding, or representing the entities and relations in a knowledge graph as points in a vector space. To make knowledge graph embeddings easier to use, our team has released a set of tools called DGL-KE, for deep-graph-learning knowledge embeddings.

Last week at SIGIR, the Association for Computing Machinery’s annual conference on information retrieval, we presented a paper describing a set of optimizations that enable DGL-KE to run more efficiently in parallel-computing environments.

In experiments, our optimizations enabled DGL-KE to compute knowledge graph embeddings two to five times more efficiently than previous parallel-computing approaches did.

Geometric interpretation

The data encoded in a knowledge graph can be represented as sets of triplets; each triplet consists of a head entity, a relation, and a tail entity. For instance, in the triplet [Nile River | is in | Africa], “Nile River” is the head, “is in” the relation, and “Africa” the tail.

The standard way to embed a knowledge graph is with a machine learning model that takes entities and relations as inputs, and outputs fixed-length vectors. During training, the model is evaluated according to some score that relates the embeddings of the heads, relations, and tails of triplets.

For instance, the score might indicate how well the sum of the head and relation vectors approximates the tail vector. The goal of training is to maximize the score for triplets that actually occur in the graph and minimize it for triplets that don’t.

Today, a knowledge graph might include millions of entities and billions of relations between them. Embedding that much data would be impractical without parallelization.

In our paper, we consider three different parallel-computing settings: (1) multicore machines, with multiple CPUs; (2) GPU machines, with multiple CPUs and multiple GPUs; and (3) computing clusters, in which the training process is broken up across machines.

In each of these settings, DGL-KE organizes data storage and access differently. In multicore training, the complete knowledge graph is stored in main memory; in GPU training, entities are stored in main memory, but relations are stored in GPU memory; and in distributed training, the knowledge graph is divided up across machines, with a key-value database (KVStore, in the figure below) managing data accesses.

DGL-KE implementations.png
DGL-KE organizes data storage and access differently for different distributed-computing settings.

DGL-KE applies different optimizations in different settings, but they all share the goal, typical of parallel-computing schemes, of minimizing communication between parallel computing resources and storing data close to the processes that will use it.

The optimizations include

1. Graph partitioning: This optimization is specific to distributed training (training across multiple machines). To split the graph up for distribution across machines, we use METIS, a software suite that implements a minimum-cut algorithm. Minimum-cut is the problem of finding the smallest number of edge deletions that will divide a graph into two parts. By minimizing the number of edges linking chunks of the graph assigned to different machines, we reduce the need for communication between machines. METIS was developed by the academic lab of our team leader, Amazon senior principal scientist George Karypis, who is also a professor of computer science at the University of Minnesota.

2. Negative sampling: For each valid triplet extracted from the knowledge graph, we artificially create multiple negative examples — typically around 200 — by substituting spurious entities for its head or tail. This is computationally costly, because retrieving the spurious substitutions requires multiple calls to main memory. To reduce that cost, we use the same set of substitute entities for each valid triplet in a group. If the group has 100 members, this reduces the number of substitute entities we have to retrieve by 99%.

3. Relation partitioning: This optimization applies to both distributed training and multiple-GPU training, but it’s particularly important in the second case. To minimize communication overhead, we attempt to assign all relations of a particular type to only one GPU. We make this assignment iteratively, in a greedy fashion. At each iteration, we assign the most common relation type remaining to the GPU with the most available memory. This approach does reduce the randomness in the training procedure, which could make the resulting model less accurate. But we avoid that consequence by ensuring different distributions of the relations during each training epoch. (Knowledge-graph-embedding models, like most machine learning models, are trained multiple times on the same data set; each pass through the data is an epoch.)

4. Overlapping computations: In the GPU machines, CPUs have two tasks: assigning batches of training data to GPUs and updating the embeddings stored in main memory. The GPUs have one task: processing their batches of training data to compute the gradients that the CPUs will use to update the embeddings. To maximize efficiency, we keep both CPUs and GPUs busy by overlapping these computations. While the CPUs are updating embeddings based on one batch of data, the GPUs have already begun computing gradients for the next batch.

In experiments, we compared a distributed-training scheme using our optimizations to two existing training schemes, using five different methods for scoring embeddings during training. On average, our method yielded a twofold speedup relative to one baseline and a fivefold speedup relative to the other.

Related content

US, WA, Seattle
We are looking for detail-oriented, organized, and responsible individuals who are eager to learn how to work with large and complicated data sets. Knowledge of econometrics, as well as basic familiarity with Python (or R, Matlab, or equivalent) is necessary, and experience with SQL 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, scientists, and product managers. 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 previous cohorts have converted to full time scientist employment at Amazon. If you are interested, please send your CV to our mailing list at econ-internship@amazon.com.
US, WA, Bellevue
As an applied scientist, you will use your experience to initiate the design, development, execution and implementation of scientific research projects. Working closely with fellow research scientists and product managers, you will use your experience in modeling, statistics, and simulation to design models of new policies, simulate their performance, and evaluate their benefits and impacts to cost, reliability, and speed of our fulfillment network. Our teams are looking for experience in network and combinatorial optimization, algorithms, data structures, statistics, and/or machine learning. This position requires superior analytical thinking, and ability to apply their technical and statistical knowledge to identify opportunities for real world applications. You should be able to mine and analyze large data, and be able to use necessary programming and statistical analysis software/tools to do so. Amazon has positions available for Research Scientists in multiple locations across the US and Canada.
US, WA, Virtual Contact Center-WA
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 Python is necessary, and experience with SQL and UNIX 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, scientists, and product managers. 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 previous cohorts have converted to full time scientist employment at Amazon. If you are interested, please send your CV to our mailing list at econ-internship@amazon.com. About the team The Selling Partner Fees team owns the end-to-end fees experience for two million active third party sellers. We own the fee strategy, fee seller experience, fee accuracy and integrity, fee science and analytics, and we provide scalable technology to monetize all services available to third-party sellers. Within the Science team, our goal is to understand the impact of changing fees on Seller (supply) and Customers (demand) behavior (e.g. price changes, advertising strategy changes, introducing new selection etc.) as well as using this information to optimize our fee structure and maximizing our long term profitability.
US, WA, Seattle
This is a unique opportunity to build technology and science that millions of people will use every day. Are you excited about working on large scale Natural Language Processing (NLP), Machine Learning (ML), and Deep Learning (DL)? We are embarking on a multi-year journey to improve the shopping experience for customers globally. Amazon Search team creates customer-focused search solutions and technologies that makes shopping delightful and effortless for our customers. Our goal is to understand what customers are looking for in whatever language happens to be their choice at the moment and help them find what they need in Amazon's vast catalog of billions of products. As Amazon expands to new geographies, we are faced with the unique challenge of maintaining the bar on Search Quality due to the diversity in user preferences, multilingual search and data scarcity in new locales. We are looking for an applied researcher to work on improving search on Amazon using NLP, ML, and DL technology. As an Applied Scientist, you will lead our efforts in query understanding, semantic matching (e.g. is a drone the same as quadcopter?), relevance ranking (what is a "funny halloween costume"?), language identification (did the customer just switch to their mother tongue?), machine translation (猫の餌を注文する). This is a highly visible role with a huge impact on Amazon customers and business. As part of this role, you will develop high precision, high recall, and low latency solutions for search. Your solutions should work for all languages that Amazon supports and will be used in all Amazon locales world-wide. You will develop scalable science and engineering solutions that work successfully in production. You will work with leaders to develop a strategic vision and long term plans to improve search globally. We are growing our collaborative group of engineers and applied scientists by expanding into new areas. This is a position on Global Search Quality team in Seattle Washington. We are moving fast to change the way Amazon search works. Together with a multi-disciplinary team you will work on building solutions with NLP/ML/DL at its core. Along the way, you’ll learn a ton, have fun and make a positive impact on millions of people. Come and join us as we invent new ways to delight Amazon customers.
US, WA, Seattle
This is a unique opportunity to build technology and science that millions of people will use every day. Are you excited about working on large scale Natural Language Processing (NLP), Machine Learning (ML), and Deep Learning (DL)? We are embarking on a multi-year journey to improve the shopping experience for customers globally. Amazon Search team creates customer-focused search solutions and technologies that makes shopping delightful and effortless for our customers. Our goal is to understand what customers are looking for in whatever language happens to be their choice at the moment and help them find what they need in Amazon's vast catalog of billions of products. As Amazon expands to new geographies, we are faced with the unique challenge of maintaining the bar on Search Quality due to the diversity in user preferences, multilingual search and data scarcity in new locales. We are looking for an applied researcher to work on improving search on Amazon using NLP, ML, and DL technology. As an Applied Scientist, you will lead our efforts in query understanding, semantic matching (e.g. is a drone the same as quadcopter?), relevance ranking (what is a "funny halloween costume"?), language identification (did the customer just switch to their mother tongue?), machine translation (猫の餌を注文する). This is a highly visible role with a huge impact on Amazon customers and business. As part of this role, you will develop high precision, high recall, and low latency solutions for search. Your solutions should work for all languages that Amazon supports and will be used in all Amazon locales world-wide. You will develop scalable science and engineering solutions that work successfully in production. You will work with leaders to develop a strategic vision and long term plans to improve search globally. We are growing our collaborative group of engineers and applied scientists by expanding into new areas. This is a position on Global Search Quality team in Seattle Washington. We are moving fast to change the way Amazon search works. Together with a multi-disciplinary team you will work on building solutions with NLP/ML/DL at its core. Along the way, you’ll learn a ton, have fun and make a positive impact on millions of people. Come and join us as we invent new ways to delight Amazon customers.
US, WA, Seattle
The retail pricing science and research group is a team of scientists and economists who design and implement the analytics powering pricing for Amazon’s on-line retail business. The team uses world-class analytics to make sure that the prices for all of Amazon’s goods and services are aligned with Amazon’s corporate goals. We are seeking an experienced high-energy Economist to help envision, design and build the next generation of retail pricing capabilities. You will work at the intersection of economic theory, statistical inference, and machine learning to design new methods and pricing strategies to deliver game changing value to our customers. Roughly 85% of previous intern cohorts have converted to full time scientist employment at Amazon. If you are interested, please send your CV to our mailing list at econ-internship@amazon.com. Key job responsibilities Amazon’s Pricing Science and Research team is seeking an Economist to help envision, design and build the next generation of pricing capabilities behind Amazon’s on-line retail business. As an economist on our team, you will work at the intersection of economic theory, statistical inference, and machine learning to design new methods and pricing strategies with the potential to deliver game changing value to our customers. This is an opportunity for a high-energy individual to work with our unprecedented retail data to bring cutting edge research into real world applications, and communicate the insights we produce to our leadership. This position is perfect for someone who has a deep and broad analytic background and is passionate about using mathematical modeling and statistical analysis to make a real difference. You should be familiar with modern tools for data science and business analysis. We are particularly interested in candidates with research background in applied microeconomics, econometrics, statistical inference and/or finance. A day in the life Discussions with business partners, as well as product managers and tech leaders to understand the business problem. Brainstorming with other scientists and economists to design the right model for the problem in hand. Present the results and new ideas for existing or forward looking problems to leadership. Deep dive into the data. Modeling and creating working prototypes. Analyze the results and review with partners. Partnering with other scientists for research problems. About the team The retail pricing science and research group is a team of scientists and economists who design and implement the analytics powering pricing for Amazon’s on-line retail business. The team uses world-class analytics to make sure that the prices for all of Amazon’s goods and services are aligned with Amazon’s corporate goals.
US, CA, San Francisco
The retail pricing science and research group is a team of scientists and economists who design and implement the analytics powering pricing for Amazon's on-line retail business. The team uses world-class analytics to make sure that the prices for all of Amazon's goods and services are aligned with Amazon's corporate goals. We are seeking an experienced high-energy Economist to help envision, design and build the next generation of retail pricing capabilities. You will work at the intersection of statistical inference, experimentation design, economic theory and machine learning to design new methods and pricing strategies for assessing pricing innovations. Roughly 85% of previous intern cohorts have converted to full time scientist employment at Amazon. If you are interested, please send your CV to our mailing list at econ-internship@amazon.com. Key job responsibilities Amazon's Pricing Science and Research team is seeking an Economist to help envision, design and build the next generation of pricing capabilities behind Amazon's on-line retail business. As an economist on our team, you will will have the opportunity to work with our unprecedented retail data to bring cutting edge research into real world applications, and communicate the insights we produce to our leadership. This position is perfect for someone who has a deep and broad analytic background and is passionate about using mathematical modeling and statistical analysis to make a real difference. You should be familiar with modern tools for data science and business analysis. We are particularly interested in candidates with research background in experimentation design, applied microeconomics, econometrics, statistical inference and/or finance. A day in the life Discussions with business partners, as well as product managers and tech leaders to understand the business problem. Brainstorming with other scientists and economists to design the right model for the problem in hand. Present the results and new ideas for existing or forward looking problems to leadership. Deep dive into the data. Modeling and creating working prototypes. Analyze the results and review with partners. Partnering with other scientists for research problems. About the team The retail pricing science and research group is a team of scientists and economists who design and implement the analytics powering pricing for Amazon's on-line retail business. The team uses world-class analytics to make sure that the prices for all of Amazon's goods and services are aligned with Amazon's corporate goals.
US, WA, Bellevue
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 Python is necessary, and experience with SQL and UNIX 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, scientists, and product managers. 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
The Amazon Supply Chain Optimization Technology (SCOT) organization is looking for an Intern in Economics to work on exciting and challenging problems related to Amazon's worldwide inventory planning. SCOT provides unique opportunities to both create and see the direct impact of your work on billions of dollars’ worth of inventory, in one of the world’s most advanced supply chains, and at massive scale. We are looking for detail-oriented, organized, and responsible individuals who are eager to learn how to work with large and complicated data sets. We are looking for a PhD candidate with exposure to Program Evaluation/Causal Inference. Knowledge of econometrics and Stata/R/or Python is necessary, and experience with SQL, Hadoop, 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, scientists, and product managers. 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 previous cohorts have converted to full time scientist employment at Amazon. If you are interested, please send your CV to our mailing list at econ-internship@amazon.com.
US, WA, Seattle
The Selling Partner Fees team owns the end-to-end fees experience for two million active third party sellers. We own the fee strategy, fee seller experience, fee accuracy and integrity, fee science and analytics, and we provide scalable technology to monetize all services available to third-party sellers. We are looking for an Intern Economist with excellent coding skills to design and develop rigorous models to assess the causal impact of fees on third party sellers’ behavior and business performance. As a Science Intern, you will have access to large datasets with billions of transactions and will translate ambiguous fee related business problems into rigorous scientific models. You will work on real world problems which will help to inform strategic direction and have the opportunity to make an impact for both Amazon and our Selling Partners.