Using supervised learning to train models for image clustering

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

Most machine learning models use supervised learning, meaning they’re trained on annotated data, which is costly and time consuming to acquire.

The chief method for doing unsupervised learning, which doesn’t require annotated data, is clustering, or grouping data points together by salient characteristics. The idea is that each cluster represents some category, such as photos of the same person or the same species of animal.

To decide where to draw boundaries between clusters, clustering algorithms typically rely on heuristics, such as a threshold distance between cluster centers or the shape of the clusters’ distributions. In a paper we’re presenting at the International Conference on Computer Vision (ICCV), we propose, instead, to learn from data how to draw boundaries.

We first represent visual data using a graph, then use a graph neural network (GNN) to produce vector representations of the graph’s nodes. So far, we follow on previous work.

Instead of relying on heuristics, however, we use labeled data to learn how to cluster the vectors and, crucially, to decide how fine-grained those clusters should be. We call the labeled data meta-training data, since the goal is to learn a general clustering technique, not a specific classification model. 

In particular, we propose a hierarchical GNN, meaning that it creates clusters by adding edges between nodes of a graph, then adds edges between the clusters to create still larger clusters, and so on, iterating until it decides that no more edges should be added.

Hierarchical clustering.png
A schematic of our graph-based hierarchical clustering approach. The colors of the image borders and of the graph nodes indicate data types (in this case, photos of the same actor). Our approach is hierarchical, iteratively treating small clusters generated at one level as the units of clustering for the next level. We call our base model LANDER, for link approximation and density estimation refinement, and our hierarchical clustering method Hi-LANDER.

Finally, we apply our hierarchical clustering technique to test sets whose classification categories are disjoint with those of the meta-training data. In our experiments we found that, compared to previous GNN-based supervised and unsupervised approaches, ours increased the F-score — which factors in both false positives and false negatives — by an average of 49% and 47%, respectively.

Constructing the graph

In our paper, we investigate the case in which we are training a model to cluster visual data that is similar to the meta-training data but has no class overlaps with it. For instance, the meta-training data might be faces of movie stars, while the target application is to cluster faces of politicians, athletes, or other public figures.

The first step in our process is to use the meta-training data to build a supervised classifier: if the meta-training data is faces of movie stars, the classifier labels input images with names of movie stars.

The classifier is an encoder-decoder model: the encoder produces a fixed-length vector representation of the input, or feature vector, and the decoder uses that vector to predict a label. Once we’ve trained the classifier, however, we use only the encoder for the rest of the process.

The feature vectors define points in a multidimensional space. On the basis of the vectors’ locations, we construct a graph, in which each node represents an image, and each image’s k nearest neighbors in the feature space are connected to it (share edges with it) in the graph.

This graph will serve as the input to the clustering model, which is also an encoder-decoder model. The encoder is a GNN, which produces a vector representation of each node in the graph, based on that node’s feature vector and those of the nodes it’s connected to. Call this vector the node embedding.

The clustering model

We adopt a hierarchical approach to clustering. Based on the node embeddings, the clustering model predicts edges between nodes. A cluster is defined as a group of nodes each of which shares an edge with at least one other node in the group and none of which shares an edge with any node outside the group.

Note that the goal of the clustering model is not just to reproduce the nearest-neighbor graph but to link nodes that represent data of the same type. The nearest-neighbor linkages are useful for predicting clustering linkages, but they are not identical with them.

After the first pass through the data, we aggregate each cluster into a single, representative “supernode” and repeat the whole process. That is, we create edges between each supernode and its k nearest neighbors, pass the resulting graph through the same GNN, and predict edges based on the supernode embeddings. We repeat this process until the clustering model predicts no edges between nodes.

We train our clustering model on two different objectives. One is to correctly predict links between nodes, where a correct link is one that picks out two representatives of the same data type in the meta-training data (say, two photos of the same actor).

We also train the model to correctly predict the density of a given data type in a given graph neighborhood. That is, for each node, the model should predict the proportion of nearby neighbors of the same data type.

