Amazon open-sources library for prediction over large output spaces

Framework improves efficiency, accuracy of applications that search for a handful of solutions in a huge space of candidates.

In the Internet age, many computational tasks involve finding a handful of solutions in an enormous space of candidates. Question-answering systems, for instance, can pull answers from anywhere on the web, while the Wikipedia taxonomy for classifying article topic classification has 500,000 terms. And of course, a product query at the Amazon Store has millions of potential matches.

Such extreme multilabel ranking (XMR) problems pose two major challenges. The first is one of scale, but the second is one of scarcity. The items in these large search spaces tend to have long-tailed distributions: most sentences rarely serve as answers to questions; most topics in the Wikipedia taxonomy rarely apply to texts; most products are rarely purchased; and so on. That means that attempts to use machine learning to solve XMR problems rarely have enough data to go on.

At Amazon, we have developed a general framework for meeting both these challenges, which we call PECOS, for prediction for enormous and correlated output spaces. After successfully using PECOS internally for key projects in product search and recommendation, we have publicly released the code to help stimulate further research on this important topic.

In the XMR context, the items retrieved from the search space are known as labels. If the task is document retrieval, the documents themselves are interpreted as candidate labels for a search string; the search string is the input. The “multilabel” in XMR indicates that a given input may have multiple labels; several different topics from the Wikipedia taxonomy, for instance, might apply to the same document.

PECOS decomposes the XMR problem into three stages:

  1. semantic label indexing, or grouping labels together according to semantic content;
  2. matching, or associating the input instance with a label group;
  3. ranking, or finding the labels in each group that best fit the input.
PECOS-framework.png
The three-stage PECOS model.
Credit: Stacy Reilly

PECOS lets users create their own algorithms to implement any of these stages, but the code release comes with a library of standard algorithms for each stage, including both a recursive linear model and a trained deep-learning model for matching.

The three-stage framework helps with both the scaling and long-tail problems. By enabling matching with groups of labels rather than individual labels, label indexing drastically reduces the search space for the matching step. It also helps with the long-tail problem, since it enables the ranking model to exploit semantic similarities between common labels and less common labels.

For machine-learning-based implementations of the ranking stage, label indexing aids in the selection of hard negatives. Machine learning models must be trained on both positive examples and negative examples; in the XMR context, most negative examples are so irrelevant as to impart little information to the model. Selecting negative examples from the same groups as the positive examples ensures that they’ll be challenging enough to improve the quality of the model.

The initial release of PECOS includes two models that implement the entire PECOS framework. One is a recursive linear model, the other a deep-learning model. In tests involving a dataset with 2.8 million labels, the deep-learning model improved the precision of the top-ranked result (precision@1) by 10% relative to the recursive linear model, but it took 265 times as long to train. It’s up to the individual users to evaluate that trade-off for their own use cases.

Semantic label indexing

Semantic label indexing has two components: a representation scheme and a grouping algorithm. For text-based inputs, the representation scheme might take advantage of pre-trained text embeddings such as Word2Vec or ELMo; for graph-based inputs, it might use information about the input’s relationships with its neighbors in the graph. PECOS includes efficient implementations of representation schemes such as positive instance indices (PII), positive instance feature aggregation (PIFA), and the graph spectrum representation.

For grouping, we’ve concentrated on clustering algorithms, but users could implement other approaches, such as approximate nearest-neighbor search. PECOS includes our implementations of the k-means and spherical k-means clustering algorithms, which feature recursive B-ary partitioning. For some value of B (usually between 2 and 16), the algorithm first partitions the label set into B clusters, then partitions each of those into B clusters, and so on.

B-ary partitioning.png
A simple example of our B-ary partitioning scheme.

In a paper about PECOS that we’ve published to the arXiv, we show that B-ary partitioning can significantly reduce the time required for semantic-label indexing, an important consideration given that we’re dealing with enormous label spaces. We also use the B-ary partitioning to implement the recursive linear model.

Built-in models

For text inputs, PECOS includes X-Transformer, which leverages pretrained transformer models from Huggingface to improve performance on extreme multilabel text classification applications. At the 2020 Conference on Knowledge Discovery and Data Mining (KDD), we presented a paper about the PECOS deep-learning model, which we also described in a related blog post on Amazon Science.

PECOS also includes a linear model, XR-Linear, which learns its matching algorithm recursively. First, it learns a B-ary partition of the label space. Then, to implement a matcher for that partition, it learns a new B-ary partition for each of the existing groups. To implement matchers for those, it learns a new B-ary partition for each, and so on, until it reaches the desired recursive depth. At that point, it learns a simple linear one-versus-all ranker for the labels in each partition.

Then, for each level of recursion, it learns a ranker for the outputs of the layer below.

Recursive matcher.png
A diagram of the recursive linear matcher.

This makes training very efficient, as the full set of weights for each recursive layer can fit in memory at once, saving time on inefficient retrieval from storage.

At inference time, XR-Linear works through the same recursion tree to identify relevant labels. For efficiency, we use beam search to restrict the search space. For instance, if the beam width is two, then at each layer of the recursion tree, the model will pursue only the two highest-weight connections to the next layer.

Beam search.gif
An example of linear ranking with a beam width of two. At each level of the tree, two nodes (green) are selected for further exploration. Each of their descendant nodes is evaluated (orange), and two of those are selected for further exploration.
Credit: Giana Bucchino

Our PECOS software has benefited from open research that has been conducted at Amazon and at other universities and companies. By open-sourcing the PECOS software, we are thrilled to contribute back to the open-research community. Our hope is to spur further research on problems where the output spaces are very large. These include zero-shot learning for extreme multilabel problems, extreme contextual bandits, and deep reinforcement learning.

For more information about the optimizations we’ve incorporated into the PECOS code release, please see our arXiv paper. The code itself can be downloaded at GitHub.

Research areas

Related content

US, CA, Santa Clara
About Amazon Health Amazon Health’s mission is to make it dramatically easier for customers to access the healthcare products and services they need to get and stay healthy. Towards this mission, we (Health Storefront and Shared Tech) are building the technology, products and services, that help customers find, buy, and engage with the healthcare solutions they need. Job summary We are seeking an exceptional Applied Scientist to join a team of experts in the field of machine learning, and work together to break new ground in the world of healthcare to make personalized and empathetic care accessible, convenient, and cost-effective. We leverage and train state-of-the-art large-language-models (LLMs) and develop entirely new experiences to help customers find the right products and services to address their health needs. We work on machine learning problems for intent detection, dialogue systems, and information retrieval. You will work in a highly collaborative environment where you can pursue both near-term productization opportunities to make immediate, meaningful customer impacts while pursuing ambitious, long-term research. You will work on hard science problems that have not been solved before, conduct rapid prototyping to validate your hypothesis, and deploy your algorithmic ideas at scale. You will get the opportunity to pursue work that makes people's lives better and pushes the envelop of science. Key job responsibilities - Translate product and CX requirements into science metrics and rigorous testing methodologies. - Invent and develop scalable methodologies to evaluate LLM outputs against metrics and guardrails. - Design and implement the best-in-class semantic retrieval system by creating high-quality knowledge base and optimizing embedding models and similarity measures. - Conduct tuning, training, and optimization of LLMs to achieve a compelling CX while reducing operational cost to be scalable. A day in the life In a fast-paced innovation environment, you work closely with product, UX, and business teams to understand customer's challenges. You translate product and business requirements into science problems. You dive deep into challenging science problems, enabling entirely new ML and LLM-driven customer experiences. You identify hypothesis and conduct rapid prototyping to learn quickly. You develop and deploy models at scale to pursue productizations. You mentor junior science team members and help influence our org in scientific best practices. About the team We are the ML Science and Engineering team, with a strong focus on Generative AI. The team consists of top-notch ML Scientists with diverse background in healthcare, robotics, customer analytics, and communication. We are committed to building and deploying the most advanced scientific capabilities and solutions for the products and services at Amazon Health. We are open to hiring candidates to work out of one of the following locations: Santa Clara, CA, USA
US, WA, Seattle
We are designing the future. If you are in quest of an iterative fast-paced environment, where you can drive innovation through scientific inquiry, and provide tangible benefit to hundreds of thousands of our associates worldwide, this is your opportunity. Come work on the Amazon Worldwide Fulfillment Design & Engineering Team! We are looking for an experienced and senior Research Scientist with background in Ergonomics and Industrial Human Factors, someone that is excited to work on complex real-world challenges for which a comprehensive scientific approach is necessary to drive solutions. Your investigations will define human factor / ergonomic thresholds resulting in design and implementation of safe and efficient workspaces and processes for our associates. Your role will entail assessment and design of manual material handling tasks throughout the entire Amazon network. You will identify fundamental questions pertaining to the human capabilities and tolerances in a myriad of work environments, and will initiate and lead studies that will drive decision making on an extreme scale. .You will provide definitive human factors/ ergonomics input and participate in design with every single design group in our network, including Amazon Robotics, Engineering R&D, and Operations Engineering. You will work closely with our Worldwide Health and Safety organization to gain feedback on designs and work tenaciously to continuously improve our associate’s experience. Key job responsibilities - Collaborating and designing work processes and workspaces that adhere to human factors / ergonomics standards worldwide. - Producing comprehensive and assessments of workstations and processes covering biomechanical, physiological, and psychophysical demands. - Effectively communicate your design rationale to multiple engineering and operations entities. - Identifying gaps in current human factors standards and guidelines, and lead comprehensive studies to redefine “industry best practices” based on solid scientific foundations. - Continuously strive to gain in-depth knowledge of your profession, as well as branch out to learn about intersecting fields, such as robotics and mechatronics. - Travelling to our various sites to perform thorough assessments and gain in-depth operational feedback, approximately 25%-50% of the time. We are open to hiring candidates to work out of one of the following locations: Seattle, WA, USA
GB, London
Amazon Advertising is looking for a Data Scientist to join its brand new initiative that powers Amazon’s contextual advertising products. Advertising at Amazon is a fast-growing multi-billion dollar business that spans across desktop, mobile and connected devices; encompasses ads on Amazon and a vast network of hundreds of thousands of third party publishers; and extends across US, EU and an increasing number of international geographies. The Supply Quality organization has the charter to solve optimization problems for ad-programs in Amazon and ensure high-quality ad-impressions. We develop advanced algorithms and infrastructure systems to optimize performance for our advertisers and publishers. We are focused on solving a wide variety of problems in computational advertising like traffic quality prediction (robot and fraud detection), Security forensics and research, Viewability prediction, Brand Safety, Contextual data processing and classification. Our team includes experts in the areas of distributed computing, machine learning, statistics, optimization, text mining, information theory and big data systems. We are looking for a dynamic, innovative and accomplished Data Scientist to work on data science initiatives for contextual data processing and classification that power our contextual advertising solutions. Are you an experienced user of sophisticated analytical techniques that can be applied to answer business questions and chart a sustainable vision? Are you exited by the prospect of communicating insights and recommendations to audiences of varying levels of technical sophistication? Above all, are you an innovator at heart and have a track record of resolving ambiguity to deliver result? As a data scientist, you help our data science team build cutting edge models and measurement solutions to power our contextual classification technology. As this is a new initiative, you will get an opportunity to act as a thought leader, work backwards from the customer needs, dive deep into data to understand the issues, define metrics, conceptualize and build algorithms and collaborate with multiple cross-functional teams. Key job responsibilities * Define a long-term science vision for contextual-classification tech, driven fundamentally from the needs of our advertisers and publishers, translating that direction into specific plans for the science team. Interpret complex and interrelated data points and anecdotes to build and communicate this vision. * Collaborate with software engineering teams to Identify and implement elegant statistical and machine learning solutions * Oversee the design, development, and implementation of production level code that handles billions of ad requests. Own the full development cycle: idea, design, prototype, impact assessment, A/B testing (including interpretation of results) and production deployment. * Promote the culture of experimentation and applied science at Amazon. * Demonstrated ability to meet deadlines while managing multiple projects. * Excellent communication and presentation skills working with multiple peer groups and different levels of management * Influence and continuously improve a sustainable team culture that exemplifies Amazon’s leadership principles. We are open to hiring candidates to work out of one of the following locations: London, GBR
JP, 13, Tokyo
We are seeking a Principal Economist to be the science leader in Amazon's customer growth and engagement. The wide remit covers Prime, delivery experiences, loyalty program (Amazon Points), and marketing. We look forward to partnering with you to advance our innovation on customers’ behalf. Amazon has a trailblazing track record of working with Ph.D. economists in the tech industry and offers a unique environment for economists to thrive. As an economist at Amazon, you will apply the frontier of econometric and economic methods to Amazon’s terabytes of data and intriguing customer problems. Your expertise in building reduced-form or structural causal inference models is exemplary in Amazon. Your strategic thinking in designing mechanisms and products influences how Amazon evolves. In this role, you will build ground-breaking, state-of-the-art econometric models to guide multi-billion-dollar investment decisions around the global Amazon marketplaces. You will own, execute, and expand a research roadmap that connects science, business, and engineering and contributes to Amazon's long term success. As one of the first economists outside North America/EU, you will make an outsized impact to our international marketplaces and pioneer in expanding Amazon’s economist community in Asia. The ideal candidate will be an experienced economist in empirical industrial organization, labour economics, or related structural/reduced-form causal inference fields. You are a self-starter who enjoys ambiguity in a fast-paced and ever-changing environment. You think big on the next game-changing opportunity but also dive deep into every detail that matters. You insist on the highest standards and are consistent in delivering results. Key job responsibilities - Work with Product, Finance, Data Science, and Data Engineering teams across the globe to deliver data-driven insights and products for regional and world-wide launches. - Innovate on how Amazon can leverage data analytics to better serve our customers through selection and pricing. - Contribute to building a strong data science community in Amazon Asia. We are open to hiring candidates to work out of one of the following locations: Tokyo, 13, JPN
DE, BE, Berlin
Ops Integration: Concessions team is looking for a motivated, creative and customer obsessed Snr. Applied Scientist with a strong machine learning background, to develop advanced analytics models (Computer Vision, LLMs, etc.) that improve customer experiences We are the voice of the customer in Amazon’s operations, and we take that role very seriously. If you join this team, you will be a key contributor to delivering the Factory of the Future: leveraging Internet of Things (IoT) and advanced analytics to drive tangible, operational change on the ground. You will collaborate with a wide range of stakeholders (You will partner with Research and Applied Scientists, SDEs, Technical Program Managers, Product Managers and Business Leaders) across the business to develop and refine new ways of assessing challenges within Amazon operations. This role will combine Amazon’s oldest Leadership Principle, with the latest analytical innovations, to deliver business change at scale and efficiently The ideal candidate will have deep and broad experience with theoretical approaches and practical implementations of vision techniques for task automation. They will be a motivated self-starter who can thrive in a fast-paced environment. They will be passionate about staying current with sensing technologies and algorithms in the broader machine vision industry. They will enjoy working in a multi-disciplinary team of engineers, scientists and business leaders. They will seek to understand processes behind data so their recommendations are grounded. Key job responsibilities Your solutions will drive new system capabilities with global impact. You will design highly scalable, large enterprise software solutions involving computer vision. You will develop complex perception algorithms integrating across multiple sensing devices. You will develop metrics to quantify the benefits of a solution and influence project resources. You will validate system performance and use insights from your live models to drive the next generation of model development. Common tasks include: • Research, design, implement and evaluate complex perception and decision making algorithms integrating across multiple disciplines • Work closely with software engineering teams to drive scalable, real-time implementations • Collaborate closely with team members on developing systems from prototyping to production level • Collaborate with teams spread all over the world • Track general business activity and provide clear, compelling management reports on a regular basis We are open to hiring candidates to work out of one of the following locations: Berlin, BE, DEU | Berlin, DEU
DE, BY, Munich
Ops Integration: Concessions team is looking for a motivated, creative and customer obsessed Snr. Applied Scientist with a strong machine learning background, to develop advanced analytics models (Computer Vision, LLMs, etc.) that improve customer experiences We are the voice of the customer in Amazon’s operations, and we take that role very seriously. If you join this team, you will be a key contributor to delivering the Factory of the Future: leveraging Internet of Things (IoT) and advanced analytics to drive tangible, operational change on the ground. You will collaborate with a wide range of stakeholders (You will partner with Research and Applied Scientists, SDEs, Technical Program Managers, Product Managers and Business Leaders) across the business to develop and refine new ways of assessing challenges within Amazon operations. This role will combine Amazon’s oldest Leadership Principle, with the latest analytical innovations, to deliver business change at scale and efficiently The ideal candidate will have deep and broad experience with theoretical approaches and practical implementations of vision techniques for task automation. They will be a motivated self-starter who can thrive in a fast-paced environment. They will be passionate about staying current with sensing technologies and algorithms in the broader machine vision industry. They will enjoy working in a multi-disciplinary team of engineers, scientists and business leaders. They will seek to understand processes behind data so their recommendations are grounded. Key job responsibilities Your solutions will drive new system capabilities with global impact. You will design highly scalable, large enterprise software solutions involving computer vision. You will develop complex perception algorithms integrating across multiple sensing devices. You will develop metrics to quantify the benefits of a solution and influence project resources. You will validate system performance and use insights from your live models to drive the next generation of model development. Common tasks include: • Research, design, implement and evaluate complex perception and decision making algorithms integrating across multiple disciplines • Work closely with software engineering teams to drive scalable, real-time implementations • Collaborate closely with team members on developing systems from prototyping to production level • Collaborate with teams spread all over the world • Track general business activity and provide clear, compelling management reports on a regular basis We are open to hiring candidates to work out of one of the following locations: Munich, BE, DEU | Munich, BY, DEU | Munich, DEU
IT, MI, Milan
Ops Integration: Concessions team is looking for a motivated, creative and customer obsessed Snr. Applied Scientist with a strong machine learning background, to develop advanced analytics models (Computer Vision, LLMs, etc.) that improve customer experiences We are the voice of the customer in Amazon’s operations, and we take that role very seriously. If you join this team, you will be a key contributor to delivering the Factory of the Future: leveraging Internet of Things (IoT) and advanced analytics to drive tangible, operational change on the ground. You will collaborate with a wide range of stakeholders (You will partner with Research and Applied Scientists, SDEs, Technical Program Managers, Product Managers and Business Leaders) across the business to develop and refine new ways of assessing challenges within Amazon operations. This role will combine Amazon’s oldest Leadership Principle, with the latest analytical innovations, to deliver business change at scale and efficiently The ideal candidate will have deep and broad experience with theoretical approaches and practical implementations of vision techniques for task automation. They will be a motivated self-starter who can thrive in a fast-paced environment. They will be passionate about staying current with sensing technologies and algorithms in the broader machine vision industry. They will enjoy working in a multi-disciplinary team of engineers, scientists and business leaders. They will seek to understand processes behind data so their recommendations are grounded. Key job responsibilities Your solutions will drive new system capabilities with global impact. You will design highly scalable, large enterprise software solutions involving computer vision. You will develop complex perception algorithms integrating across multiple sensing devices. You will develop metrics to quantify the benefits of a solution and influence project resources. You will validate system performance and use insights from your live models to drive the next generation of model development. Common tasks include: • Research, design, implement and evaluate complex perception and decision making algorithms integrating across multiple disciplines • Work closely with software engineering teams to drive scalable, real-time implementations • Collaborate closely with team members on developing systems from prototyping to production level • Collaborate with teams spread all over the world • Track general business activity and provide clear, compelling management reports on a regular basis We are open to hiring candidates to work out of one of the following locations: Milan, MI, ITA
ES, M, Madrid
Ops Integration: Concessions team is looking for a motivated, creative and customer obsessed Snr. Applied Scientist with a strong machine learning background, to develop advanced analytics models (Computer Vision, LLMs, etc.) that improve customer experiences We are the voice of the customer in Amazon’s operations, and we take that role very seriously. If you join this team, you will be a key contributor to delivering the Factory of the Future: leveraging Internet of Things (IoT) and advanced analytics to drive tangible, operational change on the ground. You will collaborate with a wide range of stakeholders (You will partner with Research and Applied Scientists, SDEs, Technical Program Managers, Product Managers and Business Leaders) across the business to develop and refine new ways of assessing challenges within Amazon operations. This role will combine Amazon’s oldest Leadership Principle, with the latest analytical innovations, to deliver business change at scale and efficiently The ideal candidate will have deep and broad experience with theoretical approaches and practical implementations of vision techniques for task automation. They will be a motivated self-starter who can thrive in a fast-paced environment. They will be passionate about staying current with sensing technologies and algorithms in the broader machine vision industry. They will enjoy working in a multi-disciplinary team of engineers, scientists and business leaders. They will seek to understand processes behind data so their recommendations are grounded. Key job responsibilities Your solutions will drive new system capabilities with global impact. You will design highly scalable, large enterprise software solutions involving computer vision. You will develop complex perception algorithms integrating across multiple sensing devices. You will develop metrics to quantify the benefits of a solution and influence project resources. You will validate system performance and use insights from your live models to drive the next generation of model development. Common tasks include: • Research, design, implement and evaluate complex perception and decision making algorithms integrating across multiple disciplines • Work closely with software engineering teams to drive scalable, real-time implementations • Collaborate closely with team members on developing systems from prototyping to production level • Collaborate with teams spread all over the world • Track general business activity and provide clear, compelling management reports on a regular basis We are open to hiring candidates to work out of one of the following locations: Madrid, ESP | Madrid, M, ESP
US, CA, San Diego
The Private Brands team is looking for an Applied Scientist to join the team in building science solutions at scale. Our team applies Optimization, Machine Learning, Statistics, Causal Inference, and Econometrics/Economics to derive actionable insights. We are an interdisciplinary team of Scientists, Engineers, and Economists and primary focus on building optimization and machine learning solutions in supply chain domain with specific focus on Amazon private brand products. Key job responsibilities You will work with business leaders, scientists, and economists to translate business and functional requirements into concrete deliverables, including the design, development, testing, and deployment of highly scalable optimization solutions and ML models. This is a unique, high visibility opportunity for someone who wants to have business impact, dive deep into large-scale problems, enable measurable actions on the consumer economy, and work closely with scientists and economists. As a scientist, you bring business and industry context to science and technology decisions. You set the standard for scientific excellence and make decisions that affect the way we build and integrate algorithms. Your solutions are exemplary in terms of algorithm design, clarity, model structure, efficiency, and extensibility. You tackle intrinsically hard problems, acquiring expertise as needed. You decompose complex problems into straightforward solutions. We are particularly interested in candidates with experience in predictive and machine learning models and working with distributed systems. Academic and/or practical background in Machine Learning are particularly relevant for this position. Familiarity and experience in applying Operations Research techniques to supply chain problems is a plus. To know more about Amazon science, Please visit https://www.amazon.science We are open to hiring candidates to work out of one of the following locations: San Diego, CA, USA | Seattle, WA, USA
LU, Luxembourg
Ops Integration: Concessions team is looking for a motivated, creative and customer obsessed Snr. Applied Scientist with a strong machine learning background, to develop advanced analytics models (Computer Vision, LLMs, etc.) that improve customer experiences We are the voice of the customer in Amazon’s operations, and we take that role very seriously. If you join this team, you will be a key contributor to delivering the Factory of the Future: leveraging Internet of Things (IoT) and advanced analytics to drive tangible, operational change on the ground. You will collaborate with a wide range of stakeholders (You will partner with Research and Applied Scientists, SDEs, Technical Program Managers, Product Managers and Business Leaders) across the business to develop and refine new ways of assessing challenges within Amazon operations. This role will combine Amazon’s oldest Leadership Principle, with the latest analytical innovations, to deliver business change at scale and efficiently The ideal candidate will have deep and broad experience with theoretical approaches and practical implementations of vision techniques for task automation. They will be a motivated self-starter who can thrive in a fast-paced environment. They will be passionate about staying current with sensing technologies and algorithms in the broader machine vision industry. They will enjoy working in a multi-disciplinary team of engineers, scientists and business leaders. They will seek to understand processes behind data so their recommendations are grounded. Key job responsibilities Your solutions will drive new system capabilities with global impact. You will design highly scalable, large enterprise software solutions involving computer vision. You will develop complex perception algorithms integrating across multiple sensing devices. You will develop metrics to quantify the benefits of a solution and influence project resources. You will validate system performance and use insights from your live models to drive the next generation of model development. Common tasks include: • Research, design, implement and evaluate complex perception and decision making algorithms integrating across multiple disciplines • Work closely with software engineering teams to drive scalable, real-time implementations • Collaborate closely with team members on developing systems from prototyping to production level • Collaborate with teams spread all over the world • Track general business activity and provide clear, compelling management reports on a regular basis We are open to hiring candidates to work out of one of the following locations: Luxembourg, LUX