Past research on clustering has shown that factoring in data density improves results. Previously, however, link prediction and data density prediction were handled by separate models. By using a single model to jointly predict both, we significantly increase computational efficiency. We believe that the combination also contributes to our increase in accuracy.

The other novelty of our approach is that, because of our hierarchical processing scheme, we optimize clustering across the entire input graph. Previous approaches would first divide the graph into subgraphs, then perform inference within subgraphs. This prevents natural parallelization, which is runtime efficient, and limits the effectiveness of information flow through the graph. The full graph-wide processing is another reason for our model’s improved efficiency.

In experiments, we considered two different sets of meta-training data. One consisted of closeups of human faces, the other of images of particular animal species. We tested the model trained on human faces on two other datasets, whose data categories had zero or very little overlap with those of the meta-training set — 0% and less than 2%. We tested the model trained on animal species on a dataset of previously unseen species. Across both models and the three test sets, our average improvements over previous GNN-based clustering models and unsupervised clustering methods were 49% and 47%, respectively.

In ongoing work, we are investigating the possibility training a more general clustering model, whose performance at inference time will be more transferrable across different data types — accurately clustering both faces and animal species, for instance.

Acknowledgements: Tianjun Xiao, Yongxin Wang, Yuanjun Xiong, Wei Xia, David Wipf, Zhang Zheng, Stefano Soatto

Related content

DE, Aachen
The Artificial General Intelligence (AGI) team is looking for a passionate, talented, and inventive Senior 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 Senior 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.
CA, BC, Vancouver
We are open to candidates located in: Seattle and Bellevue, Washington Atlanta, GA As a Senior Data Scientist, you will be on the ground floor with your team, shaping the way performance is measured, defining what questions should be asked, and scaling analytics methods and tools to support our growing business. You will work closely with Data Engineers, Product Managers, Business Intelligence Engineers, and Software Engineers to develop statistical models, design and run experiments, and find new ways to to optimize the customer experience. A successful candidate is highly analytical, able to work effectively in a matrix organization, and adept at synthesizing a variety of technologies and capabilities into products that enhances the PXF experience across multiple products. You must engage with customers to deeply understand their current and emerging needs. PXF applications are rapidly evolving and our user base is rapidly expanding, as a DS on the team you will own diving into the different users personas and inventing on behalf of the app users to meet their needs. We are looking for someone who's customer obsessed and technology savvy - with a passion for app development work. The ideal candidate will have a well-rounded technical background as well as a history of leading complex, ambiguous projects end-to-end. Key job responsibilities - Partner with business stakeholders in formulating the business problem and providing recommendations on the approach - Understanding customer behavior to personalize customer experience, build recommendation engines to provide relevant results to customers, customer lifecycle analysis and usage behavior - Conduct large scale A/B testing and offline/online experiments to evaluate performance of various programs and drive product improvements across partner teams - Process large scale datasets using distributed computing platform to build models, mining insights from data and prototyping models that optimize towards various business goals and metrics - Interact with cross-functional teams and make business recommendations i.e cost-benefit, forecasting, experiment analysis and present findings to leadership team - Driving product roadmap and prioritizations of science projects with the PMs to improve customer experience About the team PXF builds the employee experiences that connect Amazonians, support them through their employment journey, and make Amazon Earth's Best Employer. Our products include A to Z mobile application directly impacts the lives of associates by helping them identify the best shifts for their schedule, opportunities to pick up additional work, and choose when they get paid. We enable Amazon employees to easily find and access high-quality and authoritative content throughout their employment lifecycle through content management and Search capabilities. We also provide employees with a dynamic and ever-evolving learning experience to protect, prepare, and advance their careers.
IN, KA, Bengaluru
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. One of the key focus areas is Traffic Quality where we endeavour to identify non-human and invalid traffic within programmatic ad sources, and weed them out to ensure a high quality advertising marketplace. We do this by building machine learning and optimization algorithms that operate at scale, and leverage nuanced features about user, context, and creative engagement to determine the validity of traffic. The challenge is to stay one step ahead by investing in deep analytics and developing new algorithms that address emergent attack vectors in a structured and scalable fashion. We are committed to building a long-term traffic quality solution that encompasses all Amazon advertising channels and provides state-of-the-art traffic filtering that preserves advertiser trust and saves them hundreds of millions of dollars of wasted spend. We are looking for talented applied scientists who enjoy working on creative machine learning algorithms and thrive in a fast-paced, fun environment. An Applied Scientist is responsible for solving inherently hard problems in advertising fraud detection using deep learning, self-supervised techniques, representation learning and advanced clustering. An ideal candidate should have strong depth and breadth knowledge in machine learning, data mining and statistics. Traffic quality systems process billions of ad-impressions and clicks per day, by leveraging cutting-edge open source technologies like Hadoop, Spark, Redis and Amazon's cloud services like EC2, S3, EMR, DynamoDB and RedShift. The candidate should have reasonable programming and design skills to manipulate unstructured and big data and build prototypes that work on massive datasets. The candidate should be able to apply business knowledge to perform broad data analysis as a precursor to modeling and to provide valuable business intelligence.
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
About Amazon Regulatory Intelligence, Safety, and Compliance (RISC). Amazon RISC’s vision is to make Amazon Earth’s most trusted shopping destination for safe and compliant products. Towards this mission, we take a science-first approach to building technology, products and services, that protect customers from unsafe, illegal, controversial, or policy-violating products. 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 tackle challenging problems across diverse compliance domains. We leverage and train state-of-the-art multi-modal and large-language-models (LLMs) to detect illegal and unsafe products across the Amazon catalog. We work on machine learning problems for multi-modal classification, intent detection, information retrieval, anomaly and fraud detection, and generative AI. This is an exciting and challenging position to deliver scientific innovations into production systems at Amazon-scale to make immediate, meaningful customer impacts while also pursuing ambitious, long-term research. You will work in a highly collaborative environment where you can analyze and process large amounts of image, text and tabular data. 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. There will be something new to learn every day as we work in an environment with rapidly evolving regulations and adversarial actors looking to outwit your best ideas. Key job responsibilities • Design and evaluate state-of-the-art algorithms and approaches in multi-modal classification, large language models (LLMs), intent detection, information retrieval, anomaly and fraud detection, and generative AI • Translate product and CX requirements into measurable science problems and metrics. • Collaborate with product and tech partners and customers to validate hypothesis, drive adoption, and increase business impact • Key author in writing high quality scientific papers in internal and external peer-reviewed conferences. A day in the life - Understanding customer problems, project timelines, and team/project mechanisms - Proposing science formulations and brainstorming ideas with team to solve business problems - Writing code, and running experiments with re-usable science libraries - Reviewing labels and audit results with investigators and operations associates - Sharing science results with science, product and tech partners and customers - Writing science papers for submission to peer-review venues, and reviewing science papers from other scientists in the team. - Contributing to team retrospectives for continuous improvements - Driving science research collaborations and attending study groups with scientists across Amazon About the team We are a team of applied scientists building AI/ML solutions to make Amazon Earth’s most trusted shopping destination for safe and compliant products.
US, WA, Seattle
Amazon is looking for talented Postdoctoral Scientists to join our global Science teams for a one-year, full-time research position. Postdoctoral Scientists will innovate as members of Amazon’s key global Science teams. Postdoctoral Scientists will join one of our global science teams focused on solving research-intense business problems by exploring new research ideas, accelerating scientific innovation and impact, and publishing their work in peer-reviewed scientific venues. Postdocs will raise the scientific bar across Amazon by diving deep into exploratory areas of research to enhance the customer experience and improve efficiencies. Key job responsibilities * Work closely with a senior science advisor, collaborate with other scientists and engineers, and be part of Amazon’s vibrant and diverse global science community. * Publish your innovation in top-tier academic venues and hone your presentation skills. * Be inspired by challenges and opportunities to invent cutting-edge techniques in your area(s) of expertise.
US, WA, Seattle
Are you interested in applying your background in machine learning, data science, and statistical analysis to massive real-world data and impact the lives of hundreds of millions of Amazon customers? Do you like to own end-to-end business problems and 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 for our Promotion Science Team at Amazon. At Promotion Science, our vision is simple - we want to run the best promotions in the industry. Through technology, we seek to constantly increase the quality of promotions to delight customers and provide ROI for Selling Partners. In this role, you will build scalable, accurate, and interpretable statistical and machine learning models to forecast and rank promotions as well as measure their incremental business values. Key job responsibilities As a Senior Applied Scientist on this team, you will: • Tackle ambiguous and complex science problems in time-series analysis, causal inference, statistical modeling, recommendation, and reinforcement learning. • Build statistical and machine learning models, perform proof-of-concept, experiment, optimize, and deploy your models into production. • Run A/B experiments, gather data, and perform statistical tests. * Interpret and present scientific results to business stakeholders and leadership and drive decision making. • Guide and mentor other scientists on the team.
CA, BC, Vancouver
Alexa Daily Essentials is hiring a Research Scientist to research and implement large language model innovations to enhance Alexa's language understanding, knowledge representation, reasoning and generation capabilities. The Alexa Daily Essentials team delivers experiences critical to how customers interact with Alexa as part of daily life. We drive over 40 billion+ actions annually across 60 million+ monthly customers, who engage with our products across experiences connected to Timers, Alarms, Calendars, Food, and News. Our experiences include critical time saving techniques, ad-supported news audio and video, and in-depth kitchen guidance aimed at serving the needs of the family from sunset to sundown. Our upcoming launches are at the forefront of innovation, delivering step-function improvements in experiences that stretch across the customer journey, and new AI technologies that will enable customers to send Alexa information for future recall and conversation. We collaborate closely with partners such as Amazon.com, Whole Foods, Spotify, CNN, Fox, NPR, BBC, Discovery, and Food Network to deliver our vision. If you are passionate about redefining the personal assistant experience and leveraging innovative technology to improve daily life, we’d love to hear from you. This is an opportunity to make a tangible impact at the heart of the Alexa ecosystem. As a Research Scientist, you will advance state of the art techniques in ML and LLM, and work closely with product and engineering teams to build the next generation of the Alexa smart assistant. Key job responsibilities - Rapidly prototype ML/LLM solutions, evaluate feasibility, and drive projects to production deployment - Continuously monitor and improve model performance through retraining, parameter tuning, and architecture refinements - Develop new training and inference techniques to improve model performance - Work cross-functionally across engineering, product, and business teams to understand customer needs, scope science work, and drive science solutions from conception to customer delivery - Research and develop LLM innovations, and lead paper publications. - Code proficiently in Python (required) and Java (preferred); optimize systems for high performance at scale; contribute code directly into production services - Innovate and develop science and engineering solutions that optimize team operations and increase team effectiveness. - Clearly communicate complex technical concepts to non-technical stakeholders and leadership
CA, ON, Toronto
Looking for your next challenge? North America Sort Centers (NASC) are experiencing growth and looking for a skilled, highly motivated Data Scientist to join the NASC Engineering Data, Product and Simulation Team. The Sort Center network is the critical Middle-Mile solution in the Amazon Transportation Services (ATS) group, linking Fulfillment Centers to the Last Mile. The experience of our customers is dependent on our ability to efficiently execute volume flow through the middle-mile network. Key job responsibilities The Data Scientist will design and implement solutions to address complex business questions using simulation. In this role, you will apply advanced analysis techniques and statistical concepts to draw insights from massive datasets, and create intuitive simulations and data visualizations. You can contribute to each layer of a data solution – you work closely with process design engineers, business intelligence engineers and technical product managers to obtain relevant datasets and create simulation models, and review key results with business leaders and stakeholders. Your work exhibits a balance between scientific validity and business practicality. On this team, you will have a large impact on the entire NASC organization, with lots of opportunity to learn and grow within the NASC Engineering team. This role will be the first dedicated simulation expert, so you will have an exceptional opportunity to define and drive vision for simulation best practices on our team. To be successful in this role, you must be able to turn ambiguous business questions into clearly defined problems, develop quantifiable metrics and deliver results that meet high standards of data quality, security, and privacy. About the team NASC Engineering’s Product and Analytics Team’s sole objective is to develop tools for under the roof simulation and optimization, supporting the needs of our internal and external stakeholders (i.e Process Design Engineering, NASC Engineering, ACES, Finance, Safety and Operations). We develop data science tools to evaluate what-if design and operations scenarios for new and existing sort centers to understand their robustness, stability, scalability, and cost-effectiveness. We conceptualize new data science solutions, using optimization and machine learning platforms, to analyze new and existing process, identify and reduce non-value added steps, and increase overall performance and rate. We work by interfacing with various functional teams to test and pilot new hardware/software solutions.
US, TX, Dallas
At AWS, we use Artificial Intelligence to be able to identify every need of a customer across all AWS services before they have to tell us about it, and then find and seamlessly connect them to the most appropriate resolution for their need, eventually fulfilling the vision of a self-healing cloud. We are looking for Machine Learning Scientists / Applied Scientists / ML Scientists with unfettered curiosity and drive to help build “best in the world” support (contact center) experience that customers will love! You will have an opportunity to lead, invent, and design tech that will directly impact every customer across all AWS services. We are building industry-leading technology that cuts across a wide range of ML techniques from Natural Language Processing to Deep Learning and Generative Artificial Intelligence. You will be a key driver in taking something from an idea to an experiment to a prototype and finally to a live production system. Our team packs a punch with principal level engineering, science, product, and leadership talent. We are a results focused team and you have the opportunity to lead and establish a culture for the big things to come. We combine the culture of a startup, the innovation and creativity of a R&D Lab, the work-life balance of a mature organization, and technical challenges at the scale of AWS. We offer a playground of opportunities for builders to build, have fun, and make history! Key job responsibilities - Deliver real world production systems at AWS scale. - Work closely with the business to understand the problem space, identify the opportunities and formulate the problems. - Use machine learning, data mining, statistical techniques, Generative AI and others to create actionable, meaningful, and scalable solutions for the business problems. - Analyze and extract relevant information from large amounts of data and derive useful insights. - Work with software engineering teams to deliver production systems with your ML models - Establish scalable, efficient, automated processes for large scale data analyses, model development, model validation and model implementation A day in the life About the team We are working to achieve our business goals by deriving insights from a wealth of datasets like the AWS service metrics and logs, chat and call audio logs, email transcripts, support agent and support case data, customers context and sentiment, and AWS knowledge articles, tools and workflows. Diverse Experiences AWS 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 AWS? Amazon Web Services (AWS) is the world’s most comprehensive and broadly adopted cloud platform. We pioneered cloud computing and never stopped innovating — that’s why customers from the most successful startups to Global 500 companies trust our robust suite of products and services to power their businesses. Inclusive Team Culture Here at AWS, it’s in our nature to learn and be curious. Our employee-led affinity groups foster a culture of inclusion that empower us to be proud of our differences. Ongoing events and learning experiences, including our Conversations on Race and Ethnicity (CORE) and AmazeCon (gender diversity) conferences, inspire us to never stop embracing our uniqueness. Mentorship & 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, mentorship 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 we strive for flexibility as part of our working culture. When we feel supported in the workplace and at home, there’s nothing we can’t achieve in the cloud. Hybrid Work We value innovation and recognize this sometimes requires uninterrupted time to focus on a build. We also value in-person collaboration and time spent face-to-face. Our team affords employees options to work in the office every day or in a flexible, hybrid work model near one of our U.S. Amazon offices